"""Self-checks for the Mohebi lab rotation challenge.

Each notebook asks you to write a function, then calls a check from here.
A check either prints PASS or raises with a hint about what went wrong.

These are a safety net, not a grade. Passing every check does not mean your
analysis is correct — it means your building blocks behave the way ours do.
The scientific judgement is the part we actually care about.

If you think a check is wrong, say so in your write-up. That has happened.
"""

from __future__ import annotations

import numpy as np
import pandas as pd

TOL = 1e-4


class CheckFailed(AssertionError):
    """Raised when a check fails, with a hint attached."""


def _close(a, b, tol=TOL):
    return np.isfinite(a) and np.isfinite(b) and abs(a - b) <= tol * max(1.0, abs(b))


def _fail(msg, hint=""):
    raise CheckFailed(f"{msg}\n  hint: {hint}" if hint else msg)


def _ok(name, extra=""):
    print(f"PASS  {name}" + (f"  ({extra})" if extra else ""))


# The numeric expectations below were computed on the four sessions handed out
# with this challenge. If the data folder holds a different set of sessions,
# those exact values no longer apply, so the value-specific assertions are
# skipped and only the structural ones run.
REFERENCE_SESSIONS = {
    "IM-1056_M1_2025-12-14",
    "IM-1056_M1_2025-12-16",
    "IM-1056_PFC_2025-12-13",
    "IM-1056_PFC_2025-12-15",
}

_announced = set()


def _is_reference_data(df):
    if "session_id" not in df.columns:
        return True
    match = set(df["session_id"].unique()) == REFERENCE_SESSIONS
    if not match and "notice" not in _announced:
        _announced.add("notice")
        print(
            "NOTE: this is not the dataset the checks were calibrated on.\n"
            "      Structural checks will run; exact-value checks are skipped.\n"
        )
    return match


# ── Notebook 02 ─────────────────────────────────────────────────────────────
def check_causal_smooth(fn):
    """fn(signal, alpha) -> smoothed array of the same length."""
    x = np.sin(np.linspace(0, 4 * np.pi, 100))
    y = np.asarray(fn(x, 0.05), dtype=float)

    if y.shape != x.shape:
        _fail(
            f"causal_smooth returned shape {y.shape}, expected {x.shape}",
            "return one value per input sample",
        )
    if not _close(y[0], 0.0, 1e-9):
        _fail(
            f"first sample is {y[0]:.6f}, expected the raw first sample (0.0)",
            "seed the filter with signal[0]; there is no earlier sample to mix in",
        )
    for i, want in [(49, -0.32029943), (99, -0.33597860)]:
        if not _close(y[i], want):
            _fail(
                f"sample {i} is {y[i]:.6f}, expected {want:.6f}",
                "y[i] = alpha*x[i] + (1-alpha)*y[i-1]; a symmetric or centred "
                "filter will not match, and neither will scipy.ndimage defaults",
            )

    const = np.full(50, 3.0)
    if not np.allclose(fn(const, 0.05), 3.0):
        _fail(
            "smoothing a constant signal did not return that constant",
            "the filter weights must sum to 1",
        )

    step = np.concatenate([np.zeros(50), np.ones(50)])
    sm = np.asarray(fn(step, 0.05), dtype=float)
    if sm[49] > 1e-9:
        _fail(
            "the filter responded before the step arrived",
            "this must be CAUSAL: output at time i may not depend on i+1",
        )
    _ok("causal_smooth")


def check_zscore(fn, df):
    """fn(df) -> df with Ch470_dFF / Ch560_dFF z-scored per trial."""
    z = fn(df)

    if len(z) != len(df):
        _fail(
            f"z-scored frame has {len(z)} rows, input had {len(df)}",
            "z-scoring rescales values; it should not add or drop rows",
        )

    keys = ["session_id", "Trial"] if "session_id" in z.columns else ["Animal", "Trial"]
    if "session_id" not in z.columns:
        print("  note: no session_id column — checking within a single session only")

    base = z[(z["Time_s"] >= -2.0) & (z["Time_s"] <= 0.0)]
    for col in ("Ch470_dFF", "Ch560_dFF"):
        g = base.groupby(keys)[col].agg(["mean", "std"])
        if abs(g["mean"].mean()) > 1e-6:
            _fail(
                f"{col}: baseline mean across trials is {g['mean'].mean():.3e}, expected ~0",
                "subtract each trial's OWN baseline mean, not a global mean",
            )
        if abs(g["std"].mean() - 1.0) > 1e-6:
            _fail(
                f"{col}: baseline SD across trials is {g['std'].mean():.6f}, expected ~1",
                "divide by each trial's own baseline SD over the (-2, 0) s window",
            )

    n_expected = df.groupby(keys).ngroups
    if z.groupby(keys).ngroups != n_expected:
        _fail("trial grouping changed during z-scoring", "do not reindex or sort away keys")

    if "session_id" in df.columns and df["session_id"].nunique() > 1:
        n_animal_trial = df.groupby(["Animal", "Trial"]).ngroups
        if n_animal_trial < n_expected:
            print(
                f"  note: ('Animal','Trial') would have given {n_animal_trial} groups "
                f"instead of {n_expected} — you correctly used session_id"
            )
    _ok("zscore_trials", f"{n_expected} trials")


