"""Utilities for comparing MEGqc metrics across multiple samples.

This script expects the TSV files generated by the MEGqc pipeline for each
sample and produces a set of violin plots and a linear regression table
comparing the samples.

Example
-------
Assuming you have the metrics exported for two datasets, run::

    python -m meg_qc.calculation.between_sample_analysis \
        --tsv sample1/group_metrics.tsv sample2/group_metrics.tsv \
        --names sample1 sample2 \
        --output-dir results

All figures and the regression table will be stored in ``results``.
"""
import argparse
import os
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
from matplotlib import cm
from scipy import stats
import statsmodels.api as sm
from sklearn.feature_selection import mutual_info_regression
import seaborn as sns
from joblib import Parallel, delayed
from statsmodels.stats.multitest import multipletests

LABEL_MAP = {
    "GQI": "GQI",
    "GQI_std_pct": "STD noise",
    "GQI_ptp_pct": "PtP noise",
    "GQI_ecg_pct": "ECG noise",
    "GQI_eog_pct": "EOG noise",
    "GQI_muscle_pct": "Muscle noise",
    "GQI_psd_noise_pct": "PSD noise",
    "GQI_penalty_ch": "Variability penalty",
    "GQI_penalty_corr": "Correlational penalty",
    "GQI_penalty_mus": "Muscle penalty",
    "GQI_penalty_psd": "PSD penalty",
}


def _get_star(p):
    """Return asterisks representing the p-value significance."""
    if p < 0.001:
        return "***"
    if p < 0.01:
        return "**"
    if p < 0.05:
        return "*"
    return ""


def _load_tables(paths):
    """Load TSV tables into pandas DataFrames."""
    tables = []
    for p in paths:
        df = pd.read_csv(p, sep="\t")
        tables.append(df)
    return tables


def _plot_heatmap(matrix, title, out_png, mask=None):
    """Plot and save a heatmap for the given matrix."""
    plt.figure(figsize=(8, 6))
    sns.heatmap(
        matrix.astype(float), annot=True, fmt=".2f", cmap="viridis", mask=mask
    )
    plt.title(title)
    plt.tight_layout()
    plt.savefig(out_png, dpi=300)
    plt.close()


def estimate_entropy(x, bins="fd"):
    """Estimate the entropy of a 1D array using histogram binning."""
    hist, _ = np.histogram(x, bins=bins, density=False)
    probs = hist / hist.sum()
    probs = probs[probs > 0]
    return stats.entropy(probs, base=2)


def _mi_symmetric(x, y, n_neighbors):
    """Compute symmetric mutual information between two 1D arrays."""
    x = np.asarray(x).ravel()
    y = np.asarray(y).ravel()
    mi_xy = mutual_info_regression(
        x.reshape(-1, 1), y, discrete_features=False, n_neighbors=n_neighbors
    )[0]
    mi_yx = mutual_info_regression(
        y.reshape(-1, 1), x, discrete_features=False, n_neighbors=n_neighbors
    )[0]
    return 0.5 * (mi_xy + mi_yx)


def _make_violin(data, names, title, ylabel, out_png):
    """Create violin plot comparing samples."""
    violin_data = [d.dropna() for d in data]
    palette = cm.get_cmap("tab10", len(violin_data))
    plt.figure(figsize=(10, 6))
    parts = plt.violinplot(violin_data, showmeans=True, showextrema=True,
                            showmedians=False, widths=0.8)
    for i, pc in enumerate(parts["bodies"]):
        pc.set_facecolor(palette(i))
        pc.set_edgecolor("black")
        pc.set_alpha(0.5)
    parts["cmeans"].set_linewidth(2)
    parts["cmeans"].set_color("black")
    parts["cbars"].set_color("black")
    for i, y in enumerate(violin_data, start=1):
        x = np.random.normal(i, 0.05, size=len(y))
        plt.scatter(x, y, s=20, alpha=0.4, edgecolor="black", linewidth=0.5,
                    facecolor=palette(i - 1))
    plt.xticks(range(1, len(names) + 1), names, rotation=0, ha="center", fontsize=10)
    plt.ylabel(ylabel)
    plt.title(title)
    plt.tight_layout()
    plt.savefig(out_png, dpi=300)
    plt.close()


