"""
Build optimized ONNX model for Task 319.

Algorithm (267/267 verified):
1. bg = channel with most pixels
2. Template = largest non-bg channel
3. Downsample template at 4 (2,2)-block offsets, pick first with uniform blocks
4. Conv(candidate, pattern, pads=14) for cross-correlation
5. max(conv) == sum(pattern) → match (fewest-pixel candidate wins ties)
6. Pass 2: second-largest template if pass 1 fails
7. Output = matched candidate shifted to (0,0) + bg fill within bbox

Score: 11.580 pts (267/267 pass, memory=661K, params=12K)

Usage:
    python build_task319_onnx.py --task-data-dir ../task-data
"""
from onnx import TensorProto
import numpy as np
from onnx_builder import OnnxBuilder, build_and_validate

H, W, C = 30, 30, 10


def build_task319():
    b = OnnxBuilder()
    const, nd = b.const, b.nd

    # === CONSTANTS ===
    const('c_half', [0.5])
    const('c_one', [1.0])
    const('c_eps', [0.01])
    const('c_099', [0.99])
    const('c_big', [1e9])
    const('c_neg_big', [-1e9])
    const('c_30f', [30.0])
    const('c_899f', [899.0])
    const('c_zerof', [0.0])
    const('c_big_score', [1000.0])

    const('axes23', [2, 3], 'i')
    const('axes1', [1], 'i')
    const('shape_1_10', [1, 10], 'i')
    const('shape_1_1', [1, 1], 'i')
    const('shape_1', [1], 'i')
    const('shape_1_1_1_1', [1, 1, 1, 1], 'i')
    const('shape_1_10_1_1', [1, 10, 1, 1], 'i')
    const('shape_30_30', [30, 30], 'i')
    const('shape_900', [900], 'i')
    const('shape_1_1_30_30', [1, 1, 30, 30], 'i')

    const('depth_10', [10.0])
    const('oh_vals', [0.0, 1.0])

    # Slice/Pad for offsets
    const('sl_10_s', [0, 0, 1, 0], 'i')
    const('sl_10_e', [1, 1, 30, 30], 'i')
    const('sl_axes', [0, 1, 2, 3], 'i')
    const('pd_10', [0, 0, 0, 0, 0, 0, 1, 0], 'i')
    const('sl_01_s', [0, 0, 0, 1], 'i')
    const('sl_01_e', [1, 1, 30, 30], 'i')
    const('pd_01', [0, 0, 0, 0, 0, 0, 0, 1], 'i')
    const('sl_11_s', [0, 0, 1, 1], 'i')
    const('sl_11_e', [1, 1, 30, 30], 'i')
    const('pd_11', [0, 0, 0, 0, 0, 0, 1, 1], 'i')
    const('pad_val', [0.0])

    # Row/col grids
    const('row_idx_30_1', np.arange(30, dtype=np.float32).reshape(30, 1))
    const('col_idx_1_30', np.arange(30, dtype=np.float32).reshape(1, 30))
    const('ones_1_30', np.ones((1, 30), dtype=np.float32))
    const('ones_30_1', np.ones((30, 1), dtype=np.float32))

    for i in range(10):
        const(f'idx_{i}', [i], 'i')

    # === STEP 1: Channel pixel counts ===
    ch_sums = nd('ReduceSum', ['input', 'axes23'], [[1, 10, 1, 1]], keepdims=1)
    ch_sums_2d = nd('Reshape', [ch_sums, 'shape_1_10'], [[1, 10]])

    # === STEP 2: Find bg, template, template2 ===
    bg_idx = nd('ArgMax', [ch_sums_2d], [([1, 1], TensorProto.INT64)], axis=1, keepdims=1)
    bg_idx_1d = nd('Reshape', [bg_idx, 'shape_1'], [([1], TensorProto.INT64)])
    bg_oh = nd('OneHot', [bg_idx_1d, 'depth_10', 'oh_vals'], [[1, 10]], axis=1)

    not_bg = nd('Sub', ['c_one', bg_oh], [[1, 10]])
    sums_no_bg = nd('Mul', [ch_sums_2d, not_bg], [[1, 10]])

    tmpl_idx = nd('ArgMax', [sums_no_bg], [([1, 1], TensorProto.INT64)], axis=1, keepdims=1)
    tmpl_idx_1d = nd('Reshape', [tmpl_idx, 'shape_1'], [([1], TensorProto.INT64)])
    tmpl_oh = nd('OneHot', [tmpl_idx_1d, 'depth_10', 'oh_vals'], [[1, 10]], axis=1)

    not_tmpl = nd('Sub', ['c_one', tmpl_oh], [[1, 10]])
    sums_no_bg_tmpl = nd('Mul', [sums_no_bg, not_tmpl], [[1, 10]])

    tmpl2_idx = nd('ArgMax', [sums_no_bg_tmpl], [([1, 1], TensorProto.INT64)], axis=1, keepdims=1)
    tmpl2_idx_1d = nd('Reshape', [tmpl2_idx, 'shape_1'], [([1], TensorProto.INT64)])
    tmpl2_oh = nd('OneHot', [tmpl2_idx_1d, 'depth_10', 'oh_vals'], [[1, 10]], axis=1)

    # === STEP 3: Extract template masks ===
    tmpl_mask = nd('Gather', ['input', tmpl_idx_1d], [[1, 1, 30, 30]], axis=1)
    tmpl2_mask = nd('Gather', ['input', tmpl2_idx_1d], [[1, 1, 30, 30]], axis=1)

    # === Helpers ===
    def pool_at_offset(mask, offset):
        if offset == (0, 0):
            return nd('AveragePool', [mask], [[1, 1, 15, 15]],
                     kernel_shape=[2, 2], strides=[2, 2])
        elif offset == (1, 0):
            s = nd('Slice', [mask, 'sl_10_s', 'sl_10_e', 'sl_axes'], [[1, 1, 29, 30]])
            p = nd('Pad', [s, 'pd_10', 'pad_val'], [[1, 1, 30, 30]])
            return nd('AveragePool', [p], [[1, 1, 15, 15]], kernel_shape=[2, 2], strides=[2, 2])
        elif offset == (0, 1):
            s = nd('Slice', [mask, 'sl_01_s', 'sl_01_e', 'sl_axes'], [[1, 1, 30, 29]])
            p = nd('Pad', [s, 'pd_01', 'pad_val'], [[1, 1, 30, 30]])
            return nd('AveragePool', [p], [[1, 1, 15, 15]], kernel_shape=[2, 2], strides=[2, 2])
        else:
            s = nd('Slice', [mask, 'sl_11_s', 'sl_11_e', 'sl_axes'], [[1, 1, 29, 29]])
            p = nd('Pad', [s, 'pd_11', 'pad_val'], [[1, 1, 30, 30]])
            return nd('AveragePool', [p], [[1, 1, 15, 15]], kernel_shape=[2, 2], strides=[2, 2])

    def downsample_and_select(mask, label):
        """Try 4 offsets, pick first valid. Returns (pattern [1,1,15,15], sum, any_valid)."""
        patterns, valids, sums = [], [], []
        for off in [(0,0), (1,0), (0,1), (1,1)]:
            ds = pool_at_offset(mask, off)
            one_m = nd('Sub', ['c_one', ds], [[1, 1, 15, 15]])
            nu = nd('Mul', [ds, one_m], [[1, 1, 15, 15]])
            nu_sum = nd('ReduceSum', [nu, 'axes23'], [[1, 1, 1, 1]], keepdims=1)
            v_b = nd('Less', [nu_sum, 'c_eps'], [([1, 1, 1, 1], TensorProto.BOOL)])
            v_f = nd('Cast', [v_b], [[1, 1, 1, 1]], to=1)
            ds_b = nd('Greater', [ds, 'c_half'], [([1, 1, 15, 15], TensorProto.BOOL)])
            ds_bin = nd('Cast', [ds_b], [[1, 1, 15, 15]], to=1)
            p_sum = nd('ReduceSum', [ds_bin, 'axes23'], [[1, 1, 1, 1]], keepdims=1)
            nz_b = nd('Greater', [p_sum, 'c_half'], [([1, 1, 1, 1], TensorProto.BOOL)])
            nz_f = nd('Cast', [nz_b], [[1, 1, 1, 1]], to=1)
            patterns.append(ds_bin)
            valids.append(nd('Mul', [v_f, nz_f], [[1, 1, 1, 1]]))
            sums.append(p_sum)

        w = [None] * 4
        w[0] = valids[0]
        remaining = nd('Sub', ['c_one', w[0]], [[1, 1, 1, 1]])
        for i in range(1, 4):
            w[i] = nd('Mul', [remaining, valids[i]], [[1, 1, 1, 1]])
            if i < 3:
                remaining = nd('Mul', [remaining, nd('Sub', ['c_one', valids[i]], [[1, 1, 1, 1]])], [[1, 1, 1, 1]])

        sel = nd('Mul', [w[0], patterns[0]], [[1, 1, 15, 15]])
        for i in range(1, 4):
            sel = nd('Add', [sel, nd('Mul', [w[i], patterns[i]], [[1, 1, 15, 15]])], [[1, 1, 15, 15]])
        s_sel = nd('Mul', [w[0], sums[0]], [[1, 1, 1, 1]])
        for i in range(1, 4):
            s_sel = nd('Add', [s_sel, nd('Mul', [w[i], sums[i]], [[1, 1, 1, 1]])], [[1, 1, 1, 1]])
        av = w[0]
        for i in range(1, 4):
            av = nd('Add', [av, w[i]], [[1, 1, 1, 1]])
        return sel, s_sel, av

    ds1, sum1, valid1 = downsample_and_select(tmpl_mask, 't1')
    ds2, sum2, valid2 = downsample_and_select(tmpl2_mask, 't2')

    # === STEP 4: Cross-correlate with candidates ===
    def find_matches(ds_pattern, ds_sum, any_valid, exclude_oh, label):
        scores = []
        thresh = nd('Sub', [ds_sum, 'c_099'], [[1, 1, 1, 1]])
        for c in range(10):
            cand = nd('Gather', ['input', f'idx_{c}'], [[1, 1, 30, 30]], axis=1)
            conv = nd('Conv', [cand, ds_pattern], [[1, 1, 44, 44]], pads=[14, 14, 14, 14])
            mx = nd('ReduceMax', [conv, 'axes23'], [[1, 1, 1, 1]], keepdims=1)
            match_b = nd('Greater', [mx, thresh], [([1, 1, 1, 1], TensorProto.BOOL)])
            match_f = nd('Cast', [match_b], [[1, 1, 1, 1]], to=1)
            c_oh = nd('OneHot', [f'idx_{c}', 'depth_10', 'oh_vals'], [[1, 10]], axis=1)
            is_bg = nd('ReduceSum', [nd('Mul', [c_oh, bg_oh], [[1, 10]]), 'axes1'], [[1, 1]], keepdims=1)
            is_bg_4d = nd('Reshape', [is_bg, 'shape_1_1_1_1'], [[1, 1, 1, 1]])
            is_excl = nd('ReduceSum', [nd('Mul', [c_oh, exclude_oh], [[1, 10]]), 'axes1'], [[1, 1]], keepdims=1)
            is_excl_4d = nd('Reshape', [is_excl, 'shape_1_1_1_1'], [[1, 1, 1, 1]])
            not_bg_c = nd('Sub', ['c_one', is_bg_4d], [[1, 1, 1, 1]])
            not_excl_c = nd('Sub', ['c_one', is_excl_4d], [[1, 1, 1, 1]])
            cand_sum = nd('ReduceSum', [cand, 'axes23'], [[1, 1, 1, 1]], keepdims=1)
            const(f'bias_{label}_{c}', [(10 - c) * 0.01])
            priority = nd('Add', [nd('Sub', ['c_big_score', cand_sum], [[1, 1, 1, 1]]),
                                  f'bias_{label}_{c}'], [[1, 1, 1, 1]])
            sc = nd('Mul', [match_f, not_bg_c], [[1, 1, 1, 1]])
            sc = nd('Mul', [sc, not_excl_c], [[1, 1, 1, 1]])
            sc = nd('Mul', [sc, any_valid], [[1, 1, 1, 1]])
            sc = nd('Mul', [sc, priority], [[1, 1, 1, 1]])
            scores.append(sc)
        return scores

    scores1 = find_matches(ds1, sum1, valid1, tmpl_oh, 'p1')
    scores2 = find_matches(ds2, sum2, valid2, tmpl2_oh, 'p2')

    # === Winner selection ===
    def select_winner(scores, label):
        concat = nd('Concat', scores, [[1, 10, 1, 1]], axis=1)
        mx = nd('ReduceMax', [concat, 'axes1'], [[1, 1, 1, 1]], keepdims=1)
        any_b = nd('Greater', [mx, 'c_half'], [([1, 1, 1, 1], TensorProto.BOOL)])
        any_f = nd('Cast', [any_b], [[1, 1, 1, 1]], to=1)
        win_b = nd('Equal', [concat, mx], [([1, 10, 1, 1], TensorProto.BOOL)])
        win_f = nd('Cast', [win_b], [[1, 10, 1, 1]], to=1)
        return nd('Mul', [win_f, any_f], [[1, 10, 1, 1]])

    p1_winner = select_winner(scores1, 'p1')
    p1_any = nd('ReduceMax', [p1_winner, 'axes1'], [[1, 1, 1, 1]], keepdims=1)
    not_pass1 = nd('Sub', ['c_one', p1_any], [[1, 1, 1, 1]])

    scores2_gated = [nd('Mul', [s, not_pass1], [[1, 1, 1, 1]]) for s in scores2]
    p2_winner = select_winner(scores2_gated, 'p2')

    combined_winner = nd('Add', [p1_winner, p2_winner], [[1, 10, 1, 1]])

    # === STEP 5: Build matched mask ===
    const('zeros_1_1_30_30', np.zeros((1, 1, 30, 30), dtype=np.float32))
    const('zeros_1_10', np.zeros((1, 10), dtype=np.float32))
    matched_mask = 'zeros_1_1_30_30'
    matched_color_oh = 'zeros_1_10'

    for c in range(10):
        wc = nd('Gather', [combined_winner, f'idx_{c}'], [[1, 1, 1, 1]], axis=1)
        ch = nd('Gather', ['input', f'idx_{c}'], [[1, 1, 30, 30]], axis=1)
        matched_mask = nd('Add', [matched_mask, nd('Mul', [wc, ch], [[1, 1, 30, 30]])], [[1, 1, 30, 30]])
        c_oh = nd('OneHot', [f'idx_{c}', 'depth_10', 'oh_vals'], [[1, 10]], axis=1)
        wc2 = nd('Reshape', [wc, 'shape_1_1'], [[1, 1]])
        matched_color_oh = nd('Add', [matched_color_oh, nd('Mul', [wc2, c_oh], [[1, 10]])], [[1, 10]])

    # === STEP 6: Shift to origin ===
    mm_2d = nd('Reshape', [matched_mask, 'shape_30_30'], [[30, 30]])
    not_mm = nd('Sub', ['c_one', mm_2d], [[30, 30]])

    row_grid = nd('MatMul', ['row_idx_30_1', 'ones_1_30'], [[30, 30]])
    col_grid = nd('MatMul', ['ones_30_1', 'col_idx_1_30'], [[30, 30]])

    const('axes01', [0, 1], 'i')
    rmin_s = nd('Add', [nd('Mul', [mm_2d, row_grid], [[30, 30]]), nd('Mul', [not_mm, 'c_big'], [[30, 30]])], [[30, 30]])
    r_min = nd('Reshape', [nd('ReduceMin', [rmin_s, 'axes01'], [[1, 1]], keepdims=1), 'shape_1'], [[1]])
    cmin_s = nd('Add', [nd('Mul', [mm_2d, col_grid], [[30, 30]]), nd('Mul', [not_mm, 'c_big'], [[30, 30]])], [[30, 30]])
    c_min = nd('Reshape', [nd('ReduceMin', [cmin_s, 'axes01'], [[1, 1]], keepdims=1), 'shape_1'], [[1]])
    rmax_s = nd('Add', [nd('Mul', [mm_2d, row_grid], [[30, 30]]), nd('Mul', [not_mm, 'c_neg_big'], [[30, 30]])], [[30, 30]])
    r_max = nd('Reshape', [nd('ReduceMax', [rmax_s, 'axes01'], [[1, 1]], keepdims=1), 'shape_1'], [[1]])
    cmax_s = nd('Add', [nd('Mul', [mm_2d, col_grid], [[30, 30]]), nd('Mul', [not_mm, 'c_neg_big'], [[30, 30]])], [[30, 30]])
    c_max = nd('Reshape', [nd('ReduceMax', [cmax_s, 'axes01'], [[1, 1]], keepdims=1), 'shape_1'], [[1]])

    bbox_h = nd('Add', [nd('Sub', [r_max, r_min], [[1]]), 'c_one'], [[1]])
    bbox_w = nd('Add', [nd('Sub', [c_max, c_min], [[1]]), 'c_one'], [[1]])

    src_row = nd('Add', ['row_idx_30_1', r_min], [[30, 1]])
    src_col = nd('Add', ['col_idx_1_30', c_min], [[1, 30]])
    const('c_29f', [29.0])
    src_row_clip = nd('Clip', [src_row, 'c_zerof', 'c_29f'], [[30, 1]])
    src_col_clip = nd('Clip', [src_col, 'c_zerof', 'c_29f'], [[1, 30]])
    src_row_x30 = nd('Mul', [src_row_clip, 'c_30f'], [[30, 1]])
    flat_idx_f = nd('Add', [nd('MatMul', [src_row_x30, 'ones_1_30'], [[30, 30]]),
                            nd('MatMul', ['ones_30_1', src_col_clip], [[30, 30]])], [[30, 30]])
    flat_idx_clip = nd('Clip', [flat_idx_f, 'c_zerof', 'c_899f'], [[30, 30]])
    flat_idx_i = nd('Cast', [flat_idx_clip], [([30, 30], TensorProto.INT64)], to=7)
    flat_idx_1d = nd('Reshape', [flat_idx_i, 'shape_900'], [([900], TensorProto.INT64)])

    mm_flat = nd('Reshape', [mm_2d, 'shape_900'], [[900]])
    shifted_flat = nd('Gather', [mm_flat, flat_idx_1d], [[900]], axis=0)
    shifted_2d = nd('Reshape', [shifted_flat, 'shape_30_30'], [[30, 30]])

    bbox_h_m = nd('Sub', [bbox_h, 'c_half'], [[1]])
    bbox_w_m = nd('Sub', [bbox_w, 'c_half'], [[1]])
    rv_b = nd('Less', ['row_idx_30_1', bbox_h_m], [([30, 1], TensorProto.BOOL)])
    rv = nd('Cast', [rv_b], [[30, 1]], to=1)
    cv_b = nd('Less', ['col_idx_1_30', bbox_w_m], [([1, 30], TensorProto.BOOL)])
    cv = nd('Cast', [cv_b], [[1, 30]], to=1)
    bbox_valid = nd('MatMul', [rv, cv], [[30, 30]])
    shifted_masked = nd('Mul', [shifted_2d, bbox_valid], [[30, 30]])

    # === STEP 7: Output one-hot ===
    sm_4d = nd('Reshape', [shifted_masked, 'shape_1_1_30_30'], [[1, 1, 30, 30]])
    bv_4d = nd('Reshape', [bbox_valid, 'shape_1_1_30_30'], [[1, 1, 30, 30]])
    not_sm_bv = nd('Mul', [nd('Sub', ['c_one', sm_4d], [[1, 1, 30, 30]]), bv_4d], [[1, 1, 30, 30]])

    color_oh_4d = nd('Reshape', [matched_color_oh, 'shape_1_10_1_1'], [[1, 10, 1, 1]])
    bg_oh_4d = nd('Reshape', [bg_oh, 'shape_1_10_1_1'], [[1, 10, 1, 1]])

    out_color = nd('Mul', [sm_4d, color_oh_4d], [[1, 10, 30, 30]])
    out_bg = nd('Mul', [not_sm_bv, bg_oh_4d], [[1, 10, 30, 30]])
    final = nd('Add', [out_color, out_bg], [[1, 10, 30, 30]])

    return b.finish('task319', last_tensor=final)


if __name__ == '__main__':
    build_and_validate(build_task319, task_num=319)