"""Tests for giant/training/amp.py (gitea #47)."""

import tempfile
from pathlib import Path

import pytest
import torch

from giant.model.routers import EnergyRouter
from giant.model.wgan import gradient_penalty
from giant.training.amp import resolve_autocast
from giant.training.stage2_inputs import _remaining_energy_fraction
from test_train import _base_cfg, _run_train

# ---------------------------------------------------------------------------
# resolve_autocast
# ---------------------------------------------------------------------------


def test_resolve_autocast_fp32_is_disabled():
    device_type, dtype, enabled = resolve_autocast("fp32", torch.device("cpu"))
    assert device_type == "cpu"
    assert dtype is torch.float32
    assert enabled is False


def test_resolve_autocast_bf16_on_cpu_is_enabled():
    """CPU bf16 autocast is what lets the mixed-precision path be tested
    without a GPU (torch 2.3 supports it)."""
    device_type, dtype, enabled = resolve_autocast("bf16", torch.device("cpu"))
    assert device_type == "cpu"
    assert dtype is torch.bfloat16
    assert enabled is True


def test_resolve_autocast_bf16_on_unsupported_cuda_raises(monkeypatch):
    monkeypatch.setattr(torch.cuda, "is_bf16_supported", lambda: False)
    monkeypatch.setattr(torch.cuda, "get_device_capability", lambda device=None: (7, 0))
    monkeypatch.setattr(torch.cuda, "get_device_name", lambda device=None: "Tesla V100")
    with pytest.raises(ValueError, match="bf16"):
        resolve_autocast("bf16", torch.device("cuda"))


def test_resolve_autocast_bf16_on_mps_raises():
    with pytest.raises(ValueError, match="bf16"):
        resolve_autocast("bf16", torch.device("mps"))


def test_resolve_autocast_unknown_precision_raises():
    with pytest.raises(ValueError, match="fp32.*bf16"):
        resolve_autocast("fp16", torch.device("cpu"))


# ---------------------------------------------------------------------------
# End-to-end: train() under bf16 on CPU
# ---------------------------------------------------------------------------


def test_train_end_to_end_bf16_cpu_completes_and_stores_fp32_params():
    """Reuses tests/test_train.py's synthetic-batch harness — train() itself
    is device-agnostic, and CPU bf16 autocast is real (not mocked) in torch
    2.3, so this is a genuine exercise of the autocast region added to
    FlowDDPMStageTrainer.step/WGANStageTrainer.step, not just a config
    passthrough check.

    Also asserts the checkpoint's stored parameters are fp32: autocast only
    changes the dtype of intermediate activations, never the model's own
    stored weights — a regression here would mean something accidentally
    cast the model itself (e.g. `model.to(dtype=torch.bfloat16)`) rather than
    using autocast."""
    cfg = _base_cfg()
    cfg["train"]["precision"] = "bf16"
    with tempfile.TemporaryDirectory() as tmp:
        out_dir = Path(tmp) / "run"
        _run_train(cfg, out_dir)
        assert (out_dir / "last.pt").exists()
        assert (out_dir / "metrics.csv").exists()
        ckpt = torch.load(out_dir / "last.pt", weights_only=False)
        for stage_key in ("model", "sec_decoder"):
            if stage_key not in ckpt:
                continue
            for name, tensor in ckpt[stage_key].items():
                if tensor.is_floating_point():
                    assert tensor.dtype == torch.float32, f"{stage_key}.{name} is {tensor.dtype}, expected fp32"


@pytest.mark.parametrize("generator", ["wgan", "flow"])
def test_train_end_to_end_bf16_cpu_stage2_generators(generator):
    """bf16 covers both trainer subclasses (FlowDDPMStageTrainer and
    WGANStageTrainer) — the wgan default in _base_cfg exercises the
    generator-forward/critic-scoring autocast region added to
    WGANStageTrainer.step, and flow exercises the plain _compute wrap."""
    cfg = _base_cfg()
    cfg["train"]["precision"] = "bf16"
    cfg["stage2_model"]["generator"] = generator
    with tempfile.TemporaryDirectory() as tmp:
        _run_train(cfg, Path(tmp) / "run")


# ---------------------------------------------------------------------------
# fp32 guards: correct in fp32, quietly degrade in bf16 — stay fp32 even
# under an active bf16 autocast region.
# ---------------------------------------------------------------------------


def test_remaining_energy_fraction_stays_fp32_under_bf16_autocast():
    fraction = torch.rand(4, 5).to(torch.bfloat16)
    with torch.autocast("cpu", dtype=torch.bfloat16, enabled=True):
        out = _remaining_energy_fraction(fraction)
    assert out.dtype == torch.float32


def test_gradient_penalty_stays_fp32_under_bf16_autocast():
    critic = torch.nn.Linear(6, 1)

    def critic_fn(x):
        return critic(x)

    real = torch.randn(4, 6)
    fake = torch.randn(4, 6)
    with torch.autocast("cpu", dtype=torch.bfloat16, enabled=True):
        gp = gradient_penalty(critic_fn, real, fake)
    assert gp.dtype == torch.float32


def test_router_balance_and_entropy_loss_stay_fp32_under_bf16_autocast():
    router = EnergyRouter(n_experts=3)
    cond_cont = torch.randn(8, 15)
    cond_cat = torch.zeros(8, 2, dtype=torch.long)
    with torch.autocast("cpu", dtype=torch.bfloat16, enabled=True):
        balance = router.balance_loss(cond_cont, cond_cat)
        entropy = router.entropy_loss(cond_cont, cond_cat)
        weights = router.combine_weights(cond_cont, cond_cat)
    assert balance.dtype == torch.float32
    assert entropy.dtype == torch.float32
    assert weights.dtype == torch.float32