"""Regression tests for SMPLFX export coordinate-frame handling."""

from __future__ import annotations

import sys
from pathlib import Path
from types import ModuleType

import h5py
import numpy as np
import pytest

torch = pytest.importorskip("torch")

from sltk.extraction.smplfx_fitting import SMPLFXFitter


def _mock_params() -> dict[str, torch.Tensor]:
    """Build minimal parameter payload used by the exporter."""
    return {
        "transl": torch.randn(2, 3),
        "global_orient": torch.randn(2, 1, 3, 3),
        "body_pose": torch.randn(2, 21, 3, 3),
        "betas": torch.randn(1, 10),
    }


def _install_fake_utils_module(
    monkeypatch: pytest.MonkeyPatch,
    fake_get_smplx_out,
) -> None:
    """Inject a lightweight ``utils_smplx`` module used by the exporter."""
    fake_module = ModuleType("sltk.extraction.smplfx_deps.utils_smplx")
    setattr(fake_module, "getSMPLXOut", fake_get_smplx_out)
    monkeypatch.setitem(
        sys.modules,
        "sltk.extraction.smplfx_deps.utils_smplx",
        fake_module,
    )


def test_save_keypoints_uses_body_centred_features(
    monkeypatch: pytest.MonkeyPatch,
    tmp_path: Path,
) -> None:
    """H5 export stays camera-space while location features use body-centred params."""
    calls: list[dict[str, object]] = []

    def fake_get_smplx_out(
        smplx_kwargs: dict,
        device: str,
        params: dict,
        BS: int = 512,
        run_with_verts: bool = True,
    ) -> tuple[dict[str, torch.Tensor], object]:
        del smplx_kwargs, device, BS
        calls.append(
            {
                "run_with_verts": run_with_verts,
                "keys": set(params.keys()),
            }
        )
        if run_with_verts:
            return (
                {
                    "joints": torch.full((2, 127, 3), 3.0),
                    "vertices": torch.zeros((2, 10475, 3)),
                },
                object(),
            )
        return ({"joints": torch.full((2, 127, 3), 7.0)}, object())

    _install_fake_utils_module(monkeypatch, fake_get_smplx_out)

    h5_path = tmp_path / "body_smplfx.h5"
    loc_path = SMPLFXFitter._save_keypoints_h5(
        smplx_kwargs={"dummy": 1},
        smplx_params=_mock_params(),
        output_path=h5_path,
        device="cpu",
    )

    assert len(calls) == 2
    assert calls[0]["run_with_verts"] is False
    assert calls[1]["run_with_verts"] is True
    assert {"transl", "global_orient"}.issubset(calls[0]["keys"])
    assert "transl" not in calls[1]["keys"]
    assert "global_orient" not in calls[1]["keys"]

    assert h5_path.exists()
    with h5py.File(str(h5_path), "r") as hf:
        np.testing.assert_allclose(hf["joints3d"][:], 7.0)

    assert loc_path is not None
    assert loc_path.exists()
    try:
        features = torch.load(str(loc_path), map_location="cpu", weights_only=True)
    except TypeError:
        features = torch.load(str(loc_path), map_location="cpu")
    assert isinstance(features, torch.Tensor)
    assert tuple(features.shape) == (2, 1268, 3)
    torch.testing.assert_close(features[:, :127, :], torch.full((2, 127, 3), 3.0))


def test_save_keypoints_still_writes_h5_when_body_centred_export_fails(
    monkeypatch: pytest.MonkeyPatch,
    tmp_path: Path,
) -> None:
    """Failure in body-centred vertex pass should not break camera-space H5 output."""
    calls: list[bool] = []

    def fake_get_smplx_out(
        smplx_kwargs: dict,
        device: str,
        params: dict,
        BS: int = 512,
        run_with_verts: bool = True,
    ) -> tuple[dict[str, torch.Tensor], object]:
        del smplx_kwargs, device, params, BS
        calls.append(run_with_verts)
        if run_with_verts:
            raise RuntimeError("vertex pass failed")
        return ({"joints": torch.full((2, 127, 3), 11.0)}, object())

    _install_fake_utils_module(monkeypatch, fake_get_smplx_out)

    h5_path = tmp_path / "body_smplfx.h5"
    loc_path = SMPLFXFitter._save_keypoints_h5(
        smplx_kwargs={"dummy": 1},
        smplx_params=_mock_params(),
        output_path=h5_path,
        device="cpu",
    )

    assert calls == [False, True]
    assert loc_path is None
    assert h5_path.exists()
    with h5py.File(str(h5_path), "r") as hf:
        np.testing.assert_allclose(hf["joints3d"][:], 11.0)
