"""Encoder fine-tuning for multi-label ATT&CK technique classification.

A plain PyTorch loop rather than ``Trainer``: this repo's whole claim is that
its numbers are reproducible, and a visible loop with an explicit loss, an
explicit schedule and an explicit best-checkpoint rule is easier to audit than
a config dict spread across a framework whose argument names drift between
releases.

Class imbalance is handled with per-class ``pos_weight`` in the BCE loss. 79%
of sentences carry no label at all and the rarest retained technique has ~20
examples, so unweighted BCE converges to predicting nothing.
"""

from __future__ import annotations

import json
import math
from pathlib import Path

import numpy as np
import torch
from torch.utils.data import DataLoader, Dataset
from transformers import AutoModelForSequenceClassification, AutoTokenizer, get_linear_schedule_with_warmup

from . import config, evaluate


def device() -> torch.device:
    return torch.device("cuda" if torch.cuda.is_available() else "cpu")


class SentenceDataset(Dataset):
    def __init__(self, records: list[dict], labels: list[str], tokenizer, max_length: int):
        self.texts = [r["sentence"] for r in records]
        self.Y = evaluate.to_matrix(records, labels).astype(np.float32)
        self.tok = tokenizer
        self.max_length = max_length

    def __len__(self) -> int:
        return len(self.texts)

    def __getitem__(self, i: int) -> dict:
        enc = self.tok(
            self.texts[i],
            truncation=True,
            max_length=self.max_length,
            padding="max_length",
            return_tensors="pt",
        )
        item = {k: v.squeeze(0) for k, v in enc.items()}
        item["labels"] = torch.from_numpy(self.Y[i])
        return item


def compute_pos_weight(Y: np.ndarray, cap: float = 50.0) -> torch.Tensor:
    """``(negatives / positives)`` per class, capped so rare classes don't explode."""
    pos = Y.sum(axis=0)
    neg = Y.shape[0] - pos
    with np.errstate(divide="ignore", invalid="ignore"):
        w = np.where(pos > 0, neg / np.maximum(pos, 1), 1.0)
    return torch.tensor(np.clip(w, 1.0, cap), dtype=torch.float32)


@torch.no_grad()
def predict_scores(model, loader, dev, amp: bool = True) -> np.ndarray:
    model.eval()
    out = []
    for batch in loader:
        labels = batch.pop("labels", None)
        batch = {k: v.to(dev) for k, v in batch.items()}
        with torch.autocast("cuda", dtype=torch.bfloat16,
                            enabled=amp and dev.type == "cuda"):
            logits = model(**batch).logits
        out.append(torch.sigmoid(logits.float()).cpu().numpy())
    return np.concatenate(out, axis=0)


