ui/pipeline.py

from __future__ import annotations

import collections
import glob
import json
import math
import os
import pathlib
import sys

import numpy as np
import joblib
import torch
import torch.nn as nn

_PROJECT_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
if _PROJECT_ROOT not in sys.path:
    sys.path.insert(0, _PROJECT_ROOT)

from data_preparation.prepare_dataset import SELECTED_FEATURES
from models.face_mesh import FaceMeshDetector
from models.head_pose import HeadPoseEstimator
from models.eye_scorer import EyeBehaviourScorer, compute_mar, MAR_YAWN_THRESHOLD
from models.collect_features import FEATURE_NAMES, TemporalTracker, extract_features
from models.eye_scorer import compute_avg_ear

# Same 10 features used for MLP training (prepare_dataset) and inference
MLP_FEATURE_NAMES = SELECTED_FEATURES["face_orientation"]

_FEAT_IDX = {name: i for i, name in enumerate(FEATURE_NAMES)}


def _clip_features(vec):
    out = vec.copy()
    _i = _FEAT_IDX

    out[_i["yaw"]] = np.clip(out[_i["yaw"]], -45, 45)
    out[_i["pitch"]] = np.clip(out[_i["pitch"]], -30, 30)
    out[_i["roll"]] = np.clip(out[_i["roll"]], -30, 30)

    out[_i["head_deviation"]] = math.sqrt(
        float(out[_i["yaw"]]) ** 2 + float(out[_i["pitch"]]) ** 2
    )

    for f in ("ear_left", "ear_right", "ear_avg"):
        out[_i[f]] = np.clip(out[_i[f]], 0, 0.85)

    out[_i["mar"]] = np.clip(out[_i["mar"]], 0, 1.0)
    out[_i["gaze_offset"]] = np.clip(out[_i["gaze_offset"]], 0, 0.50)
    out[_i["perclos"]] = np.clip(out[_i["perclos"]], 0, 0.80)
    out[_i["blink_rate"]] = np.clip(out[_i["blink_rate"]], 0, 30.0)
    out[_i["closure_duration"]] = np.clip(out[_i["closure_duration"]], 0, 10.0)
    out[_i["yawn_duration"]] = np.clip(out[_i["yawn_duration"]], 0, 10.0)

    return out


class _OutputSmoother:
    # Asymmetric EMA: rises fast (recognise focus), falls slower (avoid flicker).
    # Grace period holds score steady for a few frames when face is lost.

    def __init__(self, alpha_up=0.55, alpha_down=0.45, grace_frames=10):
        self._alpha_up = alpha_up
        self._alpha_down = alpha_down
        self._grace = grace_frames
        self._score = 0.5
        self._no_face = 0

    def reset(self):
        self._score = 0.5
        self._no_face = 0

    def update(self, raw_score, face_detected):
        if face_detected:
            self._no_face = 0
            alpha = self._alpha_up if raw_score > self._score else self._alpha_down
            self._score += alpha * (raw_score - self._score)
        else:
            self._no_face += 1
            if self._no_face > self._grace:
                self._score *= 0.80
        return self._score


DEFAULT_HYBRID_CONFIG = {
    "use_xgb": False,
    "w_mlp": 0.3,
    "w_xgb": 0.0,
    "w_geo": 0.7,
    "threshold": 0.35,
    "use_yawn_veto": True,
    "geo_face_weight": 0.7,
    "geo_eye_weight": 0.3,
    "mar_yawn_threshold": float(MAR_YAWN_THRESHOLD),
    "combiner": None,
    "combiner_path": None,
}


