"""
End-to-end integration tests for CRISPR gRNA design pipeline.

Tests the full workflow from guide extraction to output writing.
"""

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

from crisprdesigner2 import (
    extract_guides,
    apply_filters,
    get_passing_guides,
    score_guides_rs3,
    compute_specificity_mock,
    write_outputs,
    read_design_guides_txt,
    read_design_guides_bed,
    load_predesigned_guides,
    filter_predesigned_regions,
    merge_with_predesigned,
)


class TestIntegrationWorkflow:
    """End-to-end integration tests for the full workflow."""

    @pytest.fixture
    def test_data_dir(self):
        """Return the path to test data directory."""
        return os.path.join(os.path.dirname(__file__), 'data')

    @pytest.fixture
    def temp_output_dir(self):
        """Create a temporary output directory."""
        tmpdir = tempfile.mkdtemp()
        yield tmpdir
        shutil.rmtree(tmpdir)

    @pytest.fixture
    def test_fasta(self, temp_output_dir):
        """Create a test FASTA file with known sequences."""
        # Create sequences with:
        # - Valid NGG PAM sites with good guides
        # - Some sequences that will fail filters
        # - Reverse strand PAM sites (CCN)

        # Region 1 (1000-1500): Contains good guides
        region1 = (
            "AAAAA" +                         # Padding
            "ACGTACGATCGATCGACGAA" + "AGG" +  # Good forward guide
            "CGATCGATCGATCGATCGAC" + "TGG" +  # Another forward guide
            "A" * 50 +                         # Padding
            "CC" + "A" + "ACGTACGATCGATCGACGAA" +  # Reverse strand guide (CCA = TGG on reverse)
            "A" * (500 - 23*3 - 5 - 50 - 23)  # Remaining padding
        )

        # Region 2 (5000-5500): Contains guides that will fail filters
        region2 = (
            "AAAAA" +
            "ATATATATATATATATACGA" + "AGG" +  # Will fail low_complexity
            "TTTTTTTTTTTTTTTTTTTT" + "TGG" +  # Will fail multiple filters
            "GCGCGCGCGCGCGCGCGCGC" + "CGG" +  # Will fail high_gc and low_complexity
            "ACGTACGATCGATCGACGAA" + "AGG" +  # Good guide
            "A" * (500 - 23*4 - 5)            # Padding
        )

        # Region 3 (10000-10500): More good guides
        region3 = (
            "AAAAA" +
            "TACGAACGTTACGTAACGAC" + "AGG" +  # Good forward guide
            "ACGATCAGCTACGATCAGCA" + "TGG" +  # Another good guide
            "CCA" + "ACGATCGACGATCGACGACA" +  # Reverse strand guide
            "A" * (500 - 23*3 - 5)            # Padding
        )

        # Pad to create full chromosome
        padding = "N" * 1000
        sequence = (
            padding +           # 0-1000
            region1 +           # 1000-1500
            "N" * 3500 +        # 1500-5000
            region2 +           # 5000-5500
            "N" * 4500 +        # 5500-10000
            region3 +           # 10000-10500
            "N" * 1000          # 10500-11500
        )

        fasta_path = os.path.join(temp_output_dir, "test_genome.fa")
        with open(fasta_path, 'w') as f:
            f.write(">chr1\n")
            # Write in 100-char lines
            for i in range(0, len(sequence), 100):
                f.write(sequence[i:i+100] + "\n")

        return fasta_path

    @pytest.fixture
    def test_regions(self, temp_output_dir):
        """Create a test BED file with regions."""
        bed_path = os.path.join(temp_output_dir, "test_regions.bed")
        with open(bed_path, 'w') as f:
            f.write("chr1\t1000\t1500\tregion1\n")
            f.write("chr1\t5000\t5500\tregion2\n")
            f.write("chr1\t10000\t10500\tregion3\n")
        return bed_path

    @pytest.fixture
    def mock_rs3_model(self, temp_output_dir):
        """Create a mock RS3 model for testing."""
        from sklearn.linear_model import LogisticRegression

        # Create a simple trained model
        X = np.random.rand(100, 80)
        y = np.random.randint(0, 2, 100)

        model = LogisticRegression(max_iter=100, random_state=42)
        model.fit(X, y)

        model_path = os.path.join(temp_output_dir, "mock_rs3_model.pkl")
        with open(model_path, 'wb') as f:
            pickle.dump(model, f)

        return model_path

    def test_guide_extraction_from_regions(self, test_fasta, test_regions):
        """Test that guides are correctly extracted from regions."""
        # Extract guides from region1
        guides = extract_guides('chr1', 1000, 1500, 'region1', test_fasta)

        # Should find at least some guides
        assert len(guides) > 0

        # All guides should have required fields
        for guide in guides:
            assert 'chr' in guide
            assert 'start' in guide
            assert 'end' in guide
            assert 'protospacer' in guide
            assert 'GuideSequenceWithPAM' in guide
            assert 'strand' in guide
            assert 'guideSet' in guide

            # Protospacer should be 20 nt
            assert len(guide['protospacer']) == 20

            # GuideSequenceWithPAM should be 23 nt
            assert len(guide['GuideSequenceWithPAM']) == 23

    def test_filter_removes_bad_guides(self, test_fasta):
        """Test that filters correctly remove bad guides."""
        # Extract guides from region2 (contains bad guides)
        guides = extract_guides('chr1', 5000, 5500, 'region2', test_fasta)

        if len(guides) == 0:
            pytest.skip("No guides found in test region")

        guides_df = pd.DataFrame(guides)

        # Apply filters
        filtered_df = apply_filters(guides_df)

        # Some guides should fail filters
        failed_count = (filtered_df['filter_reason'] != '').sum()
        passed_count = (filtered_df['filter_reason'] == '').sum()

        # Not all guides should pass (region2 has bad sequences)
        # This might not always be true depending on exact sequence
        assert 'filter_reason' in filtered_df.columns

    def test_rs3_scoring_integration(self, test_fasta, mock_rs3_model):
        """Test RS3 scoring on extracted guides."""
        # Extract and filter guides
        guides = extract_guides('chr1', 1000, 1500, 'region1', test_fasta)

        if len(guides) == 0:
            pytest.skip("No guides found in test region")

        guides_df = pd.DataFrame(guides)
        passing_df = get_passing_guides(guides_df)

        if len(passing_df) == 0:
            pytest.skip("No guides passed filters")

        # Score with RS3
        scored_df = score_guides_rs3(passing_df, mock_rs3_model)

        # Should have RS3_score column
        assert 'RS3_score' in scored_df.columns

        # Scores should be in valid range
        assert all(scored_df['RS3_score'] >= 0)
        assert all(scored_df['RS3_score'] <= 1)

    def test_full_pipeline_integration(self, test_fasta, mock_rs3_model, temp_output_dir):
        """Test the full pipeline from extraction to output."""
        # Step 1: Extract guides from all regions
        all_guides = []
        for region in [('chr1', 1000, 1500, 'region1'),
                       ('chr1', 5000, 5500, 'region2'),
                       ('chr1', 10000, 10500, 'region3')]:
            guides = extract_guides(region[0], region[1], region[2], region[3], test_fasta)
            all_guides.extend(guides)

        if len(all_guides) == 0:
            pytest.skip("No guides found in test regions")

        guides_df = pd.DataFrame(all_guides)

        # Step 2: Apply filters
        filtered_df = apply_filters(guides_df)
        passing_df = get_passing_guides(guides_df)

        if len(passing_df) == 0:
            pytest.skip("No guides passed filters")

        # Step 3: RS3 scoring
        scored_df = score_guides_rs3(passing_df, mock_rs3_model)

        # Step 4: Mock specificity scoring
        scored_df = compute_specificity_mock(scored_df)

        # Step 5: Write outputs
        write_outputs(scored_df, temp_output_dir)

        # Verify outputs exist
        txt_file = os.path.join(temp_output_dir, 'designGuides.txt')
        bed_file = os.path.join(temp_output_dir, 'designGuides.bed')

        assert os.path.exists(txt_file)
        assert os.path.exists(bed_file)

        # Verify output content
        result_txt = read_design_guides_txt(txt_file)
        result_bed = read_design_guides_bed(bed_file)

        assert len(result_txt) == len(scored_df)
        assert len(result_bed) == len(scored_df)

    def test_predesigned_guide_caching(self, test_fasta, mock_rs3_model, temp_output_dir):
        """Test that pre-designed guides are correctly cached and merged."""
        # Create pre-designed guides for region1
        predesigned_df = pd.DataFrame({
            'chr': ['chr1'],
            'start': [1010],
            'end': [1030],
            'locus': ['chr1:1010-1030:+'],
            'strand': ['+'],
            'GuideSequenceWithPAM': ['ACGTACGATCGATCGACGAAAGG'],
            'guideSet': ['region1'],
            'RS3_score': [0.85],
            'specificity_score': [0.92]
        })

        # Save pre-designed guides
        predesigned_path = os.path.join(temp_output_dir, 'predesigned.txt')
        predesigned_df.to_csv(predesigned_path, sep='\t', index=False)

        # Load pre-designed guides
        loaded_df = load_predesigned_guides(predesigned_path)
        assert len(loaded_df) == 1

        # Create regions DataFrame
        regions_df = pd.DataFrame({
            'chr': ['chr1', 'chr1', 'chr1'],
            'start': [1000, 5000, 10000],
            'end': [1500, 5500, 10500],
            'name': ['region1', 'region2', 'region3']
        })

        # Filter regions (region1 should be excluded)
        filtered_regions = filter_predesigned_regions(regions_df, loaded_df)
        assert len(filtered_regions) == 2
        assert 'region1' not in filtered_regions['name'].values

        # Design guides for remaining regions
        new_guides = []
        for _, row in filtered_regions.iterrows():
            guides = extract_guides(row['chr'], row['start'], row['end'],
                                    row['name'], test_fasta)
            new_guides.extend(guides)

        if len(new_guides) == 0:
            pytest.skip("No new guides found")

        new_guides_df = pd.DataFrame(new_guides)
        passing_df = get_passing_guides(new_guides_df)

        if len(passing_df) == 0:
            pytest.skip("No guides passed filters")

        # Score new guides
        scored_df = score_guides_rs3(passing_df, mock_rs3_model)
        scored_df = compute_specificity_mock(scored_df)

        # Merge with pre-designed
        final_df = merge_with_predesigned(scored_df, loaded_df)

        # Should include guides from region1 (pre-designed) and new regions
        assert 'region1' in final_df['guideSet'].values


