# SPDX-License-Identifier: Apache-2.0
"""The TT device dataflow (``reference/rewrites.py``) equals the literal graph (``reference/model.py``) on the CPU
(torch + the weights; no device, no onnxruntime):

- the whole network on the staged PandaSet sample (QT128, the faster variant; OT128 on the same frame in
  ``test_reference_cpu.py``) and on a small synthetic frame with tight capacities: every module tap PCC >= 0.9999999
  and ``pred_probs`` within 1e-4, argmax identical -- the plan's sorted / padded points, zero-sentinel gathers,
  last-voxel-wins pixel tables, log-step segmented max, K-split convs and point layers and interpolation-matrix
  up-sampling change nothing but the summation order;
- padding is harmless: two capacities give the same outputs on the real rows (to the last bits of BLAS blocking);
- the log-step segmented max equals ``scatter_reduce(amax)`` for any rounds >= the frame's need, and fails with one
  round fewer when a segment is longer;
- the align-corners interpolation matrices equal ``F.interpolate``; the emulation modes round as documented.

    TT_VISIBLE_DEVICES=none python -m pytest -q code/tt_frnet/tests/test_reference_rewrites_host.py
"""
from __future__ import annotations

import numpy as np
import pytest

torch = pytest.importorskip("torch")
pytest.importorskip("onnx")

import torch.nn.functional as F  # noqa: E402

from tt_frnet.host.frame import FRNetPreprocessor  # noqa: E402
from tt_frnet.host.preprocess import autoware_preprocess, build_plan, segmax_rounds  # noqa: E402
from tt_frnet.reference.config import FRNetConfig  # noqa: E402
from tt_frnet.reference.model import FRNetTorch, scatter_max  # noqa: E402
from tt_frnet.reference.rewrites import (Emulation, FRNetPlanTorch, plan_taps_to_literal, round_mode,  # noqa: E402
                                         segmax_logstep)
from tt_frnet.reference.weights import MODULE_TAPS, load_weights  # noqa: E402
from tt_frnet.tests import _data  # noqa: E402
from tt_frnet.ttaw.golden import TapRegistry  # noqa: E402
from tt_frnet.ttaw.metrics import pcc  # noqa: E402
from tt_frnet.ttaw.ops.upsample import interp_matrix  # noqa: E402

THREADS = 4
PATTERNS = list(MODULE_TAPS) + ["backbone.segmax*", "decoder.up*"]


@pytest.fixture(scope="module")
def qt128():
    torch.set_num_threads(THREADS)
    w = _data.weights_dir()
    return load_weights(w, "qt128"), FRNetConfig.from_model_dir(w, "qt128")


def _compare(weights, cfg, inputs, n_cap=None, m_cap=None):
    plan = build_plan(inputs, cfg, pre_norm=(weights.pre_norm_scale, weights.pre_norm_shift), n_cap=n_cap,
                      m_cap=m_cap)
    tl, tr = TapRegistry(include=PATTERNS), TapRegistry(include=PATTERNS)
    lit = FRNetTorch(weights)(inputs.points, inputs.coors, inputs.voxel_coors, inputs.inverse_map, taps=tl)
    out = FRNetPlanTorch(weights)(plan, taps=tr)
    worst = min((pcc(plan_taps_to_literal(n, tr[n], plan), tl[n]), n) for n in tl.names())
    return plan, lit, out, worst, tr


def test_plan_equals_literal_on_the_sample(qt128):
    weights, cfg = qt128
    _data.need_pandaset()
    prep = FRNetPreprocessor(cfg, plan=False)(_data.SAMPLE, _data.SAMPLE_CALIB)
    plan, lit, out, worst, _ = _compare(weights, cfg, prep.inputs)
    assert plan.n_cap == cfg.n_cap == 120000 and plan.m_cap == 60000
    assert worst[0] >= 0.9999999, worst
    got = plan.unsort(out)
    assert np.abs(got - lit).max() < 1e-4
    np.testing.assert_array_equal(got.argmax(1), lit.argmax(1))
    assert np.allclose(out[plan.n:].sum(1), 1.0, atol=1e-5)          # padded rows: finite, ignored


