"""
Tests for the Data Preprocessing and Quality Control section of
All_together_analysis_w_explanation.ipynb.

Covers functions from:
 - qc_functions.py  : filter_mirnas, run_pca_and_plot_samples,
                       get_correlation_matrix, run_outlier_detection
 - pipeline_utils.py: prepare_core_data, get_elbow_threshold,
                       calculate_group_consistency,
                       run_deseq2_pipeline (TODO), get_dds (TODO),
                       get_DEA (TODO)
"""

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 qc_functions import (
    filter_mirnas,
    run_pca_and_plot_samples,
    get_correlation_matrix,
    run_outlier_detection,
)
from pipeline_utils import (
    prepare_core_data,
    get_elbow_threshold,
    calculate_group_consistency,
)


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


# ===========================================================================
# filter_mirnas
# ===========================================================================


class TestFilterMirnas:
    def setup_method(self):
        self.cpm_df = make_cpm_df()

    def test_returns_dataframe(self):
        result = filter_mirnas(self.cpm_df, threshold=10, min_samples=2)
        assert isinstance(result, pd.DataFrame)

    def test_columns_preserved(self):
        result = filter_mirnas(self.cpm_df, threshold=10, min_samples=2)
        assert list(result.columns) == list(self.cpm_df.columns)

    def test_high_threshold_removes_all(self):
        result = filter_mirnas(self.cpm_df, threshold=1000, min_samples=1)
        assert len(result) == 0

    def test_zero_threshold_keeps_all(self):
        result = filter_mirnas(self.cpm_df, threshold=0, min_samples=1)
        assert len(result) == len(self.cpm_df)

    def test_min_samples_filter_works(self):
        # miR-1 has values [15, 18, 0.5, 1], only 2 above 10
        # With min_samples=3 it should be excluded; with min_samples=2 included
        result_strict = filter_mirnas(self.cpm_df, threshold=10, min_samples=3)
        result_loose = filter_mirnas(self.cpm_df, threshold=10, min_samples=2)
        assert len(result_loose) >= len(result_strict)

    def test_returned_mirnas_meet_threshold(self):
        threshold = 10
        min_samples = 2
        result = filter_mirnas(self.cpm_df, threshold=threshold, min_samples=min_samples)
        for idx in result.index:
            count = (self.cpm_df.loc[idx] >= threshold).sum()
            assert count >= min_samples


# ===========================================================================
# run_pca_and_plot_samples
# ===========================================================================


class TestRunPcaAndPlotSamples:
    """run_pca_and_plot_samples only produces a plot (no return value).

    Note: run_pca_and_plot_samples uses integer loop variable i to index
    pca_df columns (``pca_df.PC1[i]``).  This only works when the DataFrame
    columns – and therefore pca_df's index – are also integers.  The test
    data below deliberately uses integer column names so that label-based
    Series indexing ``series[0]`` resolves correctly.
    """

    def setup_method(self):
        np.random.seed(0)
        # Column names must be integers so that pca_df.PC1[i] works inside
        # run_pca_and_plot_samples (label-based indexing with integer i).
        self.raw_counts_df = pd.DataFrame(
            np.random.randint(1, 200, size=(20, 4)),
            columns=[0, 1, 2, 3],
            index=[f"miR-{i}" for i in range(20)],
        )
        self.samples_df = pd.DataFrame(
            {
                "Sample": [0, 1, 2, 3],
                "Individual": ["M1", "M2", "M3", "M4"],
                "IP-Nat": ["IP", "IP", "NAT", "NAT"],
                "Treatment": ["Veh", "KO", "Veh", "KO"],
            }
        )

    @patch("matplotlib.pyplot.show")
    def test_runs_without_error(self, mock_show):
        run_pca_and_plot_samples(self.raw_counts_df, self.samples_df, group="IP-Nat")
        assert mock_show.called


# ===========================================================================
# get_correlation_matrix
# ===========================================================================


