"""
GPU integration tests for the SignRep inference pipeline.

These tests require:
- A CUDA-capable GPU
- The SignRep checkpoint at sltk/weights/signrep/ckpt.pt
- Example video at /home/ef0036/Projects/SignRepID/example_video.mp4
- (Some tests) pre-built dictionary .pt files

Run with:
    pytest tests/test_signrep_integration.py -v -m gpu
    pytest tests/test_signrep_integration.py -v -m "gpu and not slow"
"""

from __future__ import annotations

import json
from pathlib import Path

import numpy as np
import pytest

torch = pytest.importorskip("torch")

EXAMPLE_VIDEO = Path("/home/ef0036/Projects/SignRepID/example_video.mp4")
CHECKPOINT = Path("/home/ef0036/Projects/VLT/sltk/weights/signrep/ckpt.pt")
DICT_PT_RACHEL = Path("/mnt/fast/nobackup/scratch4weeks/ef0036/SignRepID/bsl_dicts/rachel.pt")
DICT_PT_SIGNBANK = Path("/mnt/fast/nobackup/scratch4weeks/ef0036/SignRepID/bsl_dicts/signbank.pt")

# ---------------------------------------------------------------------------
# Module-level skip conditions
# ---------------------------------------------------------------------------

pytestmark = [
    pytest.mark.gpu,
    pytest.mark.slow,
    pytest.mark.skipif(
        not torch.cuda.is_available(),
        reason="GPU required",
    ),
    pytest.mark.skipif(
        not CHECKPOINT.exists(),
        reason=f"Checkpoint not found: {CHECKPOINT}",
    ),
    pytest.mark.skipif(
        not EXAMPLE_VIDEO.exists(),
        reason=f"Example video not found: {EXAMPLE_VIDEO}",
    ),
]


# ---------------------------------------------------------------------------
# Module-scoped pipeline fixture (expensive -- load once)
# ---------------------------------------------------------------------------


@pytest.fixture(scope="module")
def pipeline():
    """Load the real SignRep pipeline on GPU (once per module)."""
    from sltk.embedding.pipeline import SignRepPipeline

    return SignRepPipeline(
        checkpoint=str(CHECKPOINT),
        device="cuda",
        amp=True,
        batch_size=4,
        num_workers=0,  # 0 workers for test reliability
    )


# ---------------------------------------------------------------------------
# Dictionary extraction tests
# ---------------------------------------------------------------------------


class TestDictionaryExtraction:
    """Test dictionary-mode extraction on a real video."""

    def test_extract_single_video(self, pipeline):
        result = pipeline.extract_dictionary(str(EXAMPLE_VIDEO), method="middle", stride=1)
        assert result.feature.shape == (768,), f"Expected (768,), got {result.feature.shape}"
        assert result.all_features.ndim == 2
        assert result.all_features.shape[1] == 768
        assert len(result.active_probs) == len(result.all_features)
        assert len(result.frame_indices) == len(result.all_features)
        assert result.fps > 0
        assert result.total_frames > 0
        assert result.pooling_method == "middle"

    def test_feature_is_l2_normalized(self, pipeline):
        result = pipeline.extract_dictionary(str(EXAMPLE_VIDEO), method="middle", stride=1)
        norm = np.linalg.norm(result.feature)
        assert abs(norm - 1.0) < 1e-5, f"Feature norm should be ~1.0, got {norm}"

    def test_all_pooling_methods(self, pipeline):
        from sltk.embedding.pipeline import POOLING_METHODS

        for method in POOLING_METHODS:
            result = pipeline.extract_dictionary(str(EXAMPLE_VIDEO), method=method, stride=1)
            assert result.feature.shape == (
                768,
            ), f"Method '{method}' produced wrong shape: {result.feature.shape}"
            assert result.pooling_method == method

    def test_save_load_npz_roundtrip(self, pipeline, tmp_path):
        result = pipeline.extract_dictionary(str(EXAMPLE_VIDEO), method="middle", stride=1)
        out = tmp_path / "test.npz"
        result.save_npz(out)

        loaded = np.load(str(out), allow_pickle=True)
        np.testing.assert_array_almost_equal(loaded["best_feature"], result.feature, decimal=5)
        np.testing.assert_array_almost_equal(loaded["features"], result.all_features, decimal=5)

    def test_stride_affects_window_count(self, pipeline):
        r1 = pipeline.extract_dictionary(str(EXAMPLE_VIDEO), method="middle", stride=1)
        r4 = pipeline.extract_dictionary(str(EXAMPLE_VIDEO), method="middle", stride=4)
        assert len(r1.all_features) > len(r4.all_features), (
            f"stride=1 ({len(r1.all_features)} windows) should have more windows "
            f"than stride=4 ({len(r4.all_features)} windows)"
        )

    def test_active_probs_in_valid_range(self, pipeline):
        result = pipeline.extract_dictionary(str(EXAMPLE_VIDEO), method="middle", stride=1)
        assert np.all(np.isfinite(result.active_probs)), "Active probs contain non-finite values"


