models/engine.py

"""VXOneEngine — contrato de inferencia (spec F2 §8.1).

API tipada sobre un checkpoint self-contained (.pt, §8.2): la inferencia no
necesita datos de train (feature stage, vocabularios, ω, temperaturas y τ
viajan dentro del artefacto).

Invariantes congelados:
- candidate == "unknown" ⟺ (modo abstención y max_prob < τ_abs) o (modo
  trained_unknown y unknown ganó el argmax).
- probs (Choice) suma 1 sobre los candidatos NATIVOS de la corrida (en EMBER
  unknown ES candidato nativo; en UNSW/CIC nunca aparece como masa).
- policy se deriva de las salidas calibradas (recomputable, sin estado).
"""

from __future__ import annotations

from dataclasses import dataclass
from pathlib import Path

import numpy as np
import torch

from .feature_stage import FeatureStage
from .losses import corn_pmf, corn_expected
from .network import VXOneNet, HeadSpec
from .vocab import TextVocab

TAU_DET = 0.5   # g_attack (controlador; defaults serializados, §3.1 D9)
TAU_DROP = 0.5
TAU_RECON = 0.5
TAU_ABS = 0.5

SOURCE_KEY = {"unsw": "unsw_nb15_v1", "cic": "cicids2017_v1", "ember": "ember2024_v1"}
SOURCE_ALIAS = {v: k for k, v in SOURCE_KEY.items()}
LEVELS = ("benign", "low", "high", "critical")


@dataclass(frozen=True)
class ChoiceResult:
    candidate: str                 # ∈ candidatos ∪ {"unknown"}
    probs: dict[str, float]        # calibradas + de-tilt; suma 1 sobre nativos
    max_prob: float                # max_k p_cal (auditoría de abstención)


@dataclass(frozen=True)
class ScoreResult:
    expected: float                # ŝ = Σ k·π_k ∈ [0, 3]
    probs: dict[str, float]        # π calibrada sobre benign/low/high/critical


@dataclass(frozen=True)
class Decision:
    source: str
    noul: dict[str, float]         # {"is_attack"|"is_malicious": p_cal}
    choice: dict[str, ChoiceResult]
    score: ScoreResult | None      # None en EMBER (S6)
    policy: dict[str, bool]        # gates derivados


