"""
Tests for the miRNA Count Analysis section of
All_together_analysis_w_explanation.ipynb.

Covers:
 - pipeline_utils.py: get_exclusion_lists
 - plotting.py      : plot_barplot_difference_mirs, plot_boxplot_difference_mirs
"""

import sys
import types
from pathlib import Path

# ---------------------------------------------------------------------------
# Stub out heavy optional dependencies before importing pipeline_utils.
# ---------------------------------------------------------------------------
for _mod_name in [
    "anndata",
    "pydeseq2",
    "pydeseq2.dds",
    "pydeseq2.ds",
    "decoupler",
]:
    if _mod_name not in sys.modules:
        sys.modules[_mod_name] = types.ModuleType(_mod_name)

sys.modules["anndata"].AnnData = object  # type: ignore[attr-defined]
_dds_mod = sys.modules["pydeseq2.dds"]
_dds_mod.DeseqDataSet = object  # type: ignore[attr-defined]
_dds_mod.DefaultInference = object  # type: ignore[attr-defined]
sys.modules["pydeseq2.ds"].DeseqStats = object  # type: ignore[attr-defined]

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 pipeline_utils import get_exclusion_lists
from plotting import (
    plot_barplot_difference_mirs,
    plot_boxplot_difference_mirs,
)


# ---------------------------------------------------------------------------
# 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"],
        }
    )


# ===========================================================================
# get_exclusion_lists
# ===========================================================================


class TestGetExclusionLists:
    def setup_method(self):
        self.cpm_df = pd.DataFrame(
            {
                "G1A": [5.0, 0.0, 3.0],
                "G1B": [6.0, 0.0, 4.0],
                "G2A": [0.0, 8.0, 3.0],
                "G2B": [0.0, 9.0, 4.0],
            },
            index=["miR-1", "miR-2", "miR-3"],
        )

    def test_returns_two_lists(self):
        u1, u2 = get_exclusion_lists(
            self.cpm_df, ["G1A", "G1B"], ["G2A", "G2B"], threshold=0
        )
        assert isinstance(u1, list)
        assert isinstance(u2, list)

    def test_exclusive_mirnas_identified_correctly(self):
        u1, u2 = get_exclusion_lists(
            self.cpm_df, ["G1A", "G1B"], ["G2A", "G2B"], threshold=0
        )
        # miR-1 is only in group 1; miR-2 is only in group 2
        assert "miR-1" in u1
        assert "miR-2" in u2

    def test_shared_mirna_not_in_exclusive_lists(self):
        u1, u2 = get_exclusion_lists(
            self.cpm_df, ["G1A", "G1B"], ["G2A", "G2B"], threshold=0
        )
        # miR-3 is expressed in both groups
        assert "miR-3" not in u1
        assert "miR-3" not in u2

    def test_threshold_affects_results(self):
        u1_high, u2_high = get_exclusion_lists(
            self.cpm_df, ["G1A", "G1B"], ["G2A", "G2B"], threshold=5
        )
        # miR-3 has max 4 in G1 and 4 in G2, so at threshold 5 it is inactive in both groups
        assert "miR-3" not in u1_high
        assert "miR-3" not in u2_high


# ===========================================================================
# plot_barplot_difference_mirs
# ===========================================================================


class TestPlotBarplotDifferenceMirs:
    def setup_method(self):
        self.cpm_df = make_cpm_df()
        self.samples_df = make_samples_df()

    @patch("matplotlib.pyplot.show")
    def test_returns_dict(self, mock_show):
        result = plot_barplot_difference_mirs(
            self.cpm_df, self.samples_df, group="IP-Nat", threshold=5
        )
        assert isinstance(result, dict)

    @patch("matplotlib.pyplot.show")
    def test_dict_keys_are_group_values(self, mock_show):
        result = plot_barplot_difference_mirs(
            self.cpm_df, self.samples_df, group="IP-Nat", threshold=5
        )
        unique_groups = self.samples_df["IP-Nat"].unique()
        for g in unique_groups:
            assert g in result

    @patch("matplotlib.pyplot.show")
    def test_raises_when_group_not_in_samples_df(self, mock_show):
        with pytest.raises(ValueError):
            plot_barplot_difference_mirs(
                self.cpm_df, self.samples_df, group="NonExistent"
            )

    @patch("matplotlib.pyplot.show")
    def test_raises_when_no_common_samples(self, mock_show):
        other_df = pd.DataFrame({"Sample": ["X1", "X2"], "IP-Nat": ["IP", "NAT"]})
        with pytest.raises(ValueError):
            plot_barplot_difference_mirs(self.cpm_df, other_df, group="IP-Nat")


# ===========================================================================
# plot_boxplot_difference_mirs
# ===========================================================================


class TestPlotBoxplotDifferenceMirs:
    def setup_method(self):
        self.cpm_df = make_cpm_df()
        self.samples_df = make_samples_df()

    @patch("matplotlib.pyplot.show")
    def test_returns_dict(self, mock_show):
        result = plot_boxplot_difference_mirs(
            self.cpm_df, self.samples_df, group="IP-Nat", threshold=5
        )
        assert isinstance(result, dict)

    @patch("matplotlib.pyplot.show")
    def test_dict_keys_are_group_values(self, mock_show):
        result = plot_boxplot_difference_mirs(
            self.cpm_df, self.samples_df, group="IP-Nat", threshold=5
        )
        unique_groups = self.samples_df["IP-Nat"].unique()
        for g in unique_groups:
            assert g in result

    @patch("matplotlib.pyplot.show")
    def test_raises_when_group_not_in_samples_df(self, mock_show):
        with pytest.raises(ValueError):
            plot_boxplot_difference_mirs(
                self.cpm_df, self.samples_df, group="NonExistent"
            )

    @patch("matplotlib.pyplot.show")
    def test_raises_when_no_common_samples(self, mock_show):
        other_df = pd.DataFrame({"Sample": ["X1", "X2"], "IP-Nat": ["IP", "NAT"]})
        with pytest.raises(ValueError):
            plot_boxplot_difference_mirs(self.cpm_df, other_df, group="IP-Nat")