class TestGetCorrelationMatrix:
    def setup_method(self):
        np.random.seed(3)
        self.raw_counts_df = pd.DataFrame(
            np.random.randint(1, 500, size=(30, 4)),
            columns=["S1", "S2", "S3", "S4"],
            index=[f"miR-{i}" for i in range(30)],
        )

    @patch("matplotlib.pyplot.show")
    def test_returns_dataframe(self, mock_show):
        result = get_correlation_matrix(self.raw_counts_df)
        assert isinstance(result, pd.DataFrame)

    @patch("matplotlib.pyplot.show")
    def test_matrix_is_square(self, mock_show):
        result = get_correlation_matrix(self.raw_counts_df)
        assert result.shape[0] == result.shape[1]

    @patch("matplotlib.pyplot.show")
    def test_diagonal_is_one(self, mock_show):
        result = get_correlation_matrix(self.raw_counts_df)
        for col in result.columns:
            assert abs(result.loc[col, col] - 1.0) < 1e-9

    @patch("matplotlib.pyplot.show")
    def test_values_between_minus_one_and_one(self, mock_show):
        result = get_correlation_matrix(self.raw_counts_df)
        assert (result >= -1.0).all().all()
        assert (result <= 1.0).all().all()


# ===========================================================================
# run_outlier_detection
# ===========================================================================


class TestRunOutlierDetection:
    def setup_method(self):
        np.random.seed(4)
        self.cpm_df = pd.DataFrame(
            np.random.randint(1, 300, size=(30, 5)),
            columns=["S1", "S2", "S3", "S4", "S5"],
            index=[f"miR-{i}" for i in range(30)],
        )

    def test_returns_dataframe(self):
        result = run_outlier_detection(self.cpm_df, ["S1", "S2", "S3", "S4", "S5"])
        assert isinstance(result, pd.DataFrame)

    def test_output_has_expected_columns(self):
        result = run_outlier_detection(self.cpm_df, ["S1", "S2", "S3", "S4", "S5"])
        for col in ["Euclidean_Distance", "Z_Score", "Is_Outlier"]:
            assert col in result.columns

    def test_index_is_sample_names(self):
        result = run_outlier_detection(self.cpm_df, ["S1", "S2", "S3"])
        assert list(result.index) == ["S1", "S2", "S3"]

    def test_is_outlier_is_boolean(self):
        result = run_outlier_detection(self.cpm_df, ["S1", "S2", "S3", "S4", "S5"])
        assert result["Is_Outlier"].dtype == bool

    def test_extreme_outlier_flagged(self):
        # Make S5 a clear outlier with very different expression
        self.cpm_df["S5"] = 100000
        result = run_outlier_detection(self.cpm_df, ["S1", "S2", "S3", "S4", "S5"], threshold=1.0)
        assert result.loc["S5", "Is_Outlier"]


# ===========================================================================
# prepare_core_data
# ===========================================================================


class TestPrepareCorData:
    """prepare_core_data loads and standardises three CSV files."""

    def test_standardises_column_names(self, tmp_path):
        # Write minimal CPM CSV
        cpm_csv = tmp_path / "cpm.csv"
        cpm_csv.write_text("Name,S1 - CPM,S2 - CPM\nmiR-1,5.0,6.0\nmiR-2,7.0,8.0\n")

        # Write minimal raw-counts CSV
        counts_csv = tmp_path / "counts.csv"
        counts_csv.write_text(
            "Name,Chromosome,Region,Identifier,S1 - Total counts,S2 - Total counts\n"
            "miR-1,chr1,reg1,id1,100,110\n"
            "miR-2,chr1,reg2,id2,200,210\n"
        )

        # Write minimal metadata CSV
        meta_csv = tmp_path / "meta.csv"
        meta_csv.write_text("Sample,Group\nS1,A\nS2,B\n")

        cpm_df, raw_df, meta_df = prepare_core_data(
            str(cpm_csv), str(counts_csv), str(meta_csv)
        )

        assert isinstance(cpm_df, pd.DataFrame)
        assert isinstance(raw_df, pd.DataFrame)
        assert isinstance(meta_df, pd.DataFrame)

        # ' - CPM' suffix should have been stripped from CPM columns
        assert "S1" in cpm_df.columns
        assert "S2" in cpm_df.columns

        # ' - Total counts' suffix and auxiliary columns should be gone
        assert "S1" in raw_df.columns
        assert "Chromosome" not in raw_df.columns

        # CPM index should be the miRNA names
        assert cpm_df.index.name == "miRNA"


