"""
Tests for the G>T Transversion Analysis, Selective Vulnerability,
Differential miRNA Targeting, and Comparative Enrichment of Mutation
Proportions sections of All_together_analysis_w_explanation.ipynb.

Covers:
 - analysis_functions.py: calculate_gt_rates, extract_gt_positions,
                           run_diff_mutation_test, extract_first_position,
                           find_differentially_oxidized_mirs,
                           run_mutation_rate_burden
 - plotting.py           : plot_differential_heatmap, plot_mir
"""

import sys
from pathlib import Path

sys.path.insert(0, str(Path(__file__).parent.parent))

import numpy as np
import pandas as pd
import pytest
from unittest.mock import patch

from analysis_functions import (
    calculate_gt_rates,
    extract_gt_positions,
    run_diff_mutation_test,
    extract_first_position,
    find_differentially_oxidized_mirs,
    run_mutation_rate_burden,
)
from plotting import (
    plot_differential_heatmap,
    plot_mir,
)


# ---------------------------------------------------------------------------
# Shared helpers
# ---------------------------------------------------------------------------

def make_mismatch_df():
    """Minimal mismatch DataFrame as expected by G>T functions."""
    return pd.DataFrame(
        {
            "miRNA name": ["miR-1", "miR-1", "miR-2", "miR-2"],
            "pos:mut": ["PM", "5:G>T", "PM", "10:G>T"],
            "S1": [100, 5, 80, 3],
            "S2": [110, 4, 90, 2],
            "S3": [95, 6, 85, 4],
            "S4": [105, 3, 75, 5],
        }
    )


# ===========================================================================
# calculate_gt_rates
# ===========================================================================


class TestCalculateGtRates:
    def setup_method(self):
        self.mismatch_df = make_mismatch_df()
        self.sample_cols = ["S1", "S2", "S3", "S4"]

    def test_returns_dataframe(self):
        result = calculate_gt_rates(self.mismatch_df, self.sample_cols)
        assert isinstance(result, pd.DataFrame)

    def test_rates_are_between_0_and_1(self):
        result = calculate_gt_rates(self.mismatch_df, self.sample_cols)
        assert (result >= 0).all().all()
        assert (result <= 1).all().all()

    def test_no_gt_mutations_gives_zero_rate(self):
        df = pd.DataFrame(
            {
                "miRNA name": ["miR-X"],
                "pos:mut": ["PM"],
                "S1": [100],
            }
        )
        result = calculate_gt_rates(df, ["S1"])
        assert result.empty or (result == 0).all().all()

    def test_index_is_mirna_name(self):
        result = calculate_gt_rates(self.mismatch_df, self.sample_cols)
        assert result.index.name == "miRNA name"

    def test_columns_match_sample_cols(self):
        result = calculate_gt_rates(self.mismatch_df, self.sample_cols)
        assert list(result.columns) == self.sample_cols


# ===========================================================================
# extract_gt_positions
# ===========================================================================


class TestExtractGtPositions:
    def setup_method(self):
        self.mismatch_df = make_mismatch_df()

    def test_returns_dataframe(self):
        result = extract_gt_positions(self.mismatch_df)
        assert isinstance(result, pd.DataFrame)

    def test_adds_gt_position_column(self):
        result = extract_gt_positions(self.mismatch_df)
        assert "gt_position" in result.columns

    def test_pm_rows_have_nan_position(self):
        result = extract_gt_positions(self.mismatch_df)
        pm_rows = result[result["pos:mut"] == "PM"]
        assert pm_rows["gt_position"].isna().all()

    def test_gt_rows_have_integer_position(self):
        result = extract_gt_positions(self.mismatch_df)
        gt_rows = result[result["pos:mut"].str.contains("G>T", na=False)]
        assert gt_rows["gt_position"].notna().all()
        # positions should be positive integers
        assert (gt_rows["gt_position"] > 0).all()

    def test_correct_position_extracted(self):
        result = extract_gt_positions(self.mismatch_df)
        # Row with '5:G>T' should have position 5
        row_5 = result[result["pos:mut"] == "5:G>T"]
        assert row_5["gt_position"].iloc[0] == 5


# ===========================================================================
# extract_first_position
# ===========================================================================


