"""
Tests for output writer module.
"""

import os
import tempfile
import pytest
import pandas as pd

from crisprdesigner2 import (
    write_outputs, validate_dataframe, scale_score_to_bed,
    read_design_guides_txt, read_design_guides_bed,
    TXT_COLUMNS, BED_COLUMNS,
)


class TestScaleScoreToBed:
    """Tests for the scale_score_to_bed function."""

    def test_zero_score(self):
        """Score of 0 should map to 0."""
        assert scale_score_to_bed(0.0) == 0

    def test_one_score(self):
        """Score of 1 should map to 1000."""
        assert scale_score_to_bed(1.0) == 1000

    def test_half_score(self):
        """Score of 0.5 should map to 500."""
        assert scale_score_to_bed(0.5) == 500

    def test_nan_score(self):
        """NaN score should map to 0."""
        assert scale_score_to_bed(float('nan')) == 0

    def test_out_of_range_high(self):
        """Score > 1 should be clipped to 1000."""
        assert scale_score_to_bed(1.5) == 1000

    def test_out_of_range_low(self):
        """Score < 0 should be clipped to 0."""
        assert scale_score_to_bed(-0.5) == 0


class TestValidateDataframe:
    """Tests for the validate_dataframe function."""

    def test_valid_dataframe(self):
        """DataFrame with all required columns should pass validation."""
        df = pd.DataFrame({
            'chr': ['chr1'],
            'start': [100],
            'end': [120],
            'locus': ['chr1:100-120:+'],
            'strand': ['+'],
            'GuideSequenceWithPAM': ['ACGTACGATCGATCGACGAAAGG'],
            'guideSet': ['region1'],
            'RS3_score': [0.85]
        })
        # Should not raise
        validate_dataframe(df)

    def test_missing_column(self):
        """DataFrame missing required column should raise ValueError."""
        df = pd.DataFrame({
            'chr': ['chr1'],
            'start': [100],
            'end': [120],
            # Missing locus, strand, etc.
        })
        with pytest.raises(ValueError) as excinfo:
            validate_dataframe(df)
        assert "Missing required columns" in str(excinfo.value)