class _RuntimeFeatureEngine:
    _MAG_FEATURES = ["pitch", "yaw", "head_deviation", "gaze_offset", "v_gaze", "h_gaze"]
    _VEL_FEATURES = ["pitch", "yaw", "h_gaze", "v_gaze", "head_deviation", "gaze_offset"]
    _VAR_FEATURES = ["h_gaze", "v_gaze", "pitch"]
    _VAR_WINDOW = 30
    _WARMUP = 15

    def __init__(self, base_feature_names, norm_features=None):
        self._base_names = list(base_feature_names)
        self._norm_features = list(norm_features) if norm_features else []

        tracked = set(self._MAG_FEATURES) | set(self._norm_features)
        self._ema_mean = {f: 0.0 for f in tracked}
        self._ema_var = {f: 1.0 for f in tracked}
        self._n = 0
        self._prev = None
        self._var_bufs = {
            f: collections.deque(maxlen=self._VAR_WINDOW) for f in self._VAR_FEATURES
        }

        self._ext_names = (
            list(self._base_names)
            + [f"{f}_mag" for f in self._MAG_FEATURES]
            + [f"{f}_vel" for f in self._VEL_FEATURES]
            + [f"{f}_var" for f in self._VAR_FEATURES]
        )

    @property
    def extended_names(self):
        return list(self._ext_names)

    def transform(self, base_vec):
        self._n += 1
        raw = {name: float(base_vec[i]) for i, name in enumerate(self._base_names)}

        alpha = 2.0 / (min(self._n, 120) + 1)
        for feat in self._ema_mean:
            if feat not in raw:
                continue
            v = raw[feat]
            if self._n == 1:
                self._ema_mean[feat] = v
                self._ema_var[feat] = 0.0
            else:
                self._ema_mean[feat] += alpha * (v - self._ema_mean[feat])
                self._ema_var[feat] += alpha * (
                    (v - self._ema_mean[feat]) ** 2 - self._ema_var[feat]
                )

        out = base_vec.copy().astype(np.float32)
        if self._n > self._WARMUP:
            for feat in self._norm_features:
                if feat in raw:
                    idx = self._base_names.index(feat)
                    std = max(math.sqrt(self._ema_var[feat]), 1e-6)
                    out[idx] = (raw[feat] - self._ema_mean[feat]) / std

        mag = np.zeros(len(self._MAG_FEATURES), dtype=np.float32)
        for i, feat in enumerate(self._MAG_FEATURES):
            if feat in raw:
                mag[i] = abs(raw[feat] - self._ema_mean.get(feat, raw[feat]))

        vel = np.zeros(len(self._VEL_FEATURES), dtype=np.float32)
        if self._prev is not None:
            for i, feat in enumerate(self._VEL_FEATURES):
                if feat in raw and feat in self._prev:
                    vel[i] = abs(raw[feat] - self._prev[feat])
        self._prev = dict(raw)

        for feat in self._VAR_FEATURES:
            if feat in raw:
                self._var_bufs[feat].append(raw[feat])
        var = np.zeros(len(self._VAR_FEATURES), dtype=np.float32)
        for i, feat in enumerate(self._VAR_FEATURES):
            buf = self._var_bufs[feat]
            if len(buf) >= 2:
                arr = np.array(buf)
                var[i] = float(arr.var())

        return np.concatenate([out, mag, vel, var])


class FaceMeshPipeline:
    def __init__(
        self,
        max_angle: float = 22.0,
        alpha: float = 0.7,
        beta: float = 0.3,
        threshold: float = 0.55,
        detector=None,
    ):
        self.detector = detector or FaceMeshDetector()
        self._owns_detector = detector is None
        self.head_pose = HeadPoseEstimator(max_angle=max_angle)
        self.eye_scorer = EyeBehaviourScorer()
        self.alpha = alpha
        self.beta = beta
        self.threshold = threshold
        self._smoother = _OutputSmoother()

    def process_frame(self, bgr_frame: np.ndarray) -> dict:
        landmarks = self.detector.process(bgr_frame)
        h, w = bgr_frame.shape[:2]

        out = {
            "landmarks": landmarks,
            "s_face": 0.0,
            "s_eye": 0.0,
            "raw_score": 0.0,
            "is_focused": False,
            "yaw": None,
            "pitch": None,
            "roll": None,
            "mar": None,
            "is_yawning": False,
            "left_bbox": None,
            "right_bbox": None,
        }

        if landmarks is None:
            smoothed = self._smoother.update(0.0, False)
            out["raw_score"] = smoothed
            out["is_focused"] = smoothed >= self.threshold
            return out

        angles = self.head_pose.estimate(landmarks, w, h)
        if angles is not None:
            out["yaw"], out["pitch"], out["roll"] = angles
        out["s_face"] = self.head_pose.score(landmarks, w, h)
        out["s_eye"] = self.eye_scorer.score(landmarks)
        out["mar"] = compute_mar(landmarks)
        out["is_yawning"] = out["mar"] > MAR_YAWN_THRESHOLD

        raw = self.alpha * out["s_face"] + self.beta * out["s_eye"]
        if out["is_yawning"]:
            raw = 0.0
        out["raw_score"] = self._smoother.update(raw, True)
        out["is_focused"] = out["raw_score"] >= self.threshold

        return out

    def reset_session(self):
        self._smoother.reset()

    def close(self):
        if self._owns_detector:
            self.detector.close()

    def __enter__(self):
        return self

    def __exit__(self, *args):
        self.close()


