"""Task 153 Solver — Complementary shape tiling (265/265 PASS).

Rule:
1. Input has two colored shapes on bg=0 in a 10x10 grid
2. Output is always 3x3
3. Both shapes fit within a 3x3 bounding box
4. They are COMPLEMENTARY: placed together they exactly fill a 3x3 grid
5. Find the unique placement of both shapes that tiles the 3x3 grid perfectly

Base model: 14,744 nodes, score 4.91 pts
Potential optimized: ~50 nodes, score 13+ pts
Gain: +8 pts (ENOUGH FOR BRONZE FROM V20!)

Pattern shapes observed: 2x3 (146), 2x2 (136), 3x3 (136), 3x2 (112)
All fit within 3x3.

ONNX approach:
- Extract each color's binary pattern (crop to bbox)
- Try all placements in 3x3 (max 4 per shape: offsets (0,0),(0,1),(1,0),(1,1))
- Check which pair of placements sums to all-ones (complementary)
- Output the valid tiling with appropriate colors
"""
import json
import numpy as np
from pathlib import Path


def solve_153(inp_grid):
    inp = np.array(inp_grid)
    colors = sorted(set(inp.flatten()) - {0})
    if len(colors) != 2:
        return None

    # Extract patterns for both colors
    patterns = {}
    for c in colors:
        pos = np.argwhere(inp == c)
        r_min, c_min = pos.min(axis=0)
        r_max, c_max = pos.max(axis=0)
        patterns[c] = (inp[r_min:r_max+1, c_min:c_max+1] == c).astype(int)

    c1, c2 = colors
    p1, p2 = patterns[c1], patterns[c2]
    h1, w1 = p1.shape
    h2, w2 = p2.shape

    # Try all placements that tile 3x3 perfectly
    for dr1 in range(4 - h1):
        for dc1 in range(4 - w1):
            placed1 = np.zeros((3, 3), dtype=int)
            placed1[dr1:dr1+h1, dc1:dc1+w1] = p1

            for dr2 in range(4 - h2):
                for dc2 in range(4 - w2):
                    placed2 = np.zeros((3, 3), dtype=int)
                    placed2[dr2:dr2+h2, dc2:dc2+w2] = p2

                    # Check complementarity
                    if np.all((placed1 + placed2) == 1):
                        output = np.zeros((3, 3), dtype=int)
                        output[placed1 == 1] = c1
                        output[placed2 == 1] = c2
                        return output.tolist()

    return None


def main():
    task_data_dir = Path(__file__).parent.parent / 'task-data'
    with open(task_data_dir / 'task153.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_153(ex['input'])
        if result is not None and result == ex['output']:
            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()