# ---------------------------------------------------------------------------
# Continuous extraction tests
# ---------------------------------------------------------------------------


class TestContinuousExtraction:
    """Test continuous-mode extraction on a real video."""

    def test_extract_continuous(self, pipeline):
        result = pipeline.extract_continuous(str(EXAMPLE_VIDEO), stride=4)
        assert result.features.ndim == 2
        assert result.features.shape[1] == 768
        assert result.stride == 4
        assert result.fps > 0
        assert result.total_frames > 0

    def test_features_are_l2_normalized(self, pipeline):
        result = pipeline.extract_continuous(str(EXAMPLE_VIDEO), stride=4)
        norms = np.linalg.norm(result.features, axis=1)
        np.testing.assert_allclose(
            norms, 1.0, atol=1e-5, err_msg="Continuous features should be L2-normalized"
        )

    def test_frame_indices_monotonically_increasing(self, pipeline):
        result = pipeline.extract_continuous(str(EXAMPLE_VIDEO), stride=4)
        diffs = np.diff(result.frame_indices)
        assert np.all(diffs > 0), "Frame indices should be strictly increasing"

    def test_save_load_npz_roundtrip(self, pipeline, tmp_path):
        result = pipeline.extract_continuous(str(EXAMPLE_VIDEO), stride=4)
        out = tmp_path / "continuous.npz"
        result.save_npz(out)

        loaded = np.load(str(out), allow_pickle=True)
        np.testing.assert_array_almost_equal(loaded["features"], result.features, decimal=5)
        np.testing.assert_array_almost_equal(loaded["latents"], result.features, decimal=5)
        assert int(loaded["stride"]) == 4
        assert str(loaded["mode"]) == "continuous"

    def test_stride_1_vs_4(self, pipeline):
        r1 = pipeline.extract_continuous(str(EXAMPLE_VIDEO), stride=1)
        r4 = pipeline.extract_continuous(str(EXAMPLE_VIDEO), stride=4)
        assert r1.features.shape[0] > r4.features.shape[0], (
            f"stride=1 ({r1.features.shape[0]}) should yield more windows "
            f"than stride=4 ({r4.features.shape[0]})"
        )


# ---------------------------------------------------------------------------
# Dictionary loading tests
# ---------------------------------------------------------------------------


class TestDictionaryLoading:
    """Test loading pre-built dictionaries."""

    @pytest.mark.skipif(
        not DICT_PT_RACHEL.exists(),
        reason=f"Rachel dictionary not found: {DICT_PT_RACHEL}",
    )
    def test_load_rachel_pt(self, pipeline):
        dictionary = pipeline.load_dictionary([str(DICT_PT_RACHEL)])
        assert len(dictionary) > 0
        assert dictionary.features.shape[1] == 768
        # Rachel dictionary should have ~28k entries
        assert len(dictionary) > 20000, f"Rachel dictionary unexpectedly small: {len(dictionary)}"

    @pytest.mark.skipif(
        not DICT_PT_SIGNBANK.exists(),
        reason=f"Signbank dictionary not found: {DICT_PT_SIGNBANK}",
    )
    def test_load_signbank_pt(self, pipeline):
        dictionary = pipeline.load_dictionary([str(DICT_PT_SIGNBANK)])
        assert len(dictionary) > 0
        assert dictionary.features.shape[1] == 768

    @pytest.mark.skipif(
        not (DICT_PT_RACHEL.exists() and DICT_PT_SIGNBANK.exists()),
        reason="Both dictionaries required for combined test",
    )
    def test_load_combined_dictionaries(self, pipeline):
        dict_r = pipeline.load_dictionary([str(DICT_PT_RACHEL)])
        dict_s = pipeline.load_dictionary([str(DICT_PT_SIGNBANK)])
        dict_combined = pipeline.load_dictionary([str(DICT_PT_RACHEL), str(DICT_PT_SIGNBANK)])
        assert len(dict_combined) == len(dict_r) + len(dict_s), (
            f"Combined ({len(dict_combined)}) should be "
            f"Rachel ({len(dict_r)}) + Signbank ({len(dict_s)})"
        )

    @pytest.mark.skipif(
        not DICT_PT_SIGNBANK.exists(),
        reason=f"Signbank dictionary not found: {DICT_PT_SIGNBANK}",
    )
    def test_dictionary_features_normalized(self, pipeline):
        dictionary = pipeline.load_dictionary([str(DICT_PT_SIGNBANK)])
        norms = torch.norm(dictionary.features, dim=1).cpu()
        assert torch.allclose(
            norms, torch.ones_like(norms), atol=1e-5
        ), "Dictionary features should be L2-normalized"

    @pytest.mark.skipif(
        not DICT_PT_SIGNBANK.exists(),
        reason=f"Signbank dictionary not found: {DICT_PT_SIGNBANK}",
    )
    def test_dictionary_on_correct_device(self, pipeline):
        dictionary = pipeline.load_dictionary([str(DICT_PT_SIGNBANK)])
        assert (
            dictionary.features.device.type == "cuda"
        ), f"Expected features on cuda, got {dictionary.features.device}"


