models/data.py

"""Datasets y loaders del motor VX-One (spec F2 §0.1, §6.2, D12).

- UNSW / CIC: tensores completos en RAM (UNSW 149K×40×4B ≈ 24 MB;
  CIC train 1.98M×78×4B ≈ 620 MB).
- EMBER: streaming por shards (46 train parquets de 60K filas; ~616 MB
  decodificados c/u): por epoch se permutan los shards (seed 42+epoch),
  shuffle intra-shard, batches con drop_last; prefetch de 1 shard en thread.
  Shuffle de dos niveles (shard + intra) — aproximación documentada (§6.2).

Candidate set por corrida (D12): familias nativas con ≥1 fila en TRAIN bajo
el esquema de la corrida, en el ORDEN CONGELADO del registro. Pesos ω
(§3.5) derivados de los conteos de train DE LA CORRIDA (LOAO los recalcula).

φ (feature stage) se ajusta SOLO con train:
- UNSW/CIC: train completo en RAM (float64).
- EMBER: submuestra uniforme de 60,000 filas de train (1,304 por shard,
  rng seed 42, una pasada por los 46 shards) — la SE de la mediana con
  n=60K es despreciable y evita 28 GB en RAM; decisión documentada.
"""

from __future__ import annotations

import threading
from dataclasses import dataclass, field
from pathlib import Path

import numpy as np
import pyarrow.parquet as pq
import torch

from src.state.state import SOURCE_REGISTRY

from .feature_stage import FeatureStage
from .losses import class_weights
from .network import HeadSpec
from .vocab import TextVocab

ROOT = Path(__file__).resolve().parents[2]
DATA = ROOT / "data" / "processed"

REGISTRY = SOURCE_REGISTRY
ALIAS = {"unsw": "unsw_nb15_v1", "cic": "cicids2017_v1", "ember": "ember2024_v1"}


def _registry(source: str) -> dict:
    """Alias corto (unsw/cic/ember) o clave de registro → entrada del registro."""
    return REGISTRY[ALIAS.get(source, source)]

DDOS_DAYS = {
    "Wednesday-workingHours.pcap_ISCX",
    "Friday-WorkingHours-Afternoon-DDos.pcap_ISCX",
}
DDOS_FULL = {"BENIGN", "DDoS", "DoS Hulk", "DoS GoldenEye", "DoS slowloris", "DoS Slowhttptest"}
DDOS_CANDIDATES = ("normal", "tcp_flood", "http_dos", "slowloris")  # orden congelado F2 D8
_DDOS_MAP = {"BENIGN": "normal", "DDoS": "tcp_flood", "DoS Hulk": "http_dos",
             "DoS GoldenEye": "http_dos", "DoS slowloris": "slowloris",
             "DoS Slowhttptest": "slowloris"}

EMBER_STATS_ROWS_PER_SHARD = 1_304  # × 46 shards = 60,024 filas de stats


@dataclass
class Labels:
    y_bin: torch.Tensor        # int64 [N] 1[family != normal] (o label_binary en EMBER)
    y_fam: torch.Tensor        # int64 [N] índice en candidates
    y_sev: torch.Tensor | None # int64 [N] 0..3 (None en EMBER)
    ddos_mask: torch.Tensor | None  # bool [N] filas del subconjunto (solo CIC strat)
    y_ddos: torch.Tensor | None     # int64 [N] índice en DDOS_CANDIDATES


