"""Tests for reproducibility of InstaNovo predictions.

These tests verify that predictions are consistent between runs, addressing
the variability issue reported in: https://github.com/instadeepai/InstaNovo/issues/92

The key issue was that running InstaNovo locally vs on HuggingFace Spaces
produced different results. These tests ensure:
1. Predictions are deterministic with fixed seeds
2. Multiple runs produce identical results
3. The sample spectra produce consistent outputs
"""

import pytest
import os
import sys
import json
from pathlib import Path

# Add parent directory to path for imports
sys.path.insert(0, str(Path(__file__).parent.parent))


@pytest.fixture(scope="module")
def load_app_components():
    """Load the app components once for all tests."""
    # Import as a module so we read live globals (populated by load_models_and_knapsack)
    # rather than capturing None values that existed at import time.
    import app

    app.set_seed(42)
    app.load_models_and_knapsack()

    return {
        "set_seed": app.set_seed,
        "create_inference_config": app.create_inference_config,
        "run_transformer_prediction": app.run_transformer_prediction,
        "run_diffusion_prediction": app.run_diffusion_prediction,
        "INSTANOVO": app.INSTANOVO,
        "INSTANOVOPLUS": app.INSTANOVOPLUS,
        "RESIDUE_SET": app.RESIDUE_SET,
        "DEVICE": app.DEVICE,
    }


@pytest.fixture
def sample_spectra_path():
    """Path to sample spectra file."""
    return Path(__file__).parent.parent / "assets" / "sample_spectra.mgf"


@pytest.fixture
def expected_results_path():
    """Path to expected results file for comparison."""
    return Path(__file__).parent / "expected_results.json"


class TestSeedSetting:
    """Test that random seed setting works correctly."""
    
    def test_seed_produces_consistent_random_numbers(self):
        """Verify that setting seed produces the same random sequence."""
        import random
        import numpy as np
        import torch
        from app import set_seed
        
        # First run
        set_seed(42)
        rand1 = random.random()
        np1 = np.random.random()
        torch1 = torch.rand(1).item()
        
        # Second run with same seed
        set_seed(42)
        rand2 = random.random()
        np2 = np.random.random()
        torch2 = torch.rand(1).item()
        
        assert rand1 == rand2, "Python random not reproducible"
        assert np1 == np2, "NumPy random not reproducible"
        assert torch1 == torch2, "PyTorch random not reproducible"
    
    def test_different_seeds_produce_different_results(self):
        """Verify that different seeds produce different results."""
        import random
        import numpy as np
        import torch
        from app import set_seed
        
        set_seed(42)
        results_42 = (random.random(), np.random.random(), torch.rand(1).item())
        
        set_seed(123)
        results_123 = (random.random(), np.random.random(), torch.rand(1).item())
        
        assert results_42 != results_123, "Different seeds should produce different results"


class TestPredictionReproducibility:
    """Test that predictions are reproducible."""
    
    def test_transformer_prediction_reproducibility(
        self, load_app_components, sample_spectra_path, tmp_path
    ):
        """Test that transformer predictions are reproducible across multiple runs."""
        if not sample_spectra_path.exists():
            pytest.skip(f"Sample spectra not found at {sample_spectra_path}")
        
        components = load_app_components
        if components["INSTANOVO"] is None:
            pytest.skip("InstaNovo model not loaded")
        
        from instanovo.utils.data_handler import SpectrumDataFrame
        from instanovo.transformer.data import TransformerDataProcessor
        from torch.utils.data import DataLoader

        set_seed = components["set_seed"]
        create_inference_config = components["create_inference_config"]
        run_transformer_prediction = components["run_transformer_prediction"]

        # Create config
        output_path = tmp_path / "output.csv"
        config = create_inference_config(str(sample_spectra_path), str(output_path))

        # Load data
        sdf = SpectrumDataFrame.load(
            str(sample_spectra_path),
            lazy=False,
            is_annotated=False,
            shuffle=False,
        )

        import app
        RESIDUE_SET = app.RESIDUE_SET
        INSTANOVO_CONFIG = app.INSTANOVO_CONFIG

        data_processor = TransformerDataProcessor(
            residue_set=RESIDUE_SET,
            n_peaks=INSTANOVO_CONFIG.get("n_peaks", 200),
            annotated=False,
            reverse_peptide=True,
            add_eos=True,
        )
        dataset = sdf.to_dataset()

        def _make_dl():
            processed = data_processor.process_dataset(dataset)
            return DataLoader(
                processed,
                batch_size=config.batch_size,
                shuffle=False,
                collate_fn=data_processor.collate_fn,
            )

        # First run
        set_seed(42)
        results1 = run_transformer_prediction(_make_dl(), config, "Greedy Search (Fast)")
        predictions1 = ["".join(r.sequence) if r.sequence else "" for r in results1]
        log_probs1 = [r.sequence_log_probability for r in results1]
        
        # Second run
        set_seed(42)
        results2 = run_transformer_prediction(_make_dl(), config, "Greedy Search (Fast)")
        predictions2 = ["".join(r.sequence) if r.sequence else "" for r in results2]
        log_probs2 = [r.sequence_log_probability for r in results2]
        
        # Verify identical results
        assert predictions1 == predictions2, (
            f"Predictions differ between runs:\n"
            f"Run 1: {predictions1[:5]}\n"
            f"Run 2: {predictions2[:5]}"
        )
        
        assert log_probs1 == log_probs2, "Log probabilities differ between runs"


