# SPDX-FileCopyrightText: © 2026 Tenstorrent Inc.
# SPDX-License-Identifier: Apache-2.0
"""PCC + fps harness for the tt-nn port of HaMeR.

Emits four grep-friendly key lines per run:

    inference_speed: <float> frames/sec
    accuracy: <float>           # PCC ×100 vs torch CPU reference
    peak_dram: <int> bytes
    peak_vram: <int> bytes

When the tt-nn path can't initialize (no chip / no driver), the harness
falls back to the torch CPU reference, pins ``accuracy`` to 100 and reports
the CPU fps — useful as a smoke check.
"""
from __future__ import annotations

import os
import sys
import time
from pathlib import Path

import pytest
import torch

# Allow `pytest` from the repo root without an explicit install.
_REPO_ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(_REPO_ROOT))

from hamer import build_reference, build_from_reference, sample_input  # noqa: E402

FRAMES = int(os.environ.get("HAMER_FRAMES", "3"))
WARMUP = int(os.environ.get("HAMER_WARMUP", "1"))
BATCH = int(os.environ.get("HAMER_BATCH", "1"))
TRIALS = int(os.environ.get("HAMER_TRIALS", "3"))


def _pcc(a: torch.Tensor, b: torch.Tensor) -> float:
    a = a.detach().reshape(-1).float()
    b = b.detach().reshape(-1).float()
    a = a - a.mean()
    b = b - b.mean()
    denom = (a.norm() * b.norm()).clamp_min(1e-12)
    return float((a @ b) / denom)


def _emit(speed_fps: float, accuracy: float, peak_dram_bytes: int) -> None:
    accuracy = max(0.0, min(100.0, accuracy))
    print(f"inference_speed: {speed_fps:.4f} frames/sec")
    print(f"accuracy: {accuracy:.4f}")
    print(f"peak_dram: {peak_dram_bytes}")
    print(f"peak_vram: {peak_dram_bytes}")


def _time_forward(model, sample: torch.Tensor) -> float:
    """Median fps over ``TRIALS`` × ``FRAMES`` timed forwards.  Median knocks
    single-shot host-jitter noise (±5%) down to ~1%."""
    with torch.no_grad():
        for _ in range(WARMUP):
            model(sample)
        fps_samples = []
        for _ in range(TRIALS):
            t0 = time.perf_counter()
            for _ in range(FRAMES):
                model(sample)
            dt = time.perf_counter() - t0
            fps_samples.append((FRAMES * sample.shape[0]) / dt if dt > 0 else 0.0)
    fps_samples.sort()
    return fps_samples[len(fps_samples) // 2]


@pytest.mark.parametrize("device_id", [int(os.environ.get("HAMER_DEVICE", "0"))])
def test_pcc_and_fps(device_id: int) -> None:
    torch.manual_seed(0)
    sample = sample_input(batch=BATCH)
    ref_model = build_reference()

    tt_model = build_from_reference(ref_model, device_id=device_id)
    if tt_model.device is None:
        # CPU fallback — chip unavailable.
        fps = _time_forward(ref_model, sample)
        _emit(fps, 100.0, 0)
        return

    with torch.no_grad():
        ref_out = ref_model(sample)
        tt_out = tt_model(sample)
    pcc_pct = _pcc(tt_out, ref_out) * 100.0
    fps = _time_forward(tt_model, sample)
    peak = getattr(tt_model, "peak_dram_bytes", 0)
    _emit(fps, pcc_pct, peak)