#!/usr/bin/env python3
"""Edge/boundary detection solver — Laplacian Conv.

v5.2: 0 matches in current task set (edge definition too strict).
Kept for future variants (per-color edges, interior-only edges, etc.).
"""

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


def _has_edges(inp, out, edge_color, bg_color=0):
    """Check if output is edge detection of input."""
    h, w = inp.shape
    for r in range(h):
        for c in range(w):
            pix = inp[r, c]
            is_edge = False
            if pix != bg_color:
                for dr, dc in [(-1, 0), (1, 0), (0, -1), (0, 1)]:
                    nr, nc = r + dr, c + dc
                    if 0 <= nr < h and 0 <= nc < w:
                        if inp[nr, nc] != pix:
                            is_edge = True
                            break
                    else:
                        is_edge = True
                        break
            expected = edge_color if is_edge else bg_color
            if out[r, c] != expected:
                return False
    return True


def _build_edge_model(IH, IW, edge_color, bg_color=0):
    """Build ONNX model for edge detection via Laplacian conv."""
    pad_h, pad_w = GH - IH, GW - IW

    ch_sel = np.zeros((1, 10, 1, 1), dtype=np.float32)
    for c in range(10):
        if c != bg_color:
            ch_sel[0, c, 0, 0] = 1.0

    lap_k = np.array([[0, -1, 0],
                       [-1, 4, -1],
                       [0, -1, 0]], dtype=np.float32).reshape(1, 1, 3, 3)

    edge_oh = np.zeros((1, 10, 1, 1), dtype=np.float32)
    edge_oh[0, edge_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(ch_sel, 'ch_sel'),
        numpy_helper.from_array(lap_k, 'lap_k'),
        numpy_helper.from_array(np.float32(0.5), 'thresh'),
        numpy_helper.from_array(edge_oh, 'edge_oh'),
        numpy_helper.from_array(bg_oh, 'bg_oh'),
    ]

    nodes = [
        helper.make_node('Slice', ['input', 'sl_st', 'sl_en'], ['cropped']),
        helper.make_node('Conv', ['cropped', 'ch_sel'], ['occ'], kernel_shape=[1, 1]),
        helper.make_node('Conv', ['occ', 'lap_k'], ['lap_out'], kernel_shape=[3, 3], pads=[1, 1, 1, 1]),
        helper.make_node('Abs', ['lap_out'], ['lap_abs']),
        helper.make_node('Greater', ['lap_abs', 'thresh'], ['is_edge_raw']),
        helper.make_node('Greater', ['occ', 'thresh'], ['is_occ']),
        helper.make_node('And', ['is_edge_raw', 'is_occ'], ['is_edge']),
        helper.make_node('Where', ['is_edge', 'edge_oh', 'bg_oh'], ['result_small']),
    ]
    nodes.append(_build_pad_node('result_small', 'output', pad_h, pad_w, inits))
    return mk(nodes, inits)


def s_edge_detect(td):
    """Edge detection solver: output = boundary pixels of input shapes."""
    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 [0]:
        out_colors = set()
        for _, out in exs:
            out_colors.update(out.flatten())
        for edge_color in out_colors:
            if edge_color == bg_color:
                continue
            if all(_has_edges(inp, out, edge_color, bg_color) for inp, out in exs):
                return _build_edge_model(IH, IW, edge_color, bg_color)
    return None