import math

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 group_surya_regions_geometrically(regions):
    """
    Groups regions spatially.
    Instead of a linear reading order, this associates text and images 
    with the headline that is directly above them in the same column.
    """
    if not regions:
        return []
        
    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')]
    
    # If no titles exist, return everything as one article
    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
        
        # -------------------------------------------------------------
        # Rule 1: Find titles ABOVE the item that share the same column
        # -------------------------------------------------------------
        candidates_above = []
        for t in titles:
            t_bbox = t['bbox']
            if t_bbox[1] <= item_cy:
                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 < 1500:  # MAX DISTANCE ABOVE
                        candidates_above.append(t)
                    
        if candidates_above:
            best_title = min(candidates_above, key=lambda t: abs(item_bbox[1] - t['bbox'][3]))
        else:
            # -------------------------------------------------------------
            # Rule 2: If no title is above, find titles BELOW the item
            # -------------------------------------------------------------
            candidates_below = []
            for t in titles:
                t_bbox = t['bbox']
                if t_bbox[1] > item_cy:
                    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:  # MAX DISTANCE BELOW (stricter because text usually flows down, not up)
                            candidates_below.append(t)
            
            if candidates_below:
                best_title = min(candidates_below, key=lambda t: abs(t['bbox'][1] - item_bbox[3]))
            else:
                # -------------------------------------------------------------
                # Rule 3: Absolute fallback. No column overlap at all.
                # Find the absolute closest title via Euclidean distance.
                # -------------------------------------------------------------
                closest_t = min(titles, key=lambda t: math.hypot(
                    item_cx - ((t['bbox'][0] + t['bbox'][2]) / 2),
                    item_cy - ((t['bbox'][1] + t['bbox'][3]) / 2)
                ))
                # Only attach if it's reasonably close, else leave as orphan
                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:
            # Treat as an orphan (create a new article group for it, or group orphans together)
            # For simplicity, we'll assign it a virtual title ID based on its own ID so it becomes a standalone article
            virtual_id = f"orphan_{item['id']}"
            article_groups[virtual_id] = [item]
            
    # Format output
    articles = []
    # Sort groups top-to-bottom, left-to-right based on title position
    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)
        
        # Determine if it's an orphan group
        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"
        })
        
    # We should merge orphans that are vertically adjacent into the same orphan group to prevent shattering
    # Simple vertical merge for orphans
    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']
    
    # Sort orphans top-to-bottom
    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
            
        # Check if they overlap horizontally and are close vertically
        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:
            # Merge
            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])
        else:
            merged_orphans.append(current_orphan)
            current_orphan = o
            
    if current_orphan:
        merged_orphans.append(current_orphan)
        
    final_articles = valid_articles + merged_orphans
    
    # Reassign IDs
    for idx, a in enumerate(final_articles, start=1):
        a['article_id'] = idx
        
    return final_articles
