medal-solvers/build_task200_onnx.py

"""Build optimized ONNX model for Task 200.

Rule (84/84 verified):
1. Input: 10x10 grid with single colored dot at row 9
2. Output: From dot's column rightward, alternating pattern:
   - Even offsets (0,2,4,...): full vertical line of dot color
   - Odd offsets (1,3,5,...): color 5 at alternating top(row0)/bottom(row9)

ONNX approach:
- Slice to 10x10
- Detect dot column and color
- Use column grid with modular arithmetic to create masks
- Apply colors
- Pad back to 30x30

Base: 91 nodes, 198010 params, score 10.33
Target: ~30 nodes, ~500 params → score ~15.5 (gain +5.2!)
"""
import sys, os
import numpy as np
import onnx
from onnx import helper, numpy_helper, TensorProto
import math


def build_task200():
    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

    # === CONSTANTS ===
    const('c_half', [0.5])
    const('c_one', [1.0])
    const('c_two', [2.0])
    const('c_four', [4.0])
    const('c_1_5', [1.5])
    const('c_8_5', [8.5])

    const('sl_start', [0, 0, 0, 0], 'i')
    const('sl_end', [1, 10, 10, 10], 'i')
    const('sl_axes', [0, 1, 2, 3], 'i')
    const('axes1', [1], 'i')
    const('axes23', [2, 3], 'i')
    const('axes2', [2], 'i')

    # Column grid [1,1,1,10]
    const('col_grid', np.arange(10, dtype=np.float32).reshape(1, 1, 1, 10))
    # Row grid [1,1,10,1]
    const('row_grid', np.arange(10, dtype=np.float32).reshape(1, 1, 10, 1))

    const('shape_1_1_1_1', [1, 1, 1, 1], 'i')
    const('shape_1_10_1_1', [1, 10, 1, 1], 'i')
    const('shape_1', [1], 'i')
    
    const('depth_10', [10.0])
    const('oh_vals', [0.0, 1.0])
    
    # Pad to 30x30
    const('pad_10_to_30', [0, 0, 0, 0, 0, 0, 20, 20], 'i')
    const('pad_val_zero', [0.0])

    # Channel slice for non-bg
    const('sl_ch_start', [0, 1, 0, 0], 'i')
    const('sl_ch_end', [1, 10, 10, 10], 'i')
    
    # bg mask for channel selection (exclude ch0)
    const('bg_mask_10', np.array([[[[0, 1, 1, 1, 1, 1, 1, 1, 1, 1]]]], dtype=np.float32).reshape(1, 10, 1, 1))
    
    # Color 5 index
    const('idx_5', [5], 'i')
    const('idx_1_arr', [1], 'i')

    # === STEP 1: Slice input to 10x10 ===
    inp10 = nd('Slice', ['input', 'sl_start', 'sl_end', 'sl_axes'], [[1, 10, 10, 10]])

    # === STEP 2: Find dot column and color ===
    # Sum over channels to get active mask, then find the dot
    nonbg = nd('Slice', [inp10, 'sl_ch_start', 'sl_ch_end', 'sl_axes'], [[1, 9, 10, 10]])
    active = nd('ReduceSum', [nonbg, 'axes1'], [[1, 1, 10, 10]], keepdims=1)
    
    # Sum over rows to get column indicator [1,1,1,10]
    col_sums = nd('ReduceSum', [active, 'axes2'], [[1, 1, 1, 10]], keepdims=1)
    
    # Dot column = where col_sums > 0 → use col_grid * col_sums → max gives dot col
    col_weighted = nd('Mul', ['col_grid', col_sums], [[1, 1, 1, 10]])
    dot_col_f = nd('ReduceMax', [col_weighted], [[1, 1, 1, 1]], keepdims=1)  # [1,1,1,1]
    
    # Dot color: which channel has the dot?
    ch_sums = nd('ReduceSum', [inp10, 'axes23'], [[1, 10, 1, 1]], keepdims=1)
    ch_sums_nobg = nd('Mul', [ch_sums, 'bg_mask_10'], [[1, 10, 1, 1]])
    dot_color_idx = nd('ArgMax', [ch_sums_nobg], [([1, 1, 1, 1], TensorProto.INT64)], axis=1, keepdims=1)
    dot_color_1d = nd('Reshape', [dot_color_idx, 'shape_1'], [([1], TensorProto.INT64)])
    
    # === STEP 3: Compute column masks ===
    # relative column offset from dot: col - dot_col
    rel_col = nd('Sub', ['col_grid', dot_col_f], [[1, 1, 1, 10]])
    
    # Only columns >= dot_col (rel >= 0)
    in_range_b = nd('Greater', [rel_col, nd('Sub', ['c_half', 'c_one'], [[1, 1, 1, 1]])],
                    [([1, 1, 1, 10], TensorProto.BOOL)])  # rel >= 0
    in_range = nd('Cast', [in_range_b], [[1, 1, 1, 10]], to=1)
    
    # Even offset columns (rel % 2 == 0): vertical lines
    # rel / 2, floor, * 2, compare to rel
    rel_div2 = nd('Div', [rel_col, 'c_two'], [[1, 1, 1, 10]])
    rel_floor = nd('Floor', [rel_div2], [[1, 1, 1, 10]])
    rel_x2 = nd('Mul', [rel_floor, 'c_two'], [[1, 1, 1, 10]])
    is_even_diff = nd('Abs', [nd('Sub', [rel_col, rel_x2], [[1, 1, 1, 10]])], [[1, 1, 1, 10]])
    is_even_b = nd('Less', [is_even_diff, 'c_half'], [([1, 1, 1, 10], TensorProto.BOOL)])
    is_even = nd('Cast', [is_even_b], [[1, 1, 1, 10]], to=1)
    
    # Vertical line mask: is_even AND in_range → [1,1,1,10] broadcast to [1,1,10,10]
    vert_mask_col = nd('Mul', [is_even, in_range], [[1, 1, 1, 10]])  # [1,1,1,10]
    # This broadcasts to all rows (since it's independent of row)
    
    # Odd offset columns: connectors
    is_odd = nd('Sub', ['c_one', is_even], [[1, 1, 1, 10]])
    odd_mask_col = nd('Mul', [is_odd, in_range], [[1, 1, 1, 10]])
    
    # Connector position: for the k-th odd column (k=0,1,2,...):
    # k even → row 0, k odd → row 9
    # k = (rel - 1) / 2 for odd rel
    # k % 2 == 0 → top, k % 2 == 1 → bottom
    rel_minus1 = nd('Sub', [rel_col, 'c_one'], [[1, 1, 1, 10]])
    k_val = nd('Div', [rel_minus1, 'c_two'], [[1, 1, 1, 10]])  # k = (rel-1)/2 for odd cols
    k_div2 = nd('Div', [k_val, 'c_two'], [[1, 1, 1, 10]])
    k_floor = nd('Floor', [k_div2], [[1, 1, 1, 10]])
    k_x2 = nd('Mul', [k_floor, 'c_two'], [[1, 1, 1, 10]])
    k_mod2_diff = nd('Abs', [nd('Sub', [k_val, k_x2], [[1, 1, 1, 10]])], [[1, 1, 1, 10]])
    k_is_even_b = nd('Less', [k_mod2_diff, 'c_half'], [([1, 1, 1, 10], TensorProto.BOOL)])
    k_is_even = nd('Cast', [k_is_even_b], [[1, 1, 1, 10]], to=1)  # 1 if top connector
    k_is_odd = nd('Sub', ['c_one', k_is_even], [[1, 1, 1, 10]])   # 1 if bottom connector
    
    # Top connector: odd_mask_col AND k_is_even AND row==0
    # Bottom connector: odd_mask_col AND k_is_odd AND row==9
    at_row0_b = nd('Less', ['row_grid', 'c_half'], [([1, 1, 10, 1], TensorProto.BOOL)])
    at_row0 = nd('Cast', [at_row0_b], [[1, 1, 10, 1]], to=1)
    at_row9_b = nd('Greater', ['row_grid', 'c_8_5'], [([1, 1, 10, 1], TensorProto.BOOL)])
    at_row9 = nd('Cast', [at_row9_b], [[1, 1, 10, 1]], to=1)
    
    # Top connector mask [1,1,10,10]: odd_mask_col * k_is_even * at_row0
    top_conn = nd('Mul', [nd('Mul', [odd_mask_col, k_is_even], [[1, 1, 1, 10]]), at_row0], [[1, 1, 10, 10]])
    # Bottom connector mask
    bot_conn = nd('Mul', [nd('Mul', [odd_mask_col, k_is_odd], [[1, 1, 1, 10]]), at_row9], [[1, 1, 10, 10]])
    # Combined connector mask (color 5)
    conn_mask = nd('Add', [top_conn, bot_conn], [[1, 1, 10, 10]])
    
    # === STEP 4: Apply colors ===
    # Dot color one-hot
    dot_oh = nd('OneHot', [dot_color_1d, 'depth_10', 'oh_vals'], [[1, 10]], axis=1)
    dot_oh_4d = nd('Reshape', [dot_oh, 'shape_1_10_1_1'], [[1, 10, 1, 1]])
    
    # Vertical lines: dot_color * vert_mask
    vert_out = nd('Mul', [dot_oh_4d, vert_mask_col], [[1, 10, 10, 10]])
    
    # Connectors: color 5 * conn_mask
    c5_oh = nd('OneHot', ['idx_5', 'depth_10', 'oh_vals'], [[1, 10]], axis=1)
    c5_oh_4d = nd('Reshape', [c5_oh, 'shape_1_10_1_1'], [[1, 10, 1, 1]])
    conn_out = nd('Mul', [c5_oh_4d, conn_mask], [[1, 10, 10, 10]])
    
    # Combine
    colored_10 = nd('Add', [vert_out, conn_out], [[1, 10, 10, 10]])
    
    # === STEP 5: Add background channel ===
    nonbg_out = nd('Slice', [colored_10, 'sl_ch_start', 'sl_ch_end', 'sl_axes'], [[1, 9, 10, 10]])
    any_color = nd('ReduceMax', [nonbg_out, 'axes1'], [[1, 1, 10, 10]], keepdims=1)
    bg_ch = nd('Sub', ['c_one', any_color], [[1, 1, 10, 10]])
    const('ch0_idx', [0], 'i')
    ch0_oh = nd('OneHot', ['ch0_idx', 'depth_10', 'oh_vals'], [[1, 10]], axis=1)
    ch0_oh_4d = nd('Reshape', [ch0_oh, 'shape_1_10_1_1'], [[1, 10, 1, 1]])
    bg_full = nd('Mul', [ch0_oh_4d, bg_ch], [[1, 10, 10, 10]])
    out_10 = nd('Add', [colored_10, bg_full], [[1, 10, 10, 10]])
    
    # === STEP 6: Pad to 30x30 ===
    final = nd('Pad', [out_10, 'pad_10_to_30', 'pad_val_zero'], [[1, 10, 30, 30]])

    # Build model
    x = helper.make_tensor_value_info('input', TensorProto.FLOAT, [1, 10, 30, 30])
    y = helper.make_tensor_value_info('output', TensorProto.FLOAT, [1, 10, 30, 30])
    nodes.append(helper.make_node('Identity', [final], ['output']))
    
    graph = helper.make_graph(nodes, 'task200', [x], [y], initializer=inits, value_info=vis)
    model = helper.make_model(graph, ir_version=10, opset_imports=[helper.make_opsetid('', 18)])
    return model


