"""Shared helpers for the Mohebi lab rotation challenge.

These are the boring parts — finding files, loading them, keeping figures
consistent. Nothing in here does any analysis; that part is your job.

You are welcome to read, modify, or ignore this file.
"""

from __future__ import annotations

import re
from pathlib import Path

import matplotlib.pyplot as plt
import numpy as np
import pandas as pd

# Channel names, as they appear in the data files.
NE_COL = "Ch470_dFF"  # noradrenaline sensor
DA_COL = "Ch560_dFF"  # dopamine sensor
ISO_COL = "Ch410_dFF"  # isosbestic reference channel

DATA_DIR = Path(__file__).resolve().parent / "data"


# ── Loading ─────────────────────────────────────────────────────────────────
def list_sessions(data_dir=DATA_DIR):
    """Return a DataFrame describing every session file available to you.

    Region and date are parsed from the *filename*. Pay attention to whether
    the columns inside the file agree with that.
    """
    rows = []
    for path in sorted(Path(data_dir).glob("*.parquet")):
        animal, region, date = re.match(
            r"(IM-\d+)_([A-Za-z0-9]+)_(\d{4}-\d{2}-\d{2})", path.stem
        ).groups()
        rows.append(
            dict(
                session_id=path.stem,
                animal=animal,
                region=region,
                date=date,
                path=str(path),
            )
        )
    if not rows:
        raise FileNotFoundError(f"No .parquet files found in {data_dir}")
    return pd.DataFrame(rows)


def load_session(session_id, data_dir=DATA_DIR):
    """Load one session, tagged with session_id / animal / region / date."""
    meta = list_sessions(data_dir).set_index("session_id").loc[session_id]
    df = pd.read_parquet(meta["path"])
    df["session_id"] = session_id
    df["animal"] = meta["animal"]
    df["region"] = meta["region"]
    df["date"] = meta["date"]
    return df


def load_all(data_dir=DATA_DIR, region=None):
    """Load every session (optionally filtered to one region) into one frame.

    Trial numbers restart at 1 in every session, so ALWAYS group by
    ['session_id', 'Trial'] rather than 'Trial' alone.
    """
    sessions = list_sessions(data_dir)
    if region is not None:
        sessions = sessions[sessions["region"] == region]
    return pd.concat(
        [load_session(s, data_dir) for s in sessions["session_id"]],
        ignore_index=True,
    )


def sampling_rate(df):
    """Median sampling rate of the PETH time base, in Hz."""
    t = np.sort(df["Time_s"].unique())
    return 1.0 / np.median(np.diff(t))


# ── Plotting ────────────────────────────────────────────────────────────────
NE_COLOR = "seagreen"
DA_COLOR = "firebrick"


def delay_colors(delays, channel="NE"):
    """One colour per delay, light to dark. Green for NE, red for DA."""
    cmap = plt.cm.Greens if channel.upper() == "NE" else plt.cm.Reds
    return cmap(np.linspace(0.35, 0.9, len(delays)))


def style_axis(ax, event_lines=(0.0, 1.0)):
    """Drop the top/right spines and mark cue onset and cue offset."""
    ax.spines["top"].set_visible(False)
    ax.spines["right"].set_visible(False)
    for i, x in enumerate(event_lines):
        ax.axvline(
            x, color="gray", linestyle="--" if i == 0 else ":", linewidth=0.8, zorder=0
        )
    return ax


def shade_delay(ax, delay, delay_onset=1.0, **kwargs):
    """Shade the delay period on a PETH axis."""
    if delay > 0:
        kwargs.setdefault("color", "gray")
        kwargs.setdefault("alpha", 0.10)
        ax.axvspan(delay_onset, delay_onset + delay, zorder=0, **kwargs)
    return ax


def save_figure(fig, name, folder="../figures"):
    """Save a figure as PDF into the challenge figures/ folder."""
    out = Path(folder)
    out.mkdir(parents=True, exist_ok=True)
    path = out / (name if name.endswith(".pdf") else f"{name}.pdf")
    fig.savefig(path, bbox_inches="tight")
    print(f"saved {path}")
    return path