# PyTorch MLP matching models/mlp/train.py BaseModel (10 -> 64 -> 32 -> 2)
class _FocusMLP(nn.Module):
    def __init__(self, num_features: int, num_classes: int = 2):
        super().__init__()
        self.network = nn.Sequential(
            nn.Linear(num_features, 64),
            nn.ReLU(),
            nn.Linear(64, 32),
            nn.ReLU(),
            nn.Linear(32, num_classes),
        )

    def forward(self, x):
        return self.network(x)


def _mlp_artifacts_available(model_dir: str) -> bool:
    pt_path = os.path.join(model_dir, "mlp_best.pt")
    scaler_path = os.path.join(model_dir, "scaler_mlp.joblib")
    return os.path.isfile(pt_path) and os.path.isfile(scaler_path)


def _load_mlp_artifacts(model_dir: str):
    """Load PyTorch MLP + scaler from checkpoints. Returns (model, scaler, feature_names)."""
    pt_path = os.path.join(model_dir, "mlp_best.pt")
    scaler_path = os.path.join(model_dir, "scaler_mlp.joblib")
    if not os.path.isfile(pt_path):
        raise FileNotFoundError(f"No MLP checkpoint at {pt_path}")
    if not os.path.isfile(scaler_path):
        raise FileNotFoundError(f"No scaler at {scaler_path}")

    num_features = len(MLP_FEATURE_NAMES)
    num_classes = 2
    model = _FocusMLP(num_features, num_classes)
    model.load_state_dict(torch.load(pt_path, map_location="cpu", weights_only=True))
    model.eval()

    scaler = joblib.load(scaler_path)
    return model, scaler, list(MLP_FEATURE_NAMES)


def _load_hybrid_config(model_dir: str, config_path: str | None = None):
    cfg = dict(DEFAULT_HYBRID_CONFIG)
    resolved = config_path or os.path.join(model_dir, "hybrid_focus_config.json")
    if not os.path.isfile(resolved):
        print(f"[HYBRID] No config found at {resolved}; using defaults")
        return cfg, None

    with open(resolved, "r", encoding="utf-8") as f:
        file_cfg = json.load(f)

    for key in DEFAULT_HYBRID_CONFIG:
        if key in file_cfg:
            cfg[key] = file_cfg[key]

    cfg["use_xgb"] = bool(cfg.get("use_xgb", False))
    cfg["w_mlp"] = float(cfg.get("w_mlp", 0.3))
    cfg["w_xgb"] = float(cfg.get("w_xgb", 0.0))
    cfg["w_geo"] = float(cfg["w_geo"])
    if cfg["use_xgb"]:
        weight_sum = cfg["w_xgb"] + cfg["w_geo"]
        if weight_sum <= 0:
            raise ValueError("[HYBRID] Invalid config: w_xgb + w_geo must be > 0")
        cfg["w_xgb"] /= weight_sum
        cfg["w_geo"] /= weight_sum
    else:
        weight_sum = cfg["w_mlp"] + cfg["w_geo"]
        if weight_sum <= 0:
            raise ValueError("[HYBRID] Invalid config: w_mlp + w_geo must be > 0")
        cfg["w_mlp"] /= weight_sum
        cfg["w_geo"] /= weight_sum
    cfg["threshold"] = float(cfg["threshold"])
    cfg["use_yawn_veto"] = bool(cfg["use_yawn_veto"])
    cfg["geo_face_weight"] = float(cfg["geo_face_weight"])
    cfg["geo_eye_weight"] = float(cfg["geo_eye_weight"])
    cfg["mar_yawn_threshold"] = float(cfg["mar_yawn_threshold"])
    cfg["combiner"] = cfg.get("combiner") or None
    cfg["combiner_path"] = cfg.get("combiner_path") or None

    print(f"[HYBRID] Loaded config: {resolved}")
    return cfg, resolved