@dataclass
class RunData:
    """Todo lo que una corrida necesita, materializado y verificado."""
    source: str
    scale: str
    split_scheme: str          # "oficial" (UNSW) | "strat" | "time" (CIC) | "temporal" (EMBER)
    candidates: list[str]
    ddos_candidates: list[str] | None
    omega: torch.Tensor        # float32 [K]
    omega_ddos: torch.Tensor | None
    vocab: dict[str, TextVocab] | None
    feature_stage: FeatureStage
    heads: HeadSpec
    train: torch.Tensor        # u_r float32 [Ntr, d+1]
    train_txt: dict[str, torch.Tensor] | None
    train_labels: Labels
    val: torch.Tensor
    val_txt: dict[str, torch.Tensor] | None
    val_labels: Labels
    test: torch.Tensor | None  # None en EMBER (streaming, no se materializa)
    test_txt: dict[str, torch.Tensor] | None
    test_labels: Labels | None
    val_fam_names: np.ndarray | None = None   # familia nativa por fila (E2; -1 = fuera de candidatos)
    test_fam_names: np.ndarray | None = None
    ember_train_shards: list[Path] | None = None
    ember_val_shards: list[Path] | None = None
    ember_test_shards: list[Path] | None = None
    meta: dict = field(default_factory=dict)


# ── utilidades de lectura ──────────────────────────────────────────────
def _read_x_num(path: Path) -> np.ndarray:
    tbl = pq.read_table(path, columns=["x_num"])
    arr = tbl.column("x_num").combine_chunks()
    flat = arr.flatten().to_numpy(zero_copy_only=False)
    d = arr.type.list_size
    return np.asarray(flat, dtype=np.float64).reshape(-1, d)


def _frozen_candidates(source: str, fam_counts: dict[str, int]) -> list[str]:
    """D12: familias nativas con ≥1 fila en train, en orden congelado."""
    native = _registry(source)["native_families"]
    cands = [f for f in native if fam_counts.get(f, 0) > 0]
    faltan = sorted(set(native) - set(cands))
    if faltan:
        print(f"[d12] {source}: sin filas en train → fuera del candidate set: {faltan}")
    return cands


def _idx_map(candidates: list[str]) -> dict[str, int]:
    return {c: i for i, c in enumerate(candidates)}


# ── UNSW / CIC: datasets en RAM ────────────────────────────────────────
def _load_network_split(path: Path, source: str, split_col: str, split: str,
                        candidates: list[str], need_ddos: bool,
                        exclude_family: str | None = None
                        ) -> tuple[np.ndarray, dict[str, list[str | None]], Labels]:
    txt_keys = ("proto", "service", "state") if source == "unsw" else ("protocol",)
    cols = ["x_num", split_col, "label_family", "label_severity"]
    cols += [f"x_txt_{k}" for k in txt_keys]
    if need_ddos:
        cols += ["day", "label_full"]  # el subset usa etiquetas crudas (label_full)
    tbl = pq.read_table(path, columns=cols)
    mask = np.asarray(tbl.column(split_col).to_pylist()) == split
    x = _read_x_num(path)[mask]
    fam = np.asarray(tbl.column("label_family").to_pylist())[mask]
    sev = np.asarray(tbl.column("label_severity").to_pylist())[mask]
    txt = {k: list(np.asarray(tbl.column(f"x_txt_{k}").to_pylist())[mask])
           for k in txt_keys}

    # LOAO: la familia excluida sale de train Y val (§6.7); test se evalúa completo
    if exclude_family is not None and split != "test":
        keep = fam != exclude_family
        x, fam, sev = x[keep], fam[keep], sev[keep]
        txt = {k: [v for v, kp in zip(vals, keep) if kp] for k, vals in txt.items()}

    imap = _idx_map(candidates)
    # filas con familia fuera del candidate set (D12: excluida LOAO en test,
    # o recon/botnet en val/test de split_time) → sentinela -1; los nombres
    # se conservan para E2 (evaluación de la familia excluida, §6.7)
    y_fam = torch.tensor([imap.get(f, -1) for f in fam], dtype=torch.int64)
    y_bin = torch.tensor([f != "normal" for f in fam], dtype=torch.int64)
    y_sev = torch.tensor(sev, dtype=torch.int64)

    if need_ddos:
        # subset congelado (SCHEMA §3.3): días Wednesday+Friday-DDoS; solo
        # mapean a las 4 clases las filas con etiqueta nativa en DDOS_FULL
        # (642,862 de las 918,437; el resto del día no entrena este head)
        day = np.asarray(tbl.column("day").to_pylist())[mask]
        full = np.asarray(tbl.column("label_full").to_pylist())[mask]
        if exclude_family is not None and split != "test":
            day, full = day[keep], full[keep]  # misma poda LOAO que x/fam
        in_subset = np.asarray([d in DDOS_DAYS and f in DDOS_FULL
                                for d, f in zip(day, full)])
        ddos_mask = torch.tensor(in_subset)
        y_ddos = torch.zeros(len(fam), dtype=torch.int64)  # alineado a la fila
        y_ddos[ddos_mask] = torch.tensor(
            [_idx_map(DDOS_CANDIDATES)[_DDOS_MAP[f]] for f in np.asarray(full)[in_subset]],
            dtype=torch.int64)
    else:
        ddos_mask = y_ddos = None
    return x, txt, Labels(y_bin, y_fam, y_sev, ddos_mask, y_ddos), fam


