"""
Tests for RS3 on-target scoring module.

Tests both the legacy one-hot encoding functions (for backward compatibility)
and the new rs3 library integration.
"""

import os
import tempfile
import pickle
import pytest
import numpy as np
import pandas as pd

from crisprdesigner2._rs3_scoring import encode_guide_rs3, load_rs3_model


# Check if rs3 library is available
try:
    from rs3.seq import predict_seq
    RS3_AVAILABLE = True
except ImportError:
    RS3_AVAILABLE = False


class TestEncodeGuideRS3:
    """Tests for the legacy encode_guide_rs3 function."""

    def test_correct_shape(self):
        """A valid 20nt protospacer should return an array of shape (80,)."""
        protospacer = 'ACGTACGATCGATCGACGAA'
        features = encode_guide_rs3(protospacer)
        assert features.shape == (80,)

    def test_one_hot_encoding(self):
        """Each position should have exactly one 1 in its 4-element block."""
        protospacer = 'ACGTACGATCGATCGACGAA'
        features = encode_guide_rs3(protospacer)

        for i in range(20):
            block = features[i*4:(i+1)*4]
            assert np.sum(block) == 1.0, f"Position {i} should have exactly one 1"
            assert np.max(block) == 1.0
            assert np.min(block) == 0.0

    def test_specific_encoding(self):
        """Test specific nucleotide encodings."""
        # Test sequence where we know the exact encoding
        protospacer = 'ACGT' + 'A' * 16  # A=0, C=1, G=2, T=3

        features = encode_guide_rs3(protospacer)

        # Position 0: A -> index 0
        assert features[0] == 1.0
        assert features[1] == 0.0
        assert features[2] == 0.0
        assert features[3] == 0.0

        # Position 1: C -> index 5 (1*4 + 1)
        assert features[4] == 0.0
        assert features[5] == 1.0
        assert features[6] == 0.0
        assert features[7] == 0.0

        # Position 2: G -> index 10 (2*4 + 2)
        assert features[8] == 0.0
        assert features[9] == 0.0
        assert features[10] == 1.0
        assert features[11] == 0.0

        # Position 3: T -> index 15 (3*4 + 3)
        assert features[12] == 0.0
        assert features[13] == 0.0
        assert features[14] == 0.0
        assert features[15] == 1.0

    def test_short_sequence_raises_error(self):
        """A protospacer shorter than 20nt should raise ValueError."""
        with pytest.raises(ValueError) as excinfo:
            encode_guide_rs3('ACGTACGTACGTACGTACG')  # 19 nt
        assert "20 nucleotides" in str(excinfo.value)

    def test_long_sequence_raises_error(self):
        """A protospacer longer than 20nt should raise ValueError."""
        with pytest.raises(ValueError) as excinfo:
            encode_guide_rs3('ACGTACGTACGTACGTACGTA')  # 21 nt
        assert "20 nucleotides" in str(excinfo.value)

    def test_invalid_nucleotide_raises_error(self):
        """A protospacer with invalid nucleotides should raise ValueError."""
        with pytest.raises(ValueError) as excinfo:
            encode_guide_rs3('ACGTACGTACGTACGTACGN')  # Contains N
        assert "Invalid nucleotide" in str(excinfo.value)

    def test_lowercase_works(self):
        """Lowercase nucleotides should be handled correctly."""
        upper = encode_guide_rs3('ACGTACGATCGATCGACGAA')
        lower = encode_guide_rs3('acgtacgatcgatcgacgaa')
        np.testing.assert_array_equal(upper, lower)

    def test_float32_dtype(self):
        """Features should be float32 type."""
        features = encode_guide_rs3('ACGTACGATCGATCGACGAA')
        assert features.dtype == np.float32


class TestLoadRS3Model:
    """Tests for the legacy load_rs3_model function."""

    def test_missing_file_raises_error(self):
        """A missing model file should raise FileNotFoundError."""
        with pytest.raises(FileNotFoundError) as excinfo:
            load_rs3_model('/nonexistent/path/to/model.pkl')
        assert "not found" in str(excinfo.value).lower()

    def test_load_valid_model(self):
        """Loading a valid pickle file should work."""
        from sklearn.linear_model import LogisticRegression

        # Create a simple trained model that can be pickled
        X = np.random.rand(10, 80)
        y = np.random.randint(0, 2, 10)
        mock_model = LogisticRegression(max_iter=100)
        mock_model.fit(X, y)

        with tempfile.NamedTemporaryFile(mode='wb', suffix='.pkl', delete=False) as f:
            pickle.dump(mock_model, f)
            model_path = f.name

        try:
            model = load_rs3_model(model_path)
            assert hasattr(model, 'predict')
            assert hasattr(model, 'predict_proba')
        finally:
            os.unlink(model_path)


