"""Task 076 Solver — Template rotation + decoration placement (194/266).

Rule (partially cracked):
1. Input has 3 'objects' made of color 4 (L-shaped or straight bars, 5 cells each)
2. One object (template) has decorations: colors 1, 3 adjacent to the bar, plus a 2 marker
3. Other objects (incomplete) have only the 4-bar + adjacent 2 marker
4. Apply template decorations to incomplete objects, rotated/reflected to match orientation
5. Orientation determined by matching the 4-cell shape using 8 transforms (4 rots + 4 reflections)

Status: 194/266 — 67 failures with correct shape match but wrong color placement.
Possible issues:
- Some examples may need color 1↔3 swap for certain transforms
- 2-marker adjacency detection may miss diagonal cases
- Template identification may be wrong in some examples

ONNX: Would be viable IF solver reaches 266/266. The transform is just a Gather with 
precomputed indices for 8 possible rotations. ~40-60 nodes possible.
"""
import numpy as np
import json
from scipy.ndimage import label


def solve_076(inp_grid):
    inp = np.array(inp_grid)
    H, W = inp.shape
    out = inp.copy()
    
    mask4 = (inp == 4).astype(int)
    labeled, n = label(mask4)
    
    comps = []
    for cid in range(1, n+1):
        cells = set(zip(*np.where(labeled == cid)))
        adj2, adj1, adj3 = set(), set(), set()
        for r,c in cells:
            for dr,dc in [(-1,0),(1,0),(0,-1),(0,1)]:
                nr,nc = r+dr,c+dc
                if 0<=nr<H and 0<=nc<W:
                    if inp[nr,nc]==2: adj2.add((nr,nc))
                    elif inp[nr,nc]==1: adj1.add((nr,nc))
                    elif inp[nr,nc]==3: adj3.add((nr,nc))
        comps.append({"cells": cells, "adj2": adj2, "adj1": adj1, "adj3": adj3,
                      "is_template": len(adj1)>0 and len(adj3)>0})
    
    tmpl = next((c for c in comps if c["is_template"]), None)
    if not tmpl or not tmpl["adj2"]:
        return out.tolist()
    
    t_2 = list(tmpl["adj2"])[0]
    t_4_rel = set((r-t_2[0], c-t_2[1]) for r,c in tmpl["cells"])
    t_1_rel = [(r-t_2[0], c-t_2[1]) for r,c in tmpl["adj1"]]
    t_3_rel = [(r-t_2[0], c-t_2[1]) for r,c in tmpl["adj3"]]
    
    transforms = [
        lambda dr,dc: (dr, dc),
        lambda dr,dc: (dc, -dr),
        lambda dr,dc: (-dr, -dc),
        lambda dr,dc: (-dc, dr),
        lambda dr,dc: (dr, -dc),
        lambda dr,dc: (-dr, dc),
        lambda dr,dc: (dc, dr),
        lambda dr,dc: (-dc, -dr),
    ]
    
    for comp in comps:
        if comp["is_template"] or not comp["adj2"]:
            continue
        i_2 = list(comp["adj2"])[0]
        i_4_rel = set((r-i_2[0], c-i_2[1]) for r,c in comp["cells"])
        
        best_tf = None
        for tf in transforms:
            if set(tf(dr,dc) for dr,dc in t_4_rel) == i_4_rel:
                best_tf = tf
                break
        if best_tf is None:
            continue
        
        for dr, dc in t_1_rel:
            new_dr, new_dc = best_tf(dr, dc)
            nr, nc = i_2[0] + new_dr, i_2[1] + new_dc
            if 0 <= nr < H and 0 <= nc < W and out[nr, nc] == 0:
                out[nr, nc] = 1
        for dr, dc in t_3_rel:
            new_dr, new_dc = best_tf(dr, dc)
            nr, nc = i_2[0] + new_dr, i_2[1] + new_dc
            if 0 <= nr < H and 0 <= nc < W and out[nr, nc] == 0:
                out[nr, nc] = 3
    
    return out.tolist()


if __name__ == "__main__":
    with open("task-data/task076.json") as f:
        data = json.load(f)
    all_ex = data["train"] + data["test"] + data.get("arc-gen", [])
    right = sum(1 for ex in all_ex if solve_076(ex["input"]) == ex["output"])
    print(f"task076: {right}/{len(all_ex)} PASS")