class VXOneEngine:
    def __init__(self, artifact: str | Path, device: str = "cpu") -> None:
        ckpt = torch.load(artifact, map_location=device, weights_only=False)
        self.ckpt = ckpt  # introspección (E4 latencia: config, d_num)
        self.device = device
        self.source: str = SOURCE_ALIAS[ckpt["source"]]   # alias corto: unsw|cic|ember
        self.source_key: str = ckpt["source"]             # clave del registro
        self.candidates: list[str] = ckpt["head_meta"]["choice"]["candidates"]
        self.mode: str = ckpt["head_meta"]["choice"]["mode"]  # abstain|trained_unknown
        self.tau_abs: float = ckpt["head_meta"]["choice"]["tau_abs"]
        self.stage = FeatureStage.from_dict(ckpt["feature_stage"])
        self.vocab = {k: TextVocab.from_dict(v) for k, v in ckpt.get("vocab", {}).items()}
        self.omega = torch.tensor(ckpt["head_meta"]["choice"]["class_weights"])
        self.T_noul = ckpt["head_meta"]["noul"]["temperature"]
        self.T_choice = ckpt["head_meta"]["choice"]["temperature"]
        self.T_score = ckpt["head_meta"].get("score", {}).get("temperature")
        self.has_ddos = "ddos" in ckpt["head_meta"]
        if self.has_ddos:
            self.ddos_candidates = ckpt["head_meta"]["ddos"]["candidates"]
            self.omega_ddos = torch.tensor(ckpt["head_meta"]["ddos"]["class_weights"])
            self.T_ddos = ckpt["head_meta"]["ddos"]["temperature"]
        else:
            self.ddos_candidates = self.omega_ddos = self.T_ddos = None

        d_num = ckpt["config"]["d_num"]
        heads = HeadSpec(n_candidates=len(self.candidates),
                         ddos_candidates=len(self.ddos_candidates) if self.has_ddos else 0,
                         has_score=self.T_score is not None)
        self.net = VXOneNet(self.source, ckpt["config"]["scale"], heads, d_num,
                            vocab_sizes={k: v.size for k, v in self.vocab.items()})
        self.net.load_state_dict(ckpt["state_dict"])
        self.net.to(device).eval()

    # ── preparación de entrada ───────────────────────────────────────
    def _prepare(self, x_num: np.ndarray, x_txt: dict[str, list[str | None]] | None):
        x_num = np.asarray(x_num, dtype=np.float64)
        u_r, _ = self.stage.transform(x_num)
        slots = None
        if self.vocab:
            n = x_num.shape[0]
            slots = {f: torch.tensor(
                self.vocab[f].encode((x_txt or {}).get(f, [None] * n)), dtype=torch.int64)
                for f in self.vocab}
        return torch.tensor(u_r, device=self.device), slots

    # ── predicción ───────────────────────────────────────────────────
    @torch.no_grad()
    def predict_batch(self, x_num: np.ndarray,
                      x_txt: dict[str, list[str | None]] | None = None) -> list[Decision]:
        u_r, slots = self._prepare(x_num, x_txt)
        out = self.net(u_r, slots)

        # Noul calibrado
        p_mal = torch.sigmoid(out["noul"] / self.T_noul)
        noul_key = "is_attack" if self.source in ("unsw", "cic") else "is_malicious"

        # Choice: de-tilt → T → probs sobre nativos
        z = out["choice"] - torch.log(self.omega).unsqueeze(0).to(self.device)
        p_choice = torch.softmax(z / self.T_choice, dim=1)

        # Score CORN calibrado (None en EMBER)
        pi = corn_pmf(out["corn_logits"] / self.T_score)[1] if self.T_score is not None else None
        s_exp = corn_expected(pi) if pi is not None else None

        # ddos vector (solo CIC strat)
        ddos_probs = None
        if self.has_ddos:
            z_d = out["ddos"] - torch.log(self.omega_ddos).unsqueeze(0).to(self.device)
            ddos_probs = torch.softmax(z_d / self.T_ddos, dim=1)

        decisions: list[Decision] = []
        for i in range(u_r.shape[0]):
            # Choice con unknown (D1): abstención τ_abs o unknown entrenado
            pmax = float(p_choice[i].max())
            if self.mode == "abstain":
                cand = self.candidates[int(p_choice[i].argmax())] if pmax >= self.tau_abs else "unknown"
            else:  # trained_unknown (EMBER): unknown es candidato nativo
                cand = self.candidates[int(p_choice[i].argmax())]
            probs = {c: float(p_choice[i, j]) for j, c in enumerate(self.candidates)}

            choice = {"family": ChoiceResult(cand, probs, pmax)}
            if ddos_probs is not None:
                dp = {c: float(ddos_probs[i, j]) for j, c in enumerate(self.ddos_candidates)}
                choice["ddos_vector"] = ChoiceResult(
                    self.ddos_candidates[int(ddos_probs[i].argmax())], dp,
                    float(ddos_probs[i].max()))

            score = (ScoreResult(float(s_exp[i]),
                                 {lv: float(pi[i, k]) for k, lv in enumerate(LEVELS)})
                     if pi is not None else None)

            # Gates derivados (D9): funciones exactas de salidas calibradas
            policy = {"g_attack": bool(p_mal[i] >= TAU_DET)}
            if pi is not None:
                # g_drop: P(sev ≥ 2) = C2 = π2+π3 (acumulada CORN exacta)
                policy["g_drop"] = bool(float(pi[i, 2] + pi[i, 3]) >= TAU_DROP)
            if "recon" in self.candidates:
                policy["g_recon"] = bool(probs.get("recon", 0.0) >= TAU_RECON)

            decisions.append(Decision(
                source=self.source,
                noul={noul_key: float(p_mal[i])},
                choice=choice,
                score=score,
                policy=policy,
            ))
        return decisions

    def predict(self, x_num: np.ndarray, x_txt: dict[str, str | None] | None = None) -> Decision:
        """batch-1 sin maquinaria de batch (E4)."""
        xt = {k: [v] for k, v in (x_txt or {}).items()} if x_txt else None
        return self.predict_batch(np.asarray(x_num)[None, ...], xt)[0]


def save_artifact(path: str | Path, net: VXOneNet, run_meta: dict,
                  temperatures: dict[str, float], d_num: int) -> None:
    """Escribe el checkpoint self-contained (§8.2)."""
    ckpt = {
        "source": run_meta["source_key"],
        "feature_stage": run_meta["feature_stage"].to_dict(),
        "vocab": {k: v.to_dict() for k, v in (run_meta.get("vocab") or {}).items()},
        "head_meta": {
            "noul": {"temperature": temperatures["noul"]},
            "choice": {"candidates": run_meta["candidates"],
                       "class_weights": run_meta["omega"].numpy(),
                       "mode": run_meta["mode"],
                       "tau_abs": run_meta.get("tau_abs", TAU_ABS),
                       "temperature": temperatures["choice"]},
            **({"score": {"temperature": temperatures["score"]}}
               if "score" in temperatures else {}),
            **({"ddos": {"candidates": run_meta["ddos_candidates"],
                         "class_weights": run_meta["omega_ddos"].numpy(),
                         "temperature": temperatures["ddos"]}}
               if "ddos" in temperatures else {}),
        },
        "config": {**run_meta["config"], "d_num": d_num},
        "manifest": run_meta["manifest"],
    }
    ckpt["state_dict"] = net.state_dict()
    torch.save(ckpt, path)