import math
import cv2
import numpy as np

def get_horizontal_overlap(bbox1, bbox2):
    x_left = max(bbox1[0], bbox2[0])
    x_right = min(bbox1[2], bbox2[2])
    return max(0, x_right - x_left)

def extract_barriers(img_path, regions):
    """Extracts horizontal and vertical barrier lines using OpenCV."""
    if not img_path:
        return [], []
        
    import statistics
    text_widths = [(r['bbox'][2] - r['bbox'][0]) for r in regions if r['type'] == 'text']
    median_col_width = statistics.median(text_widths) if text_widths else 300
    
    img = cv2.imread(str(img_path))
    if img is None:
        return [], []
        
    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    _, thresh = cv2.threshold(gray, 200, 255, cv2.THRESH_BINARY_INV)
    
    line_length = int(median_col_width * 0.5)
    h_kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (line_length, 1))
    h_lines_raw = cv2.morphologyEx(thresh, cv2.MORPH_OPEN, h_kernel, iterations=2)
    
    v_kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (1, line_length))
    v_lines_raw = cv2.morphologyEx(thresh, cv2.MORPH_OPEN, v_kernel, iterations=2)
    
    h_contours_raw, _ = cv2.findContours(h_lines_raw, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    v_contours_raw, _ = cv2.findContours(v_lines_raw, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    
    h_lines = []
    for cnt in h_contours_raw:
        x, y, w, h = cv2.boundingRect(cnt)
        # Ignore short decorative lines (caption underlines) which are <= 0.6 col
        if h <= 15 and w >= (median_col_width * 0.6):
            h_lines.append((x, y, x + w, y + h))
            
    v_lines = []
    for cnt in v_contours_raw:
        x, y, w, h = cv2.boundingRect(cnt)
        if w <= 15:  # Thin lines only
            v_lines.append((x, y, x + w, y + h))
            
    return h_lines, v_lines

def is_blocked_by_horizontal_barrier(top_y, bottom_y, item_left, item_right, h_lines):
    """
    Checks if a horizontal line sits between top_y and bottom_y and spans across the item.
    """
    item_width = item_right - item_left
    tolerance = item_width * 0.2
    
    for (lx1, ly1, lx2, ly2) in h_lines:
        line_y = ly1
        if top_y < line_y < bottom_y:
            # Does it span across the item?
            if lx1 <= (item_left + tolerance) and lx2 >= (item_right - tolerance):
                return True
    return False

def is_blocked_by_vertical_barrier(t_bbox, item_bbox, v_lines):
    """
    Checks if a vertical line sits between the centers of two regions.
    """
    t_cx = (t_bbox[0] + t_bbox[2]) / 2
    item_cx = (item_bbox[0] + item_bbox[2]) / 2
    
    left_x = min(t_cx, item_cx)
    right_x = max(t_cx, item_cx)
    
    # Only check if they are in different columns
    if right_x - left_x < 100:
        return False
        
    y_overlap_top = max(t_bbox[1], item_bbox[1])
    y_overlap_bottom = min(t_bbox[3], item_bbox[3])
    
    for (vx1, vy1, vx2, vy2) in v_lines:
        line_x = vx1
        if left_x < line_x < right_x:
            # Does the line overlap them vertically?
            # A true vertical column separator will be very tall
            if vy1 < y_overlap_bottom and vy2 > y_overlap_top:
                return True
    return False

def group_surya_regions_geometrically(regions, img_path=None):
    if not regions:
        return []
        
    # Extract barriers using OpenCV
    h_lines, v_lines = extract_barriers(img_path, regions)
    if img_path:
        print(f"  [grouper] Extracted {len(h_lines)} horizontal and {len(v_lines)} vertical OpenCV barriers")

    titles = [r for r in regions if r['type'] in ('doc_title', 'paragraph_title')]
    non_titles = [r for r in regions if r['type'] not in ('doc_title', 'paragraph_title')]

    region_by_id = {r['id']: r for r in regions}
    
    if not titles:
        member_ids = [r['id'] for r in regions]
        min_x = min(r['bbox'][0] for r in regions)
        min_y = min(r['bbox'][1] for r in regions)
        max_x = max(r['bbox'][2] for r in regions)
        max_y = max(r['bbox'][3] for r in regions)
        return [{
            "article_id": 1,
            "title_region": None,
            "member_regions": member_ids,
            "bbox": [min_x, min_y, max_x, max_y],
            "grouped_by": "geometric_fallback"
        }]
        
    article_groups = {t['id']: [t] for t in titles}
    
    for item in non_titles:
        item_bbox = item['bbox']
        item_cy = (item_bbox[1] + item_bbox[3]) / 2
        item_cx = (item_bbox[0] + item_bbox[2]) / 2
        item_width = item_bbox[2] - item_bbox[0]
        
        best_title = None
        
        max_dist_above = 2500 if item['type'] == 'image' else 1500
        candidates_above = []
        for t in titles:
            t_bbox = t['bbox']
            if t_bbox[1] <= item_cy:
                # Barrier Check
                if is_blocked_by_horizontal_barrier(t_bbox[3], item_bbox[1], item_bbox[0], item_bbox[2], h_lines):
                    continue
                if is_blocked_by_vertical_barrier(t_bbox, item_bbox, v_lines):
                    continue
                    
                overlap = get_horizontal_overlap(item_bbox, t_bbox)
                t_width = t_bbox[2] - t_bbox[0]
                t_cx = (t_bbox[0] + t_bbox[2]) / 2
                
                if overlap > (item_width * 0.1) or overlap > (t_width * 0.1) or abs(item_cx - t_cx) < 150:
                    distance = item_bbox[1] - t_bbox[3]
                    if distance < max_dist_above:
                        candidates_above.append(t)
                    
        if candidates_above:
            def _above_score(t):
                vert  = item_bbox[1] - t['bbox'][3]
                horiz = abs(item_cx - (t['bbox'][0] + t['bbox'][2]) / 2)
                return vert + horiz * 0.8
            best_title = min(candidates_above, key=_above_score)
        else:
            candidates_below = []
            for t in titles:
                t_bbox = t['bbox']
                if t_bbox[1] > item_cy:
                    # Barrier Check
                    if is_blocked_by_horizontal_barrier(item_bbox[3], t_bbox[1], item_bbox[0], item_bbox[2], h_lines):
                        continue
                    if is_blocked_by_vertical_barrier(t_bbox, item_bbox, v_lines):
                        continue

                    overlap = get_horizontal_overlap(item_bbox, t_bbox)
                    t_width = t_bbox[2] - t_bbox[0]
                    t_cx = (t_bbox[0] + t_bbox[2]) / 2
                    
                    if overlap > (item_width * 0.1) or overlap > (t_width * 0.1) or abs(item_cx - t_cx) < 150:
                        distance = t_bbox[1] - item_bbox[3]
                        if distance < 800:
                            candidates_below.append(t)
            
            if candidates_below:
                def _below_score(t):
                    vert  = t['bbox'][1] - item_bbox[3]
                    horiz = abs(item_cx - (t['bbox'][0] + t['bbox'][2]) / 2)
                    return vert + horiz * 0.8
                best_title = min(candidates_below, key=_below_score)
            else:
                # Absolute fallback
                valid_closest = []
                for t in titles:
                    if is_blocked_by_horizontal_barrier(min(item_bbox[3], t['bbox'][3]), max(item_bbox[1], t['bbox'][1]), item_bbox[0], item_bbox[2], h_lines):
                        continue
                    if is_blocked_by_vertical_barrier(t['bbox'], item_bbox, v_lines):
                        continue
                    valid_closest.append(t)
                
                if valid_closest:
                    closest_t = min(valid_closest, key=lambda t: math.hypot(
                        item_cx - ((t['bbox'][0] + t['bbox'][2]) / 2),
                        item_cy - ((t['bbox'][1] + t['bbox'][3]) / 2)
                    ))
                    dist = math.hypot(item_cx - ((closest_t['bbox'][0] + closest_t['bbox'][2]) / 2),
                                      item_cy - ((closest_t['bbox'][1] + closest_t['bbox'][3]) / 2))
                    if dist < 1500:
                        best_title = closest_t
                
        if best_title:
            article_groups[best_title['id']].append(item)
        else:
            virtual_id = f"orphan_{item['id']}"
            article_groups[virtual_id] = [item]
            
    # Format output
    articles = []
    sorted_titles = sorted(titles, key=lambda t: (t['bbox'][1], t['bbox'][0]))
    
    for idx, (title_id, group_regions) in enumerate(article_groups.items(), start=1):
        if not group_regions:
            continue
            
        min_x = min(r['bbox'][0] for r in group_regions)
        min_y = min(r['bbox'][1] for r in group_regions)
        max_x = max(r['bbox'][2] for r in group_regions)
        max_y = max(r['bbox'][3] for r in group_regions)
        
        is_orphan = str(title_id).startswith("orphan_")
        
        articles.append({
            "article_id": idx,
            "title_region": None if is_orphan else title_id,
            "member_regions": [r['id'] for r in group_regions],
            "bbox": [min_x, min_y, max_x, max_y],
            "grouped_by": "geometric_orphan" if is_orphan else "geometric"
        })
        
    orphan_articles = [a for a in articles if a['grouped_by'] == 'geometric_orphan']
    valid_articles = [a for a in articles if a['grouped_by'] == 'geometric']
    
    orphan_articles.sort(key=lambda a: a['bbox'][1])
    
    merged_orphans = []
    current_orphan = None
    
    for o in orphan_articles:
        if not current_orphan:
            current_orphan = o
            continue
            
        overlap = get_horizontal_overlap(current_orphan['bbox'], o['bbox'])
        vertical_dist = o['bbox'][1] - current_orphan['bbox'][3]
        
        if overlap > 0 and vertical_dist < 400:
            # Check barrier before merging orphans
            if not is_blocked_by_horizontal_barrier(current_orphan['bbox'][3], o['bbox'][1], current_orphan['bbox'][0], current_orphan['bbox'][2], h_lines):
                current_orphan['member_regions'].extend(o['member_regions'])
                current_orphan['bbox'][0] = min(current_orphan['bbox'][0], o['bbox'][0])
                current_orphan['bbox'][1] = min(current_orphan['bbox'][1], o['bbox'][1])
                current_orphan['bbox'][2] = max(current_orphan['bbox'][2], o['bbox'][2])
                current_orphan['bbox'][3] = max(current_orphan['bbox'][3], o['bbox'][3])
                continue
                
        merged_orphans.append(current_orphan)
        current_orphan = o
            
    if current_orphan:
        merged_orphans.append(current_orphan)
        
    # Orphan Adoption: If an orphan sits directly below a valid article with high overlap and small gap, adopt it (ignoring photo borders)
    adopted_orphans = set()
    for o in merged_orphans:
        best_parent = None
        min_gap = float('inf')
        for a in valid_articles:
            overlap = get_horizontal_overlap(a['bbox'], o['bbox'])
            gap = o['bbox'][1] - a['bbox'][3]
            if overlap > 0 and 0 <= gap < 200:
                width_overlap_ratio = overlap / min((a['bbox'][2] - a['bbox'][0]), (o['bbox'][2] - o['bbox'][0]))
                if width_overlap_ratio > 0.8:
                    if gap < min_gap:
                        min_gap = gap
                        best_parent = a
        if best_parent:
            best_parent['member_regions'].extend(o['member_regions'])
            best_parent['bbox'][0] = min(best_parent['bbox'][0], o['bbox'][0])
            best_parent['bbox'][1] = min(best_parent['bbox'][1], o['bbox'][1])
            best_parent['bbox'][2] = max(best_parent['bbox'][2], o['bbox'][2])
            best_parent['bbox'][3] = max(best_parent['bbox'][3], o['bbox'][3])
            adopted_orphans.add(o['article_id'])
            
    final_orphans_pass_1 = [o for o in merged_orphans if o['article_id'] not in adopted_orphans]
    
    # Strict Orphan Photo Binding: An image may never stand alone.
    remaining_orphans = []
    for o in final_orphans_pass_1:
        has_image = any(region_by_id[rid]['type'] == 'image' for rid in o['member_regions'] if rid in region_by_id)
        has_title = any(region_by_id[rid]['type'] in ('doc_title', 'paragraph_title') for rid in o['member_regions'] if rid in region_by_id)
        
        if has_image and not has_title:
            best_article = None
            min_dist = float('inf')
            
            for a in valid_articles:
                dist = min(abs(a['bbox'][3] - o['bbox'][1]), abs(o['bbox'][3] - a['bbox'][1]))
                overlap = get_horizontal_overlap(a['bbox'], o['bbox'])
                if overlap > 0 and dist < min_dist:
                    min_dist = dist
                    best_article = a
            
            if best_article:
                best_article['member_regions'].extend(o['member_regions'])
                best_article['bbox'][0] = min(best_article['bbox'][0], o['bbox'][0])
                best_article['bbox'][1] = min(best_article['bbox'][1], o['bbox'][1])
                best_article['bbox'][2] = max(best_article['bbox'][2], o['bbox'][2])
                best_article['bbox'][3] = max(best_article['bbox'][3], o['bbox'][3])
                continue # successfully bound
                
        remaining_orphans.append(o)
        
    final_articles = valid_articles + remaining_orphans

    MAX_SPLIT_HEIGHT  = 2500
    MIN_MEMBERS_SPLIT = 5
    MIN_SPLIT_GAP     = 30
    MIN_LOWER_HEIGHT  = 400

    split_result = []
    for art in final_articles:
        _, art_t, _, art_b = art['bbox']
        if (art_b - art_t) <= MAX_SPLIT_HEIGHT or len(art['member_regions']) < MIN_MEMBERS_SPLIT:
            split_result.append(art)
            continue

        members = [region_by_id[rid] for rid in art['member_regions'] if rid in region_by_id]
        members.sort(key=lambda reg: reg['bbox'][1])

        best_gap  = MIN_SPLIT_GAP
        split_idx = -1
        max_bottom_so_far = 0
        for i in range(len(members) - 1):
            cur_bottom = members[i]['bbox'][3]
            max_bottom_so_far = max(max_bottom_so_far, cur_bottom)
            
            next_top = members[i + 1]['bbox'][1]
            gap = next_top - max_bottom_so_far
            
            # If there is a barrier line crossing this gap, artificially inflate the gap score to force a split here
            is_barrier = is_blocked_by_horizontal_barrier(max_bottom_so_far, next_top, members[i]['bbox'][0], members[i]['bbox'][2], h_lines)
            effective_gap = gap + 10000 if is_barrier else gap
            
            if effective_gap > best_gap:
                lower_group = members[i + 1:]
                if not any(reg['type'] == 'image' for reg in lower_group):
                    best_gap  = effective_gap
                    split_idx = i

        if split_idx < 0:
            split_result.append(art)
            continue

        upper = members[:split_idx + 1]
        lower = members[split_idx + 1:]

        lower_bottom = max(reg['bbox'][3] for reg in lower)
        if lower_bottom - lower[0]['bbox'][1] < MIN_LOWER_HEIGHT:
            split_result.append(art)
            continue

        def _bbox(regs):
            return [min(reg['bbox'][0] for reg in regs), min(reg['bbox'][1] for reg in regs),
                    max(reg['bbox'][2] for reg in regs), max(reg['bbox'][3] for reg in regs)]

        real_gap = members[split_idx+1]['bbox'][1] - max(r['bbox'][3] for r in upper)
        print(f"  [split] Article {art['article_id']}: split at gap={real_gap}px after region "
              f"{upper[-1]['id']} (upper {len(upper)} regions, lower {len(lower)} regions)")

        upper_ids = [reg['id'] for reg in upper]
        lower_ids = [reg['id'] for reg in lower]

        split_result.append({
            'article_id':     art['article_id'],
            'title_region':   art['title_region'] if art['title_region'] in upper_ids else None,
            'member_regions': upper_ids,
            'bbox':           _bbox(upper),
            'grouped_by':     art['grouped_by'],
        })
        split_result.append({
            'article_id':     None,
            'title_region':   art['title_region'] if art['title_region'] in lower_ids else None,
            'member_regions': lower_ids,
            'bbox':           _bbox(lower),
            'grouped_by':     'geometric_split',
        })

    for idx, a in enumerate(split_result, start=1):
        a['article_id'] = idx

    return split_result