def check_trial_average(fn, z_one_session):
    """fn(df, signal_col) -> (mean Series indexed by Time_s, sem Series, n_trials).

    Pass the Z-SCORED frame for a single session (IM-1056_PFC_2025-12-13).
    """
    if float(np.abs(z_one_session["Ch560_dFF"]).max()) < 5:
        _fail(
            "this frame does not look z-scored",
            "pass the z-scored single-session frame, not raw dF/F",
        )
    out = fn(z_one_session, "Ch560_dFF")
    if not (isinstance(out, tuple) and len(out) == 3):
        _fail("expected a 3-tuple of (mean, sem, n_trials)", f"got {type(out)}")
    mean, sem, n = out
    mean = pd.Series(mean)

    expected_n = z_one_session.groupby("Trial").ngroups
    if int(n) != expected_n:
        _fail(f"n_trials is {n}, expected {expected_n}", "count unique trials, not rows")
    if len(mean) != 253:
        _fail(
            f"mean has {len(mean)} time points, expected 253",
            "restrict to the (-2.0, 7.5) s plot window, inclusive at both ends",
        )

    is_ref = (
        "session_id" in z_one_session.columns
        and z_one_session["session_id"].iloc[0] == "IM-1056_PFC_2025-12-13"
    )
    if is_ref and not _close(float(np.nanmax(mean)), 2.441752, 1e-3):
        _fail(
            f"peak of the mean trace is {np.nanmax(mean):.4f}, expected ~2.4418",
            "average across trials at each time point; do NOT smooth before averaging",
        )
    _ok("trial_average", f"n={int(n)}, {len(mean)} time points")


# ── Notebook 03 ─────────────────────────────────────────────────────────────
def check_cue_auc(fn, z_df):
    """fn(trial_df, signal_col) -> duration-normalized AUC over (0, 1) s."""
    if not _is_reference_data(z_df):
        print("SKIP  cue_auc (values are specific to the distributed dataset)")
        return
    one = z_df[
        (z_df["session_id"] == "IM-1056_PFC_2025-12-13") & (z_df["Trial"] == 7)
    ].sort_values("Time_s")
    if one.empty:
        _fail("could not find trial 7 of IM-1056_PFC_2025-12-13", "pass the z-scored frame")

    for col, want in [("Ch470_dFF", 2.5631040), ("Ch560_dFF", 1.3373113)]:
        got = float(fn(one, col))
        if _close(got, want * 1.0, 1e-3):
            continue
        if _close(got, want * (1.0 - 1.0 / 26.6674), 5e-3):
            _fail(
                f"{col}: got {got:.6f}, expected {want:.6f} — off by one sample",
                "use an inclusive window: Time_s >= 0 AND Time_s <= 1",
            )
        _fail(
            f"{col}: got {got:.6f}, expected {want:.6f}",
            "trapezoid integral over (0, 1) s, divided by the 1 s duration, "
            "computed on z-scored but UNSMOOTHED data",
        )
    _ok("cue_auc")


