"""Build optimized ONNX model for Task 084.

Algorithm (175/175 verified):
1. Input has a colored column at col 0 (N rows tall = grid is NxN)
2. Draw anti-diagonal (color 2) from (0, N-1) to (N-2, 1)
3. Fill bottom row (N-1, cols 1..N-1) with color 4
4. Keep original column 0

ONNX approach:
- Detect N from sum of non-bg channels
- Use row/col grids to compute masks for anti-diagonal and bottom row
- Combine into output using channel selection

Score: 175/175 pass, file size 10KB (~12 pts estimated)

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


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

    const('c_one', [1.0])
    const('c_zero', [0.0])
    const('c_half', [0.5])
    const('axes123', [1, 2, 3], 'i')
    const('axes23', [2, 3], 'i')
    const('axes1', [1], 'i')
    const('shape_1_1_30_30', [1, 1, 30, 30], 'i')
    const('shape_1_1_1_1', [1, 1, 1, 1], 'i')
    const('shape_1_10_1_1', [1, 10, 1, 1], 'i')
    
    row_vals = np.arange(30, dtype=np.float32).reshape(1, 1, 30, 1) * np.ones((1, 1, 1, 30), dtype=np.float32)
    col_vals = np.arange(30, dtype=np.float32).reshape(1, 1, 1, 30) * np.ones((1, 1, 30, 1), dtype=np.float32)
    const('row_grid', row_vals)
    const('col_grid', col_vals)
    
    const('idx_0', [0], 'i')
    const('idx_1', [1], 'i')
    const('idx_2', [2], 'i')
    const('idx_10', [10], 'i')
    const('depth_10', [10.0])
    const('oh_vals', [0.0, 1.0])
    const('ch2_idx', [2], 'i')
    const('ch4_idx', [4], 'i')

    nonbg = nd('Slice', ['input', 'idx_1', 'idx_10', 'idx_1'], [[1, 9, 30, 30]])
    n_total = nd('ReduceSum', [nonbg, 'axes123'], [[1, 1, 1, 1]], keepdims=1)
    
    row_plus_col = nd('Add', ['row_grid', 'col_grid'], [[1, 1, 30, 30]])
    n_minus_1 = nd('Sub', [n_total, 'c_one'], [[1, 1, 1, 1]])
    
    on_antidiag = nd('Equal', [row_plus_col, n_minus_1], [([1, 1, 30, 30], TensorProto.BOOL)])
    col_gt_0 = nd('Greater', ['col_grid', 'c_zero'], [([1, 1, 30, 30], TensorProto.BOOL)])
    row_lt_nm1 = nd('Less', ['row_grid', n_minus_1], [([1, 1, 30, 30], TensorProto.BOOL)])
    ad_and_col = nd('And', [on_antidiag, col_gt_0], [([1, 1, 30, 30], TensorProto.BOOL)])
    antidiag_mask = nd('And', [ad_and_col, row_lt_nm1], [([1, 1, 30, 30], TensorProto.BOOL)])
    
    on_bottom = nd('Equal', ['row_grid', n_minus_1], [([1, 1, 30, 30], TensorProto.BOOL)])
    bottom_and_col = nd('And', [on_bottom, col_gt_0], [([1, 1, 30, 30], TensorProto.BOOL)])
    col_lt_n_bool = nd('Less', ['col_grid', n_total], [([1, 1, 30, 30], TensorProto.BOOL)])
    bottom_mask = nd('And', [bottom_and_col, col_lt_n_bool], [([1, 1, 30, 30], TensorProto.BOOL)])
    
    antidiag_f = nd('Cast', [antidiag_mask], [[1, 1, 30, 30]], to=1)
    bottom_f = nd('Cast', [bottom_mask], [[1, 1, 30, 30]], to=1)
    
    col_eq_0 = nd('Equal', ['col_grid', 'c_zero'], [([1, 1, 30, 30], TensorProto.BOOL)])
    col_eq_0_f = nd('Cast', [col_eq_0], [[1, 1, 30, 30]], to=1)
    row_lt_n = nd('Less', ['row_grid', n_total], [([1, 1, 30, 30], TensorProto.BOOL)])
    row_lt_n_f = nd('Cast', [row_lt_n], [[1, 1, 30, 30]], to=1)
    col_lt_n = nd('Less', ['col_grid', n_total], [([1, 1, 30, 30], TensorProto.BOOL)])
    col_lt_n_f = nd('Cast', [col_lt_n], [[1, 1, 30, 30]], to=1)
    inside_grid_f = nd('Mul', [row_lt_n_f, col_lt_n_f], [[1, 1, 30, 30]])
    col0_mask = nd('Mul', [col_eq_0_f, inside_grid_f], [[1, 1, 30, 30]])
    
    col0_values = nd('Mul', ['input', col0_mask], [[1, 10, 30, 30]])
    
    ch2_onehot = nd('OneHot', ['ch2_idx', 'depth_10', 'oh_vals'], [[1, 10]], axis=1)
    ch2_onehot_4d = nd('Reshape', [ch2_onehot, 'shape_1_10_1_1'], [[1, 10, 1, 1]])
    antidiag_ch2 = nd('Mul', [ch2_onehot_4d, antidiag_f], [[1, 10, 30, 30]])
    
    ch4_onehot = nd('OneHot', ['ch4_idx', 'depth_10', 'oh_vals'], [[1, 10]], axis=1)
    ch4_onehot_4d = nd('Reshape', [ch4_onehot, 'shape_1_10_1_1'], [[1, 10, 1, 1]])
    bottom_ch4 = nd('Mul', [ch4_onehot_4d, bottom_f], [[1, 10, 30, 30]])
    
    occupied_1 = nd('Add', [col0_mask, antidiag_f], [[1, 1, 30, 30]])
    occupied = nd('Add', [occupied_1, bottom_f], [[1, 1, 30, 30]])
    ch0_mask_raw = nd('Sub', [inside_grid_f, occupied], [[1, 1, 30, 30]])
    
    const('ch0_idx', [0], 'i')
    ch0_onehot = nd('OneHot', ['ch0_idx', 'depth_10', 'oh_vals'], [[1, 10]], axis=1)
    ch0_onehot_4d = nd('Reshape', [ch0_onehot, 'shape_1_10_1_1'], [[1, 10, 1, 1]])
    bg_ch0 = nd('Mul', [ch0_onehot_4d, ch0_mask_raw], [[1, 10, 30, 30]])
    
    out_1 = nd('Add', [col0_values, antidiag_ch2], [[1, 10, 30, 30]])
    out_2 = nd('Add', [out_1, bottom_ch4], [[1, 10, 30, 30]])
    final = nd('Add', [out_2, bg_ch0], [[1, 10, 30, 30]])
    
    return b.finish('task084', last_tensor=final)


if __name__ == '__main__':
    build_and_validate(build_task084, task_num=84)