import layoutparser as lp
import numpy as np

CLASS_MAP = {
    0: "Text", 1: "Title", 2: "List",
    3: "Table", 4: "Figure", 5: "Advertisement",
}

class LayoutDetector:
    def __init__(self):
        try:
            # Detectron2 PubLayNet — swap with custom DocLayNet/YOLOv11 weights
            self.model = lp.Detectron2LayoutModel(
                config_path="lp://PubLayNet/mask_rcnn_X_101_32x8d_FPN_3x/config",
                extra_config=["MODEL.ROI_HEADS.SCORE_THRESH_TEST", 0.55],
                label_map={0:"Text",1:"Title",2:"List",3:"Table",4:"Figure"},
            )
        except Exception:
            self.model = None
        # Optional secondary ad classifier (custom-trained)
        self.ad_classifier = AdClassifier()

    def detect(self, image: np.ndarray):
        if self.model is None:
            # Fallback to single block if detectron2 isn't installed
            h, w = image.shape[:2]
            return [{
                "type": "Text",
                "bbox": [0, 0, w, h],
                "score": 1.0,
                "is_ad": False
            }]
            
        layout = self.model.detect(image)
        blocks = []
        for b in layout:
            box = b.coordinates  # (x1,y1,x2,y2)
            block = {
                "type": b.type,
                "bbox": [int(c) for c in box],
                "score": float(b.score),
            }
            block["is_ad"] = self.ad_classifier.predict(image, box)
            blocks.append(block)
        return blocks


class AdClassifier:
    """
    Ads rarely match column grids and contain logos/price patterns.
    Train a lightweight CNN (MobileNet) on labeled ad/non-ad crops,
    OR use heuristics + LLM verification. Stub below.
    """
    def predict(self, image, box) -> bool:
        # Heuristic placeholder: replace with trained model
        return False