gesture_model.py

"""
GestureKey — ML Gesture Classifier
Trains an MLP on captured keypoint samples and provides inference.
Falls back gracefully when no trained model exists.

Author: Frederick Opoku-Afriyie
"""

from __future__ import annotations

import json
import pickle
import threading
import time
from pathlib import Path
from typing import Optional

import numpy as np

try:
    from sklearn.neural_network import MLPClassifier
    from sklearn.preprocessing import StandardScaler
    from sklearn.model_selection import cross_val_score, StratifiedKFold
    from sklearn.pipeline import Pipeline
    SKLEARN_OK = True
except ImportError:
    SKLEARN_OK = False

MODEL_PATH  = Path(__file__).parent / "data" / "gesture_mlp.pkl"
DATA_PATH   = Path(__file__).parent / "gesture_training_data.json"
MIN_SAMPLES = 8   # minimum samples per class to include in training


class TrainingResult:
    def __init__(self, accuracy: float, n_samples: int, n_classes: int,
                 labels: list[str], duration: float):
        self.accuracy   = accuracy
        self.n_samples  = n_samples
        self.n_classes  = n_classes
        self.labels     = labels
        self.duration   = duration

    def __repr__(self) -> str:
        return (f"TrainingResult(acc={self.accuracy:.1%}, "
                f"samples={self.n_samples}, classes={self.n_classes})")


class GestureMLModel:
    """
    MLP-based gesture classifier.

    Thread safety: classify() is read-only and safe to call from any thread.
    train() must be called from a dedicated worker thread (it blocks).
    After training, call reload() from the main thread to activate the new model.
    """

    def __init__(
        self,
        data_path:  Path = DATA_PATH,
        model_path: Path = MODEL_PATH,
    ) -> None:
        self.data_path  = data_path
        self.model_path = model_path
        self._lock      = threading.RLock()
        self._pipeline: Optional[Pipeline] = None
        self._labels:   list[str]          = []
        self._accuracy: float              = 0.0
        self._n_samples: int               = 0
        self._trained_at: float            = 0.0

        self._try_load()

    # ── Public properties ─────────────────────────────────────────────────────

    @property
    def is_ready(self) -> bool:
        with self._lock:
            return self._pipeline is not None

    @property
    def labels(self) -> list[str]:
        with self._lock:
            return list(self._labels)

    @property
    def accuracy(self) -> float:
        with self._lock:
            return self._accuracy

    @property
    def n_samples(self) -> int:
        with self._lock:
            return self._n_samples

    @property
    def trained_at(self) -> float:
        with self._lock:
            return self._trained_at

    # ── Inference ─────────────────────────────────────────────────────────────

    def classify(self, keypoints: np.ndarray) -> tuple[str, float]:
        """
        Returns (gesture_label, confidence 0–1).
        Returns ('unknown', 0.0) if model not ready or confidence too low.
        """
        with self._lock:
            if not self.is_ready:
                return "unknown", 0.0
            try:
                x      = keypoints.reshape(1, -1)
                proba  = self._pipeline.predict_proba(x)[0]
                idx    = int(proba.argmax())
                conf   = float(proba[idx])
                label  = self._labels[idx]
                return label, conf
            except Exception:
                return "unknown", 0.0

    # ── Training ──────────────────────────────────────────────────────────────

    def train(self, data: dict[str, list]) -> Optional[TrainingResult]:
        """
        Train on {gesture_name: [[63-dim vector], ...]} dict.
        Blocking — call from a background QThread.
        Returns TrainingResult or None on failure.
        """
        if not SKLEARN_OK:
            return None

        t_start = time.time()
        X, y_raw = [], []
        for label, samples in data.items():
            if len(samples) < MIN_SAMPLES:
                continue
            for s in samples:
                X.append(s)
                y_raw.append(label)

        if len(set(y_raw)) < 2 or len(X) < 20:
            return None

        X       = np.array(X, dtype=np.float32)
        labels  = sorted(set(y_raw))
        lbl2idx = {l: i for i, l in enumerate(labels)}
        y       = np.array([lbl2idx[l] for l in y_raw])

        pipeline = Pipeline([
            ("scaler", StandardScaler()),
            ("clf", MLPClassifier(
                hidden_layer_sizes=(256, 128, 64),
                activation="relu",
                solver="adam",
                alpha=1e-4,
                max_iter=1000,
                random_state=42,
                early_stopping=True,
                validation_fraction=0.15,
                n_iter_no_change=20,
                verbose=False,
            )),
        ])

        # Cross-validated accuracy
        min_class = int(np.bincount(y).min())
        n_splits  = max(2, min(5, min_class))
        cv        = StratifiedKFold(n_splits=n_splits, shuffle=True, random_state=42)
        try:
            scores   = cross_val_score(pipeline, X, y, cv=cv, scoring="accuracy")
            cv_acc   = float(scores.mean())
        except Exception:
            cv_acc   = 0.0

        # Final fit on all data
        pipeline.fit(X, y)

        result = TrainingResult(
            accuracy  = cv_acc,
            n_samples = len(X),
            n_classes = len(labels),
            labels    = labels,
            duration  = time.time() - t_start,
        )

        with self._lock:
            self._pipeline  = pipeline
            self._labels    = labels
            self._accuracy  = cv_acc
            self._n_samples = len(X)
            self._trained_at = time.time()

        self._save(pipeline, labels, cv_acc, len(X))
        return result

    def reload(self) -> bool:
        """Re-load model from disk. Call after training thread completes."""
        return self._try_load(force=True)

    def sample_counts(self) -> dict[str, int]:
        """Return {gesture: n_samples} from the training data file."""
        if not self.data_path.exists():
            return {}
        try:
            with open(self.data_path) as f:
                data = json.load(f)
            return {k: len(v) for k, v in data.items()}
        except Exception:
            return {}

    def load_training_data(self) -> dict[str, list]:
        """Load raw training data dict from disk."""
        if not self.data_path.exists():
            return {}
        try:
            with open(self.data_path) as f:
                return json.load(f)
        except Exception:
            return {}

    def save_training_data(self, data: dict[str, list]) -> None:
        """Persist updated training data to disk."""
        self.data_path.parent.mkdir(exist_ok=True)
        with open(self.data_path, "w") as f:
            json.dump(data, f)

    # ── Internal ──────────────────────────────────────────────────────────────

    def _try_load(self, force: bool = False) -> bool:
        if not self.model_path.exists():
            return False
        if (not force
                and self.data_path.exists()
                and self.data_path.stat().st_mtime > self.model_path.stat().st_mtime):
            return False  # training data newer than model — needs retraining
        try:
            with open(self.model_path, "rb") as f:
                payload = pickle.load(f)
            with self._lock:
                self._pipeline   = payload["pipeline"]
                self._labels     = payload["labels"]
                self._accuracy   = payload.get("accuracy", 0.0)
                self._n_samples  = payload.get("n_samples", 0)
                self._trained_at = payload.get("trained_at", 0.0)
            return True
        except Exception:
            return False

    def _save(self, pipeline, labels, accuracy, n_samples) -> None:
        self.model_path.parent.mkdir(exist_ok=True)
        with open(self.model_path, "wb") as f:
            pickle.dump({
                "pipeline":   pipeline,
                "labels":     labels,
                "accuracy":   accuracy,
                "n_samples":  n_samples,
                "trained_at": time.time(),
            }, f)