class TestExtractFirstPosition:
    def test_returns_list(self):
        df = pd.DataFrame({"pos:mut": ["5:G>T", "PM", "12:G>T"]})
        result = extract_first_position(df)
        assert isinstance(result, list)

    def test_correct_positions_extracted(self):
        df = pd.DataFrame({"pos:mut": ["5:G>T", "12:G>T"]})
        result = extract_first_position(df)
        assert result[0] == 5
        assert result[1] == 12

    def test_pm_gives_nan(self):
        df = pd.DataFrame({"pos:mut": ["PM"]})
        result = extract_first_position(df)
        assert np.isnan(result[0])

    def test_multiple_mutations_first_gt_taken(self):
        # '3:A>C,8:G>T' – the G>T part is at position 8
        df = pd.DataFrame({"pos:mut": ["3:A>C,8:G>T"]})
        result = extract_first_position(df)
        assert result[0] == 8

    def test_raises_when_column_missing(self):
        df = pd.DataFrame({"other_col": ["5:G>T"]})
        with pytest.raises(ValueError):
            extract_first_position(df)


# ===========================================================================
# run_diff_mutation_test
# ===========================================================================


class TestRunDiffMutationTest:
    def setup_method(self):
        np.random.seed(0)
        self.rate_df = pd.DataFrame(
            {
                "G1A": [0.01, 0.02, 0.03],
                "G1B": [0.02, 0.01, 0.04],
                "G2A": [0.05, 0.06, 0.07],
                "G2B": [0.06, 0.07, 0.08],
            },
            index=["miR-1", "miR-2", "miR-3"],
        )

    def test_returns_dataframe(self):
        result = run_diff_mutation_test(self.rate_df, ["G1A", "G1B"], ["G2A", "G2B"])
        assert isinstance(result, pd.DataFrame)

    def test_output_has_expected_columns(self):
        result = run_diff_mutation_test(self.rate_df, ["G1A", "G1B"], ["G2A", "G2B"])
        for col in ["Rate_G1", "Rate_G2", "Delta", "p_value"]:
            assert col in result.columns

    def test_sorted_by_pvalue(self):
        result = run_diff_mutation_test(self.rate_df, ["G1A", "G1B"], ["G2A", "G2B"])
        p_vals = result["p_value"].tolist()
        assert p_vals == sorted(p_vals)

    def test_index_is_mirna_name(self):
        result = run_diff_mutation_test(self.rate_df, ["G1A", "G1B"], ["G2A", "G2B"])
        assert result.index.name == "miRNA"

    def test_min_reads_filter_reduces_rows(self):
        min_reads_df = pd.DataFrame(
            {
                "G1A": [1000, 1, 1000],
                "G1B": [1000, 1, 1000],
                "G2A": [1000, 1, 1000],
                "G2B": [1000, 1, 1000],
            },
            index=["miR-1", "miR-2", "miR-3"],
        )
        result = run_diff_mutation_test(
            self.rate_df,
            ["G1A", "G1B"],
            ["G2A", "G2B"],
            min_reads_df=min_reads_df,
            min_threshold=100,
        )
        # miR-2 should be excluded (total reads = 4, below 100)
        assert "miR-2" not in result.index


# ===========================================================================
# find_differentially_oxidized_mirs
# ===========================================================================


class TestFindDifferentiallyOxidizedMirs:
    def setup_method(self):
        np.random.seed(2)
        self.rate_df = pd.DataFrame(
            {
                "G1A": [0.01, 0.02, 0.03, 0.01],
                "G1B": [0.02, 0.01, 0.02, 0.02],
                "G2A": [0.07, 0.08, 0.03, 0.01],
                "G2B": [0.08, 0.09, 0.02, 0.02],
            },
            index=["miR-1", "miR-2", "miR-3", "miR-4"],
        )
        # Total counts high enough for all miRNAs
        self.total_df = pd.DataFrame(
            {
                "G1A": [200, 200, 200, 200],
                "G1B": [200, 200, 200, 200],
                "G2A": [200, 200, 200, 200],
                "G2B": [200, 200, 200, 200],
            },
            index=["miR-1", "miR-2", "miR-3", "miR-4"],
        )

    def test_returns_dataframe(self, capsys):
        result = find_differentially_oxidized_mirs(
            self.rate_df, self.total_df, ["G1A", "G1B"], ["G2A", "G2B"]
        )
        assert isinstance(result, pd.DataFrame)

    def test_output_has_expected_columns(self, capsys):
        result = find_differentially_oxidized_mirs(
            self.rate_df, self.total_df, ["G1A", "G1B"], ["G2A", "G2B"]
        )
        for col in ["Rate_G1", "Rate_G2", "Delta_Rate", "p_value"]:
            assert col in result.columns

    def test_sorted_by_pvalue(self, capsys):
        result = find_differentially_oxidized_mirs(
            self.rate_df, self.total_df, ["G1A", "G1B"], ["G2A", "G2B"]
        )
        p_vals = result["p_value"].tolist()
        assert p_vals == sorted(p_vals)

    def test_min_reads_filter_applied(self, capsys):
        low_total_df = pd.DataFrame(
            {
                "G1A": [1, 1, 200, 200],
                "G1B": [1, 1, 200, 200],
                "G2A": [1, 1, 200, 200],
                "G2B": [1, 1, 200, 200],
            },
            index=["miR-1", "miR-2", "miR-3", "miR-4"],
        )
        result = find_differentially_oxidized_mirs(
            self.rate_df, low_total_df, ["G1A", "G1B"], ["G2A", "G2B"], min_total_reads=50
        )
        # Only miR-3 and miR-4 pass the threshold
        assert len(result) == 2