class TestWriteOutputs:
    """Tests for the write_outputs function."""

    @pytest.fixture
    def sample_guides_df(self):
        """Create a sample DataFrame of guides."""
        return pd.DataFrame({
            'chr': ['chr1', 'chr1', 'chr2'],
            'start': [100, 200, 300],
            'end': [120, 220, 320],
            'locus': ['chr1:100-120:+', 'chr1:200-220:-', 'chr2:300-320:+'],
            'strand': ['+', '-', '+'],
            'GuideSequenceWithPAM': [
                'ACGTACGATCGATCGACGAAAGG',
                'CGATCGATCGATCGATCGACAGG',
                'TACGAACGTTACGTAACGACTGG'
            ],
            'guideSet': ['region1', 'region1', 'region2'],
            'RS3_score': [0.85, 0.72, 0.91],
            'specificity_score': [0.95, 0.88, 0.76]
        })

    def test_creates_both_files(self, sample_guides_df):
        """write_outputs should create both txt and bed files."""
        with tempfile.TemporaryDirectory() as tmpdir:
            write_outputs(sample_guides_df, tmpdir)

            txt_path = os.path.join(tmpdir, 'designGuides.txt')
            bed_path = os.path.join(tmpdir, 'designGuides.bed')

            assert os.path.exists(txt_path)
            assert os.path.exists(bed_path)

    def test_txt_has_correct_headers(self, sample_guides_df):
        """designGuides.txt should have correct column headers."""
        with tempfile.TemporaryDirectory() as tmpdir:
            write_outputs(sample_guides_df, tmpdir)

            txt_path = os.path.join(tmpdir, 'designGuides.txt')
            with open(txt_path, 'r') as f:
                header = f.readline().strip()

            expected_header = '\t'.join(TXT_COLUMNS)
            assert header == expected_header

    def test_txt_has_correct_row_count(self, sample_guides_df):
        """designGuides.txt should have correct number of rows."""
        with tempfile.TemporaryDirectory() as tmpdir:
            write_outputs(sample_guides_df, tmpdir)

            txt_path = os.path.join(tmpdir, 'designGuides.txt')
            df = pd.read_csv(txt_path, sep='\t')

            assert len(df) == 3

    def test_bed_is_valid_bed6(self, sample_guides_df):
        """designGuides.bed should be valid BED6 format."""
        with tempfile.TemporaryDirectory() as tmpdir:
            write_outputs(sample_guides_df, tmpdir)

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

            # Read BED file
            bed_df = pd.read_csv(bed_path, sep='\t', header=None)

            # Should have 6 columns
            assert len(bed_df.columns) == 6

            # Score should be integer 0-1000
            scores = bed_df.iloc[:, 4]
            assert all(isinstance(s, (int, np.integer)) or (isinstance(s, float) and s.is_integer())
                       for s in scores)
            assert all(0 <= s <= 1000 for s in scores)

            # Strand should be + or -
            strands = bed_df.iloc[:, 5]
            assert all(s in ['+', '-'] for s in strands)

    def test_bed_score_scaling(self, sample_guides_df):
        """BED score should be RS3_score scaled to 0-1000."""
        with tempfile.TemporaryDirectory() as tmpdir:
            write_outputs(sample_guides_df, tmpdir)

            bed_path = os.path.join(tmpdir, 'designGuides.bed')
            bed_df = pd.read_csv(bed_path, sep='\t', header=None)

            # First guide: RS3_score = 0.85 -> scaled = 850
            assert bed_df.iloc[0, 4] == 850

            # Second guide: RS3_score = 0.72 -> scaled = 720
            assert bed_df.iloc[1, 4] == 720

            # Third guide: RS3_score = 0.91 -> scaled = 910
            assert bed_df.iloc[2, 4] == 910

    def test_creates_output_directory(self, sample_guides_df):
        """write_outputs should create output directory if it doesn't exist."""
        with tempfile.TemporaryDirectory() as tmpdir:
            new_dir = os.path.join(tmpdir, 'new', 'nested', 'directory')

            write_outputs(sample_guides_df, new_dir)

            assert os.path.exists(new_dir)
            assert os.path.exists(os.path.join(new_dir, 'designGuides.txt'))
            assert os.path.exists(os.path.join(new_dir, 'designGuides.bed'))

    def test_handles_missing_specificity_score(self):
        """write_outputs should handle missing specificity_score column."""
        df = pd.DataFrame({
            'chr': ['chr1'],
            'start': [100],
            'end': [120],
            'locus': ['chr1:100-120:+'],
            'strand': ['+'],
            'GuideSequenceWithPAM': ['ACGTACGATCGATCGACGAAAGG'],
            'guideSet': ['region1'],
            'RS3_score': [0.85]
            # No specificity_score column
        })

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

            txt_path = os.path.join(tmpdir, 'designGuides.txt')
            result_df = pd.read_csv(txt_path, sep='\t')

            assert 'specificity_score' in result_df.columns

    def test_missing_required_column_raises_error(self):
        """write_outputs should raise error for missing required columns."""
        df = pd.DataFrame({
            'chr': ['chr1'],
            'start': [100],
            # Missing other required columns
        })

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


class TestReadOutputs:
    """Tests for reading output files back."""

    @pytest.fixture
    def written_files(self):
        """Create output files and return paths."""
        df = pd.DataFrame({
            'chr': ['chr1', 'chr2'],
            'start': [100, 200],
            'end': [120, 220],
            'locus': ['chr1:100-120:+', 'chr2:200-220:-'],
            'strand': ['+', '-'],
            'GuideSequenceWithPAM': [
                'ACGTACGATCGATCGACGAAAGG',
                'CGATCGATCGATCGATCGACAGG'
            ],
            'guideSet': ['region1', 'region2'],
            'RS3_score': [0.85, 0.72],
            'specificity_score': [0.95, 0.88]
        })

        tmpdir = tempfile.mkdtemp()
        write_outputs(df, tmpdir)

        yield tmpdir, df

        # Cleanup
        import shutil
        shutil.rmtree(tmpdir)

    def test_read_txt(self, written_files):
        """read_design_guides_txt should correctly read the txt file."""
        tmpdir, original_df = written_files
        txt_path = os.path.join(tmpdir, 'designGuides.txt')

        result = read_design_guides_txt(txt_path)

        assert len(result) == 2
        assert list(result.columns) == TXT_COLUMNS
        assert result['chr'].iloc[0] == 'chr1'
        assert result['RS3_score'].iloc[0] == 0.85

    def test_read_bed(self, written_files):
        """read_design_guides_bed should correctly read the bed file."""
        tmpdir, original_df = written_files
        bed_path = os.path.join(tmpdir, 'designGuides.bed')

        result = read_design_guides_bed(bed_path)

        assert len(result) == 2
        assert list(result.columns) == BED_COLUMNS
        assert result['chr'].iloc[0] == 'chr1'
        assert result['score'].iloc[0] == 850  # 0.85 * 1000


# Import numpy for the integer check
import numpy as np


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