# SPDX-License-Identifier: Apache-2.0
"""The CPU reference's building blocks (``reference.ops``, ``reference.weights``, ``reference.model``) on the host:
the plugin semantics against literal float64 ports of the CUDA loops, the strict weight load and the port's BatchNorm
folding, and the port's rewrites against the deployed form (rebatched vs dense spatial cross-attention, rotation as a
gather). Needs torch (tt-env has it); the whole-network comparisons with ONNX Runtime are in test_reference_cpu.py.

    TT_VISIBLE_DEVICES=none python -m pytest -q -p fake_ttnn_plugin code/tt_bevformer/tests/test_reference_host.py
"""
from __future__ import annotations

import numpy as np
import pytest

torch = pytest.importorskip("torch")

from tt_bevformer.host import geometry as hg  # noqa: E402
from tt_bevformer.reference import config as C  # noqa: E402
from tt_bevformer.reference import ops  # noqa: E402
from tt_bevformer.tests import paths  # noqa: E402
from tt_bevformer.ttaw.metrics import pcc  # noqa: E402

torch.set_num_threads(4)
Q = C.NUM_QUERY_BEV


@pytest.fixture(scope="module")
def weights():
    p = paths.weights_file()
    if p is None:
        pytest.skip("BEVFormer weights not found (set BEVFORMER_WEIGHTS_DIR)")
    from tt_bevformer.reference.weights import BEVFormerWeights

    return BEVFormerWeights(p)


# --------------------------------------------------------------------------------------------- plugin semantics

@pytest.mark.parametrize("case", [dict(N=2, shapes=[(7, 9)], H=2, Cc=4, Qn=5, Pg=1, P=4),     # TSA-like
                                  dict(N=3, shapes=[(5, 8)], H=2, Cc=4, Qn=6, Pg=4, P=8),     # SCA-like
                                  dict(N=1, shapes=[(6, 6), (3, 3)], H=2, Cc=4, Qn=4, Pg=1, P=2)])
def test_msda_matches_the_cuda_loop(case):
    torch.manual_seed(0)
    S = sum(h * w for h, w in case["shapes"])
    L = len(case["shapes"])
    value = torch.randn(case["N"], S, case["H"], case["Cc"])
    ref = torch.rand(case["N"], case["Qn"], 1, case["Pg"] * 2) * 1.4 - 0.2              # some points outside
    off = torch.randn(case["N"], case["Qn"], case["H"], L * case["P"] * 2) * 3.0
    aw = torch.randn(case["N"], case["Qn"], case["H"], L * case["P"])
    a = ops.msda(value, case["shapes"], ref, off, aw, q_chunk=3)
    b = ops.msda_naive(value, case["shapes"], ref, off, aw)
    assert float((a.double() - b).abs().max()) < 1e-5


def test_mdcn_matches_the_cuda_loop_and_conv2d():
    torch.manual_seed(1)
    x = torch.randn(2, 5, 7, 9)
    offset = torch.randn(2, 18, 7, 9) * 2.0
    mask = torch.sigmoid(torch.randn(2, 9, 7, 9))
    w = torch.randn(4, 5, 3, 3)
    bias = torch.randn(4)
    a = ops.mdcn(x, offset, mask, w, bias)
    b = ops.mdcn_naive(x, offset, mask, w) + bias.double().view(1, -1, 1, 1)
    assert float((a.double() - b).abs().max()) < 2e-5
    a0 = ops.mdcn(x, torch.zeros_like(offset), torch.ones_like(mask), w)                # zero offset, mask 1
    assert float((a0 - torch.nn.functional.conv2d(x, w, None, 1, 1)).abs().max()) < 2e-5


def test_rotate_identity_and_tie_rule():
    img = torch.randn(3, C.BEV_H, C.BEV_W)
    assert torch.equal(ops.rotate(img, torch.tensor([0.0]), C.ROTATE_CENTER), img)
    a = ops.rotate(img, torch.tensor([0.37]), C.ROTATE_CENTER)
    b = ops.rotate(img, torch.tensor([0.37]), C.ROTATE_CENTER, plugin_rounding=False)   # torchvision (training)
    assert 0 < int((a != b).any(0).sum()) <= 4                                          # ties only (SPEC 4.2)


# ---------------------------------------------------------------------------------------------------- weights

def test_strict_weight_coverage(weights):
    from tt_bevformer.reference.weights import NUM_TENSORS, expected_keys

    assert len(weights.sd) == NUM_TENSORS == len(expected_keys())
    assert len(weights.dcn_modules()) == 26                                              # layer3 23 + layer4 3
    assert len(C.backbone_conv_specs()) == 1 + 33 * 3 + 26 + 4                           # stem, blocks, offsets, ds
    assert weights.conv("img_backbone.layer2.0.conv1").stride == (2, 2)                  # caffe: stride on conv1
    assert weights.conv("img_backbone.layer2.0.conv2").stride == (1, 1)
    assert weights.conv("img_backbone.layer3.5.conv2").weight.shape == (256, 256, 3, 3)
    assert weights.conv("img_backbone.layer3.5.conv2.conv_offset").weight.shape == (27, 256, 3, 3)
    with pytest.raises(KeyError):
        weights.conv("img_neck.lateral_convs.0.conv")


