#!/usr/bin/env python3
"""Wave 1 static spatial remapping solvers.

A4: downsample_stride — strided sampling of input
A7: symmetry_complete — mirror to complete L-R or T-B symmetry
A1: extract_inner — remove border frame
A2: add_border — add constant border
A6: sparse_fill — pixel to block expansion
B1: channel_filter — keep only certain colors

Scan results (2026-04-27): 0 arc-gen validated matches.
Kept for future tasks and as building blocks.
"""

import numpy as np
from ..data_loader import get_exs, fixed_shapes
from ..gather_helpers import _build_gather_model, _build_gather_model_with_const
from ..onnx_helpers import mk, _make_int64_init, _build_pad_node, add_onehot_block
from ..constants import GH, GW


def s_downsample_stride(td):
    """out[r,c] = inp[r*sH + oH, c*sW + oW] for integer strides."""
    exs = get_exs(td)
    sp = fixed_shapes(td)
    if sp is None:
        return None
    (IH, IW), (OH, OW) = sp
    if OH >= IH or OW >= IW:
        return None

    for sh in range(2, 6):
        for sw in range(2, 6):
            for oh_off in range(sh):
                for ow_off in range(sw):
                    ok = True
                    for inp, out in exs:
                        sampled = inp[oh_off::sh, ow_off::sw]
                        if sampled.shape != out.shape or not np.array_equal(sampled, out):
                            ok = False
                            break
                    if ok:
                        idx = np.zeros((OH, OW, 2), dtype=np.int64)
                        for r in range(OH):
                            for c in range(OW):
                                idx[r, c] = [r * sh + oh_off, c * sw + ow_off]
                        return _build_gather_model(OH, OW, idx)
    return None


