# SPDX-License-Identifier: Apache-2.0
"""Structure of the weights the reference (and later the device graph) reads from ``bevfusion_lidar.onnx``: every
wiring assertion of ``reference/weights.py`` holds (it raises otherwise), the layer plan matches the host rulebook
plan, every BatchNorm fold is exact (float64 folded vs unfolded on random inputs, per layer kind), the parameter
count after the folds equals an independent count over the graph, constants (``bev_pos``, pooled classes, K, heads)
are the SPEC's. numpy + onnx only (no torch, no onnxruntime): runs in tt-env with the fake ttnn and in the research
venv; skipped without the weights.
"""
from __future__ import annotations

import numpy as np
import pytest

pytest.importorskip("onnx")
pytest.importorskip("yaml")

from tt_bevfusion.host.preprocess import STRIDED_LAYERS, SUBM_LAYERS  # noqa: E402
from tt_bevfusion.reference.config import BEVFusionConfig  # noqa: E402
from tt_bevfusion.reference.weights import ENCODER_PLAN, load_weights  # noqa: E402
from tt_bevfusion.tests import _data  # noqa: E402
from tt_bevfusion.ttaw.weights import OnnxWeights  # noqa: E402


@pytest.fixture(scope="module")
def loaded():
    d = _data.weights_dir()
    cfg = BEVFusionConfig.from_model_dir(d)
    return cfg, load_weights(d, cfg), OnnxWeights(d / _data.ONNX)


def test_config_from_the_weights(loaded):
    cfg, _, _ = loaded
    assert cfg.version == "v2.0" and cfg.class_names == ("CAR", "TRUCK", "BUS", "BICYCLE", "PEDESTRIAN")
    assert cfg.grid == (1440, 1440, 40) and cfg.num_proposals == 500 and cfg.max_voxels == 256000
    assert cfg.max_points_per_voxel == 10 and not cfg.use_intensity and cfg.out_size_factor == 8
    assert cfg.remapper.apply(np.array([1, 1, 2]), np.array([5.0, 10.0, 10.0]), np.array([3.0, 4.0, 4.0])).tolist() \
        == [2, 4, 4]


def test_layer_plan(loaded):
    _, w, _ = loaded
    assert len(w.encoder) == 21 and [x.name for x in w.encoder] == [p[0] for p in ENCODER_PLAN]
    maps = {x.map for x in w.encoder}
    assert maps == set(SUBM_LAYERS) | set(STRIDED_LAYERS)
    chans = [(x.in_channels, x.out_channels, x.kernel_volume) for x in w.encoder]
    assert chans[0] == (4, 16, 27) and chans[-1] == (128, 128, 3)
    assert [x.name for x in w.encoder if x.residual] == [f"sparse.layer{s}.{b}" for s in range(1, 5) for b in range(2)]
    assert [len(b) for b in w.blocks] == [6, 6] and w.blocks[1][0].stride == 2
    assert [(d.kind, d.kernel, d.stride) for d in w.deblocks] == [("conv", 1, 1), ("conv_transpose", 2, 2)]
    h = w.head
    assert h.pooled_classes == (0, 1, 2, 3) and h.num_proposals == 500 and h.nms_kernel == 3
    assert h.decoder.self_attn.num_heads == 8 and h.decoder.self_attn.head_dim == 16
    assert h.decoder.cross_attn.q_divisor == 4.0
    assert h.bev_pos[181].tolist() == [1.5, 1.5] and h.bev_pos[-1].tolist() == [179.5, 179.5]


def test_parameter_count_after_the_folds(loaded):
    """ONNX float32 weights (8,308,692, S:bevfusion:218) minus what each BatchNorm fold removes: its 4 C statistics
    become a C bias (-3 C) when the folded conv had no bias, or merge into the existing bias (-4 C)."""
    _, w, o = loaded
    import onnx
    from onnx import numpy_helper

    inits = onnx.load(str(_data.weights_dir() / _data.ONNX)).graph.initializer
    total = sum(numpy_helper.to_array(t).size for t in inits if t.data_type == onnx.TensorProto.FLOAT)
    removed = 0
    for bn in o.nodes("BatchNormalization"):
        prod = o.producer(bn.inputs[0])
        c = o.array(bn.inputs[1]).size
        has_bias = prod.op_type in ("Conv", "ConvTranspose") and len(prod.inputs) > 2 and prod.inputs[2]
        removed += 4 * c if has_bias else 3 * c
    assert total == 8_308_692
    assert w.num_params() == total - removed


def _bn_apply(y, o, bn_name, axis):
    bn = o.batchnorm(bn_name)
    shape = [1] * y.ndim
    shape[axis] = -1
    return (y - bn.mean.reshape(shape)) / np.sqrt(bn.var.reshape(shape) + bn.eps) * bn.gamma.reshape(shape) \
        + bn.beta.reshape(shape)


def test_bn_folds_are_exact(loaded):
    _, w, o = loaded
    rng = np.random.default_rng(0)
    # sparse KRSC: y = x_im2col @ W^T per tap, BN over Cout
    for layer in (w.encoder[0], w.encoder[6], w.encoder[-1]):
        node = o.node(layer.onnx_node)
        W = np.asarray(o.array(node.inputs[1]), np.float64)
        x = rng.normal(size=(7, layer.kernel_volume, layer.in_channels))
        y = np.einsum("nkc,okc->no", x, W.reshape(W.shape[0], -1, W.shape[-1]))
        ref = _bn_apply(y, o, layer.bn_node, 1)
        got = np.einsum("nkc,okc->no", x, layer.weight.astype(np.float64).reshape(W.shape[0], -1, W.shape[-1])) \
            + layer.bias
        np.testing.assert_allclose(got, ref, rtol=2e-6, atol=2e-5)
    # dense conv (1x1 deblock as a matmul) and conv1d heads: same fold rule along the output channels
    d0 = w.deblocks[0]
    W = np.asarray(o.conv(d0.onnx_node).weight, np.float64)[:, :, 0, 0]
    x = rng.normal(size=(5, W.shape[1]))
    np.testing.assert_allclose(x @ d0.weight[:, :, 0, 0].astype(np.float64).T + d0.bias,
                               _bn_apply(x @ W.T, o, d0.bn_node, 1), rtol=2e-6, atol=2e-5)
    # ConvTranspose k == s: per output pixel a matmul over Cin with W[:, :, i, j]; BN over Cout (axis 1 of W)
    d1 = w.deblocks[1]
    W = np.asarray(o.conv(d1.onnx_node).weight, np.float64)
    x = rng.normal(size=(5, W.shape[0]))
    np.testing.assert_allclose(x @ d1.weight[:, :, 1, 0].astype(np.float64) + d1.bias,
                               _bn_apply(x @ W[:, :, 1, 0], o, d1.bn_node, 1), rtol=2e-6, atol=2e-5)
    name, hid, _ = w.head.prediction_heads[0]
    c = o.conv(hid.onnx_node)
    x = rng.normal(size=(5, 128))
    y = x @ np.asarray(c.weight, np.float64)[:, :, 0].T + (0 if c.bias is None else c.bias)
    np.testing.assert_allclose(x @ hid.weight[:, :, 0].astype(np.float64).T + hid.bias,
                               _bn_apply(y, o, hid.bn_node, 1), rtol=2e-6, atol=2e-5)


def test_onnx_taps_name_real_tensors(loaded):
    _, w, o = loaded
    produced = {t for node in o.nodes() for t in node.outputs}
    missing = {k: v for k, v in w.onnx_taps.items() if v not in produced}
    assert not missing, missing