medal-solvers/build_task285_scatter.py

"""
Build optimized ONNX model for Task 285 using ScatterElements approach.

Key insight: compute everything at ACTIVE pixels (where bbox is correct),
then SCATTER results to destination positions. No need for bbox at destinations.

Algorithm:
1. Flood fill (MaxPool 3x3, 30 iter) for component labels
2. Adjacency [900x900] for component membership
3. Dominant detection via MatMul + ArgMax
4. Bbox at active pixels via MaxPool within active mask
5. For each color c: find minority_c position per component via adjacency
6. For each dominant pixel: compute destination index in minority_c's cell
7. ScatterElements to write color at destination
8. Overlay with original input

Results (validated with official neurogolf_utils.py):
  ARC-AGI: 4 pass, 0 fail
  ARC-GEN: 261 pass, 0 fail
  Score: 8.818 points (original: 5.98, gain: +2.84)
  Memory: 10,611,900 bytes (original: 181M, 17x reduction)

Usage:
    # From medal-solvers/ directory:
    python build_task285_scatter.py \\
        --neurogolf-utils ../own-solver/neurogolf_utils.py \\
        --task-data-dir ../task-data \\
        --output optimized/task285.onnx

    # Then create submission:
    python swap_and_submit.py \\
        --base ../submission-6043.zip \\
        --models optimized/task285.onnx \\
        --output ../submission.zip

Requirements:
    pip install onnx onnxruntime numpy matplotlib onnx_tool
    # Task data: extract own-solver/neurogolf-2026.zip → task-data/
"""
import onnx
from onnx import helper, TensorProto, numpy_helper
import numpy as np
import os

H, W, C, N = 30, 30, 10, 900

