import torch
import torch.nn as nn
from torchlitex.checkpoints import load_checkpoint, save_checkpoint
from torchlitex.utils import EMA

try:
    from torch.amp import GradScaler as _GradScaler  # type: ignore
except Exception:
    from torch.cuda.amp import GradScaler as _GradScaler  # type: ignore


def test_checkpoint_round_trip(tmp_path):
    model = nn.Linear(4, 2)
    optimizer = torch.optim.SGD(model.parameters(), lr=0.1)
    scaler = _GradScaler(enabled=False)

    ckpt_path = tmp_path / "ckpt.pt"

    save_checkpoint(ckpt_path, model, optimizer=optimizer, scaler=scaler, epoch=3, best_val=1.23)
    assert ckpt_path.exists()

    new_model = nn.Linear(4, 2)
    new_optimizer = torch.optim.SGD(new_model.parameters(), lr=0.1)
    new_scaler = _GradScaler(enabled=False)

    epoch, best_val, ema_state = load_checkpoint(ckpt_path, new_model, optimizer=new_optimizer, scaler=new_scaler)

    assert epoch == 3
    assert best_val == 1.23
    assert ema_state is None
    assert torch.allclose(model.weight, new_model.weight)


def test_checkpoint_with_ema(tmp_path):
    model = nn.Linear(2, 1)
    ema = EMA(model, decay=0.9)
    for _ in range(2):
        ema.update(model)

    ckpt_path = tmp_path / "ckpt.pt"
    save_checkpoint(ckpt_path, model, epoch=1, best_val=0.5, ema_state=ema.state_dict())

    new_model = nn.Linear(2, 1)
    epoch, best_val, ema_state = load_checkpoint(ckpt_path, new_model)

    assert epoch == 1
    assert best_val == 0.5
    assert "shadow" in ema_state