def s_symmetry_complete(td):
    """Complete partial T-B symmetry by adding mirrored + original via Gather."""
    from onnx import helper, numpy_helper

    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
    if IH < 2:
        return None

    # T-B symmetry: out[r,c] = max(inp[r,c], inp[IH-1-r,c])
    ok = True
    for inp, out in exs:
        exp = inp.copy()
        for r in range(IH // 2):
            for c in range(IW):
                v = max(int(inp[r, c]), int(inp[IH - 1 - r, c]))
                exp[r, c] = v
                exp[IH - 1 - r, c] = v
        if not np.array_equal(out, exp):
            ok = False
            break

    if ok:
        # Build: Gather(self) + Gather(mirror) → Add → ArgMax → one-hot
        pad_h, pad_w = GH - OH, GW - OW
        mirror_idx = np.zeros((GH * GW,), dtype=np.int64)
        mask = np.zeros((1, 1, GH, GW), dtype=np.float32)
        self_idx = np.zeros((GH * GW,), dtype=np.int64)
        for r in range(OH):
            for c in range(OW):
                self_idx[r * GW + c] = r * GW + c
                mirror_idx[r * GW + c] = (IH - 1 - r) * GW + c
                mask[0, 0, r, c] = 1.0

        inits = [
            numpy_helper.from_array(np.array([1, 10, GH * GW], dtype=np.int64), 'fs'),
            numpy_helper.from_array(self_idx, 'self_idx'),
            numpy_helper.from_array(mirror_idx, 'mirror_idx'),
            numpy_helper.from_array(np.array([1, 10, GH, GW], dtype=np.int64), 'os'),
            numpy_helper.from_array(mask, 'mask'),
        ]
        nodes = [
            helper.make_node('Reshape', ['input', 'fs'], ['flat']),
            helper.make_node('Gather', ['flat', 'self_idx'], ['g_self'], axis=2),
            helper.make_node('Gather', ['flat', 'mirror_idx'], ['g_mirror'], axis=2),
            helper.make_node('Add', ['g_self', 'g_mirror'], ['combined']),
            helper.make_node('Reshape', ['combined', 'os'], ['combined_2d']),
            helper.make_node('ArgMax', ['combined_2d'], ['am'], axis=1, keepdims=1),
        ]
        add_onehot_block(nodes, inits, 'am', 'oh_out')
        nodes.append(helper.make_node('Mul', ['oh_out', 'mask'], ['output']))
        return mk(nodes, inits)

    # L-R symmetry: out[r,c] = max(inp[r,c], inp[r,IW-1-c])
    if IW < 2:
        return None
    ok = True
    for inp, out in exs:
        exp = inp.copy()
        for r in range(IH):
            for c in range(IW // 2):
                v = max(int(inp[r, c]), int(inp[r, IW - 1 - c]))
                exp[r, c] = v
                exp[r, IW - 1 - c] = v
        if not np.array_equal(out, exp):
            ok = False
            break

    if ok:
        mirror_idx = np.zeros((GH * GW,), dtype=np.int64)
        mask = np.zeros((1, 1, GH, GW), dtype=np.float32)
        self_idx = np.zeros((GH * GW,), dtype=np.int64)
        for r in range(OH):
            for c in range(OW):
                self_idx[r * GW + c] = r * GW + c
                mirror_idx[r * GW + c] = r * GW + (IW - 1 - c)
                mask[0, 0, r, c] = 1.0

        inits = [
            numpy_helper.from_array(np.array([1, 10, GH * GW], dtype=np.int64), 'fs'),
            numpy_helper.from_array(self_idx, 'self_idx'),
            numpy_helper.from_array(mirror_idx, 'mirror_idx'),
            numpy_helper.from_array(np.array([1, 10, GH, GW], dtype=np.int64), 'os'),
            numpy_helper.from_array(mask, 'mask'),
        ]
        nodes = [
            helper.make_node('Reshape', ['input', 'fs'], ['flat']),
            helper.make_node('Gather', ['flat', 'self_idx'], ['g_self'], axis=2),
            helper.make_node('Gather', ['flat', 'mirror_idx'], ['g_mirror'], axis=2),
            helper.make_node('Add', ['g_self', 'g_mirror'], ['combined']),
            helper.make_node('Reshape', ['combined', 'os'], ['combined_2d']),
            helper.make_node('ArgMax', ['combined_2d'], ['am'], axis=1, keepdims=1),
        ]
        add_onehot_block(nodes, inits, 'am', 'oh_out')
        nodes.append(helper.make_node('Mul', ['oh_out', 'mask'], ['output']))
        return mk(nodes, inits)

    return None


def s_extract_inner(td):
    """Remove N-pixel border frame → smaller output."""
    exs = get_exs(td)
    sp = fixed_shapes(td)
    if sp is None:
        return None
    (IH, IW), (OH, OW) = sp

    for b in range(1, min(IH, IW) // 2):
        if OH != IH - 2 * b or OW != IW - 2 * b:
            continue
        if all(np.array_equal(inp[b:IH-b, b:IW-b], out) for inp, out in exs):
            idx = np.zeros((OH, OW, 2), dtype=np.int64)
            for r in range(OH):
                for c in range(OW):
                    idx[r, c] = [r + b, c + b]
            return _build_gather_model(OH, OW, idx)
    return None


def s_add_border(td):
    """Add constant-color border frame → larger output."""
    exs = get_exs(td)
    sp = fixed_shapes(td)
    if sp is None:
        return None
    (IH, IW), (OH, OW) = sp

    for b in range(1, 5):
        if OH != IH + 2 * b or OW != IW + 2 * b:
            continue
        if OH > 30 or OW > 30:
            continue
        for bc in range(10):
            ok = True
            for inp, out in exs:
                exp = np.full((OH, OW), bc, dtype=np.int64)
                exp[b:b+IH, b:b+IW] = inp
                if not np.array_equal(out, exp):
                    ok = False
                    break
            if ok:
                idx = np.zeros((OH, OW, 2), dtype=np.int64)
                cst = np.full((OH, OW), -1, dtype=np.int64)
                for r in range(OH):
                    for c in range(OW):
                        if b <= r < b + IH and b <= c < b + IW:
                            idx[r, c] = [r - b, c - b]
                        else:
                            idx[r, c] = [-1, -1]
                            cst[r, c] = bc
                return _build_gather_model_with_const(IH, IW, OH, OW, idx, cst)
    return None


def s_sparse_fill(td):
    """Each input pixel becomes an NxN block in output."""
    exs = get_exs(td)
    sp = fixed_shapes(td)
    if sp is None:
        return None
    (IH, IW), (OH, OW) = sp

    for bh in range(2, 10):
        for bw in range(2, 10):
            if OH != IH * bh or OW != IW * bw:
                continue
            if OH > 30 or OW > 30:
                continue
            ok = True
            for inp, out in exs:
                exp = np.zeros((OH, OW), dtype=np.int64)
                for r in range(IH):
                    for c in range(IW):
                        exp[r*bh:(r+1)*bh, c*bw:(c+1)*bw] = inp[r, c]
                if not np.array_equal(out, exp):
                    ok = False
                    break
            if ok:
                idx = np.zeros((OH, OW, 2), dtype=np.int64)
                for r in range(OH):
                    for c in range(OW):
                        idx[r, c] = [r // bh, c // bw]
                return _build_gather_model(OH, OW, idx)
    return None


def s_channel_filter(td):
    """Keep only certain colors, rest → background (0)."""
    from onnx import helper, numpy_helper

    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

    in_colors = set()
    out_colors = set()
    for inp, out in exs:
        in_colors.update(inp.flatten())
        out_colors.update(out.flatten())

    if not (out_colors < in_colors):
        return None

    keep = out_colors
    for inp, out in exs:
        exp = np.where(np.isin(inp, list(keep)), inp, 0)
        if not np.array_equal(out, exp):
            return None

    ch_mask = np.zeros((1, 10, 1, 1), dtype=np.float32)
    for c in keep:
        if 0 <= c < 10:
            ch_mask[0, c, 0, 0] = 1.0

    inits = [numpy_helper.from_array(ch_mask, 'ch_mask')]
    nodes = [helper.make_node('Mul', ['input', 'ch_mask'], ['output'])]
    return mk(nodes, inits)