def build_task285():
    nodes, inits, vis = [], [], []
    counter = [0]

    def nm():
        counter[0] += 1
        return f"t{counter[0]}"

    def const(name, val, dtype='f'):
        arr = np.array(val, dtype=np.float32 if dtype == 'f' else np.int64)
        inits.append(numpy_helper.from_array(arr, name))

    def vi(name, shape, dt=TensorProto.FLOAT):
        vis.append(helper.make_tensor_value_info(name, dt, shape))

    def nd(op, ins, outs_shapes, **kwargs):
        out_names = []
        for sd in outs_shapes:
            if isinstance(sd, tuple):
                shape, dt = sd
            else:
                shape, dt = sd, TensorProto.FLOAT
            n = nm()
            vi(n, shape, dt)
            out_names.append(n)
        nodes.append(helper.make_node(op, ins, out_names, **kwargs))
        return out_names[0] if len(out_names) == 1 else out_names

    def prop_masked(init_val, mask, iters=30):
        v = init_val
        for _ in range(iters):
            mp = nd('MaxPool', [v], [[1,1,H,W]], kernel_shape=[3,3], pads=[1,1,1,1])
            v = nd('Mul', [mp, mask], [[1,1,H,W]])
        return v

    # --- Constants ---
    const('c_half', [0.5]); const('c_one', [1.0]); const('c_two', [2.0])
    const('c_zero', [0.0]); const('c_nm1', [float(N-1)])
    const('c_hm1', [float(H-1)]); const('c_wm1', [float(W-1)]); const('c_w', [float(W)])
    const('shape_n1', [N, 1], 'i'); const('shape_1n', [1, N], 'i')
    const('shape_n', [N], 'i'); const('shape_cn', [C, N], 'i')
    const('shape_hw', [1, 1, H, W], 'i'); const('shape_1chw', [1, C, H, W], 'i')
    const('axes1', [1], 'i')
    const('coord_labels', np.arange(1, N+1, dtype=np.float32).reshape(1,1,H,W))
    row_s = np.arange(H, dtype=np.float32).reshape(1,1,H,1) * np.ones((1,1,1,W), dtype=np.float32)
    col_s = np.arange(W, dtype=np.float32).reshape(1,1,1,W) * np.ones((1,1,H,1), dtype=np.float32)
    const('row_s', row_s); const('col_s', col_s)
    const('row_flat', row_s.reshape(N)); const('col_flat', col_s.reshape(N))
    const('bg_mask', np.array([0.0]+[1.0]*9, dtype=np.float32).reshape(1, C))
    const('s1', [0,1,0,0], 'i'); const('e1', [1,C,H,W], 'i'); const('ax4', [0,1,2,3], 'i')
    const('zeros_n', np.zeros(N, dtype=np.float32))
    const('ones_n', np.ones(N, dtype=np.float32))
    for c in range(1, 10):
        const(f's_ch{c}', [0, c, 0, 0], 'i')
        const(f'e_ch{c}', [1, c+1, H, W], 'i')

    # STEP 1: Active mask
    ch19 = nd('Slice', ['input', 's1', 'e1', 'ax4'], [[1, 9, H, W]])
    asum = nd('ReduceSum', [ch19, 'axes1'], [[1, 1, H, W]], keepdims=1)
    abool = nd('Greater', [asum, 'c_half'], [([1,1,H,W], TensorProto.BOOL)])
    active = nd('Cast', [abool], [[1, 1, H, W]], to=1)

    # STEP 2: Flood fill (8-connected via MaxPool 3x3)
    label = nd('Mul', [active, 'coord_labels'], [[1,1,H,W]])
    for _ in range(30):
        mp = nd('MaxPool', [label], [[1,1,H,W]], kernel_shape=[3,3], pads=[1,1,1,1])
        label = nd('Mul', [mp, active], [[1,1,H,W]])

    # STEP 3: Adjacency + dominant detection
    l_n1 = nd('Reshape', [label, 'shape_n1'], [[N, 1]])
    l_1n = nd('Reshape', [label, 'shape_1n'], [[1, N]])
    adj_b = nd('Equal', [l_n1, l_1n], [([N, N], TensorProto.BOOL)])
    adj = nd('Cast', [adj_b], [[N, N]], to=1)
    a_flat = nd('Reshape', [active, 'shape_n'], [[N]])
    a_col = nd('Reshape', [a_flat, 'shape_n1'], [[N, 1]])
    adj_m = nd('Mul', [adj, a_col], [[N, N]])
    inp_cn = nd('Reshape', ['input', 'shape_cn'], [[C, N]])
    inp_nc = nd('Transpose', [inp_cn], [[N, C]], perm=[1, 0])
    csums = nd('MatMul', [adj_m, inp_nc], [[N, C]])
    csums_masked = nd('Mul', [csums, 'bg_mask'], [[N, C]])
    dom_idx = nd('ArgMax', [csums_masked], [([N, 1], TensorProto.INT64)], axis=1, keepdims=1)
    pixel_color = nd('ArgMax', [inp_nc], [([N, 1], TensorProto.INT64)], axis=1, keepdims=1)
    is_dom_b = nd('Equal', [dom_idx, pixel_color], [([N, 1], TensorProto.BOOL)])
    is_dom_f = nd('Cast', [is_dom_b], [[N, 1]], to=1)
    is_dom = nd('Reshape', [is_dom_f, 'shape_n'], [[N]])
    is_dom_active = nd('Mul', [is_dom, a_flat], [[N]])
    dom_sp = nd('Reshape', [is_dom_active, 'shape_hw'], [[1,1,H,W]])

    # STEP 4: Bbox at active pixels
    mr_init = nd('Mul', ['row_s', dom_sp], [[1,1,H,W]])
    max_row = prop_masked(mr_init, active)
    row_inv = nd('Sub', ['c_hm1', 'row_s'], [[1,1,H,W]])
    mnr_init = nd('Mul', [row_inv, dom_sp], [[1,1,H,W]])
    mnr_prop = prop_masked(mnr_init, active)
    min_row = nd('Sub', ['c_hm1', mnr_prop], [[1,1,H,W]])
    mc_init = nd('Mul', ['col_s', dom_sp], [[1,1,H,W]])
    max_col = prop_masked(mc_init, active)
    col_inv = nd('Sub', ['c_wm1', 'col_s'], [[1,1,H,W]])
    mnc_init = nd('Mul', [col_inv, dom_sp], [[1,1,H,W]])
    mnc_prop = prop_masked(mnc_init, active)
    min_col = nd('Sub', ['c_wm1', mnc_prop], [[1,1,H,W]])

    # Flatten bbox
    min_row_f = nd('Reshape', [min_row, 'shape_n'], [[N]])
    max_row_f = nd('Reshape', [max_row, 'shape_n'], [[N]])
    min_col_f = nd('Reshape', [min_col, 'shape_n'], [[N]])
    max_col_f = nd('Reshape', [max_col, 'shape_n'], [[N]])
    h_comp_f = nd('Sub', [max_row_f, min_row_f], [[N]])
    const('ones_n_c', np.ones(N, dtype=np.float32))
    h_comp1_f = nd('Add', [h_comp_f, 'ones_n_c'], [[N]])
    w_comp_f = nd('Sub', [max_col_f, min_col_f], [[N]])
    w_comp1_f = nd('Add', [w_comp_f, 'ones_n'], [[N]])
    local_r_f = nd('Sub', ['row_flat', min_row_f], [[N]])
    local_c_f = nd('Sub', ['col_flat', min_col_f], [[N]])

    # STEP 5: Per-color scatter
    not_dom_f = nd('Sub', ['ones_n', is_dom_active], [[N]])
    painted_per_channel = []

    for c in range(1, 10):
        ch_c = nd('Slice', ['input', f's_ch{c}', f'e_ch{c}', 'ax4'], [[1,1,H,W]])
        ch_c_flat = nd('Reshape', [ch_c, 'shape_n'], [[N]])
        minority_c = nd('Mul', [ch_c_flat, not_dom_f], [[N]])
        mc_col = nd('Reshape', [minority_c, 'shape_n1'], [[N, 1]])
        has_mc = nd('MatMul', [adj_m, mc_col], [[N, 1]])
        has_mc_flat = nd('Reshape', [has_mc, 'shape_n'], [[N]])
        has_mc_b = nd('Greater', [has_mc_flat, 'c_half'], [([N], TensorProto.BOOL)])
        has_mc_f = nd('Cast', [has_mc_b], [[N]], to=1)
        row_mc = nd('Mul', ['row_flat', minority_c], [[N]])
        row_mc_col = nd('Reshape', [row_mc, 'shape_n1'], [[N, 1]])
        mc_row_sum = nd('MatMul', [adj_m, row_mc_col], [[N, 1]])
        mc_row_f = nd('Reshape', [mc_row_sum, 'shape_n'], [[N]])
        col_mc = nd('Mul', ['col_flat', minority_c], [[N]])
        col_mc_col = nd('Reshape', [col_mc, 'shape_n1'], [[N, 1]])
        mc_col_sum = nd('MatMul', [adj_m, col_mc_col], [[N, 1]])
        mc_col_f = nd('Reshape', [mc_col_sum, 'shape_n'], [[N]])
        mc_r_off = nd('Sub', [mc_row_f, min_row_f], [[N]])
        gr_m_raw = nd('Div', [mc_r_off, h_comp1_f], [[N]])
        gr_m = nd('Floor', [gr_m_raw], [[N]])
        mc_c_off = nd('Sub', [mc_col_f, min_col_f], [[N]])
        gc_m_raw = nd('Div', [mc_c_off, w_comp1_f], [[N]])
        gc_m = nd('Floor', [gc_m_raw], [[N]])
        const(f'two_n_{c}', np.full(N, 2.0, dtype=np.float32))
        gr_d2 = nd('Div', [gr_m, f'two_n_{c}'], [[N]])
        gr_fl = nd('Floor', [gr_d2], [[N]])
        gr_2fl = nd('Mul', [gr_fl, f'two_n_{c}'], [[N]])
        v_flip = nd('Sub', [gr_m, gr_2fl], [[N]])
        gc_d2 = nd('Div', [gc_m, f'two_n_{c}'], [[N]])
        gc_fl = nd('Floor', [gc_d2], [[N]])
        gc_2fl = nd('Mul', [gc_fl, f'two_n_{c}'], [[N]])
        h_flip = nd('Sub', [gc_m, gc_2fl], [[N]])
        lr2 = nd('Mul', [local_r_f, f'two_n_{c}'], [[N]])
        dr = nd('Sub', [h_comp_f, lr2], [[N]])
        adj_r = nd('Mul', [v_flip, dr], [[N]])
        dest_lr = nd('Add', [local_r_f, adj_r], [[N]])
        lc2 = nd('Mul', [local_c_f, f'two_n_{c}'], [[N]])
        dc = nd('Sub', [w_comp_f, lc2], [[N]])
        adj_c = nd('Mul', [h_flip, dc], [[N]])
        dest_lc = nd('Add', [local_c_f, adj_c], [[N]])
        gr_h = nd('Mul', [gr_m, h_comp1_f], [[N]])
        dest_r_off = nd('Add', [gr_h, dest_lr], [[N]])
        dest_row = nd('Add', [min_row_f, dest_r_off], [[N]])
        gc_w = nd('Mul', [gc_m, w_comp1_f], [[N]])
        dest_c_off = nd('Add', [gc_w, dest_lc], [[N]])
        dest_col = nd('Add', [min_col_f, dest_c_off], [[N]])
        const(f'w_n_{c}', np.full(N, float(W), dtype=np.float32))
        dest_rw = nd('Mul', [dest_row, f'w_n_{c}'], [[N]])
        dest_idx_f = nd('Add', [dest_rw, dest_col], [[N]])
        const(f'z_n_{c}', np.zeros(N, dtype=np.float32))
        const(f'nm1_n_{c}', np.full(N, float(N-1), dtype=np.float32))
        dest_idx_lo = nd('Max', [dest_idx_f, f'z_n_{c}'], [[N]])
        dest_idx_clip = nd('Min', [dest_idx_lo, f'nm1_n_{c}'], [[N]])
        dest_idx_i = nd('Cast', [dest_idx_clip], [([N], TensorProto.INT64)], to=7)
        scatter_vals = nd('Mul', [is_dom_active, has_mc_f], [[N]])
        dest_idx_i_2d = nd('Reshape', [dest_idx_i, 'shape_n1'], [([N, 1], TensorProto.INT64)])
        scatter_vals_2d = nd('Reshape', [scatter_vals, 'shape_n1'], [[N, 1]])
        const(f'zeros_n1_{c}', np.zeros((N, 1), dtype=np.float32))
        scattered = nd('ScatterElements', [f'zeros_n1_{c}', dest_idx_i_2d, scatter_vals_2d], [[N, 1]],
                      axis=0, reduction='max')
        scattered_flat = nd('Reshape', [scattered, 'shape_n'], [[N]])
        painted_per_channel.append(scattered_flat)

    # STEP 6: Build output
    const('zero_ch_flat', np.zeros(N, dtype=np.float32))
    reshaped_channels = ['zero_ch_flat']
    for pc in painted_per_channel:
        reshaped_channels.append(pc)
    ch_2d = []
    for ch_name in reshaped_channels:
        r = nd('Reshape', [ch_name, 'shape_1n'], [[1, N]])
        ch_2d.append(r)
    const('ax0', [0], 'i')
    painted_cn = nd('Concat', ch_2d, [[C, N]], axis=0)
    painted_1chw = nd('Reshape', [painted_cn, 'shape_1chw'], [[1, C, H, W]])
    const('axes1b', [1], 'i')
    any_painted = nd('ReduceMax', [painted_1chw, 'axes1b'], [[1, 1, H, W]], keepdims=1)
    keep_mask = nd('Sub', ['c_one', any_painted], [[1, 1, H, W]])
    kept = nd('Mul', ['input', keep_mask], [[1, C, H, W]])
    nodes.append(helper.make_node('Add', [kept, painted_1chw], ['output']))

    x = helper.make_tensor_value_info('input', TensorProto.FLOAT, [1, C, H, W])
    y = helper.make_tensor_value_info('output', TensorProto.FLOAT, [1, C, H, W])
    graph = helper.make_graph(nodes, 'task285_scatter', [x], [y], initializer=inits, value_info=vis)
    model = helper.make_model(graph, ir_version=10, opset_imports=[helper.make_opsetid('', 18)])
    return model


