"""Metrics and threshold selection for multi-label ATT&CK classification.

Design notes that matter for reading the numbers:

* **macro-F1 is the headline, not micro-F1.** 79% of sentences carry no label
  and the technique distribution is heavily long-tailed (T1027 has 678 examples,
  T1072 has ~20). Micro-F1 is dominated by a handful of frequent techniques and
  flatters every model. Macro-F1 is what tells you whether the tail works.
* **Two threshold regimes are reported.** A single global threshold, and
  per-class thresholds tuned on dev. Per-class tuning fits 49 free parameters on
  a dev set with very few positives per class, so it can overfit; reporting both
  makes the size of that effect visible instead of hiding it.
"""

from __future__ import annotations

import json
from dataclasses import dataclass

import numpy as np

from . import config


def to_matrix(records: list[dict], labels: list[str]) -> np.ndarray:
    index = {l: i for i, l in enumerate(labels)}
    Y = np.zeros((len(records), len(labels)), dtype=np.int8)
    for i, rec in enumerate(records):
        for l in rec["labels"]:
            j = index.get(l)
            if j is not None:
                Y[i, j] = 1
    return Y


def _prf(tp: float, fp: float, fn: float) -> tuple[float, float, float]:
    p = tp / (tp + fp) if tp + fp else 0.0
    r = tp / (tp + fn) if tp + fn else 0.0
    f = 2 * p * r / (p + r) if p + r else 0.0
    return p, r, f


@dataclass
class Report:
    micro_f1: float
    macro_f1: float
    micro_precision: float
    micro_recall: float
    macro_precision: float
    macro_recall: float
    per_class: dict
    n_examples: int
    n_positive_labels: int

    def as_dict(self) -> dict:
        return {
            "micro_f1": round(self.micro_f1, 4),
            "macro_f1": round(self.macro_f1, 4),
            "micro_precision": round(self.micro_precision, 4),
            "micro_recall": round(self.micro_recall, 4),
            "macro_precision": round(self.macro_precision, 4),
            "macro_recall": round(self.macro_recall, 4),
            "n_examples": self.n_examples,
            "n_positive_labels": self.n_positive_labels,
            "per_class": self.per_class,
        }


def evaluate(Y_true: np.ndarray, Y_pred: np.ndarray, labels: list[str]) -> Report:
    tp = (Y_true & Y_pred).sum(axis=0).astype(float)
    fp = ((1 - Y_true) & Y_pred).sum(axis=0).astype(float)
    fn = (Y_true & (1 - Y_pred)).sum(axis=0).astype(float)

    mip, mir, mif = _prf(tp.sum(), fp.sum(), fn.sum())

    per_class, ps, rs, fs = {}, [], [], []
    for j, name in enumerate(labels):
        p, r, f = _prf(tp[j], fp[j], fn[j])
        per_class[name] = {
            "precision": round(p, 4),
            "recall": round(r, 4),
            "f1": round(f, 4),
            "support": int(Y_true[:, j].sum()),
            "predicted": int(Y_pred[:, j].sum()),
        }
        ps.append(p)
        rs.append(r)
        fs.append(f)

    return Report(
        micro_f1=mif,
        macro_f1=float(np.mean(fs)),
        micro_precision=mip,
        micro_recall=mir,
        macro_precision=float(np.mean(ps)),
        macro_recall=float(np.mean(rs)),
        per_class=per_class,
        n_examples=int(Y_true.shape[0]),
        n_positive_labels=int(Y_true.sum()),
    )


def tune_global_threshold(
    Y_true: np.ndarray, scores: np.ndarray, grid=None
) -> tuple[float, float]:
    """Pick the single threshold maximising macro-F1 on the given (dev) set."""
    grid = grid or config.THRESHOLD_GRID
    best_t, best_f = 0.5, -1.0
    for t in grid:
        Y_pred = (scores >= t).astype(np.int8)
        tp = (Y_true & Y_pred).sum(axis=0).astype(float)
        fp = ((1 - Y_true) & Y_pred).sum(axis=0).astype(float)
        fn = (Y_true & (1 - Y_pred)).sum(axis=0).astype(float)
        f = float(np.mean([_prf(tp[j], fp[j], fn[j])[2] for j in range(Y_true.shape[1])]))
        if f > best_f:
            best_t, best_f = t, f
    return best_t, best_f


def tune_per_class_thresholds(
    Y_true: np.ndarray, scores: np.ndarray, grid=None
) -> np.ndarray:
    """One threshold per technique, each maximising that technique's dev F1."""
    grid = grid or config.THRESHOLD_GRID
    out = np.full(Y_true.shape[1], 0.5)
    for j in range(Y_true.shape[1]):
        best_t, best_f = 0.5, -1.0
        yt = Y_true[:, j]
        for t in grid:
            yp = (scores[:, j] >= t).astype(np.int8)
            _, _, f = _prf(
                float((yt & yp).sum()),
                float(((1 - yt) & yp).sum()),
                float((yt & (1 - yp)).sum()),
            )
            if f > best_f:
                best_t, best_f = t, f
        out[j] = best_t
    return out


def apply_thresholds(scores: np.ndarray, thresholds) -> np.ndarray:
    return (scores >= np.asarray(thresholds)).astype(np.int8)


def save_report(name: str, scheme: str, reports: dict) -> None:
    path = config.RESULTS_DIR / f"{name}__{scheme}.json"
    path.write_text(json.dumps(reports, indent=2), encoding="utf-8")
    print(f"  -> {path.relative_to(config.REPO_ROOT)}")