def _cumulative_plot(tables, metrics, names, out_png, out_tsv=None, compute_ttest=True):
    """Create a cumulative violin plot of several metrics for each sample.

    Parameters
    ----------
    tables : list of pandas.DataFrame
        Tables with metrics for each sample.
    metrics : list of str
        Metrics to include in the plot.
    names : list of str
        Sample names.
    out_png : str
        Path where the figure will be saved.
    out_tsv : str or None
        If provided and ``compute_ttest`` is True, pairwise t-test results
        are written to this file.
    compute_ttest : bool, optional
        Whether to compute pairwise t-tests and annotate the plot.
    """
    n_samples = len(names)
    fig, ax = plt.subplots(figsize=(12, 6))
    palette = cm.get_cmap("tab10", len(metrics))

    positions = []
    labels = []
    results = []
    for j, metric in enumerate(metrics):
        metric_values = []
        for i, name in enumerate(names):
            data = tables[i][metric].dropna()
            pos = j * n_samples + i + 1
            parts = ax.violinplot(
                [data],
                positions=[pos],
                showmeans=True,
                showextrema=True,
                showmedians=False,
                widths=0.8,
            )
            color = palette(j)
            for pc in parts["bodies"]:
                pc.set_facecolor(color)
                pc.set_edgecolor("black")
                pc.set_alpha(0.5)
            parts["cmeans"].set_linewidth(2)
            parts["cmeans"].set_color("black")
            parts["cbars"].set_color("black")
            x = np.random.normal(pos, 0.05, size=len(data))
            ax.scatter(
                x,
                data,
                s=20,
                alpha=0.4,
                edgecolor="black",
                linewidth=0.5,
                facecolor=color,
            )
            metric_values.append(data)
            positions.append(pos)
            labels.append(name)

        if compute_ttest:
            # add significance annotations
            y_range = ax.get_ylim()
            y_offset = (y_range[1] - y_range[0]) * 0.05
            offset_count = 0
            for a in range(n_samples):
                for b in range(a + 1, n_samples):
                    stat, p = stats.ttest_ind(
                        metric_values[a],
                        metric_values[b],
                        equal_var=False,
                    )
                    star = _get_star(p)
                    results.append(
                        {
                            "metric": metric,
                            "sample1": names[a],
                            "sample2": names[b],
                            "t_stat": stat,
                            "p": p,
                            "asterisk": star,
                        }
                    )
                    if not star:
                        continue
                    x1 = j * n_samples + a + 1
                    x2 = j * n_samples + b + 1
                    y = max(metric_values[a].max(), metric_values[b].max()) + (
                        offset_count + 1
                    ) * y_offset
                    ax.plot(
                        [x1, x1, x2, x2],
                        [y, y + y_offset / 2, y + y_offset / 2, y],
                        color="black",
                    )
                    ax.text((x1 + x2) / 2, y + y_offset / 2, star, ha="center", va="bottom")
                    offset_count += 1

    ax.set_xticks(positions)
    ax.set_xticklabels(labels, rotation=0, ha="center", fontsize=9)
    ax.set_ylabel("Percentage of Quality and Noisy Channels")
    ax.set_title("Metrics across samples")

    secax = ax.secondary_xaxis(
        "bottom",
        functions=(lambda x: x, lambda x: x),
    )
    centers = [j * n_samples + (n_samples + 1) / 2 for j in range(len(metrics))]
    secax.set_xticks(centers)
    secax.set_xticklabels([LABEL_MAP.get(m, m) for m in metrics], fontsize=10, fontweight="bold")
    secax.tick_params(pad=20)

    fig.tight_layout()
    fig.savefig(out_png, dpi=300)
    if out_tsv is not None and compute_ttest:
        pd.DataFrame(results).to_csv(out_tsv, sep="\t", index=False)
    plt.close(fig)


def _perform_regression(df, metrics, out_tsv):
    """Run linear regression and save results."""
    X = df[metrics].copy()
    # add all pairwise interactions
    for i, m1 in enumerate(metrics):
        for m2 in metrics[i + 1:]:
            X[f"{m1}:{m2}"] = df[m1] * df[m2]
    X = sm.add_constant(X)
    model = sm.OLS(df["GQI"], X, missing="drop").fit()
    res_df = pd.DataFrame({
        "variable": model.params.index,
        "coef": model.params.values,
        "std_err": model.bse.values,
        "t": model.tvalues.values,
        "p": model.pvalues.values,
    })
    res_df["asterisk"] = res_df["p"].apply(_get_star)
    res_df.to_csv(out_tsv, sep="\t", index=False)
    return model, res_df

