# SPDX-License-Identifier: Apache-2.0
"""The CPU reference's building blocks that need torch but neither onnxruntime nor the full network run:
weights addressed by consuming node + the deployed-graph check, the three plugins against the research numpy
implementations, and the TT port's exact rewrites (``reference.port_form``) against the deployed form.

    TT_VISIBLE_DEVICES=none python -m pytest -q code/tt_bevdet/tests/test_reference_host.py

Needs the pinned ONNX (``tests/paths.py``: ``$BEVDET_WEIGHTS_DIR``, the HF cache or the workspace download); the
weight tests skip without it. The end-to-end reference against ONNX Runtime is ``test_reference_cpu.py``.
"""
from __future__ import annotations

import numpy as np
import pytest

torch = pytest.importorskip("torch")

from tt_bevdet.reference import plugins, port_form  # noqa: E402
from tt_bevdet.reference.config import HEAD_NAMES, default_config  # noqa: E402
from tt_bevdet.reference.weights import ONNX_SHA256, BEVDetWeights, module_names  # noqa: E402
from tt_bevdet.tests import paths  # noqa: E402
from tt_bevdet.ttaw.metrics import pcc  # noqa: E402

CFG = default_config()
DN = "img_view_transformer.depth_net."


@pytest.fixture(scope="module")
def weights():
    path = paths.weights_onnx()
    if path is None:
        pytest.skip("bevdet_one_lt_d.onnx not found (set BEVDET_WEIGHTS_DIR)")
    pytest.importorskip("onnx")
    return BEVDetWeights(path)


@pytest.fixture(scope="module")
def net(weights):
    from tt_bevdet.reference.model import BEVDetNet

    torch.set_num_threads(2)
    return BEVDetNet(weights, CFG)


def test_weights_by_consuming_node(weights):
    assert weights.sha256 == ONNX_SHA256
    names = module_names()
    n_conv = sum(len(v) for k, v in names.items() if k != "depth_net_linear")
    assert n_conv == 107 == len(weights.onnx.nodes("Conv"))                    # SPEC 4.2: Conv 107
    c = weights.conv("img_backbone.conv1")
    assert c.node == "Conv_1" and c.weight.shape == (64, 3, 7, 7) and c.stride == (2, 2) and c.padding == (3, 3)
    assert [weights.conv(DN + f"depth_conv.3.aspp{i}.atrous_conv").dilation for i in (2, 3, 4)] == \
        [(6, 6), (12, 12), (18, 18)]
    assert weights.conv("img_bev_encoder_backbone.layers.0.0.downsample").stride == (2, 2)
    w, b = weights.linear(DN + "context_mlp.fc1")
    assert w.shape == (512, 27) and b.shape == (512,)
    bn = weights.depth_bn()
    assert bn.channels == 27 and bn.eps == pytest.approx(1e-5)
    n_params = sum(c.weight.size + (0 if c.bias is None else c.bias.size) for _, c in weights.iter_convs())
    n_params += sum(a.size for m in names["depth_net_linear"] for a in weights.linear(m)) + 4 * 27
    assert n_params == 76_187_718                                              # SPEC 4.3


def test_align_gather_bevpool_vs_research():
    bc = paths.import_research()
    if bc is None:
        pytest.skip("research reference not present")
    rng = np.random.default_rng(2)
    adj = rng.normal(size=(1, 8, 4, 20, 24)).astype(np.float32)
    trs = np.zeros((1, 8, 6), np.float32)
    for n in range(8):
        a = rng.uniform(-0.3, 0.3)
        trs[0, n] = [np.cos(a), -np.sin(a), rng.uniform(-6, 6), np.sin(a), np.cos(a), rng.uniform(-6, 6)]
    trs[0, 3] = [1, 0, 0, 0, 1, 0]
    got = plugins.align_bev(adj, trs).numpy()
    np.testing.assert_array_equal(got, bc.align_bev(adj, trs))                 # bit for bit
    np.testing.assert_array_equal(got[0, 3], adj[0, 3])                        # identity = copy
    curr = rng.normal(size=(1, 4, 20, 24)).astype(np.float32)
    for flag in (0, 1):
        np.testing.assert_array_equal(plugins.gather_bev(got, curr, flag).numpy(), bc.gather_bev(got, curr, flag))


def test_camera_gates_host_fold_vs_graph(weights, net):
    rng = np.random.default_rng(4)
    cp = np.zeros((1, 6, 27), np.float32)
    cp[0, :, :4] = rng.uniform([1000, 1000, 780, 430], [1300, 1300, 830, 470], size=(6, 4))
    cp[0, :, 4:15] = [0.44, 0, 0, 0, 0.44, -140, 1, 0, 0, 1, 1]
    cp[0, :, 15:] = rng.normal(size=(6, 12))
    g_ctx, g_depth = port_form.camera_gates(weights, cp)
    t_ctx, t_depth = net.camera_gates(torch.from_numpy(cp))
    assert np.abs(g_ctx - t_ctx.numpy()).max() < 1e-5 and np.abs(g_depth - t_depth.numpy()).max() < 1e-5
    w, b = weights.linear(DN + "depth_mlp.fc1")
    bn = weights.depth_bn()
    wf, bf = port_form.fold_bn_into_linear_input(w, b, bn)
    x = cp[0].astype(np.float64)
    np.testing.assert_allclose(x @ wf.T + bf, bn.apply(x, axis=1) @ w.astype(np.float64).T + b, rtol=1e-9, atol=1e-9)


def test_depthnet_port_form_matches_deployed(weights, net):
    """1x3 ASPP branch 4, GAP branch as a per-frame bias, depth logits padded to 128 with -1e4, host gates."""
    rng = np.random.default_rng(6)
    x = torch.from_numpy((rng.normal(size=(2, 512, 16, 44)) * 0.5).astype(np.float32))
    gates = (rng.uniform(0, 1, (2, 512)).astype(np.float32), rng.uniform(0, 1, (2, 512)).astype(np.float32))
    with torch.no_grad():
        d_ref, f_ref = net.depth_net(x, gates=tuple(torch.from_numpy(g) for g in gates))
    d_port, f_port = port_form.forward_depth_net_port(net, x, gates, weights, CFG)
    np.testing.assert_array_equal(f_port.numpy(), f_ref.numpy())
    assert d_port.shape == d_ref.shape and np.abs(d_port.numpy() - d_ref.numpy()).max() < 2e-6
    assert pcc(d_port.numpy(), d_ref.numpy()) > 0.99999999
    pad = port_form.pad_depth_logits(weights)
    assert pad.weight.shape[0] == 128 and pad.n_valid == 118 and np.all(pad.bias[118:] == -1e4)
    assert float(np.exp(np.float32(-1e4) - np.float32(-50.0))) == 0.0          # the pad columns vanish in fp32


def test_merged_head_matches_deployed(weights, net):
    rng = np.random.default_rng(8)
    x = torch.from_numpy(rng.normal(size=(1, 256, 32, 32)).astype(np.float32))
    with torch.no_grad():
        ref = net.center_head(x)
    got = port_form.forward_center_head_port(net, x, weights)
    m = port_form.merge_center_head(weights)
    assert m.w1.shape == (384, 64, 3, 3) and m.w2.shape == (20, 384, 3, 3) and [s[0] for s in m.splits] == \
        list(HEAD_NAMES)
    for k in HEAD_NAMES:
        assert np.abs(got[k].numpy() - ref[k].numpy()).max() < 5e-5, k
        assert pcc(got[k].numpy(), ref[k].numpy()) > 0.9999999, k