"""Inference API for the page structure classifier.

Usage:
    from classifier.predict import PageClassifier

    clf = PageClassifier()  # loads model_bundle.joblib
    result = clf.predict(page)  # dict[str, bool] with 23 keys
"""

from __future__ import annotations

import logging
from pathlib import Path
from typing import Dict, List, Optional

import joblib
import numpy as np

from .constraints import enforce_constraints
from .features import FEATURE_NAMES, extract_features

logger = logging.getLogger(__name__)

_DEFAULT_BUNDLE_PATH = Path(__file__).parent / "model_bundle.joblib"


class PageClassifier:
    """Lightweight wrapper around 23 binary XGBoost classifiers."""

    def __init__(self, bundle_path: Optional[Path] = None):
        path = bundle_path or _DEFAULT_BUNDLE_PATH
        bundle = joblib.load(path)
        self._models = bundle["models"]
        self._thresholds = bundle["thresholds"]
        self._feature_names = bundle["feature_names"]
        self._label_names = bundle["label_names"]

        # Validate feature alignment
        if self._feature_names != FEATURE_NAMES:
            raise ValueError(
                f"Feature mismatch: model expects {len(self._feature_names)} features, "
                f"code provides {len(FEATURE_NAMES)}"
            )

    def predict(self, page) -> Dict[str, bool]:
        """Predict 23 structural labels for a single page.

        Returns dict with boolean values, constraint-enforced.
        """
        features = extract_features(page)
        X = np.array([[features[f] for f in self._feature_names]])

        raw = {}
        for label in self._label_names:
            if label not in self._models:
                raw[label] = False
                continue
            prob = float(self._models[label].predict_proba(X)[0, 1])
            threshold = self._thresholds.get(label, 0.5)
            raw[label] = prob >= threshold

        return enforce_constraints(raw)

    def predict_proba(self, page) -> Dict[str, float]:
        """Return raw probabilities for each label (no thresholding)."""
        features = extract_features(page)
        X = np.array([[features[f] for f in self._feature_names]])

        probs = {}
        for label in self._label_names:
            if label not in self._models:
                probs[label] = 0.0
                continue
            probs[label] = float(self._models[label].predict_proba(X)[0, 1])
        return probs

    def predict_batch(self, pages) -> List[Dict[str, bool]]:
        """Predict for multiple pages."""
        feature_dicts = [extract_features(p) for p in pages]
        X = np.array([[fd[f] for f in self._feature_names] for fd in feature_dicts])

        raw_preds = [{} for _ in range(len(pages))]
        for label in self._label_names:
            if label not in self._models:
                for r in raw_preds:
                    r[label] = False
                continue
            probs = self._models[label].predict_proba(X)[:, 1]
            threshold = self._thresholds.get(label, 0.5)
            for i, prob in enumerate(probs):
                raw_preds[i][label] = float(prob) >= threshold

        return [enforce_constraints(r) for r in raw_preds]
