"""
Copyright (C) 2025 The ONNXIFIER Authors.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
"""

# pylint: disable=missing-function-docstring

import numpy as np
import onnx
import pytest
from onnx.helper import (
    make_graph,
    make_model,
    make_node,
    make_operatorsetid,
    make_tensor_value_info,
)
from onnx.numpy_helper import from_array

from onnxifier.evaluator import Evaluator
from onnxifier.graph import OnnxGraph
from onnxifier.passes.version_converter.downgrade import downgrade_op_version


def _make_col2im_op18(
    x, image_shape, block_shape, dilations=None, pads=None, strides=None
):
    assert x.ndim == 3
    node = make_node(
        "Col2Im",
        ["input", "image_shape", "block_shape"],
        ["output"],
        "col2im",
        dilations=dilations,
        pads=pads,
        strides=strides,
    )
    graph = make_graph(
        [node],
        "graph",
        [
            make_tensor_value_info("input", onnx.TensorProto.FLOAT, list(x.shape)),
        ],
        [
            make_tensor_value_info(
                "output", onnx.TensorProto.FLOAT, ["B", "C", "H", "W"]
            )
        ],
        [
            from_array(np.array(image_shape, np.int64), "image_shape"),
            from_array(np.array(block_shape, np.int64), "block_shape"),
        ],
    )
    model = make_model(
        graph,
        ir_version=9,
        opset_imports=[make_operatorsetid("", 18)],
    )
    onnx.checker.check_model(model, True)
    return OnnxGraph(model)


