"""Task 020 Solver — 4-fold rotational symmetry fill (266/266 PASS).

Rule:
1. Input has scattered non-zero cells on a mostly-zero 10x10 grid
2. Find the centroid of all non-zero cells
3. For each non-zero cell, rotate it 90°, 180°, 270° around the centroid
4. Fill any resulting 0-cell position with the rotated color
5. Output = input with symmetry-completed cells

ONNX viability:
- Center varies per example (148 unique positions, non-integer)
- Requires data-dependent coordinate transformation
- Difficult to implement compactly — needs dynamic Gather with computed indices
- Base: 245 nodes, static score ~13.6
- Need <80 nodes to meaningfully beat base → HARD
"""
import numpy as np
import json


def solve_020(inp_grid):
    """4-fold rotational symmetry fill. 266/266 PASS."""
    inp = np.array(inp_grid)
    h, w = inp.shape
    out = inp.copy()

    nz = np.argwhere(inp != 0)
    if len(nz) == 0:
        return out

    center_r = nz[:, 0].mean()
    center_c = nz[:, 1].mean()
    cr, cc = center_r, center_c

    for r, c in nz:
        color = inp[r, c]
        dr, dc = r - cr, c - cc
        for rdr, rdc in [(dr, dc), (-dc, dr), (-dr, -dc), (dc, -dr)]:
            nr, nc = cr + rdr, cc + rdc
            nri, nci = int(round(nr)), int(round(nc))
            if 0 <= nri < h and 0 <= nci < w:
                if out[nri, nci] == 0:
                    out[nri, nci] = color

    return out


if __name__ == '__main__':
    from pathlib import Path
    task_data = Path(__file__).parent.parent / 'task-data'
    if not task_data.exists():
        task_data = Path('task-data')
    with open(task_data / 'task020.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 np.array_equal(solve_020(ex['input']), np.array(ex['output'])))
    print(f'task020: {right}/{len(all_ex)} PASS')