class MLPPipeline:
    def __init__(self, model_dir=None, detector=None, threshold=0.23):
        if model_dir is None:
            model_dir = os.path.join(_PROJECT_ROOT, "checkpoints")

        self._mlp, self._scaler, self._feature_names = _load_mlp_artifacts(model_dir)
        self._indices = [FEATURE_NAMES.index(n) for n in self._feature_names]

        self._detector = detector or FaceMeshDetector()
        self._owns_detector = detector is None
        self._head_pose = HeadPoseEstimator()
        self.head_pose = self._head_pose
        self._eye_scorer = EyeBehaviourScorer()
        self._temporal = TemporalTracker()
        self._smoother = _OutputSmoother()
        self._threshold = threshold
        print(f"[MLP] Loaded PyTorch MLP from {model_dir} | {len(self._feature_names)} features | threshold={threshold}")

    def process_frame(self, bgr_frame):
        landmarks = self._detector.process(bgr_frame)
        h, w = bgr_frame.shape[:2]
        out = {
            "landmarks": landmarks,
            "is_focused": False,
            "s_face": 0.0,
            "s_eye": 0.0,
            "raw_score": 0.0,
            "mlp_prob": 0.0,
            "mar": None,
            "yaw": None,
            "pitch": None,
            "roll": None,
        }
        if landmarks is None:
            smoothed = self._smoother.update(0.0, False)
            out["raw_score"] = smoothed
            out["is_focused"] = smoothed >= self._threshold
            return out
        vec = extract_features(landmarks, w, h, self._head_pose, self._eye_scorer, self._temporal)
        vec = _clip_features(vec)

        out["yaw"] = float(vec[_FEAT_IDX["yaw"]])
        out["pitch"] = float(vec[_FEAT_IDX["pitch"]])
        out["roll"] = float(vec[_FEAT_IDX["roll"]])
        out["s_face"] = float(vec[_FEAT_IDX["s_face"]])
        out["s_eye"] = float(vec[_FEAT_IDX["s_eye"]])
        out["mar"] = float(vec[_FEAT_IDX["mar"]])

        X = vec[self._indices].reshape(1, -1).astype(np.float32)
        X_sc = self._scaler.transform(X) if self._scaler is not None else X
        with torch.no_grad():
            x_t = torch.from_numpy(X_sc).float()
            logits = self._mlp(x_t)
            probs = torch.softmax(logits, dim=1)
            mlp_prob = float(probs[0, 1])
        out["mlp_prob"] = float(np.clip(mlp_prob, 0.0, 1.0))
        out["raw_score"] = self._smoother.update(out["mlp_prob"], True)
        out["is_focused"] = out["raw_score"] >= self._threshold
        return out

    def reset_session(self):
        self._temporal = TemporalTracker()
        self._smoother.reset()

    def close(self):
        if self._owns_detector:
            self._detector.close()

    def __enter__(self):
        return self

    def __exit__(self, *args):
        self.close()


def _resolve_xgb_path():
    return os.path.join(_PROJECT_ROOT, "checkpoints", "xgboost_face_orientation_best.json")