# ===========================================================================
# run_mutation_rate_burden
# ===========================================================================


class TestRunMutationRateBurden:
    def setup_method(self):
        np.random.seed(1)
        self.rate_df = pd.DataFrame(
            {
                "G1A": np.random.uniform(0, 0.05, 10),
                "G1B": np.random.uniform(0, 0.05, 10),
                "G2A": np.random.uniform(0.04, 0.1, 10),
                "G2B": np.random.uniform(0.04, 0.1, 10),
            },
            index=[f"miR-{i}" for i in range(10)],
        )

    @patch("matplotlib.pyplot.show")
    def test_returns_dict(self, mock_show):
        result = run_mutation_rate_burden(
            self.rate_df, ["G1A", "G1B"], ["G2A", "G2B"]
        )
        assert isinstance(result, dict)

    @patch("matplotlib.pyplot.show")
    def test_dict_has_expected_keys(self, mock_show):
        result = run_mutation_rate_burden(
            self.rate_df, ["G1A", "G1B"], ["G2A", "G2B"]
        )
        for key in ["t_statistic", "p_value", "group_1_mean", "group_2_mean"]:
            assert key in result

    @patch("matplotlib.pyplot.show")
    def test_raises_when_samples_missing(self, mock_show):
        with pytest.raises(ValueError):
            run_mutation_rate_burden(
                self.rate_df, ["G1A", "MISSING"], ["G2A", "G2B"]
            )

    @patch("matplotlib.pyplot.show")
    def test_p_value_is_valid(self, mock_show):
        result = run_mutation_rate_burden(
            self.rate_df, ["G1A", "G1B"], ["G2A", "G2B"]
        )
        assert 0.0 <= result["p_value"] <= 1.0


# ===========================================================================
# plot_differential_heatmap
# ===========================================================================


class TestPlotDifferentialHeatmap:
    def setup_method(self):
        self.enrichment_df = pd.DataFrame(
            {
                "M1": [0.5, -0.2, 1.0, 0.1],
                "M2": [0.6, -0.3, 0.9, 0.2],
                "M3": [1.5, 0.8, 0.3, 0.7],
                "M4": [1.6, 0.9, 0.2, 0.8],
            },
            index=["miR-1", "miR-2", "miR-3", "miR-4"],
        )

    @patch("matplotlib.pyplot.show")
    def test_runs_without_error(self, mock_show):
        plot_differential_heatmap(
            self.enrichment_df,
            target_mirs=["miR-1", "miR-2"],
            group_1_individuals=["M1", "M2"],
            group_2_individuals=["M3", "M4"],
        )
        assert mock_show.called


# ===========================================================================
# plot_mir
# ===========================================================================


class TestPlotMir:
    def setup_method(self):
        self.mutation_rate_df = pd.DataFrame(
            {
                "G1A": [0.01, 0.02],
                "G1B": [0.02, 0.03],
                "G2A": [0.05, 0.06],
                "G2B": [0.06, 0.07],
            },
            index=["miR-1", "miR-2"],
        )

    @patch("matplotlib.pyplot.show")
    def test_runs_without_error(self, mock_show):
        plot_mir(
            self.mutation_rate_df,
            group_1_samples=["G1A", "G1B"],
            group_2_samples=["G2A", "G2B"],
            mir_name="miR-1",
            p_value=0.04,
        )
        assert mock_show.called
