"""Task 363 Solver — Pattern stamping (264/265 PASS, 1 train failure).

Rule (partially understood):
1. Input has colors 0 (background), 5 (foreground), 2 (pattern markers)
2. The color-2 cells form a small pattern (shape)
3. Stamp copies of this pattern at positions where ALL target cells are currently 0
4. When multiple valid positions conflict (overlap), prefer furthest from original centroid
5. Non-overlapping greedy selection

Status: 264/265 (train[1] fails — has 2 non-overlapping valid positions, 
        expected picks the one in same row as original but furthest-first picks wrong one)
        ALL 261 arc-gen PASS with simpler 'all zeros' approach (263/265 total)

Base model: 260 nodes, score 11.78
Potential if fully solved: Conv-based implementation ~20 nodes, score ~16+ (gain ~+4)

NOTE: Cannot build ONNX until 265/265 is achieved.
"""
import numpy as np
import json


def solve_363(inp_grid):
    """Furthest-first greedy non-overlapping stamp. 264/265."""
    inp = np.array(inp_grid)
    h, w = inp.shape
    out = inp.copy()
    c2_pos = np.argwhere(inp == 2)
    if len(c2_pos) == 0:
        return out
    r_min, c_min = c2_pos.min(axis=0)
    offsets = [(int(r - r_min), int(c - c_min)) for r, c in c2_pos]
    centroid = c2_pos.mean(axis=0)

    valid = []
    for r in range(h):
        for c in range(w):
            ok = True
            for dr, dc in offsets:
                nr, nc = r + dr, c + dc
                if nr < 0 or nr >= h or nc < 0 or nc >= w:
                    ok = False; break
                if inp[nr, nc] != 0:
                    ok = False; break
            if ok:
                stamp_center = np.array([(r+dr, c+dc) for dr, dc in offsets]).mean(axis=0)
                dist = np.sum(np.abs(stamp_center - centroid))
                cells = frozenset((r+dr, c+dc) for dr, dc in offsets)
                valid.append((r, c, dist, cells))

    # Greedy: furthest from original first
    valid.sort(key=lambda x: -x[2])
    used = set()
    for r, c, d, cells in valid:
        if cells & used:
            continue
        used |= cells
        for dr, dc in offsets:
            out[r + dr, c + dc] = 2

    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 / 'task363.json') as f:
        data = json.load(f)

    all_examples = data['train'] + data['test'] + data.get('arc-gen', [])
    pass_count = sum(1 for ex in all_examples
                     if np.array_equal(solve_363(ex['input']), np.array(ex['output'])))
    print(f'Results: {pass_count} pass, {len(all_examples)-pass_count} fail (out of {len(all_examples)})')


if __name__ == '__main__':
    main()