class HybridFocusPipeline:
    def __init__(
        self,
        model_dir=None,
        config_path: str | None = None,
        max_angle: float = 22.0,
        detector=None,
    ):
        if model_dir is None:
            model_dir = os.path.join(_PROJECT_ROOT, "checkpoints")
        self._cfg, self._cfg_path = _load_hybrid_config(model_dir=model_dir, config_path=config_path)
        self._use_xgb = self._cfg["use_xgb"]

        self._detector = detector or FaceMeshDetector()
        self._owns_detector = detector is None
        self._head_pose = HeadPoseEstimator(max_angle=max_angle)
        self._eye_scorer = EyeBehaviourScorer()
        self._temporal = TemporalTracker()
        self.head_pose = self._head_pose
        self._smoother = _OutputSmoother()

        self._combiner = None
        combiner_path = self._cfg.get("combiner_path")
        if combiner_path and self._cfg.get("combiner") == "logistic":
            resolved_combiner = combiner_path if os.path.isabs(combiner_path) else os.path.join(model_dir, combiner_path)
            if not os.path.isfile(resolved_combiner):
                resolved_combiner = os.path.join(_PROJECT_ROOT, combiner_path)
            if os.path.isfile(resolved_combiner):
                blob = joblib.load(resolved_combiner)
                self._combiner = blob.get("combiner")
                if self._combiner is None:
                    self._combiner = blob
                print(f"[HYBRID] LR combiner loaded from {resolved_combiner}")
            else:
                print(f"[HYBRID] combiner_path not found: {resolved_combiner}, using heuristic weights")
        if self._use_xgb:
            from xgboost import XGBClassifier
            xgb_path = _resolve_xgb_path()
            if not os.path.isfile(xgb_path):
                raise FileNotFoundError(f"No XGBoost checkpoint at {xgb_path}")
            self._xgb_model = XGBClassifier()
            self._xgb_model.load_model(xgb_path)
            self._xgb_indices = [FEATURE_NAMES.index(n) for n in XGBoostPipeline.SELECTED]
            self._mlp = None
            self._scaler = None
            self._indices = None
            self._feature_names = list(XGBoostPipeline.SELECTED)
            mode = "LR combiner" if self._combiner else f"w_xgb={self._cfg['w_xgb']:.2f}, w_geo={self._cfg['w_geo']:.2f}"
            print(f"[HYBRID] XGBoost+geo | {xgb_path} | {mode}, threshold={self._cfg['threshold']:.2f}")
        else:
            self._mlp, self._scaler, self._feature_names = _load_mlp_artifacts(model_dir)
            self._indices = [FEATURE_NAMES.index(n) for n in self._feature_names]
            self._xgb_model = None
            self._xgb_indices = None
            mode = "LR combiner" if self._combiner else f"w_mlp={self._cfg['w_mlp']:.2f}, w_geo={self._cfg['w_geo']:.2f}"
            print(f"[HYBRID] MLP+geo | {len(self._feature_names)} features | {mode}, threshold={self._cfg['threshold']:.2f}")

    @property
    def config(self) -> dict:
        return dict(self._cfg)

    def process_frame(self, bgr_frame: np.ndarray) -> dict:
        landmarks = self._detector.process(bgr_frame)
        h, w = bgr_frame.shape[:2]
        out = {
            "landmarks": landmarks,
            "is_focused": False,
            "focus_score": 0.0,
            "mlp_prob": 0.0,
            "geo_score": 0.0,
            "raw_score": 0.0,
            "s_face": 0.0,
            "s_eye": 0.0,
            "mar": None,
            "is_yawning": False,
            "yaw": None,
            "pitch": None,
            "roll": None,
            "left_bbox": None,
            "right_bbox": None,
        }
        if landmarks is None:
            smoothed = self._smoother.update(0.0, False)
            out["focus_score"] = smoothed
            out["raw_score"] = smoothed
            out["is_focused"] = smoothed >= self._cfg["threshold"]
            return out

        angles = self._head_pose.estimate(landmarks, w, h)
        if angles is not None:
            out["yaw"], out["pitch"], out["roll"] = angles

        out["s_face"] = self._head_pose.score(landmarks, w, h)
        out["s_eye"] = self._eye_scorer.score(landmarks)
        s_eye_geo = out["s_eye"]

        geo_score = (
            self._cfg["geo_face_weight"] * out["s_face"] +
            self._cfg["geo_eye_weight"] * out["s_eye"]
        )
        geo_score = float(np.clip(geo_score, 0.0, 1.0))

        out["mar"] = compute_mar(landmarks)
        out["is_yawning"] = out["mar"] > self._cfg["mar_yawn_threshold"]
        if self._cfg["use_yawn_veto"] and out["is_yawning"]:
            geo_score = 0.0
        out["geo_score"] = geo_score

        pre = {
            "angles": angles,
            "s_face": out["s_face"],
            "s_eye": s_eye_geo,
            "mar": out["mar"],
        }
        vec = extract_features(landmarks, w, h, self._head_pose, self._eye_scorer, self._temporal, _pre=pre)
        vec = _clip_features(vec)

        if self._use_xgb:
            X = vec[self._xgb_indices].reshape(1, -1).astype(np.float32)
            prob = self._xgb_model.predict_proba(X)[0]
            model_prob = float(np.clip(prob[1], 0.0, 1.0))
            out["mlp_prob"] = model_prob
            if self._combiner is not None:
                meta = np.array([[model_prob, out["geo_score"]]], dtype=np.float32)
                focus_score = float(self._combiner.predict_proba(meta)[0, 1])
            else:
                focus_score = self._cfg["w_xgb"] * model_prob + self._cfg["w_geo"] * out["geo_score"]
        else:
            X = vec[self._indices].reshape(1, -1).astype(np.float32)
            X_sc = self._scaler.transform(X) if self._scaler is not None else X
            with torch.no_grad():
                x_t = torch.from_numpy(X_sc).float()
                logits = self._mlp(x_t)
                probs = torch.softmax(logits, dim=1)
                mlp_prob = float(probs[0, 1])
            out["mlp_prob"] = float(np.clip(mlp_prob, 0.0, 1.0))
            if self._combiner is not None:
                meta = np.array([[out["mlp_prob"], out["geo_score"]]], dtype=np.float32)
                focus_score = float(self._combiner.predict_proba(meta)[0, 1])
            else:
                focus_score = self._cfg["w_mlp"] * out["mlp_prob"] + self._cfg["w_geo"] * out["geo_score"]

        out["focus_score"] = self._smoother.update(float(np.clip(focus_score, 0.0, 1.0)), True)
        out["raw_score"] = out["focus_score"]
        out["is_focused"] = out["focus_score"] >= self._cfg["threshold"]
        return out

    def reset_session(self):
        self._temporal = TemporalTracker()
        self._smoother.reset()

    def close(self):
        if self._owns_detector:
            self._detector.close()

    def __enter__(self):
        return self

    def __exit__(self, *args):
        self.close()