# ---------------------------------------------------------------------------
# Spotting tests
# ---------------------------------------------------------------------------


class TestSpotting:
    """Test end-to-end spotting with real model and dictionary."""

    @pytest.mark.skipif(
        not DICT_PT_SIGNBANK.exists(),
        reason=f"Signbank dictionary not found: {DICT_PT_SIGNBANK}",
    )
    def test_spot_with_inline_segments(self, pipeline):
        continuous = pipeline.extract_continuous(str(EXAMPLE_VIDEO), stride=4)
        dictionary = pipeline.load_dictionary([str(DICT_PT_SIGNBANK)])

        half = continuous.total_frames // 2
        segments = [
            {
                "segment_id": 1,
                "start_frame": 0,
                "end_frame": half,
                "original_video_id": "test_video",
            },
            {
                "segment_id": 2,
                "start_frame": half,
                "end_frame": continuous.total_frames,
                "original_video_id": "test_video",
            },
        ]

        result = pipeline.spot(continuous, segments, dictionary, top_k=5, segment_pooling="max")
        assert len(result.segments) == 2

        for seg in result.segments:
            assert len(seg.top_glosses) <= 5
            if seg.top_glosses:
                assert seg.top_glosses[0]["rank"] == 1
                sim = seg.top_glosses[0]["similarity"]
                assert 0.0 <= sim <= 1.0, f"Similarity {sim} out of [0, 1] range"

    @pytest.mark.skipif(
        not DICT_PT_SIGNBANK.exists(),
        reason=f"Signbank dictionary not found: {DICT_PT_SIGNBANK}",
    )
    def test_spot_all_segment_pooling_methods(self, pipeline):
        from sltk.embedding.pipeline import SEGMENT_POOLING

        continuous = pipeline.extract_continuous(str(EXAMPLE_VIDEO), stride=4)
        dictionary = pipeline.load_dictionary([str(DICT_PT_SIGNBANK)])

        segments = [
            {
                "segment_id": 1,
                "start_frame": 0,
                "end_frame": continuous.total_frames,
                "original_video_id": "test_video",
            },
        ]

        for method in SEGMENT_POOLING:
            result = pipeline.spot(
                continuous, segments, dictionary, top_k=3, segment_pooling=method
            )
            assert len(result.segments) == 1, f"Segment pooling '{method}' failed"
            assert len(result.segments[0].top_glosses) <= 3

    @pytest.mark.skipif(
        not DICT_PT_SIGNBANK.exists(),
        reason=f"Signbank dictionary not found: {DICT_PT_SIGNBANK}",
    )
    def test_spot_save_json(self, pipeline, tmp_path):
        continuous = pipeline.extract_continuous(str(EXAMPLE_VIDEO), stride=4)
        dictionary = pipeline.load_dictionary([str(DICT_PT_SIGNBANK)])

        segments = [
            {
                "segment_id": 1,
                "start_frame": 0,
                "end_frame": continuous.total_frames,
                "original_video_id": "test_video",
            },
        ]

        result = pipeline.spot(continuous, segments, dictionary, top_k=10)
        out = tmp_path / "spot_results.json"
        result.save_json(str(out))

        with open(out) as f:
            loaded = json.load(f)

        assert "metadata" in loaded
        assert "segments" in loaded
        assert loaded["metadata"]["top_k"] == 10

    @pytest.mark.skipif(
        not DICT_PT_SIGNBANK.exists(),
        reason=f"Signbank dictionary not found: {DICT_PT_SIGNBANK}",
    )
    def test_spot_from_npz(self, pipeline, tmp_path):
        """Spot using pre-saved continuous features from .npz file."""
        continuous = pipeline.extract_continuous(str(EXAMPLE_VIDEO), stride=4)
        npz_path = tmp_path / "continuous.npz"
        continuous.save_npz(npz_path)

        dictionary = pipeline.load_dictionary([str(DICT_PT_SIGNBANK)])

        segments = [
            {
                "segment_id": 1,
                "start_frame": 0,
                "end_frame": continuous.total_frames,
                "original_video_id": "test_video",
            },
        ]

        result = pipeline.spot(str(npz_path), segments, dictionary, top_k=5)
        assert len(result.segments) == 1
        assert result.segments[0].top_glosses[0]["rank"] == 1


