"""Classification metrics (numpy-only) + metrics.json helpers.

No sklearn dependency — macro-F1 and the fault-vs-normal binary F1 are computed
directly so the node install stays light. ``metrics.json`` is a JSON list of
per-round records that both topologies append to and the notebook reads back.
"""

from __future__ import annotations

import json
import os
from typing import List

import numpy as np


def _prf(tp: int, fp: int, fn: int):
    precision = tp / (tp + fp) if (tp + fp) else 0.0
    recall = tp / (tp + fn) if (tp + fn) else 0.0
    f1 = 2 * precision * recall / (precision + recall) if (precision + recall) else 0.0
    return precision, recall, f1


def classification_metrics(
    y_true: np.ndarray,
    y_pred: np.ndarray,
    num_classes: int,
    fault_class: int,
    normal_class: int = 0,
) -> dict:
    """Accuracy, macro-F1, plus two headline binary detectors.

    - ``fault_*``: the specific incipient-fault class (one-vs-rest).
    - ``disturbance_*``: any-disturbance vs normal (normal_class one-vs-rest).
    """
    y_true = np.asarray(y_true)
    y_pred = np.asarray(y_pred)
    n = len(y_true)
    accuracy = float(np.mean(y_true == y_pred)) if n else 0.0

    f1s = []
    for c in range(num_classes):
        tp = int(np.sum((y_pred == c) & (y_true == c)))
        fp = int(np.sum((y_pred == c) & (y_true != c)))
        fn = int(np.sum((y_pred != c) & (y_true == c)))
        f1s.append(_prf(tp, fp, fn)[2])
    macro_f1 = float(np.mean(f1s)) if f1s else 0.0

    # Incipient-fault class, one-vs-rest.
    tp = int(np.sum((y_pred == fault_class) & (y_true == fault_class)))
    fp = int(np.sum((y_pred == fault_class) & (y_true != fault_class)))
    fn = int(np.sum((y_pred != fault_class) & (y_true == fault_class)))
    fault_p, fault_r, fault_f1 = _prf(tp, fp, fn)

    # Any-disturbance (not normal) vs normal.
    dist_true = y_true != normal_class
    dist_pred = y_pred != normal_class
    tp = int(np.sum(dist_pred & dist_true))
    fp = int(np.sum(dist_pred & ~dist_true))
    fn = int(np.sum(~dist_pred & dist_true))
    dist_p, dist_r, dist_f1 = _prf(tp, fp, fn)

    return {
        "accuracy": round(accuracy, 5),
        "macro_f1": round(macro_f1, 5),
        "fault_precision": round(fault_p, 5),
        "fault_recall": round(fault_r, 5),
        "fault_f1": round(fault_f1, 5),
        "disturbance_precision": round(dist_p, 5),
        "disturbance_recall": round(dist_r, 5),
        "disturbance_f1": round(dist_f1, 5),
        "n_eval": int(n),
    }


def reset(path: str) -> None:
    """Start a fresh metrics file (empty JSON list), creating the dir if needed.

    Call once at process startup so a re-run overwrites rather than appends —
    otherwise repeated runs concatenate into one file (round sequence 1..N, 1..N)
    and downstream per-round grouping silently averages unrelated runs.
    """
    os.makedirs(os.path.dirname(os.path.abspath(path)), exist_ok=True)
    tmp = f"{path}.tmp"
    with open(tmp, "w") as fh:
        json.dump([], fh)
    os.replace(tmp, path)


def append_record(path: str, record: dict) -> None:
    """Append one per-round record to a JSON-list file (creating it if needed)."""
    os.makedirs(os.path.dirname(os.path.abspath(path)), exist_ok=True)
    records = load_records(path)
    records.append(record)
    tmp = f"{path}.tmp"
    with open(tmp, "w") as fh:
        json.dump(records, fh, indent=2)
    os.replace(tmp, path)


def load_records(path: str) -> List[dict]:
    if not os.path.exists(path):
        return []
    try:
        with open(path) as fh:
            data = json.load(fh)
        return data if isinstance(data, list) else []
    except (json.JSONDecodeError, OSError):
        return []