class TestExpectedResults:
    """Test that predictions match expected baseline results.
    
    These tests compare against known-good results to catch any regressions
    or environment-specific differences.
    """
    
    def test_predictions_match_expected_baseline(
        self, load_app_components, sample_spectra_path, expected_results_path, tmp_path
    ):
        """Test that predictions match the expected baseline results."""
        if not sample_spectra_path.exists():
            pytest.skip(f"Sample spectra not found at {sample_spectra_path}")
        
        if not expected_results_path.exists():
            pytest.skip(
                f"Expected results file not found at {expected_results_path}. "
                "Run 'pytest tests/test_reproducibility.py::TestExpectedResults::test_generate_expected_results' first."
            )
        
        components = load_app_components
        if components["INSTANOVO"] is None:
            pytest.skip("InstaNovo model not loaded")
        
        from instanovo.utils.data_handler import SpectrumDataFrame
        from instanovo.transformer.data import TransformerDataProcessor
        from torch.utils.data import DataLoader
        
        set_seed = components["set_seed"]
        create_inference_config = components["create_inference_config"]
        run_transformer_prediction = components["run_transformer_prediction"]
        
        # Load expected results
        with open(expected_results_path) as f:
            expected = json.load(f)
        
        # Create config
        output_path = tmp_path / "output.csv"
        config = create_inference_config(str(sample_spectra_path), str(output_path))
        
        # Load data
        sdf = SpectrumDataFrame.load(
            str(sample_spectra_path), 
            lazy=False, 
            is_annotated=False,
            shuffle=False
        )
        
        import app
        RESIDUE_SET = app.RESIDUE_SET
        INSTANOVO_CONFIG = app.INSTANOVO_CONFIG
        
        # Use new v1.2.2 API
        data_processor = TransformerDataProcessor(
            residue_set=RESIDUE_SET,
            n_peaks=INSTANOVO_CONFIG.get("n_peaks", 200),
            annotated=False,
            reverse_peptide=True,
            add_eos=True,
        )
        dataset = sdf.to_dataset()
        processed_dataset = data_processor.process_dataset(dataset)
        dl = DataLoader(
            processed_dataset, 
            batch_size=config.batch_size, 
            shuffle=False, 
            collate_fn=data_processor.collate_fn
        )
        
        # Run prediction
        set_seed(42)
        results = run_transformer_prediction(dl, config, "Greedy Search (Fast)")
        predictions = ["".join(r.sequence) if r.sequence else "" for r in results]
        
        # Compare against expected
        expected_predictions = expected.get("greedy_predictions", [])
        
        assert predictions == expected_predictions, (
            f"Predictions differ from expected baseline:\n"
            f"Got: {predictions[:5]}\n"
            f"Expected: {expected_predictions[:5]}\n\n"
            "This may indicate a reproducibility issue. "
            "If the new results are correct, regenerate expected_results.json"
        )
    
    def test_generate_expected_results(
        self, load_app_components, sample_spectra_path, tmp_path
    ):
        """Generate expected results file for baseline comparison.
        
        This is not a real test - it's a utility to generate the expected results.
        Run with: pytest tests/test_reproducibility.py::TestExpectedResults::test_generate_expected_results -s
        """
        if not sample_spectra_path.exists():
            pytest.skip(f"Sample spectra not found at {sample_spectra_path}")
        
        components = load_app_components
        if components["INSTANOVO"] is None:
            pytest.skip("InstaNovo model not loaded")
        
        from instanovo.utils.data_handler import SpectrumDataFrame
        from instanovo.transformer.data import TransformerDataProcessor
        from torch.utils.data import DataLoader
        
        set_seed = components["set_seed"]
        create_inference_config = components["create_inference_config"]
        run_transformer_prediction = components["run_transformer_prediction"]
        
        # Create config
        output_path = tmp_path / "output.csv"
        config = create_inference_config(str(sample_spectra_path), str(output_path))
        
        # Load data
        sdf = SpectrumDataFrame.load(
            str(sample_spectra_path), 
            lazy=False, 
            is_annotated=False,
            shuffle=False
        )
        
        import app
        RESIDUE_SET = app.RESIDUE_SET
        INSTANOVO_CONFIG = app.INSTANOVO_CONFIG
        
        # Use new v1.2.2 API
        data_processor = TransformerDataProcessor(
            residue_set=RESIDUE_SET,
            n_peaks=INSTANOVO_CONFIG.get("n_peaks", 200),
            annotated=False,
            reverse_peptide=True,
            add_eos=True,
        )
        dataset = sdf.to_dataset()
        processed_dataset = data_processor.process_dataset(dataset)
        dl = DataLoader(
            processed_dataset, 
            batch_size=config.batch_size, 
            shuffle=False, 
            collate_fn=data_processor.collate_fn
        )
        
        # Run prediction with greedy
        set_seed(42)
        results = run_transformer_prediction(dl, config, "Greedy Search (Fast)")
        greedy_predictions = ["".join(r.sequence) if r.sequence else "" for r in results]
        greedy_log_probs = [r.sequence_log_probability for r in results]
        
        # Save expected results
        expected_results_path = Path(__file__).parent / "expected_results.json"
        expected = {
            "version": "1.2.2",
            "model": "instanovo-v1.2.0",
            "seed": 42,
            "greedy_predictions": greedy_predictions,
            "greedy_log_probs": greedy_log_probs,
        }
        
        with open(expected_results_path, "w") as f:
            json.dump(expected, f, indent=2)
        
        print(f"\nExpected results saved to {expected_results_path}")
        print(f"Predictions: {greedy_predictions}")


if __name__ == "__main__":
    pytest.main([__file__, "-v"])