"""Task 34 Solver — Diagonal trail of 2x2 block (267/267 PASS).

Rule:
1. 9x9 grid with a 2x2 block containing colored cells and 2-markers
2. The FULL 2x2 bounding box is the repeating unit (all cells get the non-2 color)
3. For each 2-cell: direction = sign(2_position - block_center)
4. Slide block diagonally in each direction, accumulate all positions
5. If two 2-cells point in opposite directions → bidirectional trail

ONNX: VIABLE — 9x9 grid, iterative diagonal shifts (8 per direction), 
direction detection via centroid. Estimated 200-350 nodes.
Base: 87 nodes, 2847 params, score 13.42.
"""
import numpy as np
import json


def solve_034(inp_grid):
    inp = np.array(inp_grid)
    H, W = inp.shape
    out = np.zeros_like(inp)
    
    # Find the non-zero, non-2 color
    colors = set(np.unique(inp).tolist()) - {0, 2}
    if not colors:
        return out.tolist()
    color = int(colors.pop())
    
    # Find all non-zero cells (both colored and 2)
    all_cells = list(zip(*np.where(inp > 0)))
    if not all_cells:
        return out.tolist()
    
    # Bounding box of the shape
    rows = [int(r) for r, c in all_cells]
    cols = [int(c) for r, c in all_cells]
    min_r, max_r = min(rows), max(rows)
    min_c, max_c = min(cols), max(cols)
    
    # The repeating block = bounding box of all non-zero cells
    block_cells = set()
    for r in range(min_r, max_r + 1):
        for c in range(min_c, max_c + 1):
            block_cells.add((r, c))
    
    # Center of block
    center_r = (min_r + max_r) / 2.0
    center_c = (min_c + max_c) / 2.0
    
    # Find directions from 2-cells
    two_cells = list(zip(*np.where(inp == 2)))
    directions = set()
    for r, c in two_cells:
        dr = 1 if r > center_r else -1
        dc = 1 if c > center_c else -1
        directions.add((dr, dc))
    
    if not directions:
        directions = {(-1, 1)}
    
    # For each direction, slide block and accumulate
    for dr, dc in directions:
        current_cells = set(block_cells)
        while True:
            valid = False
            for r, c in current_cells:
                if 0 <= r < H and 0 <= c < W:
                    out[r, c] = color
                    valid = True
            if not valid:
                break
            current_cells = set((r + dr, c + dc) for r, c in current_cells)
    
    # Also place the original block position
    for r, c in block_cells:
        if 0 <= r < H and 0 <= c < W:
            out[r, c] = color
    
    return out.tolist()


if __name__ == "__main__":
    import sys, os
    task_data = 'task-data'
    if not os.path.exists(task_data):
        task_data = '../task-data'
    with open(f"{task_data}/task034.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_034(ex["input"]) == ex["output"])
    print(f"task034: {right}/{len(all_ex)} PASS")