def validate_with_official_utils(model_path, task_num, neurogolf_utils_path, task_data_dir):
    """
    Validate and score model using the official Kaggle neurogolf_utils.py.
    This is the ONLY reliable way to verify correctness and get accurate scores.
    """
    import sys
    import types

    # Mock IPython.display (not available outside Jupyter notebooks)
    mock_ipython = types.ModuleType('IPython')
    mock_display = types.ModuleType('IPython.display')
    mock_display.display = lambda *a, **k: None
    mock_display.FileLink = lambda x: x
    mock_ipython.display = mock_display
    sys.modules['IPython'] = mock_ipython
    sys.modules['IPython.display'] = mock_display

    # Use non-interactive matplotlib backend
    import matplotlib
    matplotlib.use('Agg')

    # Import neurogolf_utils from the specified path
    utils_dir = os.path.dirname(os.path.abspath(neurogolf_utils_path))
    if utils_dir not in sys.path:
        sys.path.insert(0, utils_dir)

    import neurogolf_utils

    # Override the data directory to local task data
    neurogolf_utils._NEUROGOLF_DIR = task_data_dir.rstrip('/') + '/'

    # Load examples using official loader
    examples = neurogolf_utils.load_examples(task_num)
    print(f"  Loaded examples: train={len(examples['train'])}, "
          f"test={len(examples['test'])}, arc-gen={len(examples['arc-gen'])}")

    # Load model and run official verify_network
    network = onnx.load(model_path)
    print(f"\n{'='*60}")
    print(f"  Official neurogolf_utils.verify_network() output:")
    print(f"{'='*60}\n")
    neurogolf_utils.verify_network(network, task_num, examples)
    print(f"\n{'='*60}")

    return True


