"""
Tests for the Paired o8 Enrichment Analysis section of
All_together_analysis_w_explanation.ipynb.

Covers:
 - analysis_functions.py: get_metadata_mapping, calc_enrichment_matrix,
                           get_enrichment_score, filter_enrichment_noise,
                           run_distribution_tests
 - plotting.py           : plot_global_enrichment, plot_kde, ECDF_plot,
                           plot_enrichment_distribution
"""

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 (
    get_metadata_mapping,
    calc_enrichment_matrix,
    get_enrichment_score,
    filter_enrichment_noise,
    run_distribution_tests,
)
from plotting import (
    plot_global_enrichment,
    plot_kde,
    ECDF_plot,
    plot_enrichment_distribution,
)


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

def make_cpm_df():
    """Minimal CPM DataFrame: 5 miRNAs x 4 samples."""
    return pd.DataFrame(
        {
            "S1": [15.0, 0.5, 12.0, 8.0, 20.0],
            "S2": [18.0, 1.0, 11.0, 9.0, 22.0],
            "S3": [0.5, 5.0, 14.0, 3.0, 25.0],
            "S4": [1.0, 6.0, 13.0, 2.0, 24.0],
        },
        index=["miR-1", "miR-2", "miR-3", "miR-4", "miR-5"],
    )


def make_samples_df():
    """Minimal metadata DataFrame.

    The column name 'IP-Nat' deliberately matches the upstream data format
    used in the notebook and must stay consistent with the source code.
    """
    return pd.DataFrame(
        {
            "Sample": ["S1", "S2", "S3", "S4"],
            "Individual": ["M1", "M2", "M3", "M4"],
            "IP-Nat": ["IP", "IP", "NAT", "NAT"],
            "Treatment": ["Veh", "KO", "Veh", "KO"],
        }
    )


def make_mouse_pairs_dict():
    """Minimal mouse pairs dict matching make_cpm_df samples."""
    return {
        "M1": {"IP": "S1", "NAT": "S3", "Treatment": "Veh"},
        "M2": {"IP": "S2", "NAT": "S4", "Treatment": "KO"},
    }


# ===========================================================================
# get_metadata_mapping
# ===========================================================================


class TestGetMetadataMapping:
    def setup_method(self):
        self.samples_df = make_samples_df()

    def test_returns_dict(self):
        result = get_metadata_mapping(self.samples_df)
        assert isinstance(result, dict)

    def test_keys_are_individuals(self):
        result = get_metadata_mapping(self.samples_df)
        assert set(result.keys()) == {"M1", "M2", "M3", "M4"}

    def test_ip_and_nat_stored_in_uppercase(self):
        result = get_metadata_mapping(self.samples_df)
        # M1 is IP so its key should be 'IP'
        assert "IP" in result["M1"]

    def test_extra_categorical_columns_preserved(self):
        result = get_metadata_mapping(self.samples_df)
        # Treatment should appear as a categorical field
        assert "Treatment" in result["M1"]

    def test_sample_name_correctly_mapped(self):
        result = get_metadata_mapping(self.samples_df)
        assert result["M1"]["IP"] == "S1"
        assert result["M3"]["NAT"] == "S3"


# ===========================================================================
# calc_enrichment_matrix
# ===========================================================================


class TestCalcEnrichmentMatrix:
    def setup_method(self):
        self.cpm_df = make_cpm_df()
        self.pairs = make_mouse_pairs_dict()

    def test_returns_dataframe(self):
        result = calc_enrichment_matrix(self.cpm_df, self.pairs)
        assert isinstance(result, pd.DataFrame)

    def test_columns_are_mouse_ids(self):
        result = calc_enrichment_matrix(self.cpm_df, self.pairs)
        assert set(result.columns) == {"M1", "M2"}

    def test_index_matches_cpm_index(self):
        result = calc_enrichment_matrix(self.cpm_df, self.pairs)
        assert list(result.index) == list(self.cpm_df.index)

    def test_missing_column_is_skipped(self, capsys):
        bad_pairs = {"M99": {"IP": "MISSING_IP", "NAT": "S3"}}
        result = calc_enrichment_matrix(self.cpm_df, bad_pairs)
        captured = capsys.readouterr()
        assert "Warning" in captured.out
        assert result.empty or "M99" not in result.columns

    def test_values_are_finite_floats(self):
        result = calc_enrichment_matrix(self.cpm_df, self.pairs)
        assert result.dtypes.apply(lambda dt: np.issubdtype(dt, np.floating)).all()
        assert np.isfinite(result.values).all()


# ===========================================================================
# get_enrichment_score
# ===========================================================================