# ---------------------------------------------------------------------------
# inference_forward tests
# ---------------------------------------------------------------------------


class TestInferenceForward:
    """Test the fast inference path on the model directly."""

    def test_inference_forward_returns_latent(self, pipeline):
        dummy = torch.randn(1, 3, 16, 224, 224).to(pipeline.device)
        with torch.inference_mode():
            output = pipeline.model.inference_forward(dummy, need_active=True)

        assert "latent" in output
        assert "cls_active" in output
        assert output["latent"].shape == (
            1,
            768,
        ), f"Expected latent shape (1, 768), got {output['latent'].shape}"
        assert output["cls_active"].shape == (
            1,
            2,
        ), f"Expected cls_active shape (1, 2), got {output['cls_active'].shape}"

    def test_inference_forward_without_active(self, pipeline):
        dummy = torch.randn(1, 3, 16, 224, 224).to(pipeline.device)
        with torch.inference_mode():
            output = pipeline.model.inference_forward(dummy, need_active=False)

        assert "latent" in output
        assert "cls_active" not in output

    def test_inference_forward_batch(self, pipeline):
        dummy = torch.randn(4, 3, 16, 224, 224).to(pipeline.device)
        with torch.inference_mode():
            output = pipeline.model.inference_forward(dummy, need_active=True)

        assert output["latent"].shape == (4, 768)
        assert output["cls_active"].shape == (4, 2)

    def test_inference_forward_deterministic(self, pipeline):
        """Inference should be deterministic (model in eval mode)."""
        dummy = torch.randn(1, 3, 16, 224, 224).to(pipeline.device)
        with torch.inference_mode():
            out1 = pipeline.model.inference_forward(dummy, need_active=True)
            out2 = pipeline.model.inference_forward(dummy, need_active=True)

        assert torch.allclose(
            out1["latent"], out2["latent"], atol=1e-6
        ), "inference_forward should produce deterministic results"

    def test_inference_forward_with_amp(self, pipeline):
        """Verify AMP does not corrupt outputs."""
        dummy = torch.randn(1, 3, 16, 224, 224).to(pipeline.device)
        with torch.inference_mode(), torch.amp.autocast("cuda"):
            output = pipeline.model.inference_forward(dummy, need_active=True)

        assert torch.isfinite(output["latent"]).all(), "AMP produced non-finite latent values"
        assert torch.isfinite(
            output["cls_active"]
        ).all(), "AMP produced non-finite cls_active values"


# ---------------------------------------------------------------------------
# Backward compatibility: SignRepEmbedder with real model
# ---------------------------------------------------------------------------


class TestSignRepEmbedderIntegration:
    """Test backward-compatible embedder with real model."""

    @pytest.fixture(scope="class")
    def embedder(self):
        from sltk.embedding.signrep import SignRepConfig, SignRepEmbedder

        config = SignRepConfig(
            checkpoint_path=str(CHECKPOINT),
            device="cuda",
            batch_size=4,
            num_workers=0,
        )
        return SignRepEmbedder(config=config, lazy_load=True)

    def test_embed_video_returns_768_vector(self, embedder):
        result = embedder.embed_video(str(EXAMPLE_VIDEO), aggregation="mean")
        assert isinstance(result, np.ndarray)
        assert result.shape == (768,)
        assert np.all(np.isfinite(result))

    def test_embed_video_is_l2_normalized(self, embedder):
        result = embedder.embed_video(str(EXAMPLE_VIDEO), aggregation="mean")
        norm = np.linalg.norm(result)
        assert abs(norm - 1.0) < 1e-4, f"Expected norm ~1.0, got {norm}"

    def test_embed_video_all_aggregations(self, embedder):
        for agg in ["mean", "max", "first", "last"]:
            result = embedder.embed_video(str(EXAMPLE_VIDEO), aggregation=agg)
            assert result.shape == (768,), f"Aggregation '{agg}' produced shape {result.shape}"

    def test_embed_video_return_all_segments(self, embedder):
        result = embedder.embed_video(
            str(EXAMPLE_VIDEO), aggregation="mean", return_all_segments=True
        )
        assert isinstance(result, dict)
        assert "embedding" in result
        assert "segments" in result
        assert "latents" in result
        assert "num_segments" in result
        assert result["embedding"].shape == (768,)
        assert result["segments"].shape[1] == 768
        assert result["num_segments"] > 0

    def test_pipeline_property(self, embedder):
        from sltk.embedding.pipeline import SignRepPipeline

        p = embedder.pipeline
        assert isinstance(p, SignRepPipeline)

    def test_embed_videos_batch(self, embedder):
        # Note: dict keyed by str(path) -- same path yields one entry
        results = embedder.embed_videos(
            [str(EXAMPLE_VIDEO)],
            show_progress=False,
        )
        assert len(results) == 1
        for path, emb in results.items():
            assert emb is not None
            assert emb.shape == (768,)