if __name__ == '__main__':
    import argparse

    parser = argparse.ArgumentParser(description="Build and validate optimized Task 285 model")
    parser.add_argument('--neurogolf-utils', default='../own-solver/neurogolf_utils.py',
                       help='Path to official neurogolf_utils.py')
    parser.add_argument('--task-data-dir', default='../task-data',
                       help='Directory containing taskNNN.json files (from neurogolf-2026.zip)')
    parser.add_argument('--output', default='/app/repo/medal-solvers/optimized/task285.onnx',
                       help='Output path for the optimized model')
    parser.add_argument('--skip-official-validation', action='store_true',
                       help='Skip official utils validation (NOT recommended)')
    args = parser.parse_args()

    # Auto-detect paths
    task_data_dir = args.task_data_dir
    if not os.path.exists(os.path.join(task_data_dir, 'task285.json')):
        for candidate in ['task-data', '../task-data', '.']:
            if os.path.exists(os.path.join(candidate, 'task285.json')):
                task_data_dir = candidate
                break

    neurogolf_path = args.neurogolf_utils
    if not os.path.exists(neurogolf_path):
        for candidate in ['../own-solver/neurogolf_utils.py', 'own-solver/neurogolf_utils.py',
                         'neurogolf_utils.py']:
            if os.path.exists(candidate):
                neurogolf_path = candidate
                break

    print("Building Task 285 optimized model (ScatterElements approach)...")
    model = build_task285()

    os.makedirs(os.path.dirname(args.output) or '.', exist_ok=True)
    onnx.save(model, args.output)
    fsize = os.path.getsize(args.output)
    print(f"  Nodes: {len(model.graph.node)}")
    print(f"  File size: {fsize:,} bytes (limit: {int(1.44*1024*1024):,})")
    print(f"  Saved to: {args.output}")

    # Run official validation
    if not args.skip_official_validation:
        if not os.path.exists(neurogolf_path):
            print(f"\n  WARNING: neurogolf_utils.py not found at {neurogolf_path}")
            print(f"  Use --neurogolf-utils to specify path.")
            args.skip_official_validation = True
        elif not os.path.exists(os.path.join(task_data_dir, 'task285.json')):
            print(f"\n  WARNING: task285.json not found in {task_data_dir}")
            print(f"  Extract neurogolf-2026.zip first. Use --task-data-dir to specify.")
            args.skip_official_validation = True

    if not args.skip_official_validation:
        print(f"\n  Using official neurogolf_utils.py from: {neurogolf_path}")
        print(f"  Using task data from: {task_data_dir}")
        validate_with_official_utils(args.output, 285, neurogolf_path, task_data_dir)
    else:
        # Fallback: basic validation
        import onnxruntime
        import json
        print("\n  Running basic validation (use --neurogolf-utils for official scoring)...")
        sess = onnxruntime.InferenceSession(args.output)
        task_json = os.path.join(task_data_dir, 'task285.json')
        if os.path.exists(task_json):
            with open(task_json) as f:
                td = json.load(f)
            right, wrong = 0, 0
            for split in ['train', 'test', 'arc-gen']:
                for ex in td.get(split, []):
                    inp = np.zeros((1,10,30,30), dtype=np.float32)
                    for r, row in enumerate(ex['input']):
                        for c, v in enumerate(row):
                            if r<30 and c<30: inp[0][v][r][c] = 1.0
                    result = sess.run(['output'], {'input': inp})
                    out = (result[0] > 0.0).astype(float)
                    exp = np.zeros((1,10,30,30), dtype=np.float32)
                    for r, row in enumerate(ex['output']):
                        for c, v in enumerate(row):
                            if r<30 and c<30: exp[0][v][r][c] = 1.0
                    if np.array_equal(out, exp):
                        right += 1
                    else:
                        wrong += 1
            print(f"  Results: {right} pass, {wrong} fail")
            if wrong == 0:
                print(f"  PASSED (use official utils for reliable score)")
            else:
                print(f"  FAILED - do not submit!")