def build_network_run(source: str, scale: str, split_scheme: str,
                      exclude_family: str | None = None) -> RunData:
    path = DATA / ("unsw_nb15/unsw_nb15.parquet" if source == "unsw"
                   else "cicids2017/cicids2017.parquet")
    split_col = "split" if source == "unsw" else f"split_{split_scheme}"
    need_ddos = source == "cic" and split_scheme == "strat"

    # 1) conteos de train de la corrida → candidates (D12) + ω (§3.5)
    count_cols = [split_col, "label_family"] + (["day"] if need_ddos else [])
    tbl = pq.read_table(path, columns=count_cols)
    tr = np.asarray(tbl.column(split_col).to_pylist()) == "train"
    fam_all = np.asarray(tbl.column("label_family").to_pylist())
    if exclude_family is not None:
        tr = tr & (fam_all != exclude_family)
    counts = {f: int((fam_all[tr] == f).sum()) for f in _registry(source)["native_families"]}
    candidates = _frozen_candidates(source, counts)
    omega = torch.tensor(class_weights([counts[c] for c in candidates]))

    # 2) splits (LOAO: la familia sale de train Y val; test completo §6.7)
    x_tr, txt_tr, lab_tr, _ = _load_network_split(path, source, split_col, "train", candidates, need_ddos, exclude_family)
    ddos_counts = None
    if need_ddos:
        sub = lab_tr.ddos_mask
        ddos_counts = [int((lab_tr.y_ddos[sub] == i).sum())
                       for i in range(len(DDOS_CANDIDATES))]
        if any(c == 0 for c in ddos_counts):
            # fold LOAO dos_ddos: las clases de ataque del vector ddos viven
            # en la familia excluida (el subset conserva solo BENIGN) → sin
            # head ddos en esta corrida (E2 evalúa noul/choice en la familia)
            need_ddos = False
            ddos_counts = None
            lab_tr = Labels(lab_tr.y_bin, lab_tr.y_fam, lab_tr.y_sev, None, None)
    x_va, txt_va, lab_va, va_fams = _load_network_split(path, source, split_col, "val", candidates, need_ddos, exclude_family)
    x_te, txt_te, lab_te, te_fams = _load_network_split(path, source, split_col, "test", candidates, need_ddos, None)

    # 3) vocab (train de la corrida, §4.2) + φ (train)
    vocab = None
    if source == "unsw":
        vocab = {k: TextVocab.fit(k, txt_tr[k]) for k in ("proto", "service", "state")}
    stage = FeatureStage.fit(x_tr.astype(np.float64))

    u_tr, _ = stage.transform(x_tr)
    u_va, _ = stage.transform(x_va)
    u_te, _ = stage.transform(x_te)
    train_txt = {k: torch.tensor(vocab[k].encode(txt_tr[k]), dtype=torch.int64) for k in vocab} if vocab else None
    val_txt = {k: torch.tensor(vocab[k].encode(txt_va[k]), dtype=torch.int64) for k in vocab} if vocab else None
    test_txt = {k: torch.tensor(vocab[k].encode(txt_te[k]), dtype=torch.int64) for k in vocab} if vocab else None

    ddos_candidates = list(DDOS_CANDIDATES) if need_ddos else None
    omega_ddos = None
    if need_ddos:
        assert all(c > 0 for c in ddos_counts), f"ddos class sin filas: {ddos_counts}"
        omega_ddos = torch.tensor(class_weights(ddos_counts))

    heads = HeadSpec(n_candidates=len(candidates), ddos_candidates=4 if need_ddos else 0)
    return RunData(
        source=source, scale=scale, split_scheme=split_scheme,
        candidates=candidates, ddos_candidates=ddos_candidates,
        omega=omega, omega_ddos=omega_ddos, vocab=vocab, feature_stage=stage,
        heads=heads,
        train=torch.tensor(u_tr), train_txt=train_txt, train_labels=lab_tr,
        val=torch.tensor(u_va), val_txt=val_txt, val_labels=lab_va,
        test=torch.tensor(u_te), test_txt=test_txt, test_labels=lab_te,
        val_fam_names=va_fams, test_fam_names=te_fams,
        meta={"train_counts": {c: counts[c] for c in candidates},
              "exclusion": {"exclude": exclude_family} if exclude_family else {}},
    )