def analyze_mutual_information(
    df,
    metrics,
    n_permutations=1000,
    permutation_test=False,
    seed=None,
    save_dir=None,
    n_jobs=1,
    redundancy_threshold=None,
    n_neighbors=None,
):
    """Compute mutual information and related statistics between metrics.

    Parameters
    ----------
    df : pandas.DataFrame
        Data frame containing the metrics.
    metrics : list of str
        Metrics to include in the analysis.
    n_permutations : int, optional
        Number of permutations for significance testing.
    permutation_test : bool, optional
        Whether to run the permutation-based significance test.
    seed : int or None, optional
        Random seed for reproducibility.
    save_dir : str or None, optional
        Directory where results will be saved. If ``None`` results are not
        written to disk.
    n_jobs : int, optional
        Number of parallel jobs for the permutation test.
    redundancy_threshold : float or None, optional
        Threshold for reporting redundant metric pairs based on NMI.
    n_neighbors : int or None, optional
        Number of neighbors for k-NN MI estimation. If ``None``, uses
        ``max(3, int(np.sqrt(n_samples)))``.

    Returns
    -------
    dict
        Dictionary containing the computed matrices and entropy vector.
    """
    data = df[metrics].dropna()
    if n_neighbors is None:
        n_neighbors = max(3, int(np.sqrt(len(data))))
    n_neighbors = min(n_neighbors, len(data) - 1)

    mi_matrix = pd.DataFrame(index=metrics, columns=metrics, dtype=float)
    net_mi_matrix = pd.DataFrame(index=metrics, columns=metrics, dtype=float)
    z_matrix = pd.DataFrame(index=metrics, columns=metrics, dtype=float)
    p_matrix = pd.DataFrame(index=metrics, columns=metrics, dtype=float)
    nmi_geo_matrix = pd.DataFrame(index=metrics, columns=metrics, dtype=float)
    nmi_ari_matrix = pd.DataFrame(index=metrics, columns=metrics, dtype=float)

    # Marginal entropies
    entropy_series = pd.Series({m: estimate_entropy(data[m]) for m in metrics})

    rng = np.random.default_rng(seed)

    for i, m1 in enumerate(metrics):
        x = data[m1].values
        for j, m2 in enumerate(metrics):
            if j < i:
                mi_matrix.loc[m1, m2] = mi_matrix.loc[m2, m1]
                net_mi_matrix.loc[m1, m2] = net_mi_matrix.loc[m2, m1]
                z_matrix.loc[m1, m2] = z_matrix.loc[m2, m1]
                p_matrix.loc[m1, m2] = p_matrix.loc[m2, m1]
                nmi_geo_matrix.loc[m1, m2] = nmi_geo_matrix.loc[m2, m1]
                nmi_ari_matrix.loc[m1, m2] = nmi_ari_matrix.loc[m2, m1]
                continue
            if i == j:
                mi_matrix.loc[m1, m2] = np.nan
                net_mi_matrix.loc[m1, m2] = np.nan
                z_matrix.loc[m1, m2] = np.nan
                p_matrix.loc[m1, m2] = np.nan
                nmi_geo_matrix.loc[m1, m2] = np.nan
                nmi_ari_matrix.loc[m1, m2] = np.nan
                continue

            y = data[m2].values
            mi = _mi_symmetric(x, y, n_neighbors)
            if permutation_test and n_permutations > 0:
                seeds = rng.integers(0, np.iinfo(np.int32).max, size=n_permutations)

                def _perm_mi(s):
                    perm = np.random.default_rng(s).permutation(y)
                    return _mi_symmetric(x, perm, n_neighbors)

                perm_mi = Parallel(n_jobs=n_jobs)(delayed(_perm_mi)(s) for s in seeds)
                perm_mi = np.asarray(perm_mi)
                perm_mean = perm_mi.mean()
                perm_std = perm_mi.std(ddof=1)
                net = mi - perm_mean
                z = net / perm_std if perm_std > 0 else np.nan
                p = (np.sum(perm_mi >= mi) + 1) / (n_permutations + 1)
            else:
                net = z = p = np.nan

            mi_matrix.loc[m1, m2] = mi_matrix.loc[m2, m1] = mi
            net_mi_matrix.loc[m1, m2] = net_mi_matrix.loc[m2, m1] = net
            z_matrix.loc[m1, m2] = z_matrix.loc[m2, m1] = z
            p_matrix.loc[m1, m2] = p_matrix.loc[m2, m1] = p

            hx = entropy_series[m1]
            hy = entropy_series[m2]
            nmi_geo = mi / np.sqrt(hx * hy) if hx > 0 and hy > 0 else np.nan
            nmi_ari = (2 * mi) / (hx + hy) if (hx + hy) > 0 else np.nan
            nmi_geo_matrix.loc[m1, m2] = nmi_geo_matrix.loc[m2, m1] = nmi_geo
            nmi_ari_matrix.loc[m1, m2] = nmi_ari_matrix.loc[m2, m1] = nmi_ari

    if permutation_test and n_permutations > 0:
        tri_upper = np.triu_indices_from(p_matrix, k=1)
        flat_p = p_matrix.values[tri_upper]
        _, p_fdr_flat, _, _ = multipletests(flat_p, method="fdr_bh")
        p_fdr_matrix = pd.DataFrame(
            np.nan, index=metrics, columns=metrics, dtype=float
        )
        p_fdr_matrix.values[tri_upper] = p_fdr_flat
        p_fdr_matrix.values[(tri_upper[1], tri_upper[0])] = p_fdr_flat
    else:
        p_fdr_matrix = pd.DataFrame(np.nan, index=metrics, columns=metrics)

    results = {
        "mi": mi_matrix,
        "net_mi": net_mi_matrix,
        "z": z_matrix,
        "p": p_matrix,
        "p_fdr": p_fdr_matrix,
        "nmi_geo": nmi_geo_matrix,
        "nmi_ari": nmi_ari_matrix,
        "entropy": entropy_series,
    }

    if redundancy_threshold is not None:
        redundant = []
        for i, m1 in enumerate(metrics):
            for m2 in metrics[i + 1:]:
                nmi_val = nmi_ari_matrix.loc[m1, m2]
                if nmi_val > redundancy_threshold:
                    redundant.append((m1, m2, nmi_val))
        results["redundancy"] = pd.DataFrame(
            redundant, columns=["metric1", "metric2", "nmi_ari"]
        )

    if save_dir is not None:
        os.makedirs(save_dir, exist_ok=True)
        mask = np.eye(len(metrics), dtype=bool)
        mi_matrix.to_csv(os.path.join(save_dir, "mutual_information.tsv"), sep="\t")
        _plot_heatmap(
            mi_matrix,
            "Mutual Information between noise metrics",
            os.path.join(save_dir, "mutual_information.png"),
            mask=mask,
        )
        net_mi_matrix.to_csv(
            os.path.join(save_dir, "net_mutual_information.tsv"), sep="\t"
        )
        _plot_heatmap(
            net_mi_matrix,
            "Net Mutual Information",
            os.path.join(save_dir, "net_mutual_information.png"),
            mask=mask,
        )
        z_matrix.to_csv(os.path.join(save_dir, "mi_z_scores.tsv"), sep="\t")
        _plot_heatmap(
            z_matrix,
            "MI z-scores",
            os.path.join(save_dir, "mi_z_scores.png"),
            mask=mask,
        )
        p_matrix.to_csv(os.path.join(save_dir, "mi_p_values.tsv"), sep="\t")
        _plot_heatmap(
            p_matrix,
            "MI p-values",
            os.path.join(save_dir, "mi_p_values.png"),
            mask=mask,
        )
        p_fdr_matrix.to_csv(
            os.path.join(save_dir, "mi_p_values_fdr.tsv"), sep="\t"
        )
        _plot_heatmap(
            p_fdr_matrix,
            "MI p-values (FDR)",
            os.path.join(save_dir, "mi_p_values_fdr.png"),
            mask=mask,
        )
        nmi_geo_matrix.to_csv(
            os.path.join(save_dir, "nmi_geometric.tsv"), sep="\t"
        )
        _plot_heatmap(
            nmi_geo_matrix,
            "Normalized MI (geometric)",
            os.path.join(save_dir, "nmi_geometric.png"),
            mask=mask,
        )
        nmi_ari_matrix.to_csv(
            os.path.join(save_dir, "nmi_arithmetic.tsv"), sep="\t"
        )
        _plot_heatmap(
            nmi_ari_matrix,
            "Normalized MI (arithmetic)",
            os.path.join(save_dir, "nmi_arithmetic.png"),
            mask=mask,
        )
        entropy_series.to_csv(
            os.path.join(save_dir, "marginal_entropies.tsv"),
            sep="\t",
            header=["entropy"],
        )
        if "redundancy" in results:
            results["redundancy"].to_csv(
                os.path.join(save_dir, "redundant_pairs.tsv"),
                sep="\t",
                index=False,
            )

    return results


