"""
Aggregate RPV vs quality metrics across multiple datasets.

Run bombcell (Llobet 2ms + distance metrics + drift + raw waveforms) on each
dataset, pool all high-FR units, and analyse which metrics correlate with RPV
contamination. Suggests thresholds at RPV cutoffs of 10%, 15%, 20%.

Usage:
    python BC_aggregate_rpv_vs_metrics.py

Edit the DATASETS list below to point at your recordings.
"""

import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from pathlib import Path
from scipy import stats

import bombcell as bc
from bombcell.quality_metrics import get_quality_unit_type

# ──────────────────────────────────────────────────────────────────────
# EDIT HERE: one entry per dataset
# ──────────────────────────────────────────────────────────────────────
DATASETS = [
    {
        "ks_dir": "/path/to/dataset1/kilosort4",
        "raw_file": None,      # path to .ap.bin (None if unavailable)
        "meta_file": None,     # path to .meta or .oebin (None to auto-detect)
        "ks_version": 4,
    },
    {
        "ks_dir": "/path/to/dataset2/kilosort4",
        "raw_file": None,
        "meta_file": None,
        "ks_version": 4,
    },
    # add more entries …
]

FR_THRESHOLD = 10          # Hz – high-FR cutoff
RPV_CUTOFFS = [0.1, 0.15, 0.2]
FORCE_RERUN = False        # set True to recompute even if cached results exist
OUTPUT_DIR = Path("aggregate_rpv_vs_metrics")

# ──────────────────────────────────────────────────────────────────────
# Helpers
# ──────────────────────────────────────────────────────────────────────

METRIC_COLS = [
    "isolation_distance", "l_ratio", "amplitude",
    "snr", "max_drift", "presence_ratio",
]
METRIC_LABELS = {
    "isolation_distance": "Isolation distance",
    "l_ratio": "L-ratio",
    "amplitude": "Raw amplitude (uV)",
    "snr": "SNR",
    "max_drift": "Max drift (um)",
    "presence_ratio": "Presence ratio",
}
HIGHER_IS_CLEANER = {"isolation_distance", "amplitude", "snr", "presence_ratio"}


def run_dataset(ds, force_rerun=False):
    """Run bombcell on one dataset, return per-unit DataFrame."""
    ks_dir = ds["ks_dir"]
    save_path = Path(ks_dir) / "bombcell_llobet_2ms_full"
    qm_parquet = save_path / "templates._bc_qMetrics.parquet"
    param_parquet = save_path / "_bc_parameters._bc_qMetrics.parquet"

    # ── Try loading cached results ──
    if qm_parquet.exists() and param_parquet.exists() and not force_rerun:
        print(f"  Loading cached: {ks_dir}")
        df = pd.read_parquet(qm_parquet)
        param_df = pd.read_parquet(param_parquet)
        param_out = param_df.iloc[0].to_dict()
        qm = {col: df[col].values for col in df.columns}
        ut, ut_str = get_quality_unit_type(param_out, qm)
    else:
        print(f"  Running bombcell: {ks_dir}")
        param = bc.get_default_parameters(
            ks_dir,
            raw_file=ds.get("raw_file"),
            meta_file=ds.get("meta_file"),
            kilosort_version=ds.get("ks_version", 4),
        )
        param.update({
            "rpvMethod": "llobet",
            "tauR_valuesMin": 2 / 1000,
            "tauR_valuesMax": 2 / 1000,
            "computeDistanceMetrics": True,
            "computeDrift": True,
            "extractRaw": True,
            "plotGlobal": False,
            "verbose": True,
            "computeDCISI": False,
        })
        qm, param_out, ut, ut_str = bc.run_bombcell(ks_dir, save_path, param)

    # ── Compute firing rate ──
    fs = param_out["ephys_sample_rate"]
    (spike_times_samples, spike_clusters,
     *_rest) = bc.load_ephys_data(ks_dir)
    spike_times_sec = spike_times_samples / fs

    cluster_ids = np.array(qm["phy_clusterID"]).flatten()
    n_units = len(cluster_ids)
    firing_rate = np.full(n_units, np.nan)

    t_starts = np.array(qm["useTheseTimesStart"]).flatten()
    t_stops = np.array(qm["useTheseTimesStop"]).flatten()

    for i, uid in enumerate(cluster_ids):
        t0, t1 = t_starts[i], t_stops[i]
        if np.isnan(t0) or np.isnan(t1) or t1 <= t0:
            continue
        mask = spike_clusters == uid
        spks = spike_times_sec[mask]
        spks = spks[(spks >= t0) & (spks < t1)]
        firing_rate[i] = spks.size / (t1 - t0)

    # ── Build per-unit row ──
    def _flat(key):
        v = qm.get(key)
        if v is None:
            return np.full(n_units, np.nan)
        v = np.array(v)
        if v.ndim == 2:
            return np.nanmin(v, axis=1)
        return v.flatten()

    unit_df = pd.DataFrame({
        "dataset": Path(ks_dir).name,
        "cluster_id": cluster_ids,
        "firing_rate": firing_rate,
        "unit_type": ut,
        "rpv": _flat("fractionRPVs_estimatedTauR"),
        "isolation_distance": _flat("isolationDistance"),
        "l_ratio": _flat("Lratio"),
        "amplitude": _flat("rawAmplitude"),
        "snr": _flat("signalToNoiseRatio"),
        "max_drift": _flat("maxDriftEstimate"),
        "presence_ratio": _flat("presenceRatio"),
    })
    return unit_df


