#!/usr/bin/env python3
"""Gravity solver — unrolled bubble-sort via Conv + Where.

v5.2: Solves Task 78 (direction=up, bg=0, score 8.399).
Tries all 4 directions × 10 bg colors. Fixed-shape only.
"""

import numpy as np
from onnx import helper, numpy_helper
from ..onnx_helpers import mk, _make_int64_init, _build_pad_node, add_onehot_block
from ..data_loader import get_exs, fixed_shapes
from ..constants import GH, GW


def _gravity_np(grid, direction, bg_color=0):
    """Apply gravity in numpy for verification."""
    r = np.full_like(grid, bg_color)
    h, w = grid.shape
    if direction == 'down':
        for c in range(w):
            nz = grid[:, c][grid[:, c] != bg_color]
            r[h - len(nz):h, c] = nz
    elif direction == 'up':
        for c in range(w):
            nz = grid[:, c][grid[:, c] != bg_color]
            r[:len(nz), c] = nz
    elif direction == 'right':
        for rr in range(h):
            nz = grid[rr, :][grid[rr, :] != bg_color]
            r[rr, w - len(nz):w] = nz
    elif direction == 'left':
        for rr in range(h):
            nz = grid[rr, :][grid[rr, :] != bg_color]
            r[rr, :len(nz)] = nz
    return r


def _build_gravity_model(IH, IW, direction, bg_color=0):
    """Build ONNX model for gravity via unrolled bubble-sort.

    Each step compares adjacent cells and swaps if needed:
    - If current cell is bg AND source neighbor is non-bg → fill with source
    - If current cell is non-bg AND destination neighbor is bg → vacate to bg
    After max(IH,IW) passes, all non-bg pixels settle in the gravity direction.
    """
    pad_h, pad_w = GH - IH, GW - IW
    n_steps = max(IH, IW)

    pull_above = np.zeros((10, 10, 3, 3), dtype=np.float32)
    pull_below = np.zeros((10, 10, 3, 3), dtype=np.float32)
    for ch in range(10):
        if direction == 'down':
            pull_above[ch, ch, 0, 1] = 1.0
            pull_below[ch, ch, 2, 1] = 1.0
        elif direction == 'up':
            pull_above[ch, ch, 2, 1] = 1.0
            pull_below[ch, ch, 0, 1] = 1.0
        elif direction == 'right':
            pull_above[ch, ch, 1, 0] = 1.0
            pull_below[ch, ch, 1, 2] = 1.0
        elif direction == 'left':
            pull_above[ch, ch, 1, 2] = 1.0
            pull_below[ch, ch, 1, 0] = 1.0

    bg_sel = np.zeros((1, 10, 1, 1), dtype=np.float32)
    bg_sel[0, bg_color, 0, 0] = 1.0
    bg_oh = np.zeros((1, 10, 1, 1), dtype=np.float32)
    bg_oh[0, bg_color, 0, 0] = 1.0

    inits = [
        _make_int64_init('sl_st', [0, 0, 0, 0]),
        _make_int64_init('sl_en', [1, 10, IH, IW]),
        numpy_helper.from_array(pull_above, 'pull_src'),
        numpy_helper.from_array(pull_below, 'pull_dst'),
        numpy_helper.from_array(bg_sel, 'bg_sel'),
        numpy_helper.from_array(bg_oh, 'bg_oh'),
        numpy_helper.from_array(np.float32(0.5), 'half'),
    ]

    nodes = [
        helper.make_node('Slice', ['input', 'sl_st', 'sl_en'], ['cur_0']),
    ]

    cur = 'cur_0'
    for i in range(n_steps):
        src = f'src_{i}'
        nodes.append(helper.make_node('Conv', [cur, 'pull_src'], [src],
                                       kernel_shape=[3, 3], pads=[1, 1, 1, 1]))

        nodes.append(helper.make_node('Mul', [cur, 'bg_sel'], [f'cbg_{i}']))
        inits.append(_make_int64_init(f'ax1_{i}', [1]))
        nodes.append(helper.make_node('ReduceSum', [f'cbg_{i}', f'ax1_{i}'], [f'cbgsum_{i}'], keepdims=1))
        nodes.append(helper.make_node('Greater', [f'cbgsum_{i}', 'half'], [f'cur_is_bg_{i}']))

        nodes.append(helper.make_node('Mul', [src, 'bg_sel'], [f'sbg_{i}']))
        inits.append(_make_int64_init(f'ax2_{i}', [1]))
        nodes.append(helper.make_node('ReduceSum', [f'sbg_{i}', f'ax2_{i}'], [f'sbgsum_{i}'], keepdims=1))
        nodes.append(helper.make_node('Not', [f'cur_is_bg_{i}'], [f'cur_not_bg_{i}']))

        nodes.append(helper.make_node('Greater', [f'sbgsum_{i}', 'half'], [f'src_is_bg_{i}']))
        nodes.append(helper.make_node('Not', [f'src_is_bg_{i}'], [f'src_not_bg_{i}']))
        nodes.append(helper.make_node('And', [f'cur_is_bg_{i}', f'src_not_bg_{i}'], [f'fill_{i}']))

        dst = f'dst_{i}'
        nodes.append(helper.make_node('Conv', [cur, 'pull_dst'], [dst],
                                       kernel_shape=[3, 3], pads=[1, 1, 1, 1]))
        nodes.append(helper.make_node('Mul', [dst, 'bg_sel'], [f'dbg_{i}']))
        inits.append(_make_int64_init(f'ax3_{i}', [1]))
        nodes.append(helper.make_node('ReduceSum', [f'dbg_{i}', f'ax3_{i}'], [f'dbgsum_{i}'], keepdims=1))
        nodes.append(helper.make_node('Greater', [f'dbgsum_{i}', 'half'], [f'dst_is_bg_{i}']))
        nodes.append(helper.make_node('And', [f'cur_not_bg_{i}', f'dst_is_bg_{i}'], [f'vacate_{i}']))

        nxt = f'cur_{i+1}'
        nodes.append(helper.make_node('Where', [f'fill_{i}', src, cur], [f'tmp_{i}']))
        nodes.append(helper.make_node('Where', [f'vacate_{i}', 'bg_oh', f'tmp_{i}'], [nxt]))
        cur = nxt

    nodes.append(helper.make_node('ArgMax', [cur], ['grav_am'], axis=1, keepdims=1))
    add_onehot_block(nodes, inits, 'grav_am', 'grav_oh')
    nodes.append(_build_pad_node('grav_oh', 'output', pad_h, pad_w, inits))
    return mk(nodes, inits)


def s_gravity_unrolled(td):
    """Gravity solver with unrolled Conv+Where steps.
    Tries all 4 directions × bg colors 0-9."""
    exs = get_exs(td)
    sp = fixed_shapes(td)
    if sp is None:
        return None
    (IH, IW), (OH, OW) = sp
    if (IH, IW) != (OH, OW):
        return None

    for bg_color in range(10):
        for direction in ('down', 'up', 'left', 'right'):
            if all(np.array_equal(_gravity_np(inp, direction, bg_color), out)
                   for inp, out in exs):
                return _build_gravity_model(IH, IW, direction, bg_color)
    return None