if __name__ == '__main__':
    import onnxruntime as ort
    import json, zipfile
    from huggingface_hub import hf_hub_download
    
    print("Building Task 200 ONNX model...")
    model = build_task200()
    
    output_path = '/app/task200.onnx'
    onnx.save(model, output_path)
    fsize = os.path.getsize(output_path)
    print(f"  Nodes: {len(model.graph.node)}")
    print(f"  File size: {fsize:,} bytes")
    
    # Validate
    sess = ort.InferenceSession(output_path)
    task_path = hf_hub_download('rogermt/neurogolf-solver', 'own-solver/neurogolf-2026.zip')
    with zipfile.ZipFile(task_path, 'r') as zf:
        data = json.loads(zf.read('task200.json'))
    
    all_examples = data['train'] + data['test'] + data.get('arc-gen', [])
    right_count, wrong_count = 0, 0
    
    for i, ex in enumerate(all_examples):
        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_count += 1
        else:
            wrong_count += 1
            if wrong_count <= 3:
                diff_locs = np.where(out != exp)
                print(f"  FAIL {i}: {len(diff_locs[0])} diffs")
    
    print(f"\nResults: {right_count} pass, {wrong_count} fail out of {len(all_examples)}")
    
    if wrong_count == 0:
        params = sum(int(np.prod(init.dims)) for init in model.graph.initializer)
        mem_est = sum(int(np.prod([d.dim_value for d in vi.type.tensor_type.shape.dim])) * 4
                      for vi in model.graph.value_info
                      if vi.type.HasField('tensor_type') and vi.type.tensor_type.HasField('shape'))
        score_est = max(1.0, 25.0 - math.log(max(1.0, mem_est + params)))
        print(f"  Params: {params:,}, Memory est: {mem_est:,}")
        print(f"  Score est: {score_est:.3f} (base: 10.326)")
        print(f"  Gain est: +{score_est - 10.326:.3f}")
        print(f"  ** THIS IS A STRONG CANDIDATE FOR SUBMISSION **")