"""Tests for MultimodalTSModel wrapper (T2.4).

Verifies the end-to-end path: series → Encoder → Projector → splice → LLM.
- forward returns a finite scalar loss
- generate returns text
- freeze_llm=True (stage ①): LLM frozen, encoder+projector trainable, and the
  loss carries grad into encoder/projector (backward works)
- enable_lora(r=16) (stage ②): base LLM weights frozen, LoRA adapters trainable

GPU (RTX 3060) forward/backward + peak-memory checks live here too.
"""
import pytest
import torch

from tsmm.model.ts_encoder import TSEncoder
from tsmm.model.projector import Projector
from tsmm.model.wrapper import MultimodalTSModel

LLM_PATH = "/home/zhangzp/models/Qwen2.5-0.5B-Instruct"
pytestmark = pytest.mark.skipif(
    not torch.cuda.is_available(), reason="needs CUDA for LLM forward"
)


def make_model():
    return MultimodalTSModel(
        llm_path=LLM_PATH,
        encoder=TSEncoder(d=256, layers=2, heads=4, patch_len=8, stride=4),
        projector=Projector(in_dim=256, out_dim=896),
        dtype=torch.bfloat16,
    )


def make_raw_batch(batch_size=2, T=512, C=5):
    """Raw batch contract (what the collator T2.5 will produce)."""
    series = torch.randn(batch_size, T, C)
    return {
        "series": series,
        "attributes": ["CPU 内存"] * batch_size,
        "timestamps": ["t0..t1"] * batch_size,
        "events": [""] * batch_size,
        "question": ["是否存在异常？"] * batch_size,
        "answer": ["正常。"] * batch_size,
    }


class TestForwardGenerate:
    def test_forward_returns_finite_loss(self):
        model = make_model().cuda()
        batch = make_raw_batch(batch_size=2)
        with torch.no_grad():
            out = model(batch)
        assert "loss" in out
        assert torch.isfinite(out["loss"])
        assert out["loss"].dim() == 0

    def test_forward_backward_flows_into_encoder_projector(self):
        # stage ①: LLM frozen, so the only grad path is encoder+projector
        model = make_model().cuda()
        model.freeze_llm()
        model.train()
        batch = make_raw_batch(batch_size=1)
        out = model(batch)
        out["loss"].backward()
        # encoder/projector must have received gradients
        for p in model.encoder.parameters():
            assert p.grad is not None
        for p in model.projector.parameters():
            assert p.grad is not None

    def test_generate_returns_text(self):
        model = make_model().cuda()
        batch = make_raw_batch(batch_size=1)
        text = model.generate(batch, max_new_tokens=8)
        assert isinstance(text, list)
        assert len(text) == 1
        assert isinstance(text[0], str)


class TestStageModes:
    def test_freeze_llm_only_trains_encoder_projector(self):
        model = make_model()
        model.freeze_llm()  # stage ①
        llm_requires_grad = [p.requires_grad for p in model.llm.parameters()]
        assert all(not r for r in llm_requires_grad)
        enc_proj_requires_grad = (
            [p.requires_grad for p in model.encoder.parameters()]
            + [p.requires_grad for p in model.projector.parameters()]
        )
        assert all(enc_proj_requires_grad)

    def test_enable_lora_freezes_base_injects_adapters(self):
        model = make_model()
        model.freeze_llm()
        model.enable_lora(r=16)  # stage ②
        # base (non-LoRA) LLM weights must stay frozen
        base_frozen = all(
            not p.requires_grad
            for n, p in model.llm.named_parameters()
            if "lora_" not in n
        )
        assert base_frozen
        # LoRA adapter params must be trainable
        lora_trainable = [
            p for n, p in model.llm.named_parameters()
            if "lora_" in n and p.requires_grad
        ]
        assert len(lora_trainable) > 0
        # there must be lora-named modules
        lora_names = [n for n, _ in model.named_modules() if "lora" in n.lower()]
        assert len(lora_names) > 0


class TestGPUMemory:
    def test_forward_backward_no_oom_and_reasonable_mem(self):
        torch.cuda.reset_peak_memory_stats()
        model = make_model().cuda()
        model.freeze_llm()  # stage ① mode for memory check
        model.train()
        batch = make_raw_batch(batch_size=2)
        out = model(batch)
        out["loss"].backward()
        peak_gb = torch.cuda.max_memory_allocated() / 1e9
        # stage ① should be well under the design M2 target (~5GB); use 8GB headroom
        assert peak_gb < 8.0, f"peak {peak_gb:.2f} GB"