# ---------------------------------------------------------------------------
# End-to-end: extract -> spot pipeline
# ---------------------------------------------------------------------------


class TestEndToEnd:
    """Full pipeline: extract continuous -> build dictionary -> spot."""

    @pytest.mark.skipif(
        not DICT_PT_SIGNBANK.exists(),
        reason=f"Signbank dictionary not found: {DICT_PT_SIGNBANK}",
    )
    def test_spot_from_video(self, pipeline, tmp_path):
        """End-to-end spot_from_video with a segments JSON file."""
        # Create a segments file covering the whole video
        # First extract to know total frames
        cont = pipeline.extract_continuous(str(EXAMPLE_VIDEO), stride=4)

        segments_data = [
            {
                "segment_id": 1,
                "original_video_id": EXAMPLE_VIDEO.stem,
                "start_frame": 0,
                "end_frame": cont.total_frames // 2,
            },
            {
                "segment_id": 2,
                "original_video_id": EXAMPLE_VIDEO.stem,
                "start_frame": cont.total_frames // 2,
                "end_frame": cont.total_frames,
            },
        ]
        seg_path = tmp_path / "segments.json"
        with open(seg_path, "w") as f:
            json.dump(segments_data, f)

        result = pipeline.spot_from_video(
            str(EXAMPLE_VIDEO),
            str(seg_path),
            [str(DICT_PT_SIGNBANK)],
            top_k=10,
            stride=4,
        )
        assert len(result.segments) == 2
        for seg in result.segments:
            assert seg.num_windows > 0
            assert len(seg.top_glosses) <= 10

    def test_extract_and_spot_consistency(self, pipeline, tmp_path):
        """Verify that spot() with ContinuousResult gives same results
        as spot() with saved .npz file."""
        cont = pipeline.extract_continuous(str(EXAMPLE_VIDEO), stride=4)

        # Build a small synthetic dictionary (no external files needed)
        from sltk.embedding.pipeline import DictionaryIndex

        rng = np.random.default_rng(42)
        feats = torch.from_numpy(rng.standard_normal((20, 768)).astype(np.float32))
        feats = torch.nn.functional.normalize(feats, dim=1).to(pipeline.device)
        dictionary = DictionaryIndex(
            features=feats,
            glosses=[f"G{i}" for i in range(20)],
            filenames=[f"f{i}" for i in range(20)],
            source_dirs=["synth"] * 20,
            file_paths=[f"synth/f{i}.npz" for i in range(20)],
        )

        segments = [
            {
                "segment_id": 1,
                "start_frame": 0,
                "end_frame": cont.total_frames,
                "original_video_id": "test",
            },
        ]

        # Spot from ContinuousResult
        result_direct = pipeline.spot(cont, segments, dictionary, top_k=5, segment_pooling="max")

        # Spot from saved .npz
        npz_path = tmp_path / "continuous.npz"
        cont.save_npz(npz_path)
        result_npz = pipeline.spot(
            str(npz_path), segments, dictionary, top_k=5, segment_pooling="max"
        )

        # Results should match (same features, same computation)
        assert len(result_direct.segments) == len(result_npz.segments)
        for sd, sn in zip(result_direct.segments, result_npz.segments):
            assert sd.num_windows == sn.num_windows
            assert len(sd.top_glosses) == len(sn.top_glosses)
            for gd, gn in zip(sd.top_glosses, sn.top_glosses):
                assert gd["gloss"] == gn["gloss"]
                assert abs(gd["similarity"] - gn["similarity"]) < 1e-3
