"""Causal feature computation shared by offline training and the real-time filter.

Both `train_filter_model.py` (batch, pandas) and `adaptive_filter.py` (online,
per-sample) must derive identical features from identical inputs, or a model
trained offline will not behave the same way when driven by live samples.
This module is the single source of truth for that feature set.
"""

from __future__ import annotations

import math
from collections import deque

FEATURE_WINDOW = 5
MAX_CONTINUOUS_GAP_S = 0.5

FEATURE_COLUMNS = [
    "velocity_deg_s",
    "accel_px_s2",
    "dispersion_px",
    "confidence",
    "velocity_deg_s_mean",
    "velocity_deg_s_std",
    "dispersion_px_mean",
]

EVENT_LABELS = ("fixation", "saccade", "pursuit")


def add_rolling_features(df, session_col: str = "session_id"):
    """Add the *_mean / *_std rolling columns to a DataFrame in place, returning it.

    Rolling windows are computed per session so one recording's trailing
    samples never leak into the next. min_periods=1 keeps the first few
    samples of a session usable (with a noisier estimate) instead of NaN.
    """
    grouped = df.groupby(session_col, sort=False)
    roll = grouped["velocity_deg_s"].rolling(FEATURE_WINDOW, min_periods=1)
    df["velocity_deg_s_mean"] = roll.mean().reset_index(level=0, drop=True)
    df["velocity_deg_s_std"] = roll.std().reset_index(level=0, drop=True).fillna(0.0)
    df["dispersion_px_mean"] = (
        grouped["dispersion_px"].rolling(FEATURE_WINDOW, min_periods=1)
        .mean().reset_index(level=0, drop=True)
    )
    return df


class OnlineFeatureTracker:
    """Incrementally computes the same features as add_rolling_features, one raw sample at a time.

    Mirrors the instantaneous velocity/accel/dispersion logic in
    gaze_exercises.py's record_sample so a live filter sees the exact
    features the model was trained on.
    """

    def __init__(self, cm_per_px: float, viewing_distance_cm: float) -> None:
        self.cm_per_px = cm_per_px
        self.viewing_distance_cm = viewing_distance_cm
        self.history: deque[tuple[float, float, float]] = deque(maxlen=FEATURE_WINDOW)
        self.velocity_history: deque[float] = deque(maxlen=FEATURE_WINDOW)
        self.dispersion_history: deque[float] = deque(maxlen=FEATURE_WINDOW)
        self.prev_sample: tuple[float, float, float] | None = None
        self.prev_velocity_px_s: float | None = None

    def reset(self) -> None:
        self.history.clear()
        self.velocity_history.clear()
        self.dispersion_history.clear()
        self.prev_sample = None
        self.prev_velocity_px_s = None

    def update(self, timestamp: float, x: float, y: float, confidence: float) -> dict | None:
        velocity_px_s = accel_px_s2 = None
        if self.prev_sample is not None:
            prev_t, prev_x, prev_y = self.prev_sample
            dt = timestamp - prev_t
            if 0 < dt <= MAX_CONTINUOUS_GAP_S:
                dist_px = math.dist((prev_x, prev_y), (x, y))
                velocity_px_s = dist_px / dt
                if self.prev_velocity_px_s is not None:
                    accel_px_s2 = (velocity_px_s - self.prev_velocity_px_s) / dt
        self.prev_sample = (timestamp, x, y)
        self.prev_velocity_px_s = velocity_px_s

        velocity_deg_s = None
        if velocity_px_s is not None:
            velocity_cm_s = velocity_px_s * self.cm_per_px
            velocity_deg_s = math.degrees(
                2.0 * math.atan2(velocity_cm_s / 2.0, self.viewing_distance_cm)
            )

        self.history.append((x, y, timestamp))
        dispersion_px = None
        if len(self.history) >= 2:
            xs = [p[0] for p in self.history]
            ys = [p[1] for p in self.history]
            dispersion_px = (max(xs) - min(xs)) + (max(ys) - min(ys))

        if velocity_deg_s is None or dispersion_px is None:
            return None

        self.velocity_history.append(velocity_deg_s)
        self.dispersion_history.append(dispersion_px)

        n = len(self.velocity_history)
        velocity_deg_s_mean = sum(self.velocity_history) / n
        if n >= 2:
            variance = sum((v - velocity_deg_s_mean) ** 2 for v in self.velocity_history) / n
            velocity_deg_s_std = math.sqrt(variance)
        else:
            velocity_deg_s_std = 0.0
        dispersion_px_mean = sum(self.dispersion_history) / len(self.dispersion_history)

        return {
            "velocity_deg_s": velocity_deg_s,
            "accel_px_s2": accel_px_s2 if accel_px_s2 is not None else 0.0,
            "dispersion_px": dispersion_px,
            "confidence": confidence,
            "velocity_deg_s_mean": velocity_deg_s_mean,
            "velocity_deg_s_std": velocity_deg_s_std,
            "dispersion_px_mean": dispersion_px_mean,
        }