# ──────────────────────────────────────────────────────────────────────
# 1. Run all datasets and aggregate
# ──────────────────────────────────────────────────────────────────────
print("=" * 60)
print("Processing datasets")
print("=" * 60)

all_dfs = []
for ds in DATASETS:
    try:
        unit_df = run_dataset(ds, force_rerun=FORCE_RERUN)
        all_dfs.append(unit_df)
        print(f"    {len(unit_df)} units")
    except Exception as e:
        print(f"    FAILED: {e}")

df = pd.concat(all_dfs, ignore_index=True)
print(f"\nTotal: {len(df)} units from {df['dataset'].nunique()} datasets")

noise_mask = df["unit_type"] == 0
high_fr = df["firing_rate"] > FR_THRESHOLD
df_hfr = df[high_fr & ~noise_mask].copy()
print(f"High-FR (>{FR_THRESHOLD} Hz) non-noise units: {len(df_hfr)}")

OUTPUT_DIR.mkdir(exist_ok=True)

# ──────────────────────────────────────────────────────────────────────
# 2. Firing rate distributions per dataset
# ──────────────────────────────────────────────────────────────────────
print("\nPlotting firing rate distributions...")

datasets = df["dataset"].unique()
n_ds = len(datasets)
fig, axes = plt.subplots(1, n_ds + 1, figsize=(4 * (n_ds + 1), 4), squeeze=False)
axes = axes[0]

fr_bins = np.logspace(-1, 3, 60)  # 0.1 to 1000 Hz

for i, ds_name in enumerate(datasets):
    ax = axes[i]
    fr = df.loc[(df["dataset"] == ds_name) & ~noise_mask, "firing_rate"].dropna()
    ax.hist(fr, bins=fr_bins, color="#3498db", alpha=0.7, edgecolor="white", linewidth=0.3)
    ax.set_xscale("log")
    ax.axvline(FR_THRESHOLD, color="#e74c3c", ls="--", lw=1.2, label=f"{FR_THRESHOLD} Hz")
    ax.set_xlabel("Firing rate (Hz)")
    ax.set_ylabel("Unit count")
    ax.set_title(f"{ds_name}\n(n={len(fr)}, >{FR_THRESHOLD}Hz: {(fr > FR_THRESHOLD).sum()})")
    ax.legend(fontsize=8)