class XGBoostPipeline:
    SELECTED = [
        'head_deviation', 's_face', 's_eye', 'h_gaze', 'pitch',
        'ear_left', 'ear_avg', 'ear_right', 'gaze_offset', 'perclos',
    ]

    def __init__(self, model_path=None, threshold=0.28):
        from xgboost import XGBClassifier

        if model_path is None:
            model_path = os.path.join(_PROJECT_ROOT, "checkpoints", "xgboost_face_orientation_best.json")
        if not os.path.isfile(model_path):
            raise FileNotFoundError(f"No XGBoost checkpoint at {model_path}")

        self._model = XGBClassifier()
        self._model.load_model(model_path)
        self._threshold = threshold

        self._detector = FaceMeshDetector()
        self._head_pose = HeadPoseEstimator()
        self.head_pose = self._head_pose
        self._eye_scorer = EyeBehaviourScorer()
        self._temporal = TemporalTracker()
        self._smoother = _OutputSmoother()

        self._indices = [FEATURE_NAMES.index(n) for n in self.SELECTED]
        print(f"[XGB] Loaded {model_path} | {len(self.SELECTED)} features, threshold={threshold}")

    def process_frame(self, bgr_frame):
        landmarks = self._detector.process(bgr_frame)
        h, w = bgr_frame.shape[:2]
        out = {
            "landmarks": landmarks,
            "is_focused": False,
            "s_face": 0.0,
            "s_eye": 0.0,
            "raw_score": 0.0,
            "mar": None,
            "yaw": None,
            "pitch": None,
            "roll": None,
        }
        if landmarks is None:
            smoothed = self._smoother.update(0.0, False)
            out["raw_score"] = smoothed
            out["is_focused"] = smoothed >= self._threshold
            return out

        vec = extract_features(landmarks, w, h, self._head_pose, self._eye_scorer, self._temporal)
        vec = _clip_features(vec)

        out["yaw"] = float(vec[_FEAT_IDX["yaw"]])
        out["pitch"] = float(vec[_FEAT_IDX["pitch"]])
        out["roll"] = float(vec[_FEAT_IDX["roll"]])
        out["s_face"] = float(vec[_FEAT_IDX["s_face"]])
        out["s_eye"] = float(vec[_FEAT_IDX["s_eye"]])
        out["mar"] = float(vec[_FEAT_IDX["mar"]])

        X = vec[self._indices].reshape(1, -1).astype(np.float32)
        prob = self._model.predict_proba(X)[0]  # [prob_unfocused, prob_focused]
        out["raw_score"] = self._smoother.update(float(prob[1]), True)
        out["is_focused"] = out["raw_score"] >= self._threshold
        return out

    def reset_session(self):
        self._temporal = TemporalTracker()
        self._smoother.reset()

    def close(self):
        self._detector.close()

    def __enter__(self):
        return self

    def __exit__(self, *args):
        self.close()


def _is_git_lfs_pointer(path: str) -> bool:
    # *.pkl in repo are often LFS stubs; torch.load sees "v" from "version ..." and dies
    try:
        with open(path, "rb") as f:
            return f.read(64).startswith(b"version https://git-lfs.github.com/spec/v1")
    except OSError:
        return False


