"""Real-time gaze filters: the existing fixed median+EMA baseline, and two
ML-driven variants (hard state-switch, continuous confidence-weighted) built
on the classifier trained by train_filter_model.py.
"""

from __future__ import annotations

import json
import statistics
from collections import deque
from pathlib import Path

import joblib

from gaze_features import FEATURE_COLUMNS, OnlineFeatureTracker

FIXATION_ALPHA = 0.15
PURSUIT_ALPHA = 0.35
SACCADE_ALPHA = 0.9
SOFT_MIN_ALPHA = 0.15
SOFT_MAX_ALPHA = 0.9
MEDIAN_WINDOW = 5


class BaselineFilter:
    """Fixed median + EMA + dead-zone filter, matching visualize_surface_gaze.py."""

    def __init__(self, smoothing: float = 0.18, median_window: int = 5, dead_zone: float = 3.0) -> None:
        self.smoothing = smoothing
        self.dead_zone = dead_zone
        self.x_history: deque[float] = deque(maxlen=median_window)
        self.y_history: deque[float] = deque(maxlen=median_window)
        self.filtered: tuple[float, float] | None = None

    def reset(self) -> None:
        self.x_history.clear()
        self.y_history.clear()
        self.filtered = None

    def update(self, timestamp: float, x: float, y: float, confidence: float) -> tuple[float, float]:
        self.x_history.append(x)
        self.y_history.append(y)
        median_x = statistics.median(self.x_history)
        median_y = statistics.median(self.y_history)
        if self.filtered is None:
            self.filtered = (median_x, median_y)
        else:
            prev_x, prev_y = self.filtered
            fx = prev_x + self.smoothing * (median_x - prev_x)
            fy = prev_y + self.smoothing * (median_y - prev_y)
            if abs(fx - prev_x) < self.dead_zone:
                fx = prev_x
            if abs(fy - prev_y) < self.dead_zone:
                fy = prev_y
            self.filtered = (fx, fy)
        return self.filtered


class MLAdaptiveFilter:
    """Classifier-driven filter. mode="hard" switches behavior by predicted class;
    mode="soft" blends smoothing strength continuously with P(saccade)."""

    def __init__(
        self,
        model_path: str | Path,
        cm_per_px: float,
        viewing_distance_cm: float,
        mode: str = "soft",
    ) -> None:
        if mode not in ("hard", "soft"):
            raise ValueError("mode must be 'hard' or 'soft'")
        self.mode = mode
        self.model = joblib.load(model_path)
        meta_path = Path(model_path).with_name(Path(model_path).stem + "_meta.json")
        self.meta = json.loads(meta_path.read_text(encoding="utf-8")) if meta_path.exists() else {}
        self.saccade_index = list(self.model.classes_).index("saccade")

        self.tracker = OnlineFeatureTracker(cm_per_px, viewing_distance_cm)
        self.x_history: deque[float] = deque(maxlen=MEDIAN_WINDOW)
        self.y_history: deque[float] = deque(maxlen=MEDIAN_WINDOW)
        self.filtered: tuple[float, float] | None = None
        self.last_label: str | None = None
        self.last_prob_saccade: float | None = None

    def reset(self) -> None:
        self.tracker.reset()
        self.x_history.clear()
        self.y_history.clear()
        self.filtered = None
        self.last_label = None
        self.last_prob_saccade = None

    def update(self, timestamp: float, x: float, y: float, confidence: float) -> tuple[float, float]:
        features = self.tracker.update(timestamp, x, y, confidence)
        if features is None:
            # Not enough history yet for a prediction; pass raw through as median-of-one.
            self.x_history.append(x)
            self.y_history.append(y)
            self.filtered = (x, y)
            return self.filtered

        row = [[features[c] for c in FEATURE_COLUMNS]]
        proba = self.model.predict_proba(row)[0]
        prob_saccade = float(proba[self.saccade_index])
        label = self.model.classes_[proba.argmax()]
        self.last_label = label
        self.last_prob_saccade = prob_saccade

        if self.mode == "hard":
            if label == "saccade":
                self.x_history.clear()
                self.y_history.clear()
                self.x_history.append(x)
                self.y_history.append(y)
                self.filtered = (x, y)
                return self.filtered
            alpha = FIXATION_ALPHA if label == "fixation" else PURSUIT_ALPHA
        else:
            alpha = SOFT_MIN_ALPHA + prob_saccade * (SOFT_MAX_ALPHA - SOFT_MIN_ALPHA)

        self.x_history.append(x)
        self.y_history.append(y)
        median_x = statistics.median(self.x_history)
        median_y = statistics.median(self.y_history)
        if self.filtered is None:
            self.filtered = (median_x, median_y)
        else:
            prev_x, prev_y = self.filtered
            fx = prev_x + alpha * (median_x - prev_x)
            fy = prev_y + alpha * (median_y - prev_y)
            self.filtered = (fx, fy)
        return self.filtered