# Pooled
ax = axes[-1]
fr_all = df.loc[~noise_mask, "firing_rate"].dropna()
ax.hist(fr_all, bins=fr_bins, color="#2c3e50", alpha=0.7, edgecolor="white", linewidth=0.3)
ax.set_xscale("log")
ax.axvline(FR_THRESHOLD, color="#e74c3c", ls="--", lw=1.2, label=f"{FR_THRESHOLD} Hz")
ax.set_xlabel("Firing rate (Hz)")
ax.set_title(f"Pooled\n(n={len(fr_all)}, >{FR_THRESHOLD}Hz: {(fr_all > FR_THRESHOLD).sum()})")
ax.legend(fontsize=8)

plt.suptitle("Firing rate distributions (non-noise units)", fontsize=13)
plt.tight_layout()
plt.savefig(OUTPUT_DIR / "firing_rate_distributions.png", dpi=150, bbox_inches="tight")
plt.show()

# ──────────────────────────────────────────────────────────────────────
# 3. Correlation matrix (high-FR units)
# ──────────────────────────────────────────────────────────────────────
print("\nComputing correlations (high-FR units)...")

corr_rows = []
for met_col in METRIC_COLS:
    valid = df_hfr[["rpv", met_col]].dropna()
    if len(valid) < 10:
        corr_rows.append({"metric": METRIC_LABELS[met_col], "rho": np.nan,
                          "p": np.nan, "n": len(valid)})
        continue
    rho, p = stats.spearmanr(valid["rpv"], valid[met_col])
    corr_rows.append({"metric": METRIC_LABELS[met_col], "rho": rho, "p": p,
                      "n": len(valid)})

corr_df = pd.DataFrame(corr_rows)
print("\nSpearman rho (RPV Llobet 2ms vs metric):")
print(corr_df.to_string(index=False))
corr_df.to_csv(OUTPUT_DIR / "spearman_correlations.csv", index=False)

# Heatmap
fig, ax = plt.subplots(figsize=(5, 4))
rhos = corr_df["rho"].values
im = ax.imshow(rhos.reshape(-1, 1), cmap="RdBu_r", vmin=-1, vmax=1, aspect=0.4)
ax.set_yticks(range(len(METRIC_COLS)))
ax.set_yticklabels([METRIC_LABELS[m] for m in METRIC_COLS])
ax.set_xticks([0])
ax.set_xticklabels(["Llobet 2ms RPV"])
for i, (r, p) in enumerate(zip(corr_df["rho"], corr_df["p"])):
    stars = "***" if p < 0.001 else "**" if p < 0.01 else "*" if p < 0.05 else ""
    color = "white" if abs(r) > 0.5 else "black"
    ax.text(0, i, f"{r:.2f}{stars}", ha="center", va="center", fontsize=11, color=color)
plt.colorbar(im, ax=ax, label="Spearman rho")
ax.set_title(f"RPV vs quality metrics\n(pooled high-FR, n={len(df_hfr)})")
plt.tight_layout()
plt.savefig(OUTPUT_DIR / "correlation_heatmap.png", dpi=150, bbox_inches="tight")
plt.show()

# ──────────────────────────────────────────────────────────────────────
# 4. Scatter plots (high-FR)
# ──────────────────────────────────────────────────────────────────────
print("\nPlotting scatters...")