def check_delay_auc(fn, z_df):
    """fn(trial_df, signal_col, delay) -> duration-normalized AUC over the delay."""
    if not _is_reference_data(z_df):
        # The 0 s exclusion rule is universal even when the values are not.
        any_trial = z_df[z_df["Trial"] == z_df["Trial"].iloc[0]]
        got0 = fn(any_trial, "Ch560_dFF", 0.0)
        if got0 is not None and np.isfinite(got0):
            _fail(
                f"a 0 s delay returned {got0}, expected NaN",
                "there is no delay period on a 0 s trial",
            )
        print("SKIP  delay_period_auc values (dataset-specific); 0 s rule OK")
        return
    one = z_df[
        (z_df["session_id"] == "IM-1056_PFC_2025-12-13") & (z_df["Trial"] == 7)
    ].sort_values("Time_s")
    delay = float(one["Delay_s"].iloc[0])
    if delay != 1.0:
        _fail(f"trial 7 should have a 1 s delay, found {delay}", "pass the unmodified frame")

    # Trial 7 has a 1 s delay, so dividing by the duration is a no-op there.
    # Trial 26 has a 3 s delay and does distinguish the two.
    long_trial = z_df[
        (z_df["session_id"] == "IM-1056_PFC_2025-12-13") & (z_df["Trial"] == 26)
    ].sort_values("Time_s")
    long_delay = float(long_trial["Delay_s"].iloc[0])

    cases = [
        (one, delay, "Ch470_dFF", 5.2948625),
        (one, delay, "Ch560_dFF", 1.8970308),
        (long_trial, long_delay, "Ch470_dFF", 3.9931067),
        (long_trial, long_delay, "Ch560_dFF", 0.5368784),
    ]
    for trial_df, d, col, want in cases:
        got = float(fn(trial_df, col, d))
        if _close(got, want, 1e-3):
            continue
        if _close(got, want * d, 1e-3):
            _fail(
                f"{col} at delay {d:.0f} s: got {got:.6f}, expected {want:.6f}",
                "divide the integral by the delay duration, otherwise a 3 s "
                "delay looks 3x larger than a 1 s delay for the same mean signal",
            )
        _fail(
            f"{col} at delay {d:.0f} s: got {got:.6f}, expected {want:.6f}",
            "integrate from delay_onset (1.0 s) to delay_onset + delay, "
            "then divide by the delay duration",
        )

    got0 = fn(one, "Ch560_dFF", 0.0)
    if got0 is not None and np.isfinite(got0):
        _fail(
            f"a 0 s delay returned {got0}, expected NaN",
            "there is no delay period on a 0 s trial — it must be excluded, "
            "not integrated over a zero-width window",
        )
    _ok("delay_period_auc")


def check_session_auc(fn, z_df):
    """fn(df) -> per-session mean delay AUC with columns Session/Delay/Outcome/NE_AUC/DA_AUC."""
    tab = fn(z_df)
    need = {"Session", "Delay", "Outcome", "NE_AUC", "DA_AUC"}
    if not need.issubset(tab.columns):
        _fail(f"missing columns: {sorted(need - set(tab.columns))}", f"got {list(tab.columns)}")

    if 0.0 in set(tab["Delay"]):
        _fail("delay 0 appears in the table", "0 s trials have no delay period to integrate")

    n_sessions = tab["Session"].nunique()
    expected_sessions = z_df["session_id"].nunique()
    if n_sessions != expected_sessions:
        _fail(
            f"found {n_sessions} sessions, expected {expected_sessions}",
            "aggregate per session, not per animal",
        )

    expected_delays = sorted(float(d) for d in z_df["Delay_s"].unique() if d > 0)
    got_delays = sorted(float(d) for d in set(tab["Delay"]))
    if got_delays != expected_delays:
        _fail(
            f"table covers delays {got_delays}, but this data has {expected_delays}",
            "read the delays out of the data rather than hardcoding them",
        )

    if not _is_reference_data(z_df):
        _ok("session_delay_auc (structure only)", f"{len(tab)} rows, {n_sessions} sessions")
        return

    expected = {
        (1.0, "correct"): (4.9032, 0.9950),
        (1.0, "incorrect"): (4.1260, 1.1128),
        (2.0, "correct"): (5.6892, 1.0799),
        (2.0, "incorrect"): (4.5967, 0.8193),
        (3.0, "correct"): (4.7625, 0.6862),
        (3.0, "incorrect"): (4.0165, 0.5188),
    }
    means = tab.groupby(["Delay", "Outcome"])[["NE_AUC", "DA_AUC"]].mean()
    for key, (ne, da) in expected.items():
        if key not in means.index:
            _fail(f"no rows for delay {key[0]} / {key[1]}")
        row = means.loc[key]
        if not (_close(row["NE_AUC"], ne, 2e-3) and _close(row["DA_AUC"], da, 2e-3)):
            _fail(
                f"delay {key[0]:.0f} {key[1]}: got NE={row['NE_AUC']:.4f} "
                f"DA={row['DA_AUC']:.4f}, expected NE={ne} DA={da}",
                "average trials WITHIN a session first, then average sessions "
                "(pooling all trials weights the sessions unequally)",
            )
    _ok("session_delay_auc", f"{len(tab)} rows, {n_sessions} sessions")