def _mutual_information(df, metrics, out_png, out_tsv):
    """Wrapper for backward compatibility with previous API."""
    res = analyze_mutual_information(df, metrics)
    mi_matrix = res["mi"]
    mi_matrix.to_csv(out_tsv, sep="\t")
    mask = np.eye(len(mi_matrix), dtype=bool)
    _plot_heatmap(
        mi_matrix,
        "Mutual Information between noise metrics",
        out_png,
        mask=mask,
    )




def main():
    parser = argparse.ArgumentParser(description="Between sample analysis")
    parser.add_argument("--tsv", nargs="+", required=True,
                        help="Paths to group metrics TSV files")
    parser.add_argument("--names", nargs="+", required=True,
                        help="Sample names corresponding to TSV files")
    parser.add_argument("--output-dir", required=True, help="Directory for outputs")
    parser.add_argument("--ttest", action="store_true", help="Compute pairwise t-tests for the cumulative plot")
    parser.add_argument("--mi", action="store_true", help="Perform mutual information analysis")
    parser.add_argument(
        "--mi-permutations",
        type=int,
        default=0,
        help="Number of permutations for MI significance testing (0 disables)",
    )
    parser.add_argument("--mi-seed", type=int, default=None, help="Seed for MI permutations")
    parser.add_argument(
        "--mi-threshold",
        type=float,
        default=None,
        help="Threshold for redundancy reporting based on NMI",
    )
    parser.add_argument(
        "--mi-n-jobs",
        type=int,
        default=1,
        help="Number of parallel jobs for MI permutations",
    )
    args = parser.parse_args()

    if len(args.tsv) != len(args.names):
        raise ValueError("Number of TSV files must match number of names")

    os.makedirs(args.output_dir, exist_ok=True)
    cumulative_dir = os.path.join(args.output_dir, "cummlative_violin_plot")
    separated_dir = os.path.join(args.output_dir, "separated_violin_plot")
    regression_dir = os.path.join(args.output_dir, "regression")
    mi_dir = os.path.join(args.output_dir, "mutual_information")
    os.makedirs(cumulative_dir, exist_ok=True)
    os.makedirs(separated_dir, exist_ok=True)
    os.makedirs(regression_dir, exist_ok=True)
    if args.mi:
        os.makedirs(mi_dir, exist_ok=True)

    tables = _load_tables(args.tsv)
    metrics = [
        "GQI_std_pct",
        "GQI_ptp_pct",
        "GQI_ecg_pct",
        "GQI_eog_pct",
        "GQI_muscle_pct",
        "GQI_psd_noise_pct",
    ]

    # Violin plot for GQI
    gqi_data = [tbl["GQI"] for tbl in tables]
    _make_violin(
        gqi_data,
        args.names,
        "Global Quality Index",
        LABEL_MAP["GQI"],
        os.path.join(separated_dir, "GQI_violin.png"),
    )

    # Metric specific violin plots
    for m in metrics:
        metric_data = [t[m] for t in tables if m in t.columns]
        if not metric_data:
            continue
        _make_violin(
            metric_data,
            args.names,
            LABEL_MAP.get(m, m),
            LABEL_MAP.get(m, m),
            os.path.join(separated_dir, f"{m}_violin.png"),
        )

    _cumulative_plot(
        tables,
        ["GQI", *metrics],
        args.names,
        os.path.join(cumulative_dir, "cumulative_metrics.png"),
        os.path.join(cumulative_dir, "t_test_results.tsv") if args.ttest else None,
        compute_ttest=args.ttest,
    )

    # Combine all data for regression
    df_all = []
    for name, tbl in zip(args.names, tables):
        temp = tbl.copy()
        temp["sample"] = name
        df_all.append(temp)
    df_all = pd.concat(df_all, ignore_index=True)

    reg_model, res_df = _perform_regression(
        df_all, metrics, os.path.join(regression_dir, "linear_regression_results.tsv")
    )

    if args.mi:
        for name, tbl in zip(args.names, tables):
            analyze_mutual_information(
                tbl,
                metrics,
                n_permutations=args.mi_permutations,
                permutation_test=args.mi_permutations > 0,
                seed=args.mi_seed,
                save_dir=os.path.join(mi_dir, name),
                n_jobs=args.mi_n_jobs,
                redundancy_threshold=args.mi_threshold,
            )
        analyze_mutual_information(
            df_all,
            metrics,
            n_permutations=args.mi_permutations,
            permutation_test=args.mi_permutations > 0,
            seed=args.mi_seed,
            save_dir=os.path.join(mi_dir, "combined"),
            n_jobs=args.mi_n_jobs,
            redundancy_threshold=args.mi_threshold,
        )


if __name__ == "__main__":
    main()
