"""
Tests for the Global Burden and Richness Analysis section of
All_together_analysis_w_explanation.ipynb.

Covers:
 - analysis_functions.py: calculate_burden_richness
 - plotting.py           : plot_burden_and_richness
"""

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_burden_richness
from plotting import plot_burden_and_richness


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


# ===========================================================================
# calculate_burden_richness
# ===========================================================================


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

    def test_returns_dataframe_and_stats_dict(self):
        df, stats = calculate_burden_richness(
            self.cpm_df, group="Treatment", samples_df=self.samples_df, threshold=1.0
        )
        assert isinstance(df, pd.DataFrame)
        assert isinstance(stats, dict)

    def test_dataframe_has_expected_columns(self):
        df, _ = calculate_burden_richness(
            self.cpm_df, group="Treatment", samples_df=self.samples_df, threshold=1.0
        )
        assert "Burden" in df.columns
        assert "Richness" in df.columns
        assert "Group" in df.columns

    def test_stats_dict_has_burden_and_richness_keys(self):
        _, stats = calculate_burden_richness(
            self.cpm_df, group="Treatment", samples_df=self.samples_df, threshold=1.0
        )
        assert "Burden" in stats
        assert "Richness" in stats

    def test_stats_contains_pvalue(self):
        _, stats = calculate_burden_richness(
            self.cpm_df, group="Treatment", samples_df=self.samples_df, threshold=1.0
        )
        assert "p_value" in stats["Burden"]
        assert "p_value" in stats["Richness"]

    def test_raises_when_group_is_none(self):
        with pytest.raises(ValueError):
            calculate_burden_richness(
                self.cpm_df, group=None, samples_df=self.samples_df
            )

    def test_raises_when_group_not_in_samples_df(self):
        with pytest.raises(ValueError):
            calculate_burden_richness(
                self.cpm_df, group="NonExistentColumn", samples_df=self.samples_df
            )

    def test_richness_respects_threshold(self):
        # With a very high threshold, richness should be zero
        df_high, _ = calculate_burden_richness(
            self.cpm_df, group="Treatment", samples_df=self.samples_df, threshold=1000.0
        )
        assert df_high["Richness"].sum() == 0

    def test_group_means_in_stats(self):
        _, stats = calculate_burden_richness(
            self.cpm_df, group="Treatment", samples_df=self.samples_df
        )
        # Both group labels should appear as mean keys
        assert any("mean" in k.lower() or "_mean" in k for k in stats["Burden"])


# ===========================================================================
# plot_burden_and_richness
# ===========================================================================


class TestPlotBurdenAndRichness:
    def setup_method(self):
        self.analysis_df = pd.DataFrame(
            {
                "Burden": [10.0, 12.0, 8.0, 15.0],
                "Richness": [5, 6, 4, 7],
                "Group": ["Veh", "Veh", "KO", "KO"],
            }
        )

    @patch("matplotlib.pyplot.show")
    def test_runs_without_error(self, mock_show):
        plot_burden_and_richness(
            self.analysis_df,
            condition="Group",
            burden="Burden",
            miR_count="Richness",
        )
        assert mock_show.called

    @patch("matplotlib.pyplot.show")
    def test_raises_when_group_column_missing(self, mock_show):
        with pytest.raises((ValueError, KeyError)):
            plot_burden_and_richness(
                self.analysis_df,
                condition="NonExistentColumn",
                burden="Burden",
                miR_count="Richness",
            )
