#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
evaluate.py
End-to-end evaluation for the TACO→YOLOv8-Seg setup with Alami mapping.

Measures:
- Image-level classification (derived from detection/segmentation):
    * macro_accuracy
    * per_class_f1
    * confusion_matrix (labels x preds)
- Calibration:
    * ECE (before calibration)
    * ECE (after calibration), if calibration_temp.json exists
- Optional (if data provided):
    * MAE (kg) for weight (requires val_weight_groundtruth.json + val_weight_predictions.json)

Inputs:
- artifacts/<version>/val_predictions.json    (from train_yolov8_seg.py)
- artifacts/<version>/dataset.yaml            (copy from prepare_taco.py)
- artifacts/<version>/calibration_temp.json   (from calibrate_temp.py) [optional]
- artifacts/<version>/val_weight_groundtruth.json [optional]
- artifacts/<version>/val_weight_predictions.json [optional]

Output:
- artifacts/<version>/summary_eval.json
"""

import argparse
import json
import math
import sys
from pathlib import Path
from typing import Dict, List, Tuple, Any, Optional

import numpy as np
try:
    import yaml
except ImportError:
    print("Please `pip install pyyaml`", file=sys.stderr); raise


# -----------------------
# IO helpers
# -----------------------

def read_json(path: Path):
    with path.open("r", encoding="utf-8") as f:
        return json.load(f)

def write_json(path: Path, data: Any):
    path.parent.mkdir(parents=True, exist_ok=True)
    with path.open("w", encoding="utf-8") as f:
        json.dump(data, f, ensure_ascii=False, indent=2)

def read_yaml(path: Path):
    with path.open("r", encoding="utf-8") as f:
        return yaml.safe_load(f)

# -----------------------
# Geometry / IoU
# -----------------------

def bbox_iou_xyxy(a: np.ndarray, b: np.ndarray) -> float:
    ax1, ay1, ax2, ay2 = a
    bx1, by1, bx2, by2 = b
    ix1, iy1 = max(ax1, bx1), max(ay1, by1)
    ix2, iy2 = min(ax2, bx2), min(ay2, by2)
    iw, ih = max(0.0, ix2 - ix1), max(0.0, iy2 - iy1)
    inter = iw * ih
    aw, ah = max(0.0, ax2 - ax1), max(0.0, ay2 - ay1)
    bw, bh = max(0.0, bx2 - bx1), max(0.0, by2 - by1)
    union = aw * ah + bw * bh - inter + 1e-9
    return float(inter / union)

def poly_to_bbox_xyxy(poly: List[float], img_w: int, img_h: int) -> np.ndarray:
    xs = np.array(poly[0::2], dtype=np.float32)
    ys = np.array(poly[1::2], dtype=np.float32)
    xs = np.clip(xs, 0.0, 1.0) * img_w
    ys = np.clip(ys, 0.0, 1.0) * img_h
    x1, y1 = float(xs.min()), float(ys.min())
    x2, y2 = float(xs.max()), float(ys.max())
    return np.array([x1, y1, x2, y2], dtype=np.float32)

# -----------------------
# Read GT labels (YOLO-Seg)
# -----------------------

def load_gt_for_image(label_file: Path, img_shape: Tuple[int,int]) -> List[Tuple[int, np.ndarray]]:
    """
    Read YOLO-seg label file and return list of (class_id, bbox_xyxy).
    bbox approximated from polygon (for IoU matching).
    """
    out = []
    if not label_file.exists():
        return out
    img_w, img_h = img_shape
    with label_file.open("r", encoding="utf-8") as f:
        for line in f:
            parts = line.strip().split()
            if len(parts) < 7:  # cls + at least 3 points
                continue
            try:
                cls_id = int(float(parts[0]))
            except Exception:
                continue
            try:
                poly = [float(x) for x in parts[1:]]
            except Exception:
                continue
            box = poly_to_bbox_xyxy(poly, img_w, img_h)
            out.append((cls_id, box))
    return out

# -----------------------
# Classification metrics
# -----------------------

def f1_from_confusion(cm: np.ndarray, cls: int) -> float:
    # cm: [num_classes, num_classes], rows=GT, cols=Pred
    tp = cm[cls, cls]
    fp = cm[:, cls].sum() - tp
    fn = cm[cls, :].sum() - tp
    denom = (tp + fp + fn)
    if denom <= 0:
        return 0.0
    precision = tp / (tp + fp + 1e-9)
    recall = tp / (tp + fn + 1e-9)
    if precision + recall == 0:
        return 0.0
    return 2 * precision * recall / (precision + recall + 1e-9)

def expected_calibration_error(confs: np.ndarray, y: np.ndarray, n_bins: int = 15) -> float:
    bins = np.linspace(0.0, 1.0, n_bins + 1)
    ece = 0.0
    n = len(confs)
    for i in range(n_bins):
        lo, hi = bins[i], bins[i+1]
        mask = (confs >= lo) & (confs < hi) if i < n_bins - 1 else (confs >= lo) & (confs <= hi)
        if not np.any(mask):
            continue
        acc = y[mask].mean() if mask.sum() > 0 else 0.0
        conf = confs[mask].mean()
        ece += (mask.sum() / n) * abs(acc - conf)
    return float(ece)

def logit(p: np.ndarray, eps: float = 1e-8) -> np.ndarray:
    p = np.clip(p, eps, 1 - eps)
    return np.log(p) - np.log(1 - p)

def sigmoid(z: np.ndarray) -> np.ndarray:
    return 1.0 / (1.0 + np.exp(-z))

# -----------------------
# Weight MAE (optional)
# -----------------------

def mae_weights(gt_map: Dict[str, float], pred_map: Dict[str, float]) -> Optional[float]:
    keys = sorted(set(gt_map.keys()) & set(pred_map.keys()))
    if not keys:
        return None
    diffs = [abs(pred_map[k] - gt_map[k]) for k in keys]
    return float(np.mean(diffs)) if diffs else None

# -----------------------
# Evaluation core
# -----------------------

def _normalize_names(names):
    """names can be list or dict -> coerce to list."""
    if isinstance(names, dict):
        try:
            keys = [int(k) for k in names.keys()] if names else []
            arr = [None] * (max(keys) + 1 if keys else 0)
            for k, v in names.items():
                arr[int(k)] = v
            names = arr
        except Exception:
            names = list(names.values())
    if not isinstance(names, list):
        names = list(names) if names is not None else []
    names = [("" if n is None else str(n)) for n in names]
    return names

def evaluate(
    artifacts_dir: Path,
    dataset_yaml: Path,
    iou_thr: float = 0.5,
    bins: int = 15
) -> Dict[str, Any]:

    preds = read_json(artifacts_dir / "val_predictions.json")
    if not isinstance(preds, list):
        raise RuntimeError("val_predictions.json appears corrupted or empty.")

    ds = read_yaml(dataset_yaml) or {}
    if "path" not in ds or "val" not in ds:
        raise RuntimeError("dataset.yaml missing 'path' or 'val'.")

    # Setup
    root = Path(ds["path"])
    val_rel = Path(ds["val"])

    # Wenn ds["val"] z.B. "images/val" ist → benutze suffix nach "images"
    if "images" in val_rel.parts:
        idx = val_rel.parts.index("images")
        suffix = Path(*val_rel.parts[idx+1:])  # z.B. "val"
        cand = root / "labels" / suffix        # -> <root>/labels/val
        labels_dir = cand if cand.exists() else root / "labels"
    else:
        # Fallbacks, falls ds["val"] schon "val" ist
        labels_dir = root / "labels" / val_rel.name
        if not labels_dir.exists():
            labels_dir = root / "labels"

    names = _normalize_names(ds.get("names") or [])
    num_classes = len(names)
    if num_classes == 0:
        raise RuntimeError("names[] missing/empty in dataset.yaml.")

    # Confusion matrix (GT x Pred)
    cm = np.zeros((num_classes, num_classes), dtype=np.int64)

    # For ECE: all instance confidences + correctness
    inst_confs = []
    inst_correct = []

    # Image-level predictions/GT
    image_pred_class: Dict[str, Tuple[int, float]] = {}  # stem -> (pred_cls, conf)
    image_gt_class: Dict[str, int] = {}

    for r in preds:
        path = Path(r.get("path", ""))
        if not path.name:
            continue
        stem = path.stem
        orig_shape = r.get("orig_shape")
        if not orig_shape or len(orig_shape) < 2:
            # without shape no matching/ECE
            continue
        img_h, img_w = int(orig_shape[0]), int(orig_shape[1])

        # Load GT
        gt_label_file = labels_dir / f"{stem}.txt"
        gts = load_gt_for_image(gt_label_file, (img_w, img_h))
        if len(gts) == 0:
            continue

        # Image-level GT class = most frequent class in image
        gt_classes = [cls for cls, _ in gts]
        if len(gt_classes) == 0:
            continue
        gt_counts = np.bincount(gt_classes, minlength=num_classes)
        gt_img_class = int(np.argmax(gt_counts))
        image_gt_class[stem] = gt_img_class

        # Iterate predictions
        boxes = r.get("boxes") or []
        for b in boxes:
            try:
                pred_cls = int(b["cls"])
                pred_conf = float(b["conf"])
                pred_xyxy = np.array(b["xyxy"], dtype=np.float32)
            except Exception:
                continue

            # ECE label: correct if any GT of same class has IoU>=thr
            best_iou = 0.0
            for gt_cls, gt_xyxy in gts:
                if gt_cls != pred_cls:
                    continue
                iou = bbox_iou_xyxy(pred_xyxy, gt_xyxy)
                if iou > best_iou:
                    best_iou = iou

            inst_confs.append(pred_conf)
            inst_correct.append(1.0 if best_iou >= iou_thr else 0.0)

            # Image-level PRED = class of the box with highest conf (greedy)
            bp = image_pred_class.get(stem)
            if bp is None or pred_conf > bp[1]:
                image_pred_class[stem] = (pred_cls, pred_conf)

    # Confusion from image_pred_class vs. image_gt_class
    matched = 0
    for stem, (pred_cls, _) in image_pred_class.items():
        gt_cls = image_gt_class.get(stem)
        if gt_cls is None:
            continue
        # Safety: keep IDs in valid range
        if 0 <= gt_cls < num_classes and 0 <= pred_cls < num_classes:
            cm[gt_cls, pred_cls] += 1
            matched += 1

    # Metrics: macro-accuracy, per-class F1
    total = int(cm.sum())
    acc = float(np.trace(cm) / total) if total > 0 else 0.0
    per_class_f1 = {names[i]: float(f1_from_confusion(cm, i)) for i in range(num_classes)}

    # ECE (before/after calibration)
    if len(inst_confs) == 0:
        ece_before = None
        ece_after = None
    else:
        confs = np.array(inst_confs, dtype=np.float64)
        y = np.array(inst_correct, dtype=np.float64)
        ece_before = expected_calibration_error(confs, y, n_bins=bins)

        cal_path = artifacts_dir / "calibration_temp.json"
        if cal_path.exists():
            cal = read_json(cal_path) or {}
            T = float(cal.get("temperature", 1.0))
            # numerically stable: sigmoid(logit(p)/T)
            z = logit(confs)
            confs_cal = sigmoid(z / max(T, 1e-9))
            ece_after = expected_calibration_error(confs_cal, y, n_bins=bins)
        else:
            ece_after = None

    # Optional: MAE(kg) – only if both files exist
    mae_kg = None
    gt_w_path = artifacts_dir / "val_weight_groundtruth.json"
    pred_w_path = artifacts_dir / "val_weight_predictions.json"
    if gt_w_path.exists() and pred_w_path.exists():
        gt_map = read_json(gt_w_path)  # { "stem_or_path": weight_kg, ... }
        pred_map = read_json(pred_w_path)
        if isinstance(gt_map, dict) and isinstance(pred_map, dict):
            mae_kg = mae_weights(gt_map, pred_map)

    # Summary
    summary = {
        "images_evaluated": matched,
        "classes": names,
        "macro_accuracy": acc,
        "per_class_f1": per_class_f1,
        "confusion_matrix": cm.tolist(),
        "ece": {
            "before": ece_before,
            "after": ece_after,
            "bins": bins
        },
        "weight_mae_kg": mae_kg
    }
    return summary

# -----------------------
# CLI
# -----------------------

def main():
    ap = argparse.ArgumentParser(description="Evaluate YOLOv8-seg predictions at image level + calibration ECE.")
    ap.add_argument("--artifacts_dir", required=True, type=Path, help="Path to artifacts/<version>")
    ap.add_argument("--dataset_yaml", required=True, type=Path, help="Path to dataset.yaml (from prepare_taco.py)")
    ap.add_argument("--iou_thr", type=float, default=0.5)
    ap.add_argument("--bins", type=int, default=15)
    args = ap.parse_args()

    summary = evaluate(
        artifacts_dir=args.artifacts_dir,
        dataset_yaml=args.dataset_yaml,
        iou_thr=args.iou_thr,
        bins=args.bins
    )
    out_path = args.artifacts_dir / "summary_eval.json"
    write_json(out_path, summary)

    print("=== Evaluation Summary ===")
    print(json.dumps(summary, indent=2, ensure_ascii=False))
    print(f"Saved: {out_path.resolve()}")

if __name__ == "__main__":
    main()