import numpy as np
from scipy.spatial import distance

class ArticleSegmenter:
    def __init__(self, col_gap_ratio=0.02):
        self.col_gap_ratio = col_gap_ratio

    def _vertical_overlap(self, a, b):
        top = max(a[1], b[1]); bot = min(a[3], b[3])
        return max(0, bot - top)

    def _horizontal_gap(self, a, b):
        return b[0] - a[2]

    def segment(self, blocks, page_width, page_height):
        # 1. Remove ads
        content = [b for b in blocks if not b["is_ad"] and b["type"] != "Advertisement"]

        # 2. Sort reading order: column-aware (top-to-bottom within columns)
        content = self._reading_order(content, page_width)

        # 3. Seed articles at each Title block
        articles = []
        current = None
        for b in content:
            if b["type"] == "Title":
                if current:
                    articles.append(current)
                current = {"blocks": [b], "title_block": b}
            else:
                if current is None:
                    current = {"blocks": [b], "title_block": None}
                else:
                    if self._belongs(current, b, page_width):
                        current["blocks"].append(b)
                    else:
                        articles.append(current)
                        current = {"blocks": [b], "title_block": None}
        if current:
            articles.append(current)

        # 4. Merge cross-column continuation of same article
        articles = self._merge_columns(articles, page_width)

        # 5. Compute union bounding box per article
        for a in articles:
            a["bbox"] = self._union_bbox(a["blocks"])
            a["confidence"] = float(np.mean([blk["score"] for blk in a["blocks"]]))
        return articles

    def _reading_order(self, blocks, page_width, n_cols_guess=None):
        if not blocks:
            return blocks
        xs = np.array([b["bbox"][0] for b in blocks]).reshape(-1, 1)
        # Estimate columns via 1D clustering on x-start
        from sklearn.cluster import KMeans
        k = self._estimate_columns(xs, page_width)
        km = KMeans(n_clusters=k, n_init=5).fit(xs)
        for b, lbl in zip(blocks, km.labels_):
            b["col"] = int(lbl)
        # Order columns left→right, blocks top→bottom
        col_order = np.argsort([xs[km.labels_ == c].mean() for c in range(k)])
        ordered = []
        for c in col_order:
            col_blocks = [b for b in blocks if b["col"] == c]
            col_blocks.sort(key=lambda b: b["bbox"][1])
            ordered.extend(col_blocks)
        return ordered

    def _estimate_columns(self, xs, page_width):
        from sklearn.metrics import silhouette_score
        best_k, best_score = 1, -1
        for k in range(1, min(7, len(xs))):
            if k == 1:
                continue
            from sklearn.cluster import KMeans
            labels = KMeans(n_clusters=k, n_init=3).fit_predict(xs)
            try:
                s = silhouette_score(xs, labels)
                if s > best_score:
                    best_score, best_k = s, k
            except Exception:
                pass
        return max(best_k, 1)

    def _belongs(self, article, block, page_width):
        """Block belongs to current article if vertically continuous & aligned."""
        last = article["blocks"][-1]["bbox"]
        cur = block["bbox"]
        same_col = abs(last[0] - cur[0]) < 0.05 * page_width
        gap = cur[1] - last[3]
        return same_col and gap < 0.025 * page_width

    def _merge_columns(self, articles, page_width):
        # Merge article fragments whose headline spans wide / text flows to next col
        # Simplified: if an article has no title and sits at top of next column
        # right after a titled article, merge. Production: use LayoutLMv3 relations.
        return articles

    def _union_bbox(self, blocks):
        xs1 = min(b["bbox"][0] for b in blocks)
        ys1 = min(b["bbox"][1] for b in blocks)
        xs2 = max(b["bbox"][2] for b in blocks)
        ys2 = max(b["bbox"][3] for b in blocks)
        return [xs1, ys1, xs2 - xs1, ys2 - ys1]