class TestGetEnrichmentScore:
    def test_simple_equal_inputs(self):
        ip = pd.Series([1.0, 2.0, 4.0])
        nat = pd.Series([1.0, 2.0, 4.0])
        result = get_enrichment_score(ip, nat)
        # log2(1) = 0 for equal values with pseudocount
        assert (result == 0).all()

    def test_higher_ip_gives_positive_score(self):
        ip = pd.Series([10.0])
        nat = pd.Series([1.0])
        result = get_enrichment_score(ip, nat)
        assert result.iloc[0] > 0

    def test_higher_nat_gives_negative_score(self):
        ip = pd.Series([1.0])
        nat = pd.Series([10.0])
        result = get_enrichment_score(ip, nat)
        assert result.iloc[0] < 0

    def test_returns_series_same_length(self):
        ip = pd.Series([1.0, 2.0, 3.0])
        nat = pd.Series([3.0, 2.0, 1.0])
        result = get_enrichment_score(ip, nat)
        assert len(result) == 3

    def test_zero_inputs_handled_via_pseudocount(self):
        ip = pd.Series([0.0])
        nat = pd.Series([0.0])
        result = get_enrichment_score(ip, nat)
        # pseudocount keeps ratio at 1, so log2(1) == 0
        assert result.iloc[0] == 0.0


# ===========================================================================
# filter_enrichment_noise
# ===========================================================================


class TestFilterEnrichmentNoise:
    def setup_method(self):
        self.cpm_df = make_cpm_df()
        self.pairs = make_mouse_pairs_dict()
        self.enrichment_df = calc_enrichment_matrix(self.cpm_df, self.pairs)

    def test_returns_dataframe(self):
        result = filter_enrichment_noise(self.enrichment_df, self.cpm_df, self.pairs)
        assert isinstance(result, pd.DataFrame)

    def test_high_threshold_removes_all_rows(self):
        result = filter_enrichment_noise(
            self.enrichment_df, self.cpm_df, self.pairs, threshold=10000.0
        )
        assert len(result) == 0

    def test_zero_threshold_keeps_all_rows(self):
        result = filter_enrichment_noise(
            self.enrichment_df, self.cpm_df, self.pairs, threshold=0.0
        )
        assert len(result) == len(self.enrichment_df)

    def test_columns_unchanged(self):
        result = filter_enrichment_noise(self.enrichment_df, self.cpm_df, self.pairs)
        assert list(result.columns) == list(self.enrichment_df.columns)


# ===========================================================================
# run_distribution_tests
# ===========================================================================


class TestRunDistributionTests:
    def setup_method(self):
        np.random.seed(42)
        self.enrichment_df = pd.DataFrame(
            np.random.randn(50, 6),
            columns=["A1", "A2", "A3", "B1", "B2", "B3"],
            index=[f"miR-{i}" for i in range(50)],
        )

    def test_returns_dict(self):
        result = run_distribution_tests(self.enrichment_df, ["A1", "A2", "A3"], ["B1", "B2", "B3"])
        assert isinstance(result, dict)

    def test_dict_has_required_keys(self):
        result = run_distribution_tests(self.enrichment_df, ["A1", "A2", "A3"], ["B1", "B2", "B3"])
        for key in ["ttest_p", "ks_p", "ks_stat", "group_1_median", "group_2_median"]:
            assert key in result

    def test_p_values_are_between_0_and_1(self):
        result = run_distribution_tests(self.enrichment_df, ["A1", "A2", "A3"], ["B1", "B2", "B3"])
        assert 0.0 <= result["ttest_p"] <= 1.0
        assert 0.0 <= result["ks_p"] <= 1.0

    def test_missing_individuals_are_skipped(self):
        # Pass a non-existent individual; the function should still run
        result = run_distribution_tests(
            self.enrichment_df, ["A1", "A2", "MISSING"], ["B1", "B2", "B3"]
        )
        assert "ttest_p" in result


# ===========================================================================
# plot_global_enrichment
# ===========================================================================


class TestPlotGlobalEnrichment:
    def setup_method(self):
        np.random.seed(6)
        self.g1 = pd.DataFrame(
            np.random.randn(10, 3),
            columns=["M1", "M2", "M3"],
            index=[f"miR-{i}" for i in range(10)],
        )
        self.g2 = pd.DataFrame(
            np.random.randn(10, 3) + 0.5,
            columns=["M4", "M5", "M6"],
            index=[f"miR-{i}" for i in range(10)],
        )

    @patch("matplotlib.pyplot.show")
    def test_returns_three_values(self, mock_show):
        result = plot_global_enrichment(self.g1, self.g2)
        assert len(result) == 3

    @patch("matplotlib.pyplot.show")
    def test_first_return_is_dataframe(self, mock_show):
        plot_data, t_stat, p_val = plot_global_enrichment(self.g1, self.g2)
        assert isinstance(plot_data, pd.DataFrame)

    @patch("matplotlib.pyplot.show")
    def test_p_value_is_valid(self, mock_show):
        _, _, p_val = plot_global_enrichment(self.g1, self.g2)
        assert 0.0 <= p_val <= 1.0

    @patch("matplotlib.pyplot.show")
    def test_plot_data_has_group_column(self, mock_show):
        plot_data, _, _ = plot_global_enrichment(self.g1, self.g2)
        assert "Group" in plot_data.columns


