medal-solvers/build_task255_onnx.py

"""
Build optimized ONNX model for Task 255.
Boundary-anchored rect finding + integral image + cross extensions.
Target: ~500 nodes, ~2-3M memory → score ~10-12 → gain +3-5 pts.
"""
import onnx
from onnx import helper, TensorProto, numpy_helper
import numpy as np
import os, json, sys

H, W, C = 30, 30, 10

def build_task255():
    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_three', [3.0])
    const('c_30', [30.0])
    const('c_neg1000', [-1000.0])
    const('c_pos1000', [1000.0])
    const('c_zero_s', [0.0])
    const('c_2_5', [2.5])
    
    const('shape_30', [30], 'i')
    const('shape_30x30', [30, 30], 'i')
    const('shape_1_1_30_30', [1, 1, 30, 30], 'i')
    const('shape_1_30', [1, 30], 'i')
    const('shape_30_1', [30, 1], 'i')
    const('shape_900', [900], 'i')
    
    const('axes0', [0], 'i')
    const('axes1', [1], 'i')
    const('s_ch1', [0, 1, 0, 0], 'i')
    const('e_ch1', [1, 10, 30, 30], 'i')
    const('ax4', [0, 1, 2, 3], 'i')
    
    const('row_idx', np.arange(30, dtype=np.float32))
    const('col_idx', np.arange(30, dtype=np.float32))
    
    # Thresholds for CumSum-based depth detection
    row_plus_half = (np.arange(30, dtype=np.float32) + 0.5).reshape(30, 1)
    const('row_plus_half_col', row_plus_half)
    bot_thresh = (30.0 - np.arange(30, dtype=np.float32) - 0.5).reshape(30, 1)
    const('bot_thresh', bot_thresh)
    col_plus_half = (np.arange(30, dtype=np.float32) + 0.5).reshape(1, 30)
    const('col_plus_half_row', col_plus_half)
    right_thresh = (30.0 - np.arange(30, dtype=np.float32) - 0.5).reshape(1, 30)
    const('right_thresh', right_thresh)
    
    # Sparse table constants for rect finding
    # lower_tri mask and length matrix for [30,30]
    inf_upper = np.zeros((30, 30), dtype=np.float32)
    for i in range(30):
        for j in range(i):
            inf_upper[i, j] = 1e9
    const('inf_upper', inf_upper)
    
    length_mat = np.zeros((30, 30), dtype=np.float32)
    for i in range(30):
        for j in range(i, 30):
            length_mat[i, j] = float(j - i + 1)
    const('length_mat', length_mat)
    
    const('ones_30_1', np.ones((30, 1), dtype=np.float32))
    const('ones_1_30', np.ones((1, 30), dtype=np.float32))
    
    # Color delta for output
    color3_delta = np.zeros((1, 10, 1, 1), dtype=np.float32)
    color3_delta[0, 3, 0, 0] = 1.0
    color3_delta[0, 0, 0, 0] = -1.0
    const('color3_delta', color3_delta)

    # === STEP 1: fg_mask [30,30] ===
    ch19 = nd('Slice', ['input', 's_ch1', 'e_ch1', 'ax4'], [[1, 9, 30, 30]])
    fg_sum = nd('ReduceSum', [ch19, 'axes1'], [[1, 1, 30, 30]], keepdims=1)
    fg_bool = nd('Greater', [fg_sum, 'c_half'], [([1, 1, 30, 30], TensorProto.BOOL)])
    fg_4d = nd('Cast', [fg_bool], [[1, 1, 30, 30]], to=1)
    fg = nd('Reshape', [fg_4d, 'shape_30x30'], [[30, 30]])
    not_fg = nd('Sub', ['c_one', fg], [[30, 30]])

    # === STEP 2: Depth vectors from 4 sides ===
    # TOP: CumSum(not_fg, axis=0). If cum[r,c] > r+0.5 → all empty [0..r]
    cum_nf_top = nd('CumSum', [not_fg, 'axes0'], [[30, 30]])
    lead_top_b = nd('Greater', [cum_nf_top, 'row_plus_half_col'], [([30, 30], TensorProto.BOOL)])
    lead_top_f = nd('Cast', [lead_top_b], [[30, 30]], to=1)
    depth_top_2d = nd('ReduceSum', [lead_top_f, 'axes0'], [[1, 30]], keepdims=1)
    depth_top = nd('Reshape', [depth_top_2d, 'shape_30'], [[30]])
    
    # BOTTOM: CumSum(not_fg, axis=0, reverse=1)
    cum_nf_bot = nd('CumSum', [not_fg, 'axes0'], [[30, 30]], reverse=1)
    lead_bot_b = nd('Greater', [cum_nf_bot, 'bot_thresh'], [([30, 30], TensorProto.BOOL)])
    lead_bot_f = nd('Cast', [lead_bot_b], [[30, 30]], to=1)
    depth_bot_2d = nd('ReduceSum', [lead_bot_f, 'axes0'], [[1, 30]], keepdims=1)
    depth_bot = nd('Reshape', [depth_bot_2d, 'shape_30'], [[30]])
    
    # LEFT: CumSum(not_fg, axis=1)
    cum_nf_left = nd('CumSum', [not_fg, 'axes1'], [[30, 30]])
    lead_left_b = nd('Greater', [cum_nf_left, 'col_plus_half_row'], [([30, 30], TensorProto.BOOL)])
    lead_left_f = nd('Cast', [lead_left_b], [[30, 30]], to=1)
    depth_left_2d = nd('ReduceSum', [lead_left_f, 'axes1'], [[30, 1]], keepdims=1)
    depth_left = nd('Reshape', [depth_left_2d, 'shape_30'], [[30]])
    
    # RIGHT: CumSum(not_fg, axis=1, reverse=1)
    cum_nf_right = nd('CumSum', [not_fg, 'axes1'], [[30, 30]], reverse=1)
    lead_right_b = nd('Greater', [cum_nf_right, 'right_thresh'], [([30, 30], TensorProto.BOOL)])
    lead_right_f = nd('Cast', [lead_right_b], [[30, 30]], to=1)
    depth_right_2d = nd('ReduceSum', [lead_right_f, 'axes1'], [[30, 1]], keepdims=1)
    depth_right = nd('Reshape', [depth_right_2d, 'shape_30'], [[30]])

    # === STEP 3: Largest rect in histogram for each side ===
    # Sparse table: 5 doubling steps to compute range-min matrix
    # Then area = min_height * width, find max.
    
    def hist_rect(depth_name, label):
        """Largest rectangle in histogram h[30]. Returns (area, start, width, height) as [1] tensors."""
        # Broadcast h to [30,30]: h_mat[i,j] = h[j] for all i
        h_row = nd('Reshape', [depth_name, 'shape_1_30'], [[1, 30]])
        h_mat = nd('MatMul', ['ones_30_1', h_row], [[30, 30]])
        
        # Add inf for upper triangle (j < i): ensures min is only over valid range [i..j]
        m = nd('Add', [h_mat, 'inf_upper'], [[30, 30]])
        
        # Sparse table doubling: 5 steps with shifts 1, 2, 4, 8, 16
        # Shift m RIGHT by `shift` columns (pad left with inf, slice [0:30])
        # ONNX Pad format for [30,30]: [dim0_begin, dim1_begin, dim0_end, dim1_end]
        for step, shift in enumerate([1, 2, 4, 8, 16]):
            # Pad: add `shift` inf columns on the LEFT of dim 1
            pads = np.array([0, shift, 0, 0], dtype=np.int64)
            const(f'pad_{label}_{step}', pads, 'i')
            const(f'padval_{label}_{step}', [1e9])
            padded = nd('Pad', [m, f'pad_{label}_{step}', f'padval_{label}_{step}'],
                       [[30, 30 + shift]])
            # Slice first 30 columns: [0:30, 0:30]
            const(f'sls_{label}_{step}', [0, 0], 'i')
            const(f'sle_{label}_{step}', [30, 30], 'i')
            const(f'sla_{label}_{step}', [0, 1], 'i')
            m_shifted = nd('Slice', [padded, f'sls_{label}_{step}',
                                     f'sle_{label}_{step}', f'sla_{label}_{step}'],
                          [[30, 30]])
            m = nd('Min', [m, m_shifted], [[30, 30]])
        
        # Now m[i,j] = min(h[i..j]) for j>=i, large for j<i
        # area[i,j] = m[i,j] * length_mat[i,j]
        area_mat = nd('Mul', [m, 'length_mat'], [[30, 30]])
        
        # Flatten and find best
        area_flat = nd('Reshape', [area_mat, 'shape_900'], [[900]])
        best_idx = nd('ArgMax', [area_flat], [([1], TensorProto.INT64)], axis=0, keepdims=1)
        best_idx_f = nd('Cast', [best_idx], [[1]], to=1)
        
        # Decode: i = idx // 30, j = idx % 30
        const(f'c30_{label}', [30.0])
        ij_div = nd('Div', [best_idx_f, f'c30_{label}'], [[1]])
        best_i = nd('Floor', [ij_div], [[1]])
        i_times_30 = nd('Mul', [best_i, f'c30_{label}'], [[1]])
        best_j = nd('Sub', [best_idx_f, i_times_30], [[1]])
        
        # width = j - i + 1
        width = nd('Sub', [best_j, best_i], [[1]])
        width_p1 = nd('Add', [width, 'c_one'], [[1]])
        
        # area = max of area_flat
        best_area = nd('ReduceMax', [area_flat], [[1]], keepdims=1)
        
        # height = area / width (integer division since both are integers)
        height = nd('Div', [best_area, width_p1], [[1]])
        
        return best_area, best_i, width_p1, height
    
    area_t, start_t, width_t, height_t = hist_rect(depth_top, 'ht')
    area_b, start_b, width_b, height_b = hist_rect(depth_bot, 'hb')
    area_l, start_l, width_l, height_l = hist_rect(depth_left, 'hl')
    area_r, start_r, width_r, height_r = hist_rect(depth_right, 'hr')
    
    # Find best side (0=top, 1=bot, 2=left, 3=right)
    areas = nd('Concat', [area_t, area_b, area_l, area_r], [[4]], axis=0)
    best_side = nd('ArgMax', [areas], [([1], TensorProto.INT64)], axis=0, keepdims=1)
    best_side_f = nd('Cast', [best_side], [[1]], to=1)
    
    # Compute rect (tr, lc, rh, rw) for each side:
    # TOP: tr=0, lc=start, rh=height, rw=width
    # BOT: tr=30-height, lc=start, rh=height, rw=width
    # LEFT: tr=start, lc=0, rh=width, rw=height  (width=run_len along rows, height=depth)
    # RIGHT: tr=start, lc=30-height, rh=width, rw=height
    
    tr_bot = nd('Sub', ['c_30', height_b], [[1]])
    lc_right = nd('Sub', ['c_30', height_r], [[1]])
    
    # Build [4] vectors and Gather
    tr_vec = nd('Concat', ['c_zero_s', tr_bot, start_l, start_r], [[4]], axis=0)
    lc_vec = nd('Concat', [start_t, start_b, 'c_zero_s', lc_right], [[4]], axis=0)
    rh_vec = nd('Concat', [height_t, height_b, width_l, width_r], [[4]], axis=0)
    rw_vec = nd('Concat', [width_t, width_b, height_l, height_r], [[4]], axis=0)
    
    tr_val = nd('Gather', [tr_vec, best_side], [[1]], axis=0)
    lc_val = nd('Gather', [lc_vec, best_side], [[1]], axis=0)
    rh_val = nd('Gather', [rh_vec, best_side], [[1]], axis=0)
    rw_val = nd('Gather', [rw_vec, best_side], [[1]], axis=0)

    # === STEP 4: Erosion → core rect ===
    # core_r1 = tr if tr==0 else tr+1
    # core_r2 = tr+rh-1 if tr+rh==30 else tr+rh-2
    tr_plus_rh = nd('Add', [tr_val, rh_val], [[1]])
    lc_plus_rw = nd('Add', [lc_val, rw_val], [[1]])
    
    # core_r1
    tr_is_zero = nd('Less', [tr_val, 'c_half'], [([1], TensorProto.BOOL)])
    tr_iz_f = nd('Cast', [tr_is_zero], [[1]], to=1)
    not_tr_iz = nd('Sub', ['c_one', tr_iz_f], [[1]])
    core_r1 = nd('Add', [tr_val, not_tr_iz], [[1]])
    
    # core_r2
    bot_edge = nd('Greater', [tr_plus_rh, nd('Sub', ['c_30', 'c_half'], [[1]])], [([1], TensorProto.BOOL)])
    bot_e_f = nd('Cast', [bot_edge], [[1]], to=1)
    core_r2 = nd('Sub', [tr_plus_rh, nd('Sub', ['c_two', bot_e_f], [[1]])], [[1]])
    
    # core_c1
    lc_is_zero = nd('Less', [lc_val, 'c_half'], [([1], TensorProto.BOOL)])
    lc_iz_f = nd('Cast', [lc_is_zero], [[1]], to=1)
    not_lc_iz = nd('Sub', ['c_one', lc_iz_f], [[1]])
    core_c1 = nd('Add', [lc_val, not_lc_iz], [[1]])
    
    # core_c2
    right_edge = nd('Greater', [lc_plus_rw, nd('Sub', ['c_30', 'c_half'], [[1]])], [([1], TensorProto.BOOL)])
    right_e_f = nd('Cast', [right_edge], [[1]], to=1)
    core_c2 = nd('Sub', [lc_plus_rw, nd('Sub', ['c_two', right_e_f], [[1]])], [[1]])

    # === STEP 5: Build core mask [30,30] ===
    cr1_m_half = nd('Sub', [core_r1, 'c_half'], [[1]])
    cr2_p_half = nd('Add', [core_r2, 'c_half'], [[1]])
    cc1_m_half = nd('Sub', [core_c1, 'c_half'], [[1]])
    cc2_p_half = nd('Add', [core_c2, 'c_half'], [[1]])
    
    row_ge_r1 = nd('Greater', ['row_idx', cr1_m_half], [([30], TensorProto.BOOL)])
    row_le_r2 = nd('Less', ['row_idx', cr2_p_half], [([30], TensorProto.BOOL)])
    row_in_b = nd('And', [row_ge_r1, row_le_r2], [([30], TensorProto.BOOL)])
    row_in = nd('Cast', [row_in_b], [[30]], to=1)
    
    col_ge_c1 = nd('Greater', ['col_idx', cc1_m_half], [([30], TensorProto.BOOL)])
    col_le_c2 = nd('Less', ['col_idx', cc2_p_half], [([30], TensorProto.BOOL)])
    col_in_b = nd('And', [col_ge_c1, col_le_c2], [([30], TensorProto.BOOL)])
    col_in = nd('Cast', [col_in_b], [[30]], to=1)
    
    ri_col = nd('Reshape', [row_in, 'shape_30_1'], [[30, 1]])
    ci_row = nd('Reshape', [col_in, 'shape_1_30'], [[1, 30]])
    core_mask = nd('MatMul', [ri_col, ci_row], [[30, 30]])

    # === STEP 6: FG boundary vectors ===
    # fg_row_max[r] = max col with fg (or -1000 if none)
    col_idx_row = nd('Reshape', ['col_idx', 'shape_1_30'], [[1, 30]])
    row_idx_col = nd('Reshape', ['row_idx', 'shape_30_1'], [[30, 1]])
    
    fg_col_w = nd('Mul', [fg, col_idx_row], [[30, 30]])  # fg * col_idx
    nf_neg = nd('Mul', [not_fg, 'c_neg1000'], [[30, 30]])
    fg_col_scored = nd('Add', [fg_col_w, nf_neg], [[30, 30]])
    fg_row_max_2d = nd('ReduceMax', [fg_col_scored, 'axes1'], [[30, 1]], keepdims=1)
    fg_row_max = nd('Reshape', [fg_row_max_2d, 'shape_30'], [[30]])
    
    # fg_row_min[r] = min col with fg (or 1000 if none)
    nf_pos = nd('Mul', [not_fg, 'c_pos1000'], [[30, 30]])
    fg_col_min_scored = nd('Add', [fg_col_w, nf_pos], [[30, 30]])
    fg_row_min_2d = nd('ReduceMin', [fg_col_min_scored, 'axes1'], [[30, 1]], keepdims=1)
    fg_row_min = nd('Reshape', [fg_row_min_2d, 'shape_30'], [[30]])
    
    # fg_col_max[c] = max row with fg (or -1000)
    fg_row_w = nd('Mul', [fg, row_idx_col], [[30, 30]])
    nf_neg2 = nd('Mul', [not_fg, 'c_neg1000'], [[30, 30]])
    fg_row_scored = nd('Add', [fg_row_w, nf_neg2], [[30, 30]])
    fg_col_max_2d = nd('ReduceMax', [fg_row_scored, 'axes0'], [[1, 30]], keepdims=1)
    fg_col_max = nd('Reshape', [fg_col_max_2d, 'shape_30'], [[30]])
    
    # fg_col_min[c] = min row with fg (or 1000)
    nf_pos2 = nd('Mul', [not_fg, 'c_pos1000'], [[30, 30]])
    fg_row_min_scored = nd('Add', [fg_row_w, nf_pos2], [[30, 30]])
    fg_col_min_2d = nd('ReduceMin', [fg_row_min_scored, 'axes0'], [[1, 30]], keepdims=1)
    fg_col_min = nd('Reshape', [fg_col_min_2d, 'shape_30'], [[30]])
    
    # has_fg per col
    fg_col_sum = nd('ReduceSum', [fg, 'axes0'], [[1, 30]], keepdims=1)
    fg_col_sum_1d = nd('Reshape', [fg_col_sum, 'shape_30'], [[30]])
    has_fg_col_b = nd('Greater', [fg_col_sum_1d, 'c_half'], [([30], TensorProto.BOOL)])
    has_fg_col = nd('Cast', [has_fg_col_b], [[30]], to=1)
    no_fg_col = nd('Sub', ['c_one', has_fg_col], [[30]])

    # === Erosion helper for binary vectors ===
    def erode_vec(elig, label):
        """Erode contiguous runs in binary [30] vector.
        Remove first element of run if not at pos 0, last if not at pos 29.
        """
        # Get prev (shift right) and next (shift left)
        e_2d = nd('Reshape', [elig, 'shape_1_30'], [[1, 30]])
        
        # Prev: pad 1 on left of dim1, slice [0:30]
        const(f'pl_{label}', [0, 1, 0, 0], 'i')
        const(f'pv_{label}', [0.0])
        e_pl = nd('Pad', [e_2d, f'pl_{label}', f'pv_{label}'], [[1, 31]])
        const(f'ss_{label}', [0, 0], 'i')
        const(f'se_{label}', [1, 30], 'i')
        const(f'sa_{label}', [0, 1], 'i')
        prev_2d = nd('Slice', [e_pl, f'ss_{label}', f'se_{label}', f'sa_{label}'], [[1, 30]])
        prev = nd('Reshape', [prev_2d, 'shape_30'], [[30]])
        
        # Next: pad 1 on right of dim1, slice [1:31]
        const(f'pr_{label}', [0, 0, 0, 1], 'i')
        e_pr = nd('Pad', [e_2d, f'pr_{label}', f'pv_{label}'], [[1, 31]])
        const(f'ss2_{label}', [0, 1], 'i')
        const(f'se2_{label}', [1, 31], 'i')
        next_2d = nd('Slice', [e_pr, f'ss2_{label}', f'se2_{label}', f'sa_{label}'], [[1, 30]])
        nxt = nd('Reshape', [next_2d, 'shape_30'], [[30]])
        
        # run_start = elig * (1 - prev)
        not_prev = nd('Sub', ['c_one', prev], [[30]])
        run_start = nd('Mul', [elig, not_prev], [[30]])
        
        # run_end = elig * (1 - next)
        not_next = nd('Sub', ['c_one', nxt], [[30]])
        run_end = nd('Mul', [elig, not_next], [[30]])
        
        # Remove start at non-boundary (pos > 0), remove end at non-boundary (pos < 29)
        const(f'p0_{label}', np.array([1.0] + [0.0]*29, dtype=np.float32))
        const(f'p29_{label}', np.array([0.0]*29 + [1.0], dtype=np.float32))
        
        not_p0 = nd('Sub', ['c_one', f'p0_{label}'], [[30]])
        rm_start = nd('Mul', [run_start, not_p0], [[30]])
        
        not_p29 = nd('Sub', ['c_one', f'p29_{label}'], [[30]])
        rm_end = nd('Mul', [run_end, not_p29], [[30]])
        
        e1 = nd('Sub', [elig, rm_start], [[30]])
        e2 = nd('Sub', [e1, rm_end], [[30]])
        # Clip to [0,1]
        eroded = nd('Clip', [e2, 'c_zero_s', 'c_one'], [[30]])
        return eroded

    # === PHASE 1: RIGHT extension ===
    # Eligible: row in core AND fg_row_max < core_c1 (all fg left of core)
    fg_lt_cc1 = nd('Less', [fg_row_max, cc1_m_half], [([30], TensorProto.BOOL)])
    fg_lt_cc1_f = nd('Cast', [fg_lt_cc1], [[30]], to=1)
    right_elig = nd('Mul', [row_in, fg_lt_cc1_f], [[30]])
    right_eroded = erode_vec(right_elig, 'r1')
    
    # Fill: eroded rows, cols > core_c2
    col_gt_cc2 = nd('Greater', ['col_idx', cc2_p_half], [([30], TensorProto.BOOL)])
    col_gt_cc2_f = nd('Cast', [col_gt_cc2], [[30]], to=1)
    re_col = nd('Reshape', [right_eroded, 'shape_30_1'], [[30, 1]])
    cg_row = nd('Reshape', [col_gt_cc2_f, 'shape_1_30'], [[1, 30]])
    right_fill = nd('MatMul', [re_col, cg_row], [[30, 30]])
    result = nd('Max', [core_mask, right_fill], [[30, 30]])

    # === PHASE 1b: LEFT extension ===
    fg_gt_cc2 = nd('Greater', [fg_row_min, cc2_p_half], [([30], TensorProto.BOOL)])
    fg_gt_cc2_f = nd('Cast', [fg_gt_cc2], [[30]], to=1)
    left_elig = nd('Mul', [row_in, fg_gt_cc2_f], [[30]])
    left_eroded = erode_vec(left_elig, 'l1')
    
    col_lt_cc1 = nd('Less', ['col_idx', core_c1], [([30], TensorProto.BOOL)])
    col_lt_cc1_f = nd('Cast', [col_lt_cc1], [[30]], to=1)
    le_col = nd('Reshape', [left_eroded, 'shape_30_1'], [[30, 1]])
    cl_row = nd('Reshape', [col_lt_cc1_f, 'shape_1_30'], [[1, 30]])
    left_fill = nd('MatMul', [le_col, cl_row], [[30, 30]])
    result = nd('Max', [result, left_fill], [[30, 30]])

    # === PHASE 2: Sub-rect post-processing ===
    # For RIGHT sub-rect: eligible rows where fg_row_max < core_c1
    # Within each contiguous run of eligible rows, compute max(fg_row_max)
    # sub_left = max_fg_in_run + 1. If main_c1 - sub_left >= 4, fill from sub_left+1 (eroded)
    
    # Static RIGHT eligibility
    sr_elig = fg_lt_cc1_f  # [30] — already computed in Phase 1
    
    # Propagate max fg_row_max within eligible runs using iterative max
    # masked_val[r] = fg_row_max[r] if eligible else -1000
    sr_not_elig = nd('Sub', ['c_one', sr_elig], [[30]])
    const('c_neg1000_sr', [-1000.0])
    sr_fg_masked = nd('Mul', [fg_row_max, sr_elig], [[30]])
    sr_ne_w = nd('Mul', [sr_not_elig, 'c_neg1000_sr'], [[30]])
    sr_vals = nd('Add', [sr_fg_masked, sr_ne_w], [[30]])  # [30] fg_max where eligible, -1000 elsewhere
    
    # Reshape to [1,1,30,1] for MaxPool propagation along row dimension
    const('shape_1_1_30_1', [1, 1, 30, 1], 'i')
    sr_4d2 = nd('Reshape', [sr_vals, 'shape_1_1_30_1'], [[1, 1, 30, 1]])
    
    # MaxPool with kernel [3,1] propagates max along the row dimension
    # Do 15 iterations to propagate max across runs of length up to 30
    for it in range(15):
        sr_4d2 = nd('MaxPool', [sr_4d2], [[1, 1, 30, 1]], kernel_shape=[3, 1], pads=[1, 0, 1, 0])
        # Re-mask: only keep values in eligible positions
        sr_elig_4d = nd('Reshape', [sr_elig, 'shape_1_1_30_1'], [[1, 1, 30, 1]])
        sr_4d2 = nd('Mul', [sr_4d2, sr_elig_4d], [[1, 1, 30, 1]])
        # Add -1000 for non-eligible
        sr_ne_4d = nd('Reshape', [sr_ne_w, 'shape_1_1_30_1'], [[1, 1, 30, 1]])
        sr_4d2 = nd('Add', [sr_4d2, sr_ne_4d], [[1, 1, 30, 1]])
    
    # Now sr_4d2[r] = max(fg_row_max) across the contiguous eligible run containing r
    sr_run_max = nd('Reshape', [sr_4d2, 'shape_30'], [[30]])
    
    # sub_left per row = sr_run_max + 1
    sr_sub_left = nd('Add', [sr_run_max, 'c_two'], [[30]])  # eroded: +2 from max_fg (= sub_lc + 1)
    
    # Threshold: lc_val - (sr_run_max + 1) >= 4 (per row) — check before erosion
    const('c_3_5', [3.5])
    sr_sub_left_uneroded = nd('Add', [sr_run_max, 'c_one'], [[30]])  # before col erosion
    sr_gap = nd('Sub', [lc_val, sr_sub_left_uneroded], [[30]])
    sr_thresh_b = nd('Greater', [sr_gap, 'c_3_5'], [([30], TensorProto.BOOL)])
    sr_thresh_f = nd('Cast', [sr_thresh_b], [[30]], to=1)
    
    # Erode the eligible vector
    sr_eroded = erode_vec(sr_elig, 'sr')
    
    # Active = eroded AND threshold met
    sr_active = nd('Mul', [sr_eroded, sr_thresh_f], [[30]])
    
    # Fill: for each active row r, fill cols >= sr_sub_left[r]
    # Build [30,30] mask where mask[r,c] = (c >= sr_sub_left[r]) * sr_active[r]
    sr_sub_left_col = nd('Reshape', [sr_sub_left, 'shape_30_1'], [[30, 1]])  # [30,1]
    # col_idx [30] broadcast with sr_sub_left_col [30,1]
    col_ge_srl = nd('Greater', [col_idx_row, nd('Sub', [sr_sub_left_col, 'c_half'], [[30, 1]])],
                   [([30, 30], TensorProto.BOOL)])
    col_ge_srl_f = nd('Cast', [col_ge_srl], [[30, 30]], to=1)
    sr_act_col = nd('Reshape', [sr_active, 'shape_30_1'], [[30, 1]])
    sr_fill = nd('Mul', [nd('MatMul', [sr_act_col, 'ones_1_30'], [[30, 30]]), col_ge_srl_f], [[30, 30]])
    result = nd('Max', [result, sr_fill], [[30, 30]])
    
    # For LEFT sub-rect: eligible rows where fg_row_min > core_c2
    # Propagate min fg_row_min within eligible runs
    sl_elig = fg_gt_cc2_f  # [30]
    sl_not_elig = nd('Sub', ['c_one', sl_elig], [[30]])
    const('c_pos1000_sl', [1000.0])
    sl_fg_masked = nd('Mul', [fg_row_min, sl_elig], [[30]])
    sl_ne_w = nd('Mul', [sl_not_elig, 'c_pos1000_sl'], [[30]])
    sl_vals = nd('Add', [sl_fg_masked, sl_ne_w], [[30]])
    
    # For min propagation: negate, maxpool, negate
    sl_neg = nd('Mul', [sl_vals, nd('Sub', ['c_zero_s', 'c_one'], [[1]])], [[30]])
    sl_4d = nd('Reshape', [sl_neg, 'shape_1_1_30_1'], [[1, 1, 30, 1]])
    
    for it in range(15):
        sl_4d = nd('MaxPool', [sl_4d], [[1, 1, 30, 1]], kernel_shape=[3, 1], pads=[1, 0, 1, 0])
        sl_elig_4d = nd('Reshape', [sl_elig, 'shape_1_1_30_1'], [[1, 1, 30, 1]])
        sl_4d = nd('Mul', [sl_4d, sl_elig_4d], [[1, 1, 30, 1]])
        # Non-eligible: add -1000 (since negated, this means +1000 in original = not participating)
        sl_ne_neg = nd('Reshape', [nd('Mul', [sl_not_elig, 'c_neg1000_sr'], [[30]]), 'shape_1_1_30_1'], [[1, 1, 30, 1]])
        sl_4d = nd('Add', [sl_4d, sl_ne_neg], [[1, 1, 30, 1]])
    
    # Negate back to get min
    sl_run_min_neg = nd('Reshape', [sl_4d, 'shape_30'], [[30]])
    sl_run_min = nd('Mul', [sl_run_min_neg, nd('Sub', ['c_zero_s', 'c_one'], [[1]])], [[30]])
    
    # sub_right per row = sl_run_min - 1
    sl_sub_right_uneroded = nd('Sub', [sl_run_min, 'c_one'], [[30]])  # before erosion
    sl_sub_right = nd('Sub', [sl_run_min, 'c_two'], [[30]])  # eroded: -2 from min_fg
    
    # Threshold: sl_sub_right_uneroded - main_c2 >= 4
    main_c2_val = nd('Sub', [lc_plus_rw, 'c_one'], [[1]])
    sl_gap = nd('Sub', [sl_sub_right_uneroded, main_c2_val], [[30]])
    sl_thresh_b = nd('Greater', [sl_gap, 'c_3_5'], [([30], TensorProto.BOOL)])
    sl_thresh_f = nd('Cast', [sl_thresh_b], [[30]], to=1)
    
    sl_eroded = erode_vec(sl_elig, 'sl')
    sl_active = nd('Mul', [sl_eroded, sl_thresh_f], [[30]])
    
    # Fill cols <= sl_sub_right[r]
    sl_sub_right_col = nd('Reshape', [sl_sub_right, 'shape_30_1'], [[30, 1]])
    col_le_slr = nd('Less', [col_idx_row, nd('Add', [sl_sub_right_col, 'c_half'], [[30, 1]])],
                   [([30, 30], TensorProto.BOOL)])
    col_le_slr_f = nd('Cast', [col_le_slr], [[30, 30]], to=1)
    sl_act_col = nd('Reshape', [sl_active, 'shape_30_1'], [[30, 1]])
    sl_fill = nd('Mul', [nd('MatMul', [sl_act_col, 'ones_1_30'], [[30, 30]]), col_le_slr_f], [[30, 30]])
    result = nd('Max', [result, sl_fill], [[30, 30]])

    # === PHASE 3: DOWN extension ===
    # Eligible: col has result height >= 3 AND all fg in col is above result_top
    # Compute per-column result height and top
    
    # result_col_sum[c] = number of result cells in col c
    res_col_sum = nd('ReduceSum', [result, 'axes0'], [[1, 30]], keepdims=1)
    res_col_sum_1d = nd('Reshape', [res_col_sum, 'shape_30'], [[30]])
    col_h_ge3 = nd('Greater', [res_col_sum_1d, 'c_2_5'], [([30], TensorProto.BOOL)])
    col_h_ge3_f = nd('Cast', [col_h_ge3], [[30]], to=1)
    
    # result_top per col: min row with result
    # Use: result * row_idx + (1-result) * 1000 → ReduceMin along rows
    not_res_pre = nd('Sub', ['c_one', result], [[30, 30]])
    res_row_w = nd('Mul', [result, row_idx_col], [[30, 30]])
    nrp_big = nd('Mul', [not_res_pre, 'c_pos1000'], [[30, 30]])
    res_top_scored = nd('Add', [res_row_w, nrp_big], [[30, 30]])
    res_top_2d = nd('ReduceMin', [res_top_scored, 'axes0'], [[1, 30]], keepdims=1)
    res_top_1d = nd('Reshape', [res_top_2d, 'shape_30'], [[30]])
    
    # fg_col_max < res_top (all fg above result top in that col)
    res_top_m_half = nd('Sub', [res_top_1d, 'c_half'], [[30]])
    fg_above = nd('Less', [fg_col_max, res_top_m_half], [([30], TensorProto.BOOL)])
    fg_above_f = nd('Cast', [fg_above], [[30]], to=1)
    
    down_e1 = nd('Mul', [col_h_ge3_f, fg_above_f], [[30]])
    
    # Extend with adjacent no-fg cols
    de_2d = nd('Reshape', [down_e1, 'shape_1_30'], [[1, 30]])
    const('pl_de', [0, 1, 0, 0], 'i')
    const('pv_de', [0.0])
    de_pl = nd('Pad', [de_2d, 'pl_de', 'pv_de'], [[1, 31]])
    const('ss_de', [0, 0], 'i')
    const('se_de', [1, 30], 'i')
    const('sa_de', [0, 1], 'i')
    de_prev = nd('Slice', [de_pl, 'ss_de', 'se_de', 'sa_de'], [[1, 30]])
    de_prev_1d = nd('Reshape', [de_prev, 'shape_30'], [[30]])
    
    const('pr_de', [0, 0, 0, 1], 'i')
    de_pr = nd('Pad', [de_2d, 'pr_de', 'pv_de'], [[1, 31]])
    const('ss2_de', [0, 1], 'i')
    const('se2_de', [1, 31], 'i')
    de_next = nd('Slice', [de_pr, 'ss2_de', 'se2_de', 'sa_de'], [[1, 30]])
    de_next_1d = nd('Reshape', [de_next, 'shape_30'], [[30]])
    
    adj_p_d = nd('Mul', [de_prev_1d, no_fg_col], [[30]])
    adj_n_d = nd('Mul', [de_next_1d, no_fg_col], [[30]])
    down_ext = nd('Max', [down_e1, adj_p_d], [[30]])
    down_ext2 = nd('Max', [down_ext, adj_n_d], [[30]])
    down_eroded = erode_vec(down_ext2, 'd1')
    
    # Fill: for each eligible col, fill rows > result_bottom[c]
    # result_bottom[c] = max row with result in col c
    res_bot_scored = nd('Add', [res_row_w, nd('Mul', [not_res_pre, 'c_neg1000'], [[30, 30]])], [[30, 30]])
    res_bot_2d = nd('ReduceMax', [res_bot_scored, 'axes0'], [[1, 30]], keepdims=1)
    res_bot_1d = nd('Reshape', [res_bot_2d, 'shape_30'], [[30]])
    
    # For each col c: fill rows > res_bot_1d[c]
    # Build [30,30] mask: mask[r,c] = (row_idx[r] > res_bot_1d[c]) * down_eroded[c]
    res_bot_row = nd('Reshape', [res_bot_1d, 'shape_1_30'], [[1, 30]])  # [1,30]
    row_gt_rb = nd('Greater', [row_idx_col, res_bot_row], [([30, 30], TensorProto.BOOL)])  # [30,30]
    row_gt_rb_f = nd('Cast', [row_gt_rb], [[30, 30]], to=1)
    de_row2 = nd('Reshape', [down_eroded, 'shape_1_30'], [[1, 30]])
    down_mask = nd('Mul', [nd('MatMul', ['ones_30_1', de_row2], [[30, 30]]), row_gt_rb_f], [[30, 30]])
    result = nd('Max', [result, down_mask], [[30, 30]])

    # === PHASE 3b: UP extension ===
    # Eligible: col has result height >= 3 AND all fg below result_bottom
    # Re-compute result column info after DOWN phase
    res_col_sum2 = nd('ReduceSum', [result, 'axes0'], [[1, 30]], keepdims=1)
    res_col_sum2_1d = nd('Reshape', [res_col_sum2, 'shape_30'], [[30]])
    col_h_ge3_2 = nd('Greater', [res_col_sum2_1d, 'c_2_5'], [([30], TensorProto.BOOL)])
    col_h_ge3_2f = nd('Cast', [col_h_ge3_2], [[30]], to=1)
    
    # result_bottom per col (after DOWN)
    not_res_2 = nd('Sub', ['c_one', result], [[30, 30]])
    res_row_w2 = nd('Mul', [result, row_idx_col], [[30, 30]])
    nrp_neg2 = nd('Mul', [not_res_2, 'c_neg1000'], [[30, 30]])
    res_bot_scored2 = nd('Add', [res_row_w2, nrp_neg2], [[30, 30]])
    res_bot_2d2 = nd('ReduceMax', [res_bot_scored2, 'axes0'], [[1, 30]], keepdims=1)
    res_bot_1d2 = nd('Reshape', [res_bot_2d2, 'shape_30'], [[30]])
    
    # fg_col_min > result_bottom (all fg below result)
    res_bot_p_half = nd('Add', [res_bot_1d2, 'c_half'], [[30]])
    fg_below = nd('Greater', [fg_col_min, res_bot_p_half], [([30], TensorProto.BOOL)])
    fg_below_f = nd('Cast', [fg_below], [[30]], to=1)
    
    up_e1 = nd('Mul', [col_h_ge3_2f, fg_below_f], [[30]])
    
    ue_2d = nd('Reshape', [up_e1, 'shape_1_30'], [[1, 30]])
    const('pl_ue', [0, 1, 0, 0], 'i')
    const('pv_ue', [0.0])
    ue_pl = nd('Pad', [ue_2d, 'pl_ue', 'pv_ue'], [[1, 31]])
    const('ss_ue', [0, 0], 'i')
    const('se_ue', [1, 30], 'i')
    const('sa_ue', [0, 1], 'i')
    ue_prev = nd('Slice', [ue_pl, 'ss_ue', 'se_ue', 'sa_ue'], [[1, 30]])
    ue_prev_1d = nd('Reshape', [ue_prev, 'shape_30'], [[30]])
    
    const('pr_ue', [0, 0, 0, 1], 'i')
    ue_pr = nd('Pad', [ue_2d, 'pr_ue', 'pv_ue'], [[1, 31]])
    const('ss2_ue', [0, 1], 'i')
    const('se2_ue', [1, 31], 'i')
    ue_next = nd('Slice', [ue_pr, 'ss2_ue', 'se2_ue', 'sa_ue'], [[1, 30]])
    ue_next_1d = nd('Reshape', [ue_next, 'shape_30'], [[30]])
    
    adj_p_u = nd('Mul', [ue_prev_1d, no_fg_col], [[30]])
    adj_n_u = nd('Mul', [ue_next_1d, no_fg_col], [[30]])
    up_ext = nd('Max', [up_e1, adj_p_u], [[30]])
    up_ext2 = nd('Max', [up_ext, adj_n_u], [[30]])
    up_eroded = erode_vec(up_ext2, 'u1')
    
    # Fill: for each eligible col, fill rows < result_top[c]
    # result_top per col (after DOWN fills)
    nrp_big2 = nd('Mul', [not_res_2, 'c_pos1000'], [[30, 30]])
    res_top_scored2 = nd('Add', [res_row_w2, nrp_big2], [[30, 30]])
    res_top_2d2 = nd('ReduceMin', [res_top_scored2, 'axes0'], [[1, 30]], keepdims=1)
    res_top_1d2 = nd('Reshape', [res_top_2d2, 'shape_30'], [[30]])
    
    # mask[r,c] = (row_idx[r] < res_top_1d2[c]) * up_eroded[c]
    res_top_row2 = nd('Reshape', [res_top_1d2, 'shape_1_30'], [[1, 30]])
    row_lt_rt = nd('Less', [row_idx_col, res_top_row2], [([30, 30], TensorProto.BOOL)])
    row_lt_rt_f = nd('Cast', [row_lt_rt], [[30, 30]], to=1)
    ue_row2 = nd('Reshape', [up_eroded, 'shape_1_30'], [[1, 30]])
    up_mask = nd('Mul', [nd('MatMul', ['ones_30_1', ue_row2], [[30, 30]]), row_lt_rt_f], [[30, 30]])
    result = nd('Max', [result, up_mask], [[30, 30]])

    # === PHASE 4: Second RIGHT/LEFT (width>=3) ===
    # result_row info
    res_row_sum = nd('ReduceSum', [result, 'axes1'], [[30, 1]], keepdims=1)
    res_row_sum_1d = nd('Reshape', [res_row_sum, 'shape_30'], [[30]])
    has_res_b = nd('Greater', [res_row_sum_1d, 'c_half'], [([30], TensorProto.BOOL)])
    has_res = nd('Cast', [has_res_b], [[30]], to=1)
    w_ge3_b = nd('Greater', [res_row_sum_1d, 'c_2_5'], [([30], TensorProto.BOOL)])
    w_ge3 = nd('Cast', [w_ge3_b], [[30]], to=1)
    
    # result left boundary per row
    not_res = nd('Sub', ['c_one', result], [[30, 30]])
    res_col_w = nd('Mul', [result, col_idx_row], [[30, 30]])
    nr_big = nd('Mul', [not_res, 'c_pos1000'], [[30, 30]])
    res_min_scored = nd('Add', [res_col_w, nr_big], [[30, 30]])
    res_min_2d = nd('ReduceMin', [res_min_scored, 'axes1'], [[30, 1]], keepdims=1)
    res_min_1d = nd('Reshape', [res_min_2d, 'shape_30'], [[30]])
    
    # result right boundary per row
    nr_neg = nd('Mul', [not_res, 'c_neg1000'], [[30, 30]])
    res_max_scored = nd('Add', [res_col_w, nr_neg], [[30, 30]])
    res_max_2d = nd('ReduceMax', [res_max_scored, 'axes1'], [[30, 1]], keepdims=1)
    res_max_1d = nd('Reshape', [res_max_2d, 'shape_30'], [[30]])
    
    # RIGHT2: fg_row_max < res_min (all fg left of result)
    res_min_m_half = nd('Sub', [res_min_1d, 'c_half'], [[30]])
    fg_lt_res = nd('Less', [fg_row_max, res_min_m_half], [([30], TensorProto.BOOL)])
    fg_lt_res_f = nd('Cast', [fg_lt_res], [[30]], to=1)
    r2_elig = nd('Mul', [has_res, w_ge3], [[30]])
    r2_elig2 = nd('Mul', [r2_elig, fg_lt_res_f], [[30]])
    r2_eroded = erode_vec(r2_elig2, 'r2')
    
    # Fill cols > res_max per row
    rm_col = nd('Reshape', [res_max_1d, 'shape_30_1'], [[30, 1]])
    col_gt_rm = nd('Greater', [col_idx_row, rm_col], [([30, 30], TensorProto.BOOL)])
    col_gt_rm_f = nd('Cast', [col_gt_rm], [[30, 30]], to=1)
    r2e_col = nd('Reshape', [r2_eroded, 'shape_30_1'], [[30, 1]])
    r2_mask = nd('MatMul', [r2e_col, 'ones_1_30'], [[30, 30]])
    r2_fill = nd('Mul', [r2_mask, col_gt_rm_f], [[30, 30]])
    result = nd('Max', [result, r2_fill], [[30, 30]])
    
    # LEFT2: fg_row_min > res_max (all fg right of result)
    res_max_p_half = nd('Add', [res_max_1d, 'c_half'], [[30]])
    fg_gt_res = nd('Greater', [fg_row_min, res_max_p_half], [([30], TensorProto.BOOL)])
    fg_gt_res_f = nd('Cast', [fg_gt_res], [[30]], to=1)
    l2_elig = nd('Mul', [has_res, w_ge3], [[30]])
    l2_elig2 = nd('Mul', [l2_elig, fg_gt_res_f], [[30]])
    l2_eroded = erode_vec(l2_elig2, 'l2')
    
    # Fill cols < res_min per row
    rm2_col = nd('Reshape', [res_min_1d, 'shape_30_1'], [[30, 1]])
    col_lt_rm = nd('Less', [col_idx_row, rm2_col], [([30, 30], TensorProto.BOOL)])
    col_lt_rm_f = nd('Cast', [col_lt_rm], [[30, 30]], to=1)
    l2e_col = nd('Reshape', [l2_eroded, 'shape_30_1'], [[30, 1]])
    l2_mask = nd('MatMul', [l2e_col, 'ones_1_30'], [[30, 30]])
    l2_fill = nd('Mul', [l2_mask, col_lt_rm_f], [[30, 30]])
    result = nd('Max', [result, l2_fill], [[30, 30]])

    # === STEP 7: Mask to background + output ===
    final_mask = nd('Mul', [result, not_fg], [[30, 30]])
    fm_4d = nd('Reshape', [final_mask, 'shape_1_1_30_30'], [[1, 1, 30, 30]])
    delta = nd('Mul', [fm_4d, 'color3_delta'], [[1, 10, 30, 30]])
    nodes.append(helper.make_node('Add', ['input', delta], ['output']))

    # === Build model ===
    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, 'task255_v2', [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__':
    print("Building Task 255 ONNX model v2...")
    model = build_task255()
    
    os.makedirs('/app/repo/medal-solvers/optimized', exist_ok=True)
    output_path = '/app/repo/medal-solvers/optimized/task255.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 (limit: {int(1.44*1024*1024):,})")
    
    # Quick validation
    import onnxruntime as ort
    with open('task-data/task255.json') as f:
        data = json.load(f)
    
    sess = ort.InferenceSession(output_path)
    all_examples = data['train'] + data['test'] + data['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 = (out != exp)
                ch3_pred = set(zip(*np.where(out[0, 3] > 0)))
                ch3_exp = set(zip(*np.where(exp[0, 3] > 0)))
                extra = ch3_pred - ch3_exp
                missing = ch3_exp - ch3_pred
                print(f"  FAIL {i}: ch3 extra={len(extra)}, missing={len(missing)}")
    
    print(f"\nResults: {right_count} pass, {wrong_count} fail out of {len(all_examples)}")