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

"""
calibrate_temp.py
- Learn a global temperature T for confidence calibration (Temperature Scaling)
- Uses artifacts/<version>/val_predictions.json (from train_yolov8_seg.py)
- Matches predicted boxes against GT (val split) via IoU -> binary y in {0,1}
- Optimizes T on validation negative log loss (better calibrated probabilities)
- Writes artifacts/<version>/calibration_temp.json
"""

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

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:
    """IoU of two boxes (xyxy, 4 floats each)."""
    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:
    """YOLO-seg polygon (normalized) -> BBox in pixels."""
    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 .txt and return (class_id, bbox_xyxy) with bbox approximated from polygon.
    """
    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:
                continue
            try:
                cls_id = int(float(parts[0]))
                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

# -----------------------
# Build calibration dataset
# -----------------------

def build_calibration_pairs(
    val_predictions_json: Path,
    dataset_yaml: Path,
    iou_thr: float = 0.5
) -> Tuple[np.ndarray, np.ndarray]:
    """
    Returns (p, y):
        p = predicted box confidences (0..1)
        y = 1 if prediction is correct (IoU >= iou_thr and same class), else 0
    """
    preds = read_json(val_predictions_json)
    if not isinstance(preds, list):
        raise RuntimeError("val_predictions.json appears empty or corrupted.")
    ds = read_yaml(dataset_yaml) or {}
    if "path" not in ds or "val" not in ds:
        raise RuntimeError("dataset.yaml missing 'path' or 'val'.")

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

    # Case: ds["val"] is "images/val" or "images/<split>"
    if "images" in val_rel.parts:
        idx = val_rel.parts.index("images")
        suffix = Path(*val_rel.parts[idx+1:]) 
        cand = root / "labels" / suffix
        labels_dir = cand if cand.exists() else root / "labels"
    else:
        # Fallbacks
        labels_dir = root / "labels" / "val"
        if not labels_dir.exists():
            labels_dir = root / "labels"


    confs: List[float] = []
    ys: List[float] = []

    for r in preds:
        # Image size
        orig_shape = r.get("orig_shape")
        if not orig_shape or len(orig_shape) < 2:
            continue
        img_h, img_w = int(orig_shape[0]), int(orig_shape[1])

        # GT labels
        img_path = Path(r.get("path", ""))
        stem = img_path.stem
        if not stem:
            continue
        gt_label = labels_dir / f"{stem}.txt"
        gts = load_gt_for_image(gt_label, (img_w, img_h))
        if len(gts) == 0:
            continue

        # 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

            # best IoU with GT of the same class
            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

            confs.append(pred_conf)
            ys.append(1.0 if best_iou >= iou_thr else 0.0)

    if len(confs) == 0:
        raise RuntimeError("No validation pairs found for calibration (check val_predictions/labels).")

    # clamp in [0,1]
    confs = np.clip(np.array(confs, dtype=np.float64), 1e-8, 1 - 1e-8)
    y = np.array(ys, dtype=np.float64)
    return confs, y

# -----------------------
# Temperature scaling
# -----------------------

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))

def nll_for_temperature(confs: np.ndarray, y: np.ndarray, T: float) -> float:
    """
    Binary NLL after temperature scaling:
        pT = sigmoid(logit(confs) / T)
        NLL = - sum( y*log(pT) + (1-y)*log(1-pT) )
    """
    z = _logit(confs) / max(T, 1e-6)
    pT = _sigmoid(z)
    eps = 1e-12
    return float(-np.sum(y * np.log(pT + eps) + (1 - y) * np.log(1 - pT + eps)))

def fit_temperature(confs: np.ndarray, y: np.ndarray, t_min: float, t_max: float, steps: int) -> float:
    """
    1D search (grid + ternary fine search) for global temperature T.
    """
    best_T, best_nll = None, float("inf")
    grid = np.linspace(t_min, t_max, max(3, steps))
    for T in grid:
        nll = nll_for_temperature(confs, y, T)
        if nll < best_nll:
            best_T, best_nll = T, nll

    # local fine search around best_T
    lo = max(t_min, best_T - 0.5)
    hi = min(t_max, best_T + 0.5)
    for _ in range(50):
        mid1 = lo + (hi - lo) / 3.0
        mid2 = hi - (hi - lo) / 3.0
        n1 = nll_for_temperature(confs, y, mid1)
        n2 = nll_for_temperature(confs, y, mid2)
        if n1 < n2:
            hi = mid2
        else:
            lo = mid1
    return float((lo + hi) / 2.0)

def expected_calibration_error(confs: np.ndarray, y: np.ndarray, n_bins: int = 15) -> float:
    """ECE estimate (binning-based)."""
    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)

# -----------------------
# Main
# -----------------------

def main():
    ap = argparse.ArgumentParser(description="Calibrate confidences via temperature scaling.")
    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)
    ap.add_argument("--t_min", type=float, default=0.5, help="lower bound for T search")
    ap.add_argument("--t_max", type=float, default=5.0, help="upper bound for T search")
    ap.add_argument("--t_steps", type=int, default=200, help="grid steps for coarse search")
    args = ap.parse_args()

    preds_path = args.artifacts_dir / "val_predictions.json"
    if not preds_path.exists():
        raise FileNotFoundError(f"{preds_path} not found. Please run train_yolov8_seg.py first.")

    # build (p, y)
    confs_raw, y = build_calibration_pairs(preds_path, args.dataset_yaml, iou_thr=args.iou_thr)

    # Before: NLL/ECE
    nll_raw = nll_for_temperature(confs_raw, y, T=1.0)
    ece_raw = expected_calibration_error(confs_raw, y, n_bins=args.bins)

    # Fit T
    T = fit_temperature(confs_raw, y, t_min=float(args.t_min), t_max=float(args.t_max), steps=int(args.t_steps))

    # After: NLL/ECE
    z = _logit(confs_raw) / max(T, 1e-6)
    confs_cal = _sigmoid(z)
    nll_cal = nll_for_temperature(confs_raw, y, T=T)  # equivalent to using confs_cal in NLL
    ece_cal = expected_calibration_error(confs_cal, y, n_bins=args.bins)

    out = {
        "method": "temperature_scaling",
        "temperature": float(T),
        "metrics": {
            "val_pairs": int(confs_raw.shape[0]),
            "nll_before": float(nll_raw),
            "nll_after": float(nll_cal),
            "ece_before": float(ece_raw),
            "ece_after": float(ece_cal),
            "iou_thr": float(args.iou_thr),
            "bins": int(args.bins),
            "t_min": float(args.t_min),
            "t_max": float(args.t_max),
            "t_steps": int(args.t_steps)
        }
    }
    out_path = args.artifacts_dir / "calibration_temp.json"
    write_json(out_path, out)

    print("=== Calibration Done ===")
    print(json.dumps(out, indent=2, ensure_ascii=False))
    print(f"Saved: {out_path.resolve()}")

if __name__ == "__main__":
    main()