def _numpy_impl(
    x: np.ndarray, image_shape, block_shape, dilations=None, pads=None, strides=None
):
    """Implement Col2Im using numpy with proper handling of padding and overlaps."""
    B, C, L = x.shape
    blocks = np.prod(block_shape)
    x = x.reshape([B, C // blocks, *block_shape, L])
    if dilations is None:
        dilations = np.ones_like(image_shape)
    if pads is None:
        pads = np.repeat(np.zeros_like(image_shape), 2)
    if strides is None:
        strides = np.ones_like(image_shape)
    vpad = np.array(pads, np.int64).reshape([2, -1]).T
    spatial_shape = []
    for ch, d, s, p, b in zip(image_shape, dilations, strides, vpad, block_shape):
        spatial_shape.append(
            np.ceil((ch + p[0] + p[1] - (d * (b - 1))) / s).astype(np.int64)
        )
    assert L == np.prod(spatial_shape)

    # canvas based filling
    y = np.zeros([B, C // blocks, *image_shape], dtype=x.dtype)
    for ch in range(y.shape[1]):
        # === Below this block-wise loop is generated by AI (Trae)
        block = x[:, ch].reshape([B, *block_shape, *spatial_shape])
        # Generate all combinations of kernel indices
        kernel_indices = np.meshgrid(
            *[np.arange(d) for d in block_shape], indexing="ij"
        )
        kernel_indices = np.stack(kernel_indices, axis=-1).reshape(-1, len(block_shape))

        # Generate all combinations of spatial indices
        spatial_indices = np.meshgrid(
            *[np.arange(d) for d in spatial_shape], indexing="ij"
        )
        spatial_indices = np.stack(spatial_indices, axis=-1).reshape(
            -1, len(spatial_shape)
        )

        # Create all combinations of kernel and spatial indices
        kdx = np.repeat(kernel_indices, len(spatial_indices), axis=0)
        sdx = np.tile(spatial_indices, (len(kernel_indices), 1))

        # Calculate output positions using correct Col2Im formula
        output_positions = []
        for i in range(len(spatial_shape)):
            pos = sdx[:, i] * strides[i] + kdx[:, i] * dilations[i] - vpad[i][0]
            output_positions.append(pos)
        odx = np.stack(output_positions, axis=1)

        # Filter valid positions within output bounds
        output_shape = np.array(y.shape[2:])
        valid_mask = np.logical_and.reduce(
            [odx[:, i] >= 0 for i in range(len(spatial_shape))]
        ) & np.logical_and.reduce(
            [odx[:, i] < output_shape[i] for i in range(len(spatial_shape))]
        )

        # Apply valid mask
        valid_kernel = kdx[valid_mask]
        valid_spatial = sdx[valid_mask]
        valid_positions = odx[valid_mask]

        # Accumulate block values to output image
        if valid_positions.size > 0:
            # Create multi-dimensional indices for block
            block_idx = tuple(valid_kernel.T) + tuple(valid_spatial.T)
            # Create indices for output image y
            y_idx = (slice(None), ch) + tuple(valid_positions.T)
            # Accumulate values using numpy's scatter add
            np.add.at(y, y_idx, block[(slice(None),) + block_idx])
        # === Generation done

    return y


@pytest.fixture(
    params=[
        dict(
            shape=(1, 32, 625),
            image_shape=(24, 24),
            block_size=(2, 2),
            paddings=(1, 1, 1, 1),
            strides=(1, 1),
            output_shape=(1, 8, 24, 24),
        ),
        dict(
            shape=(1, 18, 144),
            image_shape=(12, 12),
            block_size=(3, 3),
            paddings=(1, 1, 1, 1),
            strides=(1, 1),
            output_shape=(1, 2, 12, 12),
        ),
        dict(
            shape=(1, 32, 144),
            image_shape=(24, 24),
            block_size=(2, 2),
            paddings=None,
            strides=(2, 2),
            output_shape=(1, 8, 24, 24),
        ),
    ]
)
def test_params(request):
    return request.param


def test_col2im_numpy_impl(test_params):
    shape = test_params["shape"]
    image_shape = test_params["image_shape"]
    block_size = test_params["block_size"]
    paddings = test_params["paddings"]
    strides = test_params["strides"]
    output_shape = test_params["output_shape"]

    x = np.random.randn(*shape).astype("float32")
    graph = _make_col2im_op18(
        x,
        image_shape=image_shape,
        block_shape=block_size,
        pads=paddings,
        strides=strides,
    )
    runner1 = Evaluator(graph.model)
    [y_ref] = runner1(["output"], input=x)

    # Compute our numpy implementation output
    y = _numpy_impl(
        x,
        image_shape=image_shape,
        block_shape=block_size,
        pads=paddings,
        strides=strides,
    )

    # Validate output shape and numerical accuracy
    assert y_ref.shape == output_shape
    # Temporarily relax tolerance for debugging
    assert np.allclose(
        y, y_ref, atol=1e-5
    ), f"Numerical mismatch: max diff {np.max(np.abs(y - y_ref)):.6f}"


def test_downgrade_col2im_op18(test_params):
    shape = test_params["shape"]
    image_shape = test_params["image_shape"]
    block_size = test_params["block_size"]
    paddings = test_params["paddings"]
    strides = test_params["strides"]
    output_shape = test_params["output_shape"]

    x = np.random.randn(*shape).astype("float32")
    graph = _make_col2im_op18(
        x,
        image_shape=image_shape,
        block_shape=block_size,
        pads=paddings,
        strides=strides,
    )
    runner1 = Evaluator(graph.model)
    [y_ref] = runner1(["output"], input=x)

    graph = downgrade_op_version(graph, op_version=17)
    runner2 = Evaluator(graph.model)
    [y] = runner2(["output"], input=x)

    assert y_ref.shape == output_shape
    assert y.shape == y_ref.shape
    assert np.allclose(y, y_ref)


def test_downgrade_col2im_to_depthtospace_op18():
    x = np.random.randn(1, 32, 144).astype("float32")
    graph = _make_col2im_op18(
        x, image_shape=[24, 24], block_shape=[2, 2], strides=[2, 2]
    )
    runner1 = Evaluator(graph.model)
    [y_ref] = runner1(["output"], input=x)

    graph = downgrade_op_version(graph, op_version=17)
    runner2 = Evaluator(graph.model)
    [y] = runner2(["output"], input=x)

    assert y_ref.shape == (1, 8, 24, 24)
    assert y.shape == (1, 8, 24, 24)
    assert np.allclose(y, y_ref)
