#!/usr/bin/env python3
"""ONNX model building helper functions (opset 17)."""

import numpy as np
from onnx import helper, TensorProto, numpy_helper
from .constants import DT, IR, GRID_SHAPE, INT64_MIN, GH, GW
from .config import make_opset


def _make_int64_init(name, values):
    """Create int64 initializer."""
    return numpy_helper.from_array(np.array(values, dtype=np.int64), name)


def _build_pad_node(input_name, output_name, pad_h, pad_w, inits, suffix=''):
    """Pad with tensor-based pads input (opset 11+)."""
    pads_name = f'pads{suffix}'
    cv_name = f'pad_cv{suffix}'
    pads_arr = np.array([0, 0, 0, 0, 0, 0, pad_h, pad_w], dtype=np.int64)
    inits.append(numpy_helper.from_array(pads_arr, pads_name))
    inits.append(numpy_helper.from_array(np.array(0.0, dtype=np.float32), cv_name))
    return helper.make_node('Pad', [input_name, pads_name, cv_name], [output_name], mode='constant')


def _build_slice_crop(input_name, output_name, IH, IW, inits, suffix=''):
    """Slice to crop [1,10,30,30] to [1,10,IH,IW]."""
    st_name = f'crop_st{suffix}'
    en_name = f'crop_en{suffix}'
    inits.append(_make_int64_init(st_name, [0, 0, 0, 0]))
    inits.append(_make_int64_init(en_name, [1, 10, IH, IW]))
    return helper.make_node('Slice', [input_name, st_name, en_name], [output_name])


def _build_slice_reverse(input_name, output_name, axis, dim_size, inits, suffix=''):
    """Slice(step=-1) to reverse one axis. Zero MACs."""
    st_name = f'rev_st{suffix}'
    en_name = f'rev_en{suffix}'
    ax_name = f'rev_ax{suffix}'
    sp_name = f'rev_sp{suffix}'
    inits.append(_make_int64_init(st_name, [dim_size - 1]))
    inits.append(_make_int64_init(en_name, [INT64_MIN]))
    inits.append(_make_int64_init(ax_name, [axis]))
    inits.append(_make_int64_init(sp_name, [-1]))
    return helper.make_node('Slice', [input_name, st_name, en_name, ax_name, sp_name], [output_name])


def _build_reducesum(input_name, output_name, axes_list, inits, suffix=''):
    """ReduceSum with axes as tensor input (opset 13+). keepdims=1."""
    axes_name = f'rs_axes{suffix}'
    inits.append(_make_int64_init(axes_name, axes_list))
    return helper.make_node('ReduceSum', [input_name, axes_name], [output_name], keepdims=1)


def mk(nodes, inits=None, opset_version=17):
    """Create ONNX model from nodes and initializers."""
    x = helper.make_tensor_value_info("input", DT, GRID_SHAPE)
    y = helper.make_tensor_value_info("output", DT, GRID_SHAPE)
    g = helper.make_graph(nodes, "g", [x], [y], initializer=inits or [])
    return helper.make_model(g, ir_version=IR, opset_imports=make_opset(opset_version))


def add_onehot_block(nodes, inits, am_name, oh_name):
    """Add ArgMax one-hot conversion block."""
    classes = np.arange(10, dtype=np.int64).reshape(1, 10, 1, 1)
    inits.append(numpy_helper.from_array(classes, 'classes'))
    nodes.append(helper.make_node('Equal', [am_name, 'classes'], ['eq']))
    nodes.append(helper.make_node('Cast', ['eq'], [oh_name], to=TensorProto.FLOAT))