# ── EMBER: streaming por shards ────────────────────────────────────────
def ember_shards(subset: str) -> list[Path]:
    d = DATA / "ember2024"
    if subset == "train":
        return [d / f"part_{i:02d}_train.parquet" for i in range(46)]
    if subset == "val":
        return [d / f"part_{i:02d}_val.parquet" for i in range(46, 52)]
    return [d / f"part_{i:02d}_test.parquet" for i in range(52, 64)]


def _ember_labels(candidates: list[str]):
    imap = _idx_map(candidates)

    def decode(tbl) -> tuple[np.ndarray, torch.Tensor, torch.Tensor]:
        fam = tbl.column("label_family").to_pylist()
        # A5-EMBER: 'unknown' fuera del candidate set → sentinela -1 (la
        # máscara de _batch_losses la saca del CE de choice)
        y_fam = torch.tensor([imap.get(f, -1) for f in fam], dtype=torch.int64)
        y_bin = torch.tensor(tbl.column("label_binary").to_pylist(), dtype=torch.int64)
        return y_fam, y_bin
    return decode


def build_ember_run(scale: str, exclude_family: str | None = None) -> RunData:
    tr_shards, va_shards, te_shards = ember_shards("train"), ember_shards("val"), ember_shards("test")

    # 1) conteos de train (una pasada ligera por shard, solo label_family)
    counts = {f: 0 for f in _registry("ember")["native_families"]}
    for p in tr_shards:
        for f in pq.read_table(p, columns=["label_family"]).column("label_family").to_pylist():
            counts[f] += 1
    if exclude_family is not None:
        counts[exclude_family] = 0
    candidates = _frozen_candidates("ember2024_v1", counts)
    omega = torch.tensor(class_weights([counts[c] for c in candidates]))

    # 2) φ con submuestra uniforme de train (decisión documentada arriba)
    rng = np.random.default_rng(42)
    sample_chunks: list[np.ndarray] = []
    for p in tr_shards:
        x = _read_x_num(p).astype(np.float32)
        take = rng.choice(x.shape[0], size=EMBER_STATS_ROWS_PER_SHARD, replace=False)
        sample_chunks.append(x[take].astype(np.float64))
        del x
    stage = FeatureStage.fit(np.concatenate(sample_chunks))
    del sample_chunks

    return RunData(
        source="ember2024_v1", scale=scale, split_scheme="temporal",
        candidates=candidates, ddos_candidates=None,
        omega=omega, omega_ddos=None, vocab=None, feature_stage=stage,
        heads=HeadSpec(n_candidates=len(candidates), has_score=False),
        train=torch.empty(0), train_txt=None,
        train_labels=Labels(torch.empty(0, dtype=torch.int64), torch.empty(0, dtype=torch.int64), None, None, None),
        val=torch.empty(0), val_txt=None,
        val_labels=Labels(torch.empty(0, dtype=torch.int64), torch.empty(0, dtype=torch.int64), None, None, None),
        test=None, test_txt=None, test_labels=None,
        ember_train_shards=tr_shards, ember_val_shards=va_shards, ember_test_shards=te_shards,
        meta={"train_counts": {c: counts[c] for c in candidates}, "exclusion": {"exclude": exclude_family} if exclude_family else {}},
    )


