# SPDX-FileCopyrightText: © 2026 Tenstorrent Inc.
# SPDX-License-Identifier: Apache-2.0
"""Host-only MANO checks (no device, no licensed MANO_RIGHT.pkl needed).

A synthetic pkl with the real MANO array shapes stands in for MANO_RIGHT.pkl: its outputs are meaningless, but
the shape blend / kinematic walk / skinning code paths are the same, which is all the bit-identity test needs.
"""
import pickle

import numpy as np
import pytest
import torch

from hamer.mano import MANO


@pytest.fixture(scope="module")
def mano(tmp_path_factory):
    g = np.random.default_rng(0)
    w = g.random((778, 16))
    w /= w.sum(1, keepdims=True)
    jr = np.zeros((16, 778))
    for j in range(16):
        jr[j, g.choice(778, 20, replace=False)] = 1 / 20
    d = {"v_template": g.random((778, 3)) * 0.1, "shapedirs": g.normal(0, 1e-3, (778, 3, 10)),
         "posedirs": g.normal(0, 1e-3, (778, 3, 135)), "J_regressor": jr, "weights": w,
         "kintree_table": np.array([[4294967295, 0, 1, 2, 0, 4, 5, 0, 7, 8, 0, 10, 11, 0, 13, 14], list(range(16))]),
         "f": g.integers(0, 778, (1538, 3))}
    p = tmp_path_factory.mktemp("mano") / "MANO_RIGHT.pkl"
    with open(p, "wb") as f:
        pickle.dump(d, f)
    return MANO(p).eval()


def _random_pose(g, batch):
    a = torch.randn(batch, 16, 3, 3, generator=g, dtype=torch.float64) * 0.5
    rot = torch.linalg.matrix_exp(a - a.transpose(-1, -2)).float()     # proper rotations
    return rot, torch.randn(batch, 10, generator=g)


@pytest.mark.parametrize("batch", [1, 3])
def test_forward_with_joints_bit_identical(mano, batch):
    """The server's one-walk path (default) returns exactly forward() + joints_21() (the HAMER_MANO_ONEPASS=0 path)."""
    g = torch.Generator().manual_seed(batch)
    with torch.inference_mode():
        for _ in range(100):
            rot, betas = _random_pose(g, batch)
            v0, j0 = mano(rot, betas), mano.joints_21(rot, betas)
            v1, j1 = mano.forward_with_joints(rot, betas)
            assert v1.shape == (batch, 778, 3) and j1.shape == (batch, 21, 3)
            assert torch.equal(v0, v1) and torch.equal(j0, j1)


def test_server_mano_onepass_switch(monkeypatch):
    import importlib

    import hamer.server.app as app
    try:
        monkeypatch.setenv("HAMER_MANO_ONEPASS", "0")
        assert importlib.reload(app).MANO_ONEPASS is False
        monkeypatch.delenv("HAMER_MANO_ONEPASS")
        assert importlib.reload(app).MANO_ONEPASS is True     # default on
    finally:
        monkeypatch.delenv("HAMER_MANO_ONEPASS", raising=False)
        importlib.reload(app)