# SPDX-License-Identifier: Apache-2.0
"""The CPU reference without onnxruntime (host suite: tt-env, research venv): architecture, weights loaders, the
fp32 forward against the stored goldens, teacher-forced partial runs and the stored smoke references.

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

The weights are found without downloading (``reference.assets``: ``SCENESEG_WEIGHTS_DIR``, the staged
``weights_upload/``, the workspace download); tests that need them skip elsewhere. The goldens were written in the
research venv (torch 2.14) by ``code/scripts/ref_golden.py``; this suite may run on another torch (tt-env: 2.11).
"""
from __future__ import annotations

import json

import numpy as np
import pytest

from tt_sceneseg.reference import assets, config
from tt_sceneseg.reference.weights import ARCH, ARCH_BY_KEY, WeightsError, load_weights, param_count
from tt_sceneseg.tests.research import GOLDENS, LARGE_GOLDENS, SAMPLE, SAMPLES

ONNX = assets.onnx_path(required=False)
needs_onnx = pytest.mark.skipif(ONNX is None, reason="SceneSeg_FP32.onnx not found (set SCENESEG_WEIGHTS_DIR)")


def _macs(spec, hw):
    """MACs of one layer at its input resolution ``hw``; returns (macs, output hw)."""
    h, w = hw
    if spec.kind == "linear":
        return spec.cin * spec.cout, hw
    if spec.kind == "conv_transpose":
        return spec.cin * spec.cout * spec.k * spec.k * h * w, (h * spec.stride, w * spec.stride)
    oh, ow = (h + 2 * spec.pad - spec.k) // spec.stride + 1, (w + 2 * spec.pad - spec.k) // spec.stride + 1
    return spec.cout * oh * ow * (spec.cin // spec.groups) * spec.k * spec.k, (oh, ow)


def test_architecture_counts():
    """108 layers (100 Conv, 5 ConvTranspose, 3 Gemm), 16 depthwise, 48,341,495 parameters and 183.514 GMAC at
    batch 1 (research/sceneseg/SPEC.md 4.1, verify_onnx.json)."""
    assert len(ARCH) == 108 and len(ARCH_BY_KEY) == 108
    kinds = [s.kind for s in ARCH]
    assert (kinds.count("conv"), kinds.count("conv_transpose"), kinds.count("linear")) == (100, 5, 3)
    dw = [s for s in ARCH if s.depthwise]
    assert len(dw) == 16 and sorted({(s.k, s.stride) for s in dw}) == [(3, 1), (3, 2), (5, 1), (5, 2)]
    params = sum(int(np.prod(s.weight_shape)) + s.cout for s in ARCH)
    assert params == 48_341_495
    # MACs: walk the layers at their resolutions (backbone strides, context 10x20, decoder up-sampling)
    res = {"Backbone": (320, 640)}
    total = 0
    hw = (320, 640)
    for s in ARCH:
        if s.key.startswith("Backbone"):
            if s.key.endswith((".fc1", ".fc2")):
                m, _ = _macs(s, (1, 1))
            else:
                m, hw = _macs(s, hw)
            res["Backbone"] = hw
        elif s.kind == "linear":
            m, _ = _macs(s, (1, 1))
        elif s.key.startswith("SceneContext"):
            m, _ = _macs(s, config.CONTEXT_HW)
            hw = config.CONTEXT_HW
        elif s.kind == "conv_transpose":
            m, hw = _macs(s, hw)
        else:
            m, _ = _macs(s, hw)
        total += m
    assert res["Backbone"] == (10, 20) and hw == (320, 640)
    assert total == pytest.approx(183.514e9, rel=1e-5)


@needs_onnx
def test_onnx_weights():
    w = load_weights(ONNX)
    assert len(w) == 108 and param_count(w) == 48_341_495
    assert all(lw.folded_bn is None for lw in w.values())                    # BN was folded at export
    assert all(lw.weight.dtype == np.float32 and lw.weight.flags["C_CONTIGUOUS"] for lw in w.values())
    sd = assets.state_dict_path(required=False)
    if sd is not None:                                                       # the .pth / safetensors cross-check
        from tt_sceneseg.reference.weights import load_state_dict_weights

        p = load_state_dict_weights(sd)
        for k, lw in w.items():
            ref = p[k]
            assert np.abs(lw.weight - ref.weight).max() <= 2e-7 * max(np.abs(ref.weight).max(), 1.0), k
            assert np.abs(lw.bias - ref.bias).max() <= 1e-5, k
            if lw.spec.bn is None:                                           # unfolded layers: the same numbers
                np.testing.assert_array_equal(lw.weight, ref.weight)
                np.testing.assert_array_equal(lw.bias, ref.bias)


def test_loader_refuses_other_files(tmp_path, monkeypatch):
    bad = tmp_path / config.ONNX_FILE
    bad.write_bytes(b"not an onnx file")
    with pytest.raises(Exception):
        load_weights(bad)
    if ONNX is not None:                                                     # another export: refused by its sha256
        monkeypatch.setattr(config, "ONNX_SHA256", "0" * 64)
        with pytest.raises(WeightsError, match="sha256"):
            load_weights(ONNX)
        assert len(load_weights(ONNX, check_sha256=False)) == len(ARCH)     # ... unless asked not to check


def _small_golden(name="highway_normal_1"):
    path = GOLDENS / f"{name}.npz"
    if not path.is_file():
        pytest.skip("tests/goldens not generated (code/scripts/ref_golden.py)")
    return np.load(path)


@needs_onnx
@pytest.mark.parametrize("order", config.CHANNEL_ORDERS)
def test_forward_matches_stored_goldens(order):
    """The fp32 forward on the shipped sample against the stored goldens of the research venv run: logits at the
    sampled pixels PCC >= 0.99999, class-map agreement >= 0.9999, tap statistics within 1e-4 relative."""
    torch = pytest.importorskip("torch")
    from tt_sceneseg.host.postprocess import class_map_from_logits
    from tt_sceneseg.host.preprocess import prepare_image
    from tt_sceneseg.reference.model import SceneSegReference
    from tt_sceneseg.ttaw.golden import TapRegistry
    from tt_sceneseg.ttaw.metrics import pcc

    torch.set_num_threads(4)
    g = _small_golden()
    meta = json.loads((GOLDENS / "highway_normal_1.json").read_text())
    x = prepare_image(str(SAMPLE), channel_order=order).network_input()
    taps = TapRegistry(exclude=config.LARGE_TAPS)
    logits = SceneSegReference.from_onnx(ONNX).run_numpy(x, taps=taps)["logits"]
    yx = g["sample_yx"]
    assert pcc(logits[0][:, yx[0], yx[1]], g[f"{order}.logits_samples"]) >= 0.99999
    assert (class_map_from_logits(logits) == g[f"{order}.class_map"]).mean() >= 0.9999
    for name, st in meta["variants"][order]["taps"].items():
        if name in taps.names():
            v = taps[name]
            assert list(v.shape) == st["shape"], name
            scale = max(st["absmax"], 1e-6)
            assert abs(float(v.mean()) - st["mean"]) <= 1e-4 * scale and abs(float(v.std()) - st["std"]) <= 1e-4 * scale


@needs_onnx
def test_teacher_forced_runs_equal_the_full_forward():
    """``run_from`` (the inputs of a TT module case are reference taps) reproduces the full forward."""
    pytest.importorskip("torch")
    from tt_sceneseg.host.preprocess import prepare_image
    from tt_sceneseg.reference.model import SceneSegReference
    from tt_sceneseg.ttaw.golden import TapRegistry

    ref = SceneSegReference.from_onnx(ONNX)
    taps = TapRegistry(include=("l0", "l2", "l3", "l4", "l8", "context", "neck", "logits"))
    full = ref.run_numpy(prepare_image(str(SAMPLE)).network_input(), taps=taps)["logits"]
    t = taps.to_dict()
    for start in ("context", "neck", "head"):
        got = ref.run_from(start, t)["logits"]
        np.testing.assert_allclose(got, full, rtol=0, atol=1e-4, err_msg=start)
    with pytest.raises(ValueError):
        ref.run_from("backbone", t)


@needs_onnx
@pytest.mark.parametrize("order", config.CHANNEL_ORDERS)
def test_pipeline_equals_stored_smoke_reference(order):
    """``samples/highway_normal_1.<variant>.reference.json`` (the container smoke's reference) is what the
    reference pipeline returns: the same published mask and class map (agreement >= 0.9999 across torch builds)."""
    pytest.importorskip("torch")
    from tt_sceneseg.reference.pipeline import ReferencePipeline
    from tt_sceneseg.ttaw.server import smoke

    path = SAMPLES / f"highway_normal_1.{order}.reference.json"
    if not path.is_file():
        pytest.skip("stored reference not generated (code/scripts/ref_golden.py)")
    stored = json.loads(path.read_text())
    body = ReferencePipeline(ONNX, order, threads=4).reference_body(str(SAMPLE))
    metrics, fails = smoke.compare_with_reference(body, stored, gates={"min_label_agreement": 0.9999})
    assert fails == [], fails
    assert set(metrics) == {"mask", "class_map"}
    assert body["meta"]["channel_order"] == order and body["class_names"] == list(config.CLASS_NAMES)


def test_large_goldens_index():
    """research/sceneseg/goldens/index.json (written with the goldens): the module taps of the gate samples exist
    with the recorded sha256 (sizes only checked here; the device tests read them)."""
    idx = LARGE_GOLDENS / "index.json"
    if not idx.is_file():
        pytest.skip("research/sceneseg/goldens not generated")
    meta = json.loads(idx.read_text())
    for rel, rec in meta["files"].items():
        p = LARGE_GOLDENS / rel
        assert p.is_file() and p.stat().st_size == rec["bytes"], rel


def test_comma10k_goldens_reproduce_the_spec():
    """The comma10k ground-truth metric (tests/sceneseg_eval.py) on the stored fp32 class maps (shipped copies in
    tests/goldens/comma10k/): the per-image IoUs of
    the index, and the SPEC section 3 means (training ROI crop: foreground 0.729 bgr / 0.765 rgb, road 0.907 /
    0.919): the FP32 side of the PLAN 2.2 comma10k gate."""
    from tt_sceneseg.tests.sceneseg_eval import COMMA10K_GOLDENS, comma10k_ious

    rec = json.loads((COMMA10K_GOLDENS / "index.json").read_text())   # shipped since the release docs
    for name, r in rec["images"].items():
        with np.load(COMMA10K_GOLDENS / f"{name[:4]}.npz") as z:
            for order in config.CHANNEL_ORDERS:
                assert comma10k_ious(z[f"cls_{order}"], z["gt"]) == pytest.approx(r[order], nan_ok=True), (name, order)
    spec = {"bgr": (0.729, 0.907), "rgb": (0.765, 0.919)}
    for order, (fg, road) in spec.items():
        assert rec["mean"][order]["foreground"] == pytest.approx(fg, abs=1e-3)
        assert rec["mean"][order]["road"] == pytest.approx(road, abs=1e-3)


def test_shipped_comma10k_matches_the_gate_inputs():
    """The comma10k images, masks and fp32 goldens shipped since the release docs are byte-identical to the ones the
    frozen comma10k gate was set on (sha256 of research/sceneseg/goldens/index.json, recorded with the goldens), and
    the workspace copies (when present) still match."""
    import hashlib

    from tt_sceneseg.tests.sceneseg_eval import COMMA10K_GOLDENS, comma10k_samples

    def sha(p):
        return hashlib.sha256(p.read_bytes()).hexdigest()

    rec = json.loads((COMMA10K_GOLDENS / "index.json").read_text())
    samples = comma10k_samples()
    assert sorted(samples) == sorted(rec["images"])
    for name, (img, mask) in samples.items():
        assert sha(img) == rec["images"][name]["image_sha256"], name
        assert sha(mask) == rec["images"][name]["mask_sha256"], name
    for fname, f in rec["files"].items():
        assert sha(COMMA10K_GOLDENS / fname) == f["sha256"], fname
    idx = LARGE_GOLDENS / "index.json"
    if idx.is_file():
        large = json.loads(idx.read_text())
        assert large["comma10k"]["images"] == rec["images"] and large["comma10k"]["mean"] == rec["mean"]
        for fname, f in rec["files"].items():
            assert large["files"][f"comma10k/{fname}"] == f, fname