class EmberBatcher:
    """Itera batches de train (o evalúa val) con streaming por shards.

    Train: permutación de shards seed 42+epoch → decode → shuffle intra-shard
    → batches con drop_last. Prefetch de 1 shard en thread (§6.2).
    Eval (val/test/challenge): shards en orden fijo, batches sin drop_last,
    shuffle off.
    """

    def __init__(self, run: RunData) -> None:
        self.run = run
        self.stage = run.feature_stage
        self._decode = _ember_labels(run.candidates)

    def _load_shard(self, path: Path) -> tuple[np.ndarray, torch.Tensor, torch.Tensor]:
        tbl = pq.read_table(path, columns=["x_num", "label_family", "label_binary"])
        x = _read_x_num(path)  # float64 para φ
        u, _ = self.stage.transform(x)
        del x
        y_fam, y_bin = self._decode(tbl)
        return u, y_fam, y_bin

    def _iter_shard_batches(self, path: Path, batch_size: int, rng: np.random.Generator | None):
        u, y_fam, y_bin = self._load_shard(path)
        n = u.shape[0]
        order = rng.permutation(n) if rng is not None else np.arange(n)
        n_batches = n // batch_size if rng is not None else int(np.ceil(n / batch_size))
        for b in range(n_batches):
            idx = order[b * batch_size:(b + 1) * batch_size]
            yield (torch.tensor(u[idx]), y_fam[idx], y_bin[idx])

    def _stream(self, shards: list[Path], batch_size: int, rng: np.random.Generator | None):
        """Genera batches; con prefetch de 1 shard (thread) si rng (train)."""
        if rng is None:  # eval: secuencial simple
            for p in shards:
                yield from self._iter_shard_batches(p, batch_size, None)
            return

        # train: productor en thread descarga el shard i+1 mientras se
        # consumen los batches del shard i (prefetch=1, §6.2).
        import queue

        q: queue.Queue = queue.Queue(maxsize=1)
        exc: list[BaseException] = []

        def producer() -> None:
            try:
                for p in shards:
                    q.put(self._load_shard(p))
            except BaseException as e:
                exc.append(e)
            finally:
                q.put(None)

        t = threading.Thread(target=producer, daemon=True)
        t.start()
        while True:
            item = q.get()
            if item is None:
                break
            u, y_fam, y_bin = item
            n = u.shape[0]
            order = rng.permutation(n)
            for b in range(n // batch_size):  # drop_last
                idx = order[b * batch_size:(b + 1) * batch_size]
                yield (torch.tensor(u[idx]), y_fam[idx], y_bin[idx])
        if exc:
            raise exc[0]

    def train_batches(self, epoch: int, batch_size: int):
        rng = np.random.default_rng(42 + epoch)
        return self._stream(self.run.ember_train_shards, batch_size, rng)

    def eval_batches(self, subset: str, batch_size: int):
        # "train" = streaming en orden (sin shuffle/drop_last): para
        # materializar features en los baselines tabulares
        shards = {"train": self.run.ember_train_shards,
                  "val": self.run.ember_val_shards,
                  "test": self.run.ember_test_shards}[subset]
        return self._stream(shards, batch_size, None)


# ── fábrica ────────────────────────────────────────────────────────────
def build_run(source: str, scale: str, split_scheme: str,
              exclude_family: str | None = None) -> RunData:
    if source == "unsw":
        return build_network_run("unsw", scale, "oficial", exclude_family)
    if source == "cic":
        return build_network_run("cic", scale, split_scheme, exclude_family)
    if source == "ember":
        return build_ember_run(scale, exclude_family)
    raise ValueError(source)