"""
Inference helper for HowzerRiskScorer.

Supports both PyTorch and ONNX Runtime backends.

Usage:
    from inference import RiskScorer
    scorer = RiskScorer.from_onnx("howzer_risk_scorer.onnx")
    result = scorer.predict(features)  # np.ndarray of shape (120,) or (N, 120)
"""

from __future__ import annotations

from dataclasses import dataclass
from pathlib import Path

import numpy as np

TIER_NAMES = ["low", "medium", "high", "critical"]
TIER_THRESHOLDS = {"medium": 0.38, "high": 0.72, "critical": 0.88}


@dataclass
class RiskAssessment:
    """Risk assessment output for a single sample."""
    global_risk: float
    tier: str
    tier_index: int
    escalation: float
    churn: float
    brand: float
    revenue: float
    confidence: float
    tier_probabilities: dict[str, float]

    def to_dict(self) -> dict:
        return {
            "global_risk": round(self.global_risk, 4),
            "tier": self.tier,
            "escalation": round(self.escalation, 4),
            "churn": round(self.churn, 4),
            "brand": round(self.brand, 4),
            "revenue": round(self.revenue, 4),
            "confidence": round(self.confidence, 4),
            "tier_probabilities": {k: round(v, 4) for k, v in self.tier_probabilities.items()},
        }


class RiskScorer:
    """Unified inference wrapper for HowzerRiskScorer."""

    def __init__(self, session):
        self._session = session
        self._backend = "onnx"

    @classmethod
    def from_onnx(cls, model_path: str | Path) -> "RiskScorer":
        """Load from ONNX file using ONNX Runtime."""
        import onnxruntime as ort
        session = ort.InferenceSession(
            str(model_path),
            providers=["CPUExecutionProvider"],
        )
        return cls(session)

    @classmethod
    def from_pytorch(cls, checkpoint_path: str | Path, device: str = "cpu") -> "RiskScorer":
        """Load from PyTorch checkpoint."""
        import torch
        # Import from relative location or installed package
        import sys
        sys.path.insert(0, str(Path(__file__).resolve().parent.parent.parent))
        from src.model.risk_scorer import HowzerRiskScorer

        model = HowzerRiskScorer(
            input_dim=120, hidden_dims=(256, 128, 64),
            num_factors=4, num_tiers=4, dropout=0.3,
        )
        ckpt = torch.load(checkpoint_path, map_location=device, weights_only=False)
        model.load_state_dict(ckpt["model_state_dict"])
        model.eval()
        model.to(device)

        scorer = cls(model)
        scorer._backend = "pytorch"
        scorer._device = device
        return scorer

    def predict_raw(self, features: np.ndarray) -> dict[str, np.ndarray]:
        """
        Run inference on features array.

        Args:
            features: (N, 120) or (120,) float32 array.

        Returns:
            Dict with numpy arrays for each output.
        """
        if features.ndim == 1:
            features = features[np.newaxis, :]
        features = features.astype(np.float32)

        if self._backend == "onnx":
            outputs = self._session.run(None, {"features": features})
            names = ["escalation", "churn", "brand", "revenue",
                     "global_risk", "tier_logits", "confidence"]
            return {name: out for name, out in zip(names, outputs)}
        else:
            import torch
            with torch.no_grad():
                x = torch.from_numpy(features).to(self._device)
                preds = self._session(x)
            return {
                "escalation": preds["escalation"].cpu().numpy(),
                "churn": preds["churn"].cpu().numpy(),
                "brand": preds["brand"].cpu().numpy(),
                "revenue": preds["revenue"].cpu().numpy(),
                "global_risk": preds["global_risk"].cpu().numpy(),
                "tier_logits": preds["tier_logits"].cpu().numpy(),
                "confidence": preds["confidence"].cpu().numpy(),
            }

    def predict(self, features: np.ndarray) -> list[RiskAssessment]:
        """
        Run inference and return structured RiskAssessment objects.

        Args:
            features: (N, 120) or (120,) float32 array.

        Returns:
            List of RiskAssessment objects (one per sample).
        """
        single = features.ndim == 1
        raw = self.predict_raw(features)

        n = raw["global_risk"].shape[0]
        results = []

        # Compute tier probabilities from logits
        logits = raw["tier_logits"]
        exp_logits = np.exp(logits - logits.max(axis=-1, keepdims=True))
        probs = exp_logits / exp_logits.sum(axis=-1, keepdims=True)

        for i in range(n):
            tier_idx = int(probs[i].argmax())
            tier_probs = {TIER_NAMES[j]: float(probs[i, j]) for j in range(4)}

            results.append(RiskAssessment(
                global_risk=float(raw["global_risk"][i, 0]),
                tier=TIER_NAMES[tier_idx],
                tier_index=tier_idx,
                escalation=float(raw["escalation"][i, 0]),
                churn=float(raw["churn"][i, 0]),
                brand=float(raw["brand"][i, 0]),
                revenue=float(raw["revenue"][i, 0]),
                confidence=float(raw["confidence"][i, 0]),
                tier_probabilities=tier_probs,
            ))

        return results[0] if single else results


if __name__ == "__main__":
    # Quick demo
    model_dir = Path(__file__).parent
    onnx_path = model_dir / "howzer_risk_scorer.onnx"

    if onnx_path.exists():
        scorer = RiskScorer.from_onnx(onnx_path)
        # Random input
        features = np.random.randn(3, 120).astype(np.float32)
        results = scorer.predict(features)
        for i, r in enumerate(results):
            print(f"Sample {i}: {r.to_dict()}")
    else:
        print(f"ONNX model not found at {onnx_path}")