# ===========================================================================
# get_elbow_threshold
# ===========================================================================


class TestGetElbowThreshold:
    """get_elbow_threshold uses the kneed library to find a signal/noise cutoff."""

    def test_returns_numeric_or_none(self):
        """The function returns knee_y (float) or None if no knee is detected."""
        np.random.seed(42)
        high_values = np.linspace(100, 10, num=20)
        low_values = np.ones(30) * 0.5
        values = np.concatenate([high_values, low_values])
        cpm_df = pd.DataFrame(
            {f"S{i}": values + np.random.rand(50) * 0.1 for i in range(3)},
            index=[f"miR-{i}" for i in range(50)],
        )
        result = get_elbow_threshold(cpm_df, plot=False)
        # knee_y is either a float (knee found) or None (no knee detectable)
        assert result is None or isinstance(result, float)

    def test_returns_positive_value_when_knee_detected(self):
        """When a knee is found, the threshold should be a positive number."""
        # Sharp step: 5 high values followed by near-zero flat region
        step = [100.0, 80.0, 60.0, 40.0, 20.0] + [0.1] * 45
        cpm_df = pd.DataFrame(
            {"S1": step, "S2": step},
            index=[f"miR-{i}" for i in range(50)],
        )
        result = get_elbow_threshold(cpm_df, plot=False)
        if result is not None:
            assert result > 0


# ===========================================================================
# calculate_group_consistency
# ===========================================================================


class TestCalculateGroupConsistency:
    def setup_method(self):
        np.random.seed(5)
        self.cpm_df = pd.DataFrame(
            np.random.randint(1, 200, size=(20, 3)),
            columns=["S1", "S2", "S3"],
            index=[f"miR-{i}" for i in range(20)],
        )

    def test_returns_float(self):
        result = calculate_group_consistency(self.cpm_df, ["S1", "S2", "S3"])
        assert isinstance(result, float)

    def test_single_sample_gives_one(self):
        result = calculate_group_consistency(self.cpm_df, ["S1"])
        assert result == 1.0

    def test_correlation_between_minus_one_and_one(self):
        result = calculate_group_consistency(self.cpm_df, ["S1", "S2", "S3"])
        assert -1.0 <= result <= 1.0

    def test_identical_samples_give_one(self):
        self.cpm_df["S2"] = self.cpm_df["S1"]
        result = calculate_group_consistency(self.cpm_df, ["S1", "S2"])
        assert abs(result - 1.0) < 1e-9


# ===========================================================================
# run_deseq2_pipeline / get_dds / get_DEA  (TODO — require pydeseq2/anndata)
# ===========================================================================


class TestRunDeseq2Pipeline:
    """TODO: run_deseq2_pipeline depends on pydeseq2 and anndata.
    Tests are omitted here because these heavy dependencies may not be
    available in all CI environments.
    """

    def test_placeholder(self):
        # TODO: create a minimal integer counts DataFrame and matching group_labels
        # and verify the function returns a DataFrame with expected columns
        pass


class TestGetDds:
    """TODO: get_dds depends on pydeseq2 and anndata."""

    def test_placeholder(self):
        # TODO: build a minimal AnnData object and verify a DESeq2 dataset is returned
        pass


class TestGetDea:
    """TODO: get_DEA depends on a DESeq2 dataset object (dds)."""

    def test_placeholder(self):
        # TODO: use a pre-built dds fixture once DESeq2 test infrastructure exists
        pass
