"""
Build Task 153 ONNX model — Rebuilt using ONLY patterns from task319/255.

Avoid:
- Multi-dimensional Gather indices [9,3,3] (suspected Kaggle discrepancy cause)
- Broadcast Add [9,1,3,3]+[1,9,3,3]→[9,9,3,3] (suspected cause)

Use instead:
- Flat [100] mask + scalar-index Gather operations
- MatMul for pair checking (proven in task319)
- Per-placement computation using explicit row/col offsets
- Only ops proven to work: MatMul, ReduceSum, Cast, Gather(axis=0, scalar), 
  Reshape, Add, Mul, Less, Greater, ArgMax, Concat, Pad

Strategy (keeps similar node count to v3 but avoids problematic patterns):
1. Extract masks, find bbox
2. For each placement k (0-8), compute the flat-index-into-100 for each of 9 output cells
3. Gather from flat mask → get 9 values per placement
4. Stack into [9,9] matrix per color
5. MatMul P1 @ P2^T to get overlap matrix [9,9]
6. Find valid complement pair, extract, output
"""
import sys
import os
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))

from onnx import TensorProto
import numpy as np
from onnx_builder import OnnxBuilder, build_and_validate


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

    # === CONSTANTS ===
    const('c_half', [0.5])
    const('c_one', [1.0])
    const('c_big', [1e9])
    const('c_8_5', [8.5])
    const('c_9f', [9.0])
    const('c_zero', [0.0])
    const('c_99f', [99.0])
    const('c_10f', [10.0])

    const('axes23', [2, 3], 'i')
    const('axes1', [1], 'i')
    const('axes01', [0, 1], 'i')

    const('shape_1_10', [1, 10], 'i')
    const('shape_1', [1], 'i')
    const('shape_1_10_1_1', [1, 10, 1, 1], 'i')
    const('shape_81', [81], 'i')
    const('shape_1_1_3_3', [1, 1, 3, 3], 'i')
    const('shape_10_10', [10, 10], 'i')
    const('shape_100', [100], 'i')
    const('shape_1_9', [1, 9], 'i')
    const('shape_9_1', [9, 1], 'i')
    
    const('depth_10', [10.0])
    const('oh_vals', [0.0, 1.0])
    
    # Row/col grids for 10x10
    const('row_idx_10_1', np.arange(10, dtype=np.float32).reshape(10, 1))
    const('col_idx_1_10', np.arange(10, dtype=np.float32).reshape(1, 10))
    const('ones_1_10', np.ones((1, 10), dtype=np.float32))
    const('ones_10_1', np.ones((10, 1), dtype=np.float32))

    # Background mask
    const('bg_mask_10', np.array([[0, 1, 1, 1, 1, 1, 1, 1, 1, 1]], dtype=np.float32))
    
    # Pad for output
    const('pad_10_3_to_30', [0, 0, 0, 0, 0, 0, 27, 27], 'i')
    const('pad_val_zero', [0.0])
    
    # Precompute constant row/col offsets: [9, 9] matrices
    # row_off[k, i] = (i//3) - (k//3), col_off[k, i] = (i%3) - (k%3)
    row_off = np.zeros((9, 9), dtype=np.float32)
    col_off = np.zeros((9, 9), dtype=np.float32)
    for k in range(9):
        dr = k // 3
        dc = k % 3
        for i in range(9):
            r = i // 3
            c = i % 3
            row_off[k, i] = r - dr
            col_off[k, i] = c - dc
    
    const('row_off_9_9', row_off)
    const('col_off_9_9', col_off)
    
    # Static validity: row_off[k][i] >= 0 AND col_off[k][i] >= 0
    valid_static = ((row_off >= 0) & (col_off >= 0)).astype(np.float32)
    const('valid_static_9_9', valid_static)

    # === STEP 1: Slice input to 10x10 ===
    const('sl_start', [0, 0, 0, 0], 'i')
    const('sl_end', [1, 10, 10, 10], 'i')
    const('sl_axes_4', [0, 1, 2, 3], 'i')
    inp_10 = nd('Slice', ['input', 'sl_start', 'sl_end', 'sl_axes_4'], [[1, 10, 10, 10]])
    
    # === STEP 2: Channel pixel counts ===
    ch_sums = nd('ReduceSum', [inp_10, 'axes23'], [[1, 10, 1, 1]], keepdims=1)
    ch_sums_2d = nd('Reshape', [ch_sums, 'shape_1_10'], [[1, 10]])
    ch_sums_nobg = nd('Mul', [ch_sums_2d, 'bg_mask_10'], [[1, 10]])

    # === STEP 3: Find color1 and color2 ===
    c1_idx = nd('ArgMax', [ch_sums_nobg], [([1, 1], TensorProto.INT64)], axis=1, keepdims=1)
    c1_idx_1d = nd('Reshape', [c1_idx, 'shape_1'], [([1], TensorProto.INT64)])
    c1_oh = nd('OneHot', [c1_idx_1d, 'depth_10', 'oh_vals'], [[1, 10]], axis=1)

    not_c1 = nd('Sub', ['c_one', c1_oh], [[1, 10]])
    ch_sums_noc1 = nd('Mul', [ch_sums_nobg, not_c1], [[1, 10]])

    c2_idx = nd('ArgMax', [ch_sums_noc1], [([1, 1], TensorProto.INT64)], axis=1, keepdims=1)
    c2_idx_1d = nd('Reshape', [c2_idx, 'shape_1'], [([1], TensorProto.INT64)])
    c2_oh = nd('OneHot', [c2_idx_1d, 'depth_10', 'oh_vals'], [[1, 10]], axis=1)

    # === STEP 4: Extract binary masks (10x10) ===
    mask1_4d = nd('Gather', [inp_10, c1_idx_1d], [[1, 1, 10, 10]], axis=1)
    mask2_4d = nd('Gather', [inp_10, c2_idx_1d], [[1, 1, 10, 10]], axis=1)

    # === STEP 5: Flatten masks and find bbox ===
    def get_flat_and_bbox(mask_4d):
        """Get flat [100] mask and scalar r_min, c_min."""
        m2d = nd('Reshape', [mask_4d, 'shape_10_10'], [[10, 10]])
        m_flat = nd('Reshape', [m2d, 'shape_100'], [[100]])
        not_m = nd('Sub', ['c_one', m2d], [[10, 10]])

        rgrid = nd('MatMul', ['row_idx_10_1', 'ones_1_10'], [[10, 10]])
        cgrid = nd('MatMul', ['ones_10_1', 'col_idx_1_10'], [[10, 10]])

        rmin_v = nd('Add', [nd('Mul', [m2d, rgrid], [[10, 10]]),
                            nd('Mul', [not_m, 'c_big'], [[10, 10]])], [[10, 10]])
        r_min = nd('ReduceMin', [rmin_v, 'axes01'], [[1, 1]], keepdims=1)

        cmin_v = nd('Add', [nd('Mul', [m2d, cgrid], [[10, 10]]),
                            nd('Mul', [not_m, 'c_big'], [[10, 10]])], [[10, 10]])
        c_min = nd('ReduceMin', [cmin_v, 'axes01'], [[1, 1]], keepdims=1)

        # r_min, c_min are [1,1] float scalars
        return m_flat, r_min, c_min

    m1_flat, r1_min, c1_min = get_flat_and_bbox(mask1_4d)
    m2_flat, r2_min, c2_min = get_flat_and_bbox(mask2_4d)

    # === STEP 6: Compute placement matrices [9, 9] ===
    def compute_placements(m_flat, r_min, c_min):
        """Compute [9,9] placement matrix. Row k = flattened placement at offset k."""
        # abs_rows [9, 9] = r_min + row_off
        abs_rows = nd('Add', [r_min, 'row_off_9_9'], [[9, 9]])
        abs_cols = nd('Add', [c_min, 'col_off_9_9'], [[9, 9]])
        
        # flat_idx = abs_rows * 10 + abs_cols
        flat_f = nd('Add', [nd('Mul', [abs_rows, 'c_10f'], [[9, 9]]), abs_cols], [[9, 9]])
        
        # Clip to [0, 99]
        flat_clip = nd('Clip', [flat_f, 'c_zero', 'c_99f'], [[9, 9]])
        
        # Cast to int64
        flat_i = nd('Cast', [flat_clip], [([9, 9], TensorProto.INT64)], to=7)
        
        # Gather from flat mask [100] using indices [9,9] -> [9,9]
        gathered = nd('Gather', [m_flat, flat_i], [[9, 9]], axis=0)
        
        # Apply static validity mask
        valid_gathered = nd('Mul', [gathered, 'valid_static_9_9'], [[9, 9]])
        
        return valid_gathered  # [9, 9]: row k is the 9 values for placement k
    
    P1 = compute_placements(m1_flat, r1_min, c1_min)
    P2 = compute_placements(m2_flat, r2_min, c2_min)
    
    # === STEP 7: Check all 81 pairs ===
    # dots[i,j] = P1[i] · P2[j] = overlap count
    P2_T = nd('Transpose', [P2], [[9, 9]], perm=[1, 0])
    dots = nd('MatMul', [P1, P2_T], [[9, 9]])
    
    # Sums per placement
    sum1 = nd('ReduceSum', [P1, 'axes1'], [[9, 1]], keepdims=1)  # [9,1]
    sum2 = nd('ReduceSum', [P2, 'axes1'], [[9, 1]], keepdims=1)  # [9,1]
    sum2_T = nd('Transpose', [sum2], [[1, 9]], perm=[1, 0])       # [1,9]
    total_sum = nd('Add', [sum1, sum2_T], [[9, 9]])               # [9,9]
    
    # Valid: total_sum > 8.5 AND dots < 0.5
    sum_ok_b = nd('Greater', [total_sum, 'c_8_5'], [([9, 9], TensorProto.BOOL)])
    sum_ok = nd('Cast', [sum_ok_b], [[9, 9]], to=1)
    dots_ok_b = nd('Less', [dots, 'c_half'], [([9, 9], TensorProto.BOOL)])
    dots_ok = nd('Cast', [dots_ok_b], [[9, 9]], to=1)
    is_valid = nd('Mul', [sum_ok, dots_ok], [[9, 9]])
    
    # Flatten to [81], ArgMax
    valid_flat = nd('Reshape', [is_valid, 'shape_81'], [[81]])
    winner_idx = nd('ArgMax', [valid_flat], [([1], TensorProto.INT64)], axis=0, keepdims=1)
    
    # Decode: i = winner // 9, j = winner % 9
    winner_f = nd('Cast', [winner_idx], [[1]], to=1)
    i_f = nd('Floor', [nd('Div', [winner_f, 'c_9f'], [[1]])], [[1]])
    i_idx = nd('Cast', [i_f], [([1], TensorProto.INT64)], to=7)
    j_idx_f = nd('Sub', [winner_f, nd('Mul', [i_f, 'c_9f'], [[1]])], [[1]])
    j_idx = nd('Cast', [j_idx_f], [([1], TensorProto.INT64)], to=7)
    
    # Get winning placements: P1[i] -> [1, 9], P2[j] -> [1, 9]
    best_p1_flat = nd('Gather', [P1, i_idx], [[1, 9]], axis=0)
    best_p2_flat = nd('Gather', [P2, j_idx], [[1, 9]], axis=0)
    
    # Reshape to [1, 1, 3, 3]
    best_p1 = nd('Reshape', [best_p1_flat, 'shape_1_1_3_3'], [[1, 1, 3, 3]])
    best_p2 = nd('Reshape', [best_p2_flat, 'shape_1_1_3_3'], [[1, 1, 3, 3]])
    
    # === STEP 8: Build output [1, 10, 30, 30] ===
    c1_oh_4d = nd('Reshape', [c1_oh, 'shape_1_10_1_1'], [[1, 10, 1, 1]])
    c2_oh_4d = nd('Reshape', [c2_oh, 'shape_1_10_1_1'], [[1, 10, 1, 1]])

    out_c1_small = nd('Mul', [c1_oh_4d, best_p1], [[1, 10, 3, 3]])
    out_c2_small = nd('Mul', [c2_oh_4d, best_p2], [[1, 10, 3, 3]])
    out_small = nd('Add', [out_c1_small, out_c2_small], [[1, 10, 3, 3]])

    # Pad [1,10,3,3] to [1,10,30,30]
    final = nd('Pad', [out_small, 'pad_10_3_to_30', 'pad_val_zero'], [[1, 10, 30, 30]])

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


if __name__ == '__main__':
    build_and_validate(build_task153, task_num=153)