# ── Notebook 04 ─────────────────────────────────────────────────────────────
def check_dprime(fn):
    """fn(n_correct, n_total) -> 2AFC d' with loglinear correction."""
    cases = [
        ((50, 100), 0.0, "chance performance must give d' = 0"),
        ((80, 100), 1.1752942, "d' = sqrt(2) * Phi^-1(p_adjusted)"),
        (
            (100, 100),
            3.6476360,
            "at ceiling, use p_adj = (n_correct + 0.5) / (n_total + 1); "
            "without the correction this is +inf",
        ),
        ((0, 40), -3.1832896, "at floor, the same correction keeps d' finite"),
    ]
    for (nc, nt), want, hint in cases:
        got = float(fn(nc, nt))
        if not _close(got, want, 1e-5):
            if _close(got, want / np.sqrt(2), 1e-5):
                _fail(
                    f"d'({nc}/{nt}) = {got:.6f}, expected {want:.6f}",
                    "missing the sqrt(2) factor for the 2AFC design",
                )
            _fail(f"d'({nc}/{nt}) = {got:.6f}, expected {want:.6f}", hint)

    if np.isfinite(fn(0, 0)):
        _fail("d' with zero trials should be NaN", "guard against division by zero")
    _ok("dprime_2afc")


def check_performance(fn, z_df):
    """fn(df) -> per-session table with Delay/PctCorrect/DPrime."""
    tab = fn(z_df)
    need = {"Session", "Delay", "PctCorrect", "DPrime", "N_correct", "N_total"}
    if not need.issubset(tab.columns):
        _fail(f"missing columns: {sorted(need - set(tab.columns))}", f"got {list(tab.columns)}")
    n_expected = z_df["session_id"].nunique() * z_df["Delay_s"].nunique()
    if len(tab) != n_expected:
        _fail(
            f"table has {len(tab)} rows, expected {n_expected} "
            f"({z_df['session_id'].nunique()} sessions x {z_df['Delay_s'].nunique()} delays)",
            "one row per session per delay",
        )

    # Every trial has exactly 533 samples, so percent-correct computed over ROWS
    # happens to equal percent-correct computed over TRIALS. The trial counts
    # are what actually reveal the mistake.
    n_trials = z_df.groupby(["session_id", "Trial"]).ngroups
    total = int(tab["N_total"].sum())
    if total != n_trials:
        hint = "count unique trials, not rows"
        if total == len(z_df):
            hint = (
                "you counted samples, not trials — every trial contributes 533 "
                "rows. Reduce to one row per trial first (groupby(...).first())"
            )
        _fail(f"N_total sums to {total}, expected {n_trials} trials", hint)

    if not _is_reference_data(z_df):
        _ok("performance_table (structure only)", f"{len(tab)} rows, {total} trials")
        return

    expected_pct = {0.0: 73.5642, 1.0: 65.8167, 2.0: 62.0139, 3.0: 62.3626}
    got = tab.groupby("Delay").PctCorrect.mean()
    for d, want in expected_pct.items():
        if not _close(float(got.loc[d]), want, 1e-3):
            _fail(
                f"delay {d:.0f}: mean % correct = {got.loc[d]:.4f}, expected {want}",
                "average percent-correct across the 4 sessions, weighting each "
                "session equally",
            )

    expected_dp = {0.0: 0.8792, 1.0: 0.5819, 2.0: 0.4376, 3.0: 0.4341}
    got_dp = tab.groupby("Delay").DPrime.mean()
    for d, want in expected_dp.items():
        if not _close(float(got_dp.loc[d]), want, 2e-3):
            _fail(
                f"delay {d:.0f}: mean d' = {got_dp.loc[d]:.4f}, expected {want}",
                "compute d' per session from that session's trial counts, "
                "then average across sessions",
            )
    _ok("performance_table", f"{len(tab)} rows, {total} trials")


# ── Runner ──────────────────────────────────────────────────────────────────
def run_all(namespace, df, z_df=None):
    """Run every check for which the matching function exists in `namespace`.

    Usage inside a notebook:  checks.run_all(globals(), df, z)
    """
    ran = 0
    for name, check in [
        ("causal_smooth", lambda f: check_causal_smooth(f)),
        ("zscore_trials", lambda f: check_zscore(f, df)),
        ("dprime_2afc", lambda f: check_dprime(f)),
    ]:
        if name in namespace:
            check(namespace[name])
            ran += 1
    if z_df is not None:
        for name, check in [
            ("cue_auc", lambda f: check_cue_auc(f, z_df)),
            ("delay_period_auc", lambda f: check_delay_auc(f, z_df)),
            ("session_delay_auc", lambda f: check_session_auc(f, z_df)),
            ("performance_table", lambda f: check_performance(f, z_df)),
        ]:
            if name in namespace:
                check(namespace[name])
                ran += 1
    print(f"\n{ran} check(s) passed.")
    return ran
