"""Train a per-sample eye-movement classifier (fixation/saccade/pursuit) from
gaze_exercises.py recordings, for use by adaptive_filter.py.

Usage:
    python train_filter_model.py
    python train_filter_model.py --results-dir results --test-size 0.25
"""

from __future__ import annotations

import argparse
import json
from pathlib import Path

import joblib
import numpy as np
import pandas as pd
from sklearn.ensemble import RandomForestClassifier
from sklearn.metrics import classification_report, confusion_matrix
from sklearn.model_selection import GroupShuffleSplit

from gaze_features import EVENT_LABELS, FEATURE_COLUMNS, add_rolling_features

BASELINE_VELOCITY_THRESHOLD_DEG_S = 30.0


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="Train the adaptive-filter event classifier.")
    parser.add_argument("--results-dir", default="results")
    parser.add_argument("--model-dir", default="models")
    parser.add_argument("--test-size", type=float, default=0.25,
                        help="fraction of trials held out for evaluation (default: 0.25)")
    parser.add_argument("--random-state", type=int, default=42)
    return parser.parse_args()


def load_dataset(results_dir: Path) -> pd.DataFrame:
    csv_paths = sorted(p for p in results_dir.glob("gaze_exercises_*.csv"))
    if not csv_paths:
        raise SystemExit(f"error: no gaze_exercises_*.csv files found in {results_dir}")
    frames = [pd.read_csv(p) for p in csv_paths]
    df = pd.concat(frames, ignore_index=True)
    print(f"Loaded {len(df)} rows from {len(csv_paths)} recording(s): "
          f"{', '.join(p.name for p in csv_paths)}")
    print("Sessions:", ", ".join(sorted(df['session_id'].unique())))
    return df


def prepare_features(df: pd.DataFrame) -> pd.DataFrame:
    df = df.sort_values(["session_id", "pupil_timestamp"]).reset_index(drop=True)
    df["accel_px_s2"] = df["accel_px_s2"].fillna(0.0)
    df = add_rolling_features(df)
    before = len(df)
    df = df.dropna(subset=FEATURE_COLUMNS + ["event_label"]).reset_index(drop=True)
    dropped = before - len(df)
    if dropped:
        print(f"Dropped {dropped} rows missing features (first sample(s) of each block).")
    df["trial_group"] = (
        df["session_id"] + "_" + df["block"] + "_" +
        df["round"].astype(str) + "_" + df["trial"].astype(str)
    )
    return df


def trial_level_split(df: pd.DataFrame, test_size: float, random_state: int):
    splitter = GroupShuffleSplit(n_splits=1, test_size=test_size, random_state=random_state)
    train_idx, test_idx = next(splitter.split(df, groups=df["trial_group"]))
    train_df, test_df = df.iloc[train_idx], df.iloc[test_idx]
    print(f"\nTrain: {len(train_df)} rows / {train_df['trial_group'].nunique()} trials")
    print(f"Test:  {len(test_df)} rows / {test_df['trial_group'].nunique()} trials")
    print("Train block mix:", train_df["block"].value_counts().to_dict())
    print("Test block mix: ", test_df["block"].value_counts().to_dict())
    if df["session_id"].nunique() == 1:
        print("\nNOTE: only one session/subject present. This split holds out unseen trials\n"
              "      within the same session - it does NOT test cross-subject generalization.\n"
              "      Re-run training once more subjects/sessions are recorded and switch to\n"
              "      leave-one-subject-out (group by subject_id instead of trial_group).")
    return train_df, test_df


def baseline_predict(df: pd.DataFrame) -> pd.Series:
    """Classic velocity-threshold (I-VT style) baseline: no pursuit class."""
    return np.where(df["velocity_deg_s"] > BASELINE_VELOCITY_THRESHOLD_DEG_S, "saccade", "fixation")


def detection_latency(df: pd.DataFrame, predicted: np.ndarray, label_name: str) -> None:
    """For each saccade block trial, measure ms between true and predicted onset of 'fixation'
    (i.e. how late the filter would recognize the saccade has ended and settle back down)."""
    work = df.copy()
    work["predicted"] = predicted
    deltas = []
    for _, trial in work[work["block"] == "saccade"].groupby("trial_group"):
        trial = trial.sort_values("pupil_timestamp")
        true_onset = trial.loc[trial["event_label"] == "fixation", "pupil_timestamp"]
        pred_onset = trial.loc[trial["predicted"] == "fixation", "pupil_timestamp"]
        if true_onset.empty or pred_onset.empty:
            continue
        deltas.append((pred_onset.iloc[0] - true_onset.iloc[0]) * 1000.0)
    if not deltas:
        print(f"{label_name}: no comparable saccade trials found")
        return
    deltas = np.array(deltas)
    print(f"{label_name}: fixation-recognition delay over {len(deltas)} saccade trials - "
          f"mean {deltas.mean():+.0f} ms, median {np.median(deltas):+.0f} ms "
          f"(positive = detected late)")


def main() -> None:
    args = parse_args()
    results_dir = Path(args.results_dir)
    model_dir = Path(args.model_dir)
    model_dir.mkdir(exist_ok=True)

    df = load_dataset(results_dir)
    df = prepare_features(df)
    train_df, test_df = trial_level_split(df, args.test_size, args.random_state)

    X_train, y_train = train_df[FEATURE_COLUMNS], train_df["event_label"]
    X_test, y_test = test_df[FEATURE_COLUMNS], test_df["event_label"]

    model = RandomForestClassifier(
        n_estimators=200, max_depth=12, class_weight="balanced",
        random_state=args.random_state, n_jobs=-1,
    )
    model.fit(X_train, y_train)
    pred = model.predict(X_test)

    print("\n=== Model (RandomForest) ===")
    print(classification_report(y_test, pred, labels=list(EVENT_LABELS)))
    print("Confusion matrix (rows=true, cols=pred), order", EVENT_LABELS)
    print(confusion_matrix(y_test, pred, labels=list(EVENT_LABELS)))

    print("\n=== Baseline (fixed velocity threshold, "
          f"{BASELINE_VELOCITY_THRESHOLD_DEG_S:.0f} deg/s) ===")
    baseline_pred = baseline_predict(test_df)
    print(classification_report(y_test, baseline_pred, labels=list(EVENT_LABELS), zero_division=0))

    print("\n=== Detection latency (lower is better; how late 'fixation' is recognized "
          "after a saccade) ===")
    detection_latency(test_df, pred, "Model")
    detection_latency(test_df, baseline_pred, "Baseline")

    importances = sorted(zip(FEATURE_COLUMNS, model.feature_importances_),
                          key=lambda kv: -kv[1])
    print("\nFeature importances:")
    for name, importance in importances:
        print(f"  {name:<22} {importance:.3f}")

    model_path = model_dir / "filter_model.joblib"
    joblib.dump(model, model_path)
    meta_path = model_dir / "filter_model_meta.json"
    meta_path.write_text(json.dumps({
        "feature_columns": FEATURE_COLUMNS,
        "event_labels": list(model.classes_),
        "baseline_velocity_threshold_deg_s": BASELINE_VELOCITY_THRESHOLD_DEG_S,
        "trained_on_sessions": sorted(df["session_id"].unique().tolist()),
        "train_rows": len(train_df),
        "test_rows": len(test_df),
    }, indent=2), encoding="utf-8")
    print(f"\nSaved model to {model_path}")
    print(f"Saved metadata to {meta_path}")


if __name__ == "__main__":
    main()