def _resolve_l2cs_weights():
    preferred = os.path.join(_PROJECT_ROOT, "checkpoints", "L2CSNet_gaze360.pkl")
    if os.path.isfile(preferred) and not _is_git_lfs_pointer(preferred):
        return preferred

    # Backward-compatible fallback paths for older setups.
    for p in [
        os.path.join(_PROJECT_ROOT, "models", "L2CS-Net", "models", "L2CSNet_gaze360.pkl"),
        os.path.join(_PROJECT_ROOT, "models", "L2CSNet_gaze360.pkl"),
    ]:
        if os.path.isfile(p) and not _is_git_lfs_pointer(p):
            return p
    return None


def is_l2cs_weights_available():
    return _resolve_l2cs_weights() is not None


class L2CSPipeline:
    # Uses in-tree l2cs.Pipeline (RetinaFace + ResNet50) for gaze estimation
    # and MediaPipe for head pose, EAR, MAR, and roll de-rotation.
    # L2CS inference is throttled to every Nth frame to reduce latency;
    # intermediate frames reuse the last gaze result.

    YAW_THRESHOLD = 22.0
    PITCH_THRESHOLD = 20.0
    _SKIP_CPU = 5   # run L2CS every 5th frame on CPU
    _SKIP_GPU = 1   # run every frame on GPU (fast enough)

    def __init__(self, weights_path=None, arch="ResNet50", device=None,
                 threshold=0.52, detector=None):
        resolved = weights_path or _resolve_l2cs_weights()
        if resolved is None or not os.path.isfile(resolved):
            raise FileNotFoundError(
                "L2CS weights missing or Git LFS not pulled. "
                "Run: git lfs pull  or  python download_l2cs_weights.py  "
                "(real .pkl in checkpoints/ or models/L2CS-Net/models/)"
            )

        # add in-tree L2CS-Net to import path
        l2cs_root = os.path.join(_PROJECT_ROOT, "models", "L2CS-Net")
        if l2cs_root not in sys.path:
            sys.path.insert(0, l2cs_root)
        from l2cs import Pipeline as _L2CSPipeline

        import torch

        # Auto-detect GPU if no device specified
        if device is None:
            device = "cuda" if torch.cuda.is_available() else "cpu"
        self._device_str = device
        self._on_gpu = device.startswith("cuda")

        # torch.device passed explicitly for reliable CPU/CUDA selection
        self._pipeline = _L2CSPipeline(
            weights=pathlib.Path(resolved), arch=arch, device=torch.device(device),
        )

        self._detector = detector or FaceMeshDetector()
        self._owns_detector = detector is None
        self._head_pose = HeadPoseEstimator()
        self.head_pose = self._head_pose
        self._eye_scorer = EyeBehaviourScorer()
        self._threshold = threshold
        self._smoother = _OutputSmoother()

        # Frame skipping: GPU is fast enough to run every frame
        self.L2CS_SKIP_FRAMES = self._SKIP_GPU if self._on_gpu else self._SKIP_CPU
        self._frame_count = 0
        self._last_l2cs_result = None  # cached (derotated pitch_rad, yaw_rad)
        self._calibrating = False  # set True during calibration to disable frame skipping

        # Blink tolerance: hold score steady during brief blinks
        self._blink_streak = 0
        self._BLINK_EAR = 0.18
        self._BLINK_GRACE = 5  # ignore blinks shorter than this many frames (~300ms)

        print(
            f"[L2CS] Loaded {resolved} | arch={arch} device={device} "
            f"yaw_thresh={self.YAW_THRESHOLD} pitch_thresh={self.PITCH_THRESHOLD} "
            f"threshold={threshold} skip_frames={self.L2CS_SKIP_FRAMES}"
        )

    @staticmethod
    def _derotate_gaze(pitch_rad, yaw_rad, roll_deg):
        # remove head roll so tilted-but-looking-at-screen reads as (0,0)
        roll_rad = -math.radians(roll_deg)
        cos_r, sin_r = math.cos(roll_rad), math.sin(roll_rad)
        return (yaw_rad * sin_r + pitch_rad * cos_r,
                yaw_rad * cos_r - pitch_rad * sin_r)

    def process_frame(self, bgr_frame):
        landmarks = self._detector.process(bgr_frame)
        h, w = bgr_frame.shape[:2]

        out = {
            "landmarks": landmarks, "is_focused": False, "raw_score": 0.0,
            "s_face": 0.0, "s_eye": 0.0, "gaze_pitch": None, "gaze_yaw": None,
            "yaw": None, "pitch": None, "roll": None, "mar": None, "is_yawning": False,
        }

        # MediaPipe: head pose, eye/mouth scores (runs every frame — fast)
        roll_deg = 0.0
        blinking = False
        if landmarks is not None:
            angles = self._head_pose.estimate(landmarks, w, h)
            if angles is not None:
                out["yaw"], out["pitch"], out["roll"] = angles
                roll_deg = angles[2]
            out["s_face"] = self._head_pose.score(landmarks, w, h)
            out["s_eye"] = self._eye_scorer.score(landmarks)
            out["mar"] = compute_mar(landmarks)
            out["is_yawning"] = out["mar"] > MAR_YAWN_THRESHOLD

            # Detect blink — EAR drops below threshold
            ear = compute_avg_ear(landmarks)
            if ear < self._BLINK_EAR:
                self._blink_streak += 1
                blinking = True
            else:
                self._blink_streak = 0

        # During a brief blink, L2CS gaze angles are unreliable (eyes closed).
        # Hold the previous score steady until blink ends or becomes sustained.
        if blinking and self._blink_streak < self._BLINK_GRACE:
            # Brief blink — freeze score, skip L2CS inference
            out["raw_score"] = self._smoother._score
            out["is_focused"] = out["raw_score"] >= self._threshold
            # Keep previous gaze angles for visualization continuity
            if self._last_l2cs_result is not None:
                out["gaze_pitch"] = self._last_l2cs_result[0]
                out["gaze_yaw"] = self._last_l2cs_result[1]
            return out

        # L2CS gaze — throttled: only run every Nth frame, reuse cached result otherwise.
        # During calibration, run every frame for accurate sample collection.
        self._frame_count += 1
        if self._calibrating:
            run_l2cs = True
        else:
            run_l2cs = (self._frame_count % self.L2CS_SKIP_FRAMES == 1) or self._last_l2cs_result is None

        if run_l2cs:
            results = self._pipeline.step(bgr_frame)
            if results is not None and results.pitch.shape[0] > 0:
                raw_pitch = float(results.pitch[0])
                raw_yaw = float(results.yaw[0])
                # Derotate immediately and cache the derotated result
                # so cached frames don't get re-derotated with a different roll.
                dr_pitch, dr_yaw = self._derotate_gaze(raw_pitch, raw_yaw, roll_deg)
                self._last_l2cs_result = (dr_pitch, dr_yaw)
            else:
                self._last_l2cs_result = None

        if self._last_l2cs_result is None:
            smoothed = self._smoother.update(0.0, landmarks is not None)
            out["raw_score"] = smoothed
            out["is_focused"] = smoothed >= self._threshold
            return out

        pitch_rad, yaw_rad = self._last_l2cs_result
        # Already derotated above — use directly
        out["gaze_pitch"] = pitch_rad
        out["gaze_yaw"] = yaw_rad

        yaw_deg = abs(math.degrees(yaw_rad))
        pitch_deg = abs(math.degrees(pitch_rad))

        # fall back to L2CS angles if MediaPipe didn't produce head pose
        out["yaw"] = out.get("yaw") or math.degrees(yaw_rad)
        out["pitch"] = out.get("pitch") or math.degrees(pitch_rad)

        # cosine scoring: 1.0 at centre, 0.0 at threshold
        yaw_t = min(yaw_deg / self.YAW_THRESHOLD, 1.0)
        pitch_t = min(pitch_deg / self.PITCH_THRESHOLD, 1.0)
        yaw_score = 0.5 * (1.0 + math.cos(math.pi * yaw_t))
        pitch_score = 0.5 * (1.0 + math.cos(math.pi * pitch_t))
        gaze_score = 0.55 * yaw_score + 0.45 * pitch_score

        if out["is_yawning"]:
            gaze_score = 0.0

        # Sustained closed eyes — let score drop
        if self._blink_streak >= self._BLINK_GRACE:
            gaze_score = 0.0

        out["raw_score"] = self._smoother.update(float(gaze_score), True)
        out["is_focused"] = out["raw_score"] >= self._threshold
        return out

    def reset_session(self):
        self._smoother.reset()
        self._frame_count = 0
        self._last_l2cs_result = None
        self._calibrating = False
        self._blink_streak = 0

    def close(self):
        if self._owns_detector:
            self._detector.close()

    def __enter__(self):
        return self

    def __exit__(self, *args):
        self.close()