fig, axes = plt.subplots(2, 3, figsize=(15, 9))
for idx, met_col in enumerate(METRIC_COLS):
    ax = axes[idx // 3, idx % 3]
    valid = df_hfr[["rpv", met_col, "dataset"]].dropna()
    if len(valid) < 5:
        ax.text(0.5, 0.5, "insufficient data", transform=ax.transAxes, ha="center")
        continue

    for ds_name in valid["dataset"].unique():
        sub = valid[valid["dataset"] == ds_name]
        ax.scatter(sub["rpv"], sub[met_col], s=10, alpha=0.4, label=ds_name)

    rho, p = stats.spearmanr(valid["rpv"], valid[met_col])
    stars = "***" if p < 0.001 else "**" if p < 0.01 else "*" if p < 0.05 else ""
    ax.set_title(f"{METRIC_LABELS[met_col]}  rho={rho:.2f}{stars}")
    ax.set_xlabel("RPV contamination (Llobet 2ms)")
    ax.set_ylabel(METRIC_LABELS[met_col])

    if len(valid["dataset"].unique()) <= 8:
        ax.legend(fontsize=6, markerscale=2)

plt.suptitle(f"RPV vs quality metrics (pooled, FR > {FR_THRESHOLD} Hz, n={len(df_hfr)})",
             fontsize=13)
plt.tight_layout()
plt.savefig(OUTPUT_DIR / "scatter_rpv_vs_metrics.png", dpi=150, bbox_inches="tight")
plt.show()

# ──────────────────────────────────────────────────────────────────────
# 5. Threshold suggestions
# ──────────────────────────────────────────────────────────────────────
print("\nComputing threshold suggestions...")

threshold_rows = []
for met_col in METRIC_COLS:
    higher_is_cleaner = met_col in HIGHER_IS_CLEANER
    for rpv_cut in RPV_CUTOFFS:
        sub = df_hfr[["rpv", met_col]].dropna()
        if len(sub) < 10:
            threshold_rows.append({
                "metric": METRIC_LABELS[met_col], "rpv_cutoff": rpv_cut,
                "threshold_median": np.nan, "threshold_logistic": np.nan,
                "accuracy": np.nan, "n_clean": 0, "n_dirty": 0,
            })
            continue

        is_clean = sub["rpv"] <= rpv_cut
        clean_vals = sub.loc[is_clean, met_col].values
        dirty_vals = sub.loc[~is_clean, met_col].values

        # Median-based threshold
        if len(clean_vals) >= 3 and len(dirty_vals) >= 3:
            thresh_median = (np.median(clean_vals) + np.median(dirty_vals)) / 2
        else:
            thresh_median = np.nan

        # Logistic regression threshold
        thresh_logistic = np.nan
        if len(clean_vals) >= 3 and len(dirty_vals) >= 3:
            from sklearn.linear_model import LogisticRegression
            X = sub[met_col].values.reshape(-1, 1)
            y = is_clean.astype(int).values
            lr = LogisticRegression(solver="lbfgs", max_iter=1000)
            lr.fit(X, y)
            if lr.coef_[0, 0] != 0:
                thresh_logistic = -lr.intercept_[0] / lr.coef_[0, 0]

        # Accuracy at median threshold
        if not np.isnan(thresh_median):
            if higher_is_cleaner:
                pred_clean = sub[met_col] >= thresh_median
            else:
                pred_clean = sub[met_col] <= thresh_median
            acc = np.mean(pred_clean == is_clean)
        else:
            acc = np.nan

        threshold_rows.append({
            "metric": METRIC_LABELS[met_col],
            "rpv_cutoff": rpv_cut,
            "threshold_median": thresh_median,
            "threshold_logistic": thresh_logistic,
            "accuracy": acc,
            "n_clean": int(is_clean.sum()),
            "n_dirty": int((~is_clean).sum()),
        })

thresh_df = pd.DataFrame(threshold_rows)
thresh_df.to_csv(OUTPUT_DIR / "suggested_thresholds.csv", index=False)

print(f"\nSuggested thresholds (pooled high-FR units, n={len(df_hfr)}):")
print(f"'Clean' = RPV <= cutoff, 'Dirty' = RPV > cutoff\n")
print(thresh_df.round(3).to_string(index=False))

# ──────────────────────────────────────────────────────────────────────
# 6. Threshold visualisation
# ──────────────────────────────────────────────────────────────────────
fig, axes = plt.subplots(len(METRIC_COLS), len(RPV_CUTOFFS), figsize=(14, 20))

for i, met_col in enumerate(METRIC_COLS):
    for j, rpv_cut in enumerate(RPV_CUTOFFS):
        ax = axes[i, j]
        sub = df_hfr[["rpv", met_col]].dropna()
        if len(sub) < 5:
            ax.text(0.5, 0.5, "insufficient data", transform=ax.transAxes, ha="center")
            continue

        is_clean = sub["rpv"] <= rpv_cut
        ax.scatter(sub.loc[is_clean, "rpv"], sub.loc[is_clean, met_col],
                   s=12, alpha=0.4, color="#2ecc71", label="clean")
        ax.scatter(sub.loc[~is_clean, "rpv"], sub.loc[~is_clean, met_col],
                   s=12, alpha=0.4, color="#e74c3c", label="dirty")

        ax.axvline(rpv_cut, color="gray", ls="--", lw=0.8)

        row = thresh_df[(thresh_df["metric"] == METRIC_LABELS[met_col]) &
                        (thresh_df["rpv_cutoff"] == rpv_cut)]
        if len(row) > 0:
            t_med = row["threshold_median"].values[0]
            t_log = row["threshold_logistic"].values[0]
            if not np.isnan(t_med):
                ax.axhline(t_med, color="#3498db", ls="-", lw=1.2,
                           label=f"median={t_med:.2f}")
            if not np.isnan(t_log):
                ax.axhline(t_log, color="#9b59b6", ls="-", lw=1.2,
                           label=f"logistic={t_log:.2f}")

        if i == 0:
            ax.set_title(f"RPV cutoff = {rpv_cut}", fontsize=11)
            ax.legend(fontsize=6, loc="best")
        if j == 0:
            ax.set_ylabel(METRIC_LABELS[met_col])
        if i == len(METRIC_COLS) - 1:
            ax.set_xlabel("RPV contamination (Llobet 2ms)")

plt.suptitle(f"Metric thresholds at RPV cutoffs (pooled, FR > {FR_THRESHOLD} Hz, n={len(df_hfr)})",
             fontsize=13, y=1.01)
plt.tight_layout()
plt.savefig(OUTPUT_DIR / "threshold_scatter.png", dpi=150, bbox_inches="tight")
plt.show()

# ──────────────────────────────────────────────────────────────────────
# 7. Correlation by firing rate band
# ──────────────────────────────────────────────────────────────────────
print("\nCorrelation by firing rate band...")

fr_bands = [(0, 2, "<2"), (2, 5, "2-5"), (5, 10, "5-10"),
            (10, 20, "10-20"), (20, 50, "20-50"), (50, np.inf, "50+")]

fig, axes = plt.subplots(2, 3, figsize=(14, 8))
for idx, met_col in enumerate(METRIC_COLS):
    ax = axes[idx // 3, idx % 3]
    rhos, ns, labels = [], [], []
    for lo, hi, label in fr_bands:
        band = df[(df["firing_rate"] > lo) & (df["firing_rate"] <= hi) & ~noise_mask]
        sub = band[["rpv", met_col]].dropna()
        if len(sub) >= 10:
            rho, _ = stats.spearmanr(sub["rpv"], sub[met_col])
        else:
            rho = np.nan
        rhos.append(rho)
        ns.append(len(sub))
        labels.append(label)

    colors = ["#e74c3c" if n < 10 else "#3498db" for n in ns]
    ax.bar(range(len(labels)), rhos, color=colors)
    ax.set_xticks(range(len(labels)))
    ax.set_xticklabels([f"{l}\n(n={n})" for l, n in zip(labels, ns)], fontsize=7)
    ax.set_ylabel("Spearman rho")
    ax.set_title(METRIC_LABELS[met_col])
    ax.axhline(0, color="gray", lw=0.5)
    ax.set_ylim(-1, 1)

plt.suptitle("Spearman rho (RPV vs metric) by firing rate band (pooled)", fontsize=13)
plt.tight_layout()
plt.savefig(OUTPUT_DIR / "correlation_by_fr_band.png", dpi=150, bbox_inches="tight")
plt.show()

# ──────────────────────────────────────────────────────────────────────
# Done
# ──────────────────────────────────────────────────────────────────────
df.to_csv(OUTPUT_DIR / "all_units.csv", index=False)
df_hfr.to_csv(OUTPUT_DIR / "high_fr_units.csv", index=False)

print(f"\nAll outputs saved to {OUTPUT_DIR.resolve()}")
print("Done.")
