#!/usr/bin/env python3
"""Basic analytical solvers: identity, constant, color_map, transpose."""

import numpy as np
from onnx import helper, numpy_helper, TensorProto
from ..onnx_helpers import mk, _make_int64_init
from ..data_loader import get_exs, fixed_shapes


def s_identity(td):
    """Identity solver."""
    for ex in td['train'] + td['test']:
        if ex['input'] != ex['output']:
            return None
    return mk([helper.make_node('Identity', ['input'], ['output'])])


def s_color_map(td):
    """Color mapping solver."""
    cm = {}
    for ex in td['train'] + td['test']:
        inp, out = np.array(ex['input']), np.array(ex['output'])
        if inp.shape != out.shape:
            return None
        for iv, ov in zip(inp.flat, out.flat):
            iv, ov = int(iv), int(ov)
            if iv in cm and cm[iv] != ov:
                return None
            cm[iv] = ov
    is_permutation = (set(cm.keys()) == set(cm.values()))
    if is_permutation:
        gather_ch = np.arange(10, dtype=np.int32)
        for src, dst in cm.items():
            if 0 <= src < 10 and 0 <= dst < 10:
                gather_ch[dst] = src
        inits = [numpy_helper.from_array(gather_ch, 'gi')]
        nodes = [helper.make_node('Gather', ['input', 'gi'], ['output'], axis=1)]
        return mk(nodes, inits)
    else:
        W = np.zeros((10, 10, 1, 1), dtype=np.float32)
        for ic in range(10):
            W[cm.get(ic, ic), ic, 0, 0] = 1.0
        return mk([helper.make_node('Conv', ['input', 'W'], ['output'], kernel_shape=[1, 1])],
                  [numpy_helper.from_array(W, 'W')])


def s_transpose(td):
    """Transpose solver."""
    for ex in td['train'] + td['test']:
        if not np.array_equal(np.array(ex['output']), np.array(ex['input']).T):
            return None
    return mk([helper.make_node('Transpose', ['input'], ['output'], perm=[0, 1, 3, 2])])


def s_constant(td):
    """Constant output solver using opset 17 ReduceSum."""
    sp = fixed_shapes(td)
    if sp is None:
        return None
    exs = get_exs(td)
    outs = [out for _, out in exs]
    if not all(np.array_equal(outs[0], o) for o in outs[1:]):
        return None
    const = np.zeros((1, 10, 30, 30), dtype=np.float32)
    for r, row in enumerate(outs[0]):
        for c, v in enumerate(row):
            const[0, int(v), r, c] = 1.0
    inits = [
        numpy_helper.from_array(np.array(0.0, dtype=np.float32), 'z'),
        numpy_helper.from_array(const, 'c'),
        _make_int64_init('rs_axes_cst', [1, 2, 3]),
    ]
    nodes = [
        helper.make_node('Mul', ['input', 'z'], ['zd']),
        helper.make_node('ReduceSum', ['zd', 'rs_axes_cst'], ['s'], keepdims=1),
        helper.make_node('Add', ['s', 'c'], ['output']),
    ]
    return mk(nodes, inits)