# ===========================================================================
# plot_kde
# ===========================================================================


class TestPlotKde:
    def setup_method(self):
        np.random.seed(7)
        self.g1 = pd.DataFrame(
            np.random.randn(10, 3),
            columns=["M1", "M2", "M3"],
            index=[f"miR-{i}" for i in range(10)],
        )
        self.g2 = pd.DataFrame(
            np.random.randn(10, 3) + 0.5,
            columns=["M4", "M5", "M6"],
            index=[f"miR-{i}" for i in range(10)],
        )

    @patch("matplotlib.pyplot.show")
    def test_returns_two_values(self, mock_show):
        result = plot_kde(self.g1, self.g2)
        assert len(result) == 2

    @patch("matplotlib.pyplot.show")
    def test_first_return_is_dataframe(self, mock_show):
        all_data, stats_dict = plot_kde(self.g1, self.g2)
        assert isinstance(all_data, pd.DataFrame)

    @patch("matplotlib.pyplot.show")
    def test_stats_dict_has_ks_keys(self, mock_show):
        _, stats_dict = plot_kde(self.g1, self.g2)
        assert "KS-statistic" in stats_dict
        assert "KS-p-value" in stats_dict

    @patch("matplotlib.pyplot.show")
    def test_all_data_has_enrichment_and_group(self, mock_show):
        all_data, _ = plot_kde(self.g1, self.g2)
        assert "Enrichment" in all_data.columns
        assert "Group" in all_data.columns


# ===========================================================================
# ECDF_plot
# ===========================================================================


class TestEcdfPlot:
    def setup_method(self):
        np.random.seed(8)
        self.enrichment_df = pd.DataFrame(
            np.random.randn(20, 4),
            columns=["M1", "M2", "M3", "M4"],
            index=[f"miR-{i}" for i in range(20)],
        )

    @patch("matplotlib.pyplot.show")
    def test_runs_without_error(self, mock_show):
        ECDF_plot(
            self.enrichment_df,
            group_1_individuals=["M1", "M2"],
            group_2_individuals=["M3", "M4"],
            group_1_name="Group1",
            group_2_name="Group2",
        )
        assert mock_show.called

    def test_raises_on_missing_individual(self):
        with pytest.raises(ValueError):
            ECDF_plot(
                self.enrichment_df,
                group_1_individuals=["M1", "MISSING"],
                group_2_individuals=["M3", "M4"],
                group_1_name="Group1",
                group_2_name="Group2",
            )

    @patch("matplotlib.pyplot.show")
    def test_ignore_missing_flag(self, mock_show):
        # Should not raise when ignore_missing=True
        ECDF_plot(
            self.enrichment_df,
            group_1_individuals=["M1", "MISSING"],
            group_2_individuals=["M3", "M4"],
            group_1_name="Group1",
            group_2_name="Group2",
            ignore_missing=True,
        )
        assert mock_show.called


# ===========================================================================
# plot_enrichment_distribution
# ===========================================================================


class TestPlotEnrichmentDistribution:
    def setup_method(self):
        np.random.seed(9)
        self.g1 = pd.DataFrame(
            np.random.randn(10, 2),
            columns=["M1", "M2"],
            index=[f"miR-{i}" for i in range(10)],
        )
        self.g2 = pd.DataFrame(
            np.random.randn(10, 2) + 0.5,
            columns=["M3", "M4"],
            index=[f"miR-{i}" for i in range(10)],
        )

    @patch("matplotlib.pyplot.show")
    def test_returns_two_arrays(self, mock_show):
        g1_pts, g2_pts = plot_enrichment_distribution(self.g1, self.g2)
        assert len(g1_pts) > 0
        assert len(g2_pts) > 0

    @patch("matplotlib.pyplot.show")
    def test_output_lengths_match_input(self, mock_show):
        g1_pts, g2_pts = plot_enrichment_distribution(self.g1, self.g2)
        assert len(g1_pts) == self.g1.size
        assert len(g2_pts) == self.g2.size