def train(
    model_key: str,
    scheme: str,
    train_records: list[dict],
    dev_records: list[dict],
    labels: list[str],
    output_dir: Path,
    epochs: int = config.EPOCHS,
    seed: int = config.SEED,
) -> dict:
    torch.manual_seed(seed)
    np.random.seed(seed)

    dev_ = device()
    model_name = config.BASE_MODELS[model_key]
    amp = model_key not in config.FP32_ONLY_MODELS
    print(f"\n  base model : {model_name}")
    print(f"  device     : {dev_} ({torch.cuda.get_device_name(0) if dev_.type == 'cuda' else 'cpu'})")
    print(f"  precision  : {'bf16 autocast' if amp else 'fp32 (autocast disabled for this model)'}")

    tok = AutoTokenizer.from_pretrained(model_name)
    model = AutoModelForSequenceClassification.from_pretrained(
        model_name,
        num_labels=len(labels),
        problem_type="multi_label_classification",
        id2label={i: l for i, l in enumerate(labels)},
        label2id={l: i for i, l in enumerate(labels)},
    ).to(dev_)

    ds_tr = SentenceDataset(train_records, labels, tok, config.MAX_LENGTH)
    ds_dv = SentenceDataset(dev_records, labels, tok, config.MAX_LENGTH)
    dl_tr = DataLoader(ds_tr, batch_size=config.BATCH_SIZE, shuffle=True, drop_last=False)
    dl_dv = DataLoader(ds_dv, batch_size=config.BATCH_SIZE * 2, shuffle=False)

    pos_weight = compute_pos_weight(ds_tr.Y).to(dev_)
    loss_fn = torch.nn.BCEWithLogitsLoss(pos_weight=pos_weight)

    decay = [p for n, p in model.named_parameters() if not any(x in n for x in ("bias", "LayerNorm.weight", "norm.weight"))]
    no_decay = [p for n, p in model.named_parameters() if any(x in n for x in ("bias", "LayerNorm.weight", "norm.weight"))]
    optim = torch.optim.AdamW(
        [{"params": decay, "weight_decay": config.WEIGHT_DECAY},
         {"params": no_decay, "weight_decay": 0.0}],
        lr=config.LEARNING_RATE,
    )

    steps_per_epoch = math.ceil(len(dl_tr) / config.GRAD_ACCUM)
    total_steps = steps_per_epoch * epochs
    sched = get_linear_schedule_with_warmup(
        optim, int(total_steps * config.WARMUP_RATIO), total_steps)

    Ydv = ds_dv.Y.astype(np.int8)
    best = {"macro_f1": -1.0, "epoch": -1, "threshold": 0.5}
    output_dir.mkdir(parents=True, exist_ok=True)

    for epoch in range(1, epochs + 1):
        model.train()
        running, nb = 0.0, 0
        optim.zero_grad(set_to_none=True)
        for step, batch in enumerate(dl_tr):
            y = batch.pop("labels").to(dev_)
            batch = {k: v.to(dev_) for k, v in batch.items()}
            with torch.autocast("cuda", dtype=torch.bfloat16,
                                enabled=amp and dev_.type == "cuda"):
                logits = model(**batch).logits
            loss = loss_fn(logits.float(), y) / config.GRAD_ACCUM

            # Fail loudly. A silent nan run previously trained for six full
            # epochs, saved a checkpoint and reported macro-F1 0.0000 as though
            # it were a legitimate result.
            if not torch.isfinite(loss):
                raise RuntimeError(
                    f"non-finite loss at epoch {epoch} step {step} for "
                    f"{model_name!r} (amp={'bf16' if amp else 'off'}). Training "
                    f"aborted rather than reporting a meaningless score. If this "
                    f"model is new, add it to config.FP32_ONLY_MODELS."
                )

            loss.backward()
            running += loss.item() * config.GRAD_ACCUM
            nb += 1
            if (step + 1) % config.GRAD_ACCUM == 0 or step + 1 == len(dl_tr):
                torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
                optim.step()
                sched.step()
                optim.zero_grad(set_to_none=True)

        scores = predict_scores(model, dl_dv, dev_, amp=amp)
        thr, dev_macro = evaluate.tune_global_threshold(Ydv, scores)
        print(f"  epoch {epoch}/{epochs}  loss={running / max(nb, 1):.4f}  "
              f"dev macro-F1={dev_macro:.4f} (t={thr})")

        if dev_macro > best["macro_f1"]:
            best = {"macro_f1": dev_macro, "epoch": epoch, "threshold": thr}
            model.save_pretrained(output_dir)
            tok.save_pretrained(output_dir)
            (output_dir / "labels.json").write_text(
                json.dumps(labels, indent=2), encoding="utf-8")

    (output_dir / "training_meta.json").write_text(json.dumps({
        "base_model": model_name,
        "split_scheme": scheme,
        "epochs": epochs,
        "best_epoch": best["epoch"],
        "best_dev_macro_f1": round(best["macro_f1"], 4),
        "best_dev_global_threshold": best["threshold"],
        "max_length": config.MAX_LENGTH,
        "batch_size": config.BATCH_SIZE,
        "grad_accum": config.GRAD_ACCUM,
        "learning_rate": config.LEARNING_RATE,
        "seed": seed,
        "n_train": len(train_records),
        "n_dev": len(dev_records),
        "n_labels": len(labels),
    }, indent=2), encoding="utf-8")

    print(f"  best: epoch {best['epoch']}, dev macro-F1 {best['macro_f1']:.4f}")
    return best


def load_for_inference(model_dir: str | Path):
    tok = AutoTokenizer.from_pretrained(str(model_dir))
    model = AutoModelForSequenceClassification.from_pretrained(str(model_dir))
    model.eval()
    return model, tok