def test_bn_folding_equals_conv_then_bn(weights):
    """The folded convs equal conv + frozen BN (float64 oracle) to float32 rounding, also through a DCN."""
    torch.manual_seed(2)
    sd = weights.sd
    for module, bn, cin in (("img_backbone.conv1", "img_backbone.bn1", 3),
                            ("img_backbone.layer3.0.conv2", "img_backbone.layer3.0.bn2", 256)):
        x = torch.randn(1, cin, 12, 20, dtype=torch.float64)
        w = torch.from_numpy(np.asarray(sd[f"{module}.weight"], np.float64))
        cp = weights.conv(module)
        stride, pad = cp.stride[0], cp.padding[0]
        g, b, m, v = (torch.from_numpy(np.asarray(sd[f"{bn}.{k}"], np.float64)) for k in
                      ("weight", "bias", "running_mean", "running_var"))
        if module.endswith("conv2"):                                                     # DCN with zero offsets
            ref = ops.mdcn(x, torch.zeros(1, 18, 12, 20, dtype=torch.float64),
                           torch.ones(1, 9, 12, 20, dtype=torch.float64), w)
            got = ops.mdcn(x.float(), torch.zeros(1, 18, 12, 20), torch.ones(1, 9, 12, 20),
                           torch.from_numpy(cp.weight), torch.from_numpy(cp.bias)).double()
        else:
            ref = torch.nn.functional.conv2d(x, w, None, stride, pad)
            got = torch.nn.functional.conv2d(x.float(), torch.from_numpy(cp.weight), torch.from_numpy(cp.bias),
                                             stride, pad).double()
        ref = (ref - m.view(1, -1, 1, 1)) / torch.sqrt(v.view(1, -1, 1, 1) + C.BN_EPS) * g.view(1, -1, 1, 1) \
            + b.view(1, -1, 1, 1)
        rel = float((got - ref).abs().max() / ref.abs().max())
        assert rel < 1e-5, (module, rel)


# ----------------------------------------------------------------------------------------- port rewrites

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

    return BEVFormerNet(weights)


def test_rebatched_sca_equals_dense(net):
    """Encoder layer 0's spatial cross-attention, dense over 6 x 22,500 pairs (the TRT graph) vs rebatched to the
    valid pairs with the gather-form scatter (the port), on the sample0 rig with random query / value."""
    golden = paths.research_file("golden", "sample0_golden.npz")
    if golden is None:
        pytest.skip("research/bevformer golden not in this workspace")
    from tt_bevformer.reference.model import PortGeometry, point_sampling, reference_points_3d

    with np.load(golden) as g:
        l2i = g["lidar2img"]
    geo_np = hg.frame_geometry(l2i, np.zeros(18, np.float32))
    geo = PortGeometry.from_numpy(geo_np)
    torch.manual_seed(3)
    query = torch.randn(Q, C.EMBED)
    value = torch.randn(C.NUM_CAMS, C.FEAT_LEN, C.EMBED)
    ref_cam, mask_n = point_sampling(reference_points_3d(), torch.from_numpy(l2i))
    with torch.no_grad():
        dense = net.sca_dense(0, query, value, ref_cam, mask_n)
        dense_host = net.sca_dense(0, query, value, geo.ref_cam, torch.from_numpy(geo_np["bev_mask_n"]))
        rebatched = net.sca_rebatched(0, query, value, geo)
    assert float((rebatched - dense_host).abs().max()) < 1e-4                     # same geometry: summation order
    assert pcc(rebatched.numpy(), dense.numpy()) > 0.9999999                      # host vs graph geometry


def test_rotation_gather_equals_rotate_op():
    """forward_port rotates prev_bev with the host map as a row gather; on angles without ties it is the plugin."""
    torch.manual_seed(4)
    prev = torch.randn(Q, C.EMBED)
    table = torch.cat([prev, torch.zeros(1, C.EMBED)], 0)
    for a in (0.0, 1.0, -3.0, 3.0, 12.5):
        idx = torch.from_numpy(hg.rotation_index_map(a))
        ours = table[idx]
        theirs = ops.rotate(prev.view(C.BEV_H, C.BEV_W, C.EMBED).permute(2, 0, 1), torch.tensor([a]),
                            C.ROTATE_CENTER).permute(1, 2, 0).reshape(Q, C.EMBED)
        assert int((ours != theirs).any(1).sum()) <= 2, a