#!/usr/bin/env python3
"""
Statistical Category Classifier — inference wrapper (PLAN_TO_PROD Phase 3A)
==========================================================================

Loads the supervised device-category model trained by
``scripts/train_category_classifier.py`` and predicts a device *category*
(with class probabilities) from the timing/statistical features of a RAW
Sub-GHz capture.

This is the "statistical" leg of the designed ensemble. It is used as a
category *prior* over the heuristic RAW matches (see
``pattern_decoder._apply_ml_category_boost``) — it re-ranks toward the
category the model believes the signal belongs to, but never invents a match.

Graceful degradation: if the model file is missing or joblib/sklearn are
unavailable, ``available`` is False and the decoder falls back to the pure
heuristic path unchanged.

Feature parity: the vector is built in the frozen order carried in the model
bundle (``features``), matching ``train_category_classifier.extract_features``
exactly. If the two ever diverge, retrain rather than hand-edit one side.
"""

from dataclasses import dataclass, field
from pathlib import Path
from typing import Dict, List, Optional

import numpy as np

from src.matcher.timing_analyzer import get_timing_analyzer
from src.matcher.preamble_detector import get_preamble_detector

DEFAULT_MODEL_PATH = (
    Path(__file__).parent.parent.parent / "models" / "category_classifier.joblib"
)

# Must match scripts/train_category_classifier._PREAMBLE_ID
_PREAMBLE_ID = {"none": 0, "long_burst": 1, "alternating": 2,
                "sync_word": 3, "custom": 4}


@dataclass
class MLCategoryPrediction:
    """Result of a statistical category prediction."""
    available: bool = False
    top_category: Optional[str] = None
    top_prob: float = 0.0
    probabilities: Dict[str, float] = field(default_factory=dict)


class MLCategoryClassifier:
    """Lazy-loading inference wrapper around the trained sklearn model."""

    def __init__(self, model_path: Path = DEFAULT_MODEL_PATH):
        self.model_path = Path(model_path)
        self._loaded = False
        self._model = None
        self._features: List[str] = []
        self._classes: List[str] = []
        self.timing_analyzer = get_timing_analyzer()
        self.preamble_detector = get_preamble_detector()

    @property
    def available(self) -> bool:
        self._ensure_loaded()
        return self._model is not None

    def _ensure_loaded(self):
        if self._loaded:
            return
        self._loaded = True
        try:
            import joblib
            bundle = joblib.load(self.model_path)
            self._model = bundle["model"]
            self._features = list(bundle["features"])
            self._classes = [str(c) for c in bundle["classes"]]
        except Exception:
            # Missing model, joblib, or sklearn — degrade to heuristic-only.
            self._model = None

    def _feature_vector(self, pulses, frequency) -> Optional[List[float]]:
        """Build the frozen-order feature vector for one RAW capture.

        Mirrors ``train_category_classifier.extract_features`` — keep in sync.
        """
        if not pulses:
            return None
        timing = self.timing_analyzer.extract_timing(pulses)
        short, long = timing.short_pulse_us, timing.long_pulse_us
        if short == 0 or long == 0:
            return None

        detected = self.preamble_detector.detect(pulses, short, long)
        ptype = detected.type if detected else "none"

        abs_p = np.abs(np.asarray(pulses, dtype=float))
        gap_ratio = (timing.long_gap_us / timing.short_gap_us
                     if timing.short_gap_us > 0 else 0.0)

        named = {
            "freq_mhz": (frequency or 0) / 1_000_000,
            "short_pulse_us": short,
            "long_pulse_us": long,
            "pulse_ratio": timing.pulse_ratio,
            "short_gap_us": timing.short_gap_us,
            "long_gap_us": timing.long_gap_us,
            "gap_ratio": gap_ratio,
            "duty_cycle": timing.duty_cycle,
            "pulse_count": len(pulses),
            "pulse_mean_abs": float(abs_p.mean()),
            "pulse_std_abs": float(abs_p.std()),
            "pulse_min_abs": float(abs_p.min()),
            "pulse_max_abs": float(abs_p.max()),
            "preamble_type": _PREAMBLE_ID.get(ptype, 0),
        }
        return [named[name] for name in self._features]

    def predict(self, pulses, frequency) -> MLCategoryPrediction:
        """Predict a device category from RAW pulses. Never raises."""
        self._ensure_loaded()
        if self._model is None:
            return MLCategoryPrediction(available=False)

        try:
            vec = self._feature_vector(pulses, frequency)
            if vec is None:
                return MLCategoryPrediction(available=False)
            X = np.asarray([vec], dtype=float)
            probs = self._model.predict_proba(X)[0]
            prob_map = {self._classes[i]: float(probs[i])
                        for i in range(len(self._classes))}
            top_i = int(np.argmax(probs))
            return MLCategoryPrediction(
                available=True,
                top_category=self._classes[top_i],
                top_prob=float(probs[top_i]),
                probabilities=prob_map,
            )
        except Exception:
            return MLCategoryPrediction(available=False)


# Singleton
_classifier: Optional[MLCategoryClassifier] = None


def get_category_classifier() -> MLCategoryClassifier:
    """Get singleton statistical category classifier."""
    global _classifier
    if _classifier is None:
        _classifier = MLCategoryClassifier()
    return _classifier
