"""Task 077 Solver — Bridge-based cluster merging + bbox fill (266/266 PASS).

Rule:
1. Input has 3 colors: 0 (background), fg_color (most common non-0 non-2), and 2 (markers)
2. Find 8-connected clusters of color-2 cells
3. Merge clusters that share a 'bridge' cell (a fg cell 8-adjacent to cells from 2+ different clusters)
4. For each merged group: compute bounding box of all color-2 cells in the group
5. Fill all fg-color cells within the bounding box with color 4

NOTE: Base model already scores 12.83 pts with only 72 nodes.
An ONNX implementation of this algorithm requires 400+ nodes (30-iter MaxPool flood fill)
which would score WORSE than base. NOT a viable optimization target.

Validated: 266/266 PASS (all train + test + arc-gen)
"""
import numpy as np
import json
from scipy import ndimage
from collections import defaultdict


def solve_077(inp_grid):
    inp = np.array(inp_grid)
    out = inp.copy()
    h, w = inp.shape
    
    # Find fg color (most common non-0, non-2)
    colors, counts = np.unique(inp, return_counts=True)
    fg_candidates = [(c, cnt) for c, cnt in zip(colors, counts) if c != 0 and c != 2]
    if not fg_candidates:
        return out
    fg_color = int(max(fg_candidates, key=lambda x: x[1])[0])
    
    # Find 8-connected clusters of color-2 cells
    c2_mask = (inp == 2).astype(int)
    c2_labeled, n_c2 = ndimage.label(c2_mask, structure=np.ones((3, 3)))
    
    if n_c2 == 0:
        return out
    
    # Union-Find for merging clusters via bridge cells
    parent = list(range(n_c2 + 1))
    
    def find(x):
        while parent[x] != x:
            parent[x] = parent[parent[x]]
            x = parent[x]
        return x
    
    def union(a, b):
        a, b = find(a), find(b)
        if a != b:
            parent[a] = b
    
    # For each fg cell, check which 2-clusters it's 8-adjacent to
    fg_mask = (inp == fg_color)
    for r in range(h):
        for c in range(w):
            if not fg_mask[r, c]:
                continue
            touching = set()
            for dr in [-1, 0, 1]:
                for dc in [-1, 0, 1]:
                    if dr == 0 and dc == 0:
                        continue
                    nr, nc = r + dr, c + dc
                    if 0 <= nr < h and 0 <= nc < w and c2_labeled[nr, nc] > 0:
                        touching.add(c2_labeled[nr, nc])
            # If this fg cell touches 2+ clusters, merge them
            touching = list(touching)
            for i in range(1, len(touching)):
                union(touching[0], touching[i])
    
    # Group clusters by their root
    groups = defaultdict(list)
    for i in range(1, n_c2 + 1):
        groups[find(i)].append(i)
    
    # For each merged group, compute bbox of all 2s and fill fg cells
    for root, cluster_ids in groups.items():
        all_twos = []
        for cid in cluster_ids:
            all_twos.extend(np.argwhere(c2_labeled == cid).tolist())
        all_twos = np.array(all_twos)
        if len(all_twos) < 2:
            continue
        r_min, c_min = all_twos.min(axis=0)
        r_max, c_max = all_twos.max(axis=0)
        if r_min == r_max and c_min == c_max:
            continue
        for r in range(r_min, r_max + 1):
            for c in range(c_min, c_max + 1):
                if inp[r, c] == fg_color:
                    out[r, c] = 4
    
    return out


def main():
    from pathlib import Path
    task_data_dir = Path(__file__).parent.parent / 'task-data'
    if not task_data_dir.exists():
        task_data_dir = Path('task-data')
    
    with open(task_data_dir / 'task077.json') as f:
        data = json.load(f)
    
    all_examples = data['train'] + data['test'] + data.get('arc-gen', [])
    pass_count = 0
    fail_count = 0
    
    for i, ex in enumerate(all_examples):
        result = solve_077(ex['input'])
        expected = np.array(ex['output'])
        if np.array_equal(result, expected):
            pass_count += 1
        else:
            fail_count += 1
            if fail_count <= 5:
                print(f'Example {i}: FAIL')
    
    total = pass_count + fail_count
    print(f'\nResults: {pass_count} pass, {fail_count} fail (out of {total})')


if __name__ == '__main__':
    main()