@pytest.mark.skipif(not RS3_AVAILABLE, reason="rs3 library not installed")
class TestScoreGuidesRS3WithLibrary:
    """Integration tests for score_guides_rs3 using the official rs3 library."""

    def test_adds_rs3_score_column(self):
        """score_guides_rs3 should add RS3_score column to DataFrame."""
        from crisprdesigner2 import score_guides_rs3

        # 30nt context sequences (4bp upstream + 20bp protospacer + 3bp PAM + 3bp downstream)
        df = pd.DataFrame({
            'context_sequence': [
                'AAAACGTACGATCGATCGACGAACAGGCCC',  # 30nt: AAAA + 20bp + AGG + CCC
                'TTTTCGATCGATCGATCGATCGACTGGAAA',  # 30nt: TTTT + 20bp + TGG + AAA
            ]
        })

        result = score_guides_rs3(df)

        assert 'RS3_score' in result.columns
        assert len(result) == 2

    def test_preserves_other_columns(self):
        """score_guides_rs3 should preserve other columns in DataFrame."""
        from crisprdesigner2 import score_guides_rs3

        df = pd.DataFrame({
            'context_sequence': ['AAAACGTACGATCGATCGACGAACAGGCCC'],  # 30nt
            'chr': ['chr1'],
            'start': [100],
            'end': [120]
        })

        result = score_guides_rs3(df)

        assert result['chr'].iloc[0] == 'chr1'
        assert result['start'].iloc[0] == 100
        assert result['end'].iloc[0] == 120

    def test_missing_context_sequence_raises_error(self):
        """Missing context_sequence column should raise ValueError."""
        from crisprdesigner2 import score_guides_rs3

        df = pd.DataFrame({
            'protospacer': ['ACGTACGATCGATCGACGAA']  # Old format, missing context
        })

        with pytest.raises(ValueError) as excinfo:
            score_guides_rs3(df)
        assert "context_sequence" in str(excinfo.value)

    def test_invalid_context_length_raises_error(self):
        """Invalid context sequence length should raise ValueError."""
        from crisprdesigner2 import score_guides_rs3

        df = pd.DataFrame({
            'context_sequence': ['ACGTACGTACGT']  # Only 12 nt, should be 30
        })

        with pytest.raises(ValueError) as excinfo:
            score_guides_rs3(df)
        assert "30nt" in str(excinfo.value) or "30" in str(excinfo.value)


@pytest.mark.skipif(not RS3_AVAILABLE, reason="rs3 library not installed")
class TestScoreSingleGuideRS3WithLibrary:
    """Tests for score_single_guide_rs3 using the official rs3 library."""

    def test_single_guide_score(self):
        """Scoring a single guide should return a float."""
        from crisprdesigner2 import score_single_guide_rs3

        # 30nt context sequence (4bp upstream + 20bp protospacer + 3bp PAM + 3bp downstream)
        context = 'AAAACGTACGATCGATCGACGAACAGGCCC'
        score = score_single_guide_rs3(context)

        assert isinstance(score, float)

    def test_invalid_length_raises_error(self):
        """Invalid context length should raise ValueError."""
        from crisprdesigner2 import score_single_guide_rs3

        with pytest.raises(ValueError) as excinfo:
            score_single_guide_rs3('ACGTACGATCGATCGACGAA')  # Only 20nt
        assert "30 nucleotides" in str(excinfo.value)


class TestScoreGuidesRS3MockFallback:
    """Tests using mock scoring when rs3 library is not available."""

    def test_context_sequence_validation(self):
        """Test that context_sequence column is required."""
        from crisprdesigner2 import score_guides_rs3

        df = pd.DataFrame({
            'protospacer': ['ACGTACGATCGATCGACGAA']
        })

        # Should raise error about missing context_sequence
        with pytest.raises((ValueError, ImportError)):
            score_guides_rs3(df)


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