def test_plan_on_a_small_frame_and_padding_is_harmless(qt128):
    weights, cfg = qt128
    rng = np.random.default_rng(7)
    n = 6000
    az = rng.uniform(-np.pi, np.pi, n)
    el = np.radians(rng.uniform(-50, 50, n))
    r = rng.uniform(2, 60, n)
    pts = np.c_[r * np.cos(el) * np.cos(az), r * np.cos(el) * np.sin(az), r * np.sin(el),
                rng.uniform(0, 255, n)].astype(np.float32)
    pts[:500, :3] = pts[0, :3] + rng.normal(0, 0.01, (500, 3))       # one dense cell: a long segment
    inputs = autoware_preprocess(pts, cfg, crop_bounds=(0.0,) * 6)
    p1, lit, o1, worst, t1 = _compare(weights, cfg, inputs, n_cap=inputs.num_points + 1, m_cap=inputs.num_voxels + 1)
    assert p1.max_segment_length >= 400 and segmax_rounds(p1.max_segment_length) >= 9
    assert worst[0] >= 0.9999999, worst
    assert np.abs(p1.unsort(o1) - lit).max() < 1e-4
    p2 = build_plan(inputs, cfg, pre_norm=(weights.pre_norm_scale, weights.pre_norm_shift),
                    n_cap=inputs.num_points + 4096, m_cap=inputs.num_voxels + 777)
    o2 = FRNetPlanTorch(weights)(p2)
    # padded rows never reach a real row: only BLAS blocking of the taller matrices may move the last bits
    np.testing.assert_allclose(p2.unsort(o2), p1.unsort(o1), rtol=0, atol=1e-6)


def test_segmax_logstep_equals_scatter_reduce(qt128):
    _, cfg = qt128
    rng = np.random.default_rng(3)
    pts = np.c_[rng.normal(0, 15, (3000, 3)), rng.uniform(0, 255, 3000)].astype(np.float32)
    pts[:300, :3] = pts[0, :3]
    inputs = autoware_preprocess(pts, cfg, crop_bounds=(0.0,) * 6)
    plan = build_plan(inputs, cfg, n_cap=inputs.num_points + 64, m_cap=inputs.num_voxels + 9)
    x = torch.randn(plan.n_cap, 5)
    longest = int(np.argmax(np.diff(plan.seg_offsets)))
    x[int(plan.seg_offsets[longest])] = 100.0       # the longest segment's max sits on its FIRST row
    ref = scatter_max(x[: plan.n], torch.from_numpy(plan.seg_id[: plan.n]), plan.m)
    need = segmax_rounds(plan.max_segment_length)
    for rounds in (need, need + 2):
        assert torch.equal(segmax_logstep(x, plan, rounds)[: plan.m], ref)
    assert not torch.equal(segmax_logstep(x, plan, need - 1)[: plan.m], ref)


@pytest.mark.parametrize("hw,s", [((128, 1024), 2), ((128, 256), 8), ((16, 128), 4)])
def test_interp_matrices_equal_align_corners_interpolate(hw, s):
    H, W = hw
    x = torch.randn(1, 3, H // s, W // s, dtype=torch.float64)
    ah = torch.from_numpy(interp_matrix(H // s, H, mode="linear", align_corners=True))
    aw = torch.from_numpy(interp_matrix(W // s, W, mode="linear", align_corners=True))
    y = torch.einsum("hi,bcij,wj->bchw", ah, x, aw)
    ref = F.interpolate(x, size=(H, W), mode="bilinear", align_corners=True)
    assert torch.allclose(y, ref, atol=1e-12)


def test_rounding_modes():
    x = torch.tensor([1.0, 1.0 + 2 ** -8, 1.0 + 2 ** -7 + 2 ** -9, -3.14159265, 1e-30])
    assert torch.equal(round_mode(x, "fp32"), x)
    assert torch.equal(round_mode(x, "bf16"), x.to(torch.bfloat16).float())
    t = round_mode(x, "tf32")
    assert (t.view(torch.int32) & 0x1FFF).eq(0).all() and (t.abs() <= x.abs()).all()
    e = Emulation.parse("conv_act=bf16,point_act=tf32")
    assert e.label() == "conv_act=bf16,point_act=tf32" and not e.exact and Emulation().exact
    with pytest.raises(ValueError):
        Emulation.parse("point_act=fp16")