from typing import Dict, Optional
import numpy as np
from config import NUM_CLASSES, CLASS_NAMES, IGNORE_INDEX


def compute_metrics(pred: np.ndarray, gt: np.ndarray, num_classes: int = NUM_CLASSES) -> Dict:
    pred = pred.astype(np.int64)
    gt   = gt.astype(np.int64)

    labeled  = gt != IGNORE_INDEX
    pred_l   = pred[labeled]
    gt_l     = gt[labeled]

    if len(gt_l) == 0:
        return {
            "overall_acc":      0.0,
            "miou":             0.0,
            "per_class_acc":    [None] * num_classes,
            "per_class_iou":    [None] * num_classes,
            "confusion_matrix": [[0] * num_classes] * num_classes,
        }

    cm = np.zeros((num_classes, num_classes), dtype=np.int64)
    for g, p in zip(gt_l, pred_l):
        if 0 <= g < num_classes and 0 <= p < num_classes:
            cm[g, p] += 1

    overall_acc = float((gt_l == pred_l).mean())
    per_class_acc, per_class_iou = [], []
    for c in range(num_classes):
        tp         = cm[c, c]
        gt_total   = cm[c, :].sum()
        pred_total = cm[:, c].sum()
        union      = gt_total + pred_total - tp
        per_class_acc.append(float(tp / gt_total) if gt_total > 0 else None)
        per_class_iou.append(float(tp / union)    if union > 0    else None)

    miou = float(np.nanmean([x if x is not None else np.nan for x in per_class_iou]))
    return {
        "overall_acc":      overall_acc,
        "miou":             miou,
        "per_class_acc":    per_class_acc,
        "per_class_iou":    per_class_iou,
        "confusion_matrix": cm.tolist(),
    }


def metrics_markdown(metrics: Dict, title: str = "Metrics") -> str:
    lines = [f"### {title}"]
    lines.append(f"- Précision globale : **{metrics['overall_acc'] * 100:.2f}%**")
    lines.append(f"- IoU moyen : **{metrics['miou'] * 100:.2f}%**")
    lines.append("")
    lines.append("| Classe | Précision | IoU |")
    lines.append("|---|---:|---:|")
    for name, acc, iou in zip(CLASS_NAMES, metrics["per_class_acc"], metrics["per_class_iou"]):
        acc_s = "—" if acc is None else f"{acc * 100:.1f}%"
        iou_s = "—" if iou is None else f"{iou * 100:.1f}%"
        lines.append(f"| {name} | {acc_s} | {iou_s} |")
    return "\n".join(lines)