class TestDataIntegrity:
    """Tests for data integrity throughout the pipeline."""

    def test_coordinates_are_valid(self):
        """Test that output coordinates are valid BED format."""
        df = pd.DataFrame({
            'chr': ['chr1'],
            'start': [100],
            'end': [120],
            'locus': ['chr1:100-120:+'],
            'strand': ['+'],
            'GuideSequenceWithPAM': ['ACGTACGATCGATCGACGAAAGG'],
            'guideSet': ['test'],
            'RS3_score': [0.85],
            'specificity_score': [0.92]
        })

        with tempfile.TemporaryDirectory() as tmpdir:
            write_outputs(df, tmpdir)

            bed_df = read_design_guides_bed(os.path.join(tmpdir, 'designGuides.bed'))

            # Start should be less than end
            assert all(bed_df['start'] < bed_df['end'])

            # Coordinates should be non-negative
            assert all(bed_df['start'] >= 0)
            assert all(bed_df['end'] >= 0)

    def test_scores_are_normalized(self):
        """Test that all scores are in valid range."""
        df = pd.DataFrame({
            'chr': ['chr1'],
            'start': [100],
            'end': [120],
            'locus': ['chr1:100-120:+'],
            'strand': ['+'],
            'GuideSequenceWithPAM': ['ACGTACGATCGATCGACGAAAGG'],
            'guideSet': ['test'],
            'RS3_score': [0.85],
            'specificity_score': [0.92]
        })

        # RS3 scores should be 0-1
        assert all(df['RS3_score'] >= 0)
        assert all(df['RS3_score'] <= 1)

        # Specificity scores should be 0-1
        assert all(df['specificity_score'] >= 0)
        assert all(df['specificity_score'] <= 1)


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