#!/usr/bin/env python3
"""Evaluate deployable multi-prompt mask stability as an offline QC signal."""

from __future__ import annotations

import argparse
import itertools
from pathlib import Path

import numpy as np
import pandas as pd
from PIL import Image


def parse_args():
    p = argparse.ArgumentParser()
    p.add_argument("--predictions_csv", default="outputs/sam3_multiprompt_lora_qc/multiprompt_lora_eval/all_prompt_predictions.csv")
    p.add_argument("--index_csv", default="data/sam3_multiprompt_lora_qc/index/test_index.csv")
    p.add_argument("--output_dir", default="outputs/sam3_multiprompt_lora_qc/stability_qc")
    p.add_argument("--default_prompt", default="breast lesion")
    p.add_argument("--overwrite", action="store_true")
    return p.parse_args()


def read(path):
    return np.asarray(Image.open(path).convert("L")) > 0


def align(a, shape):
    return a if a.shape == shape else np.asarray(Image.fromarray(a.astype(np.uint8) * 255).resize((shape[1], shape[0]), Image.Resampling.NEAREST)) > 0


def overlap(a, b):
    b = align(b, a.shape)
    inter = int(np.logical_and(a, b).sum())
    den = int(a.sum() + b.sum())
    union = int(np.logical_or(a, b).sum())
    return (2 * inter / den if den else 1.0, inter / union if union else 1.0)


def bbox(mask):
    y, x = np.where(mask)
    return None if not len(x) else (x.min(), y.min(), x.max() + 1, y.max() + 1)


def bbox_iou(a, b):
    if a is None or b is None:
        return 0.0
    ix1, iy1, ix2, iy2 = max(a[0], b[0]), max(a[1], b[1]), min(a[2], b[2]), min(a[3], b[3])
    inter = max(0, ix2 - ix1) * max(0, iy2 - iy1)
    aa, bb = (a[2] - a[0]) * (a[3] - a[1]), (b[2] - b[0]) * (b[3] - b[1])
    return inter / (aa + bb - inter) if aa + bb > inter else 0.0


def auc(y, score):
    y = np.asarray(y, dtype=bool)
    score = np.asarray(score, dtype=float)
    keep = np.isfinite(score)
    y, score = y[keep], score[keep]
    positives, negatives = int(y.sum()), int((~y).sum())
    if not positives or not negatives:
        return np.nan
    ranks = pd.Series(score).rank(method="average").to_numpy()
    return float((ranks[y].sum() - positives * (positives + 1) / 2) / (positives * negatives))


def save_stability_overlay(group, output_path):
    valid = group[group["valid_mask"].astype(bool)]
    if valid.empty:
        return
    image = Image.open(valid.iloc[0]["original_image_path"]).convert("RGB")
    masks = [align(read(path), (image.height, image.width)) for path in valid["pred_mask_path"]]
    agreement = np.mean(np.stack(masks), axis=0)
    base = np.asarray(image).copy()
    high, uncertain = agreement >= 0.8, (agreement > 0.2) & (agreement < 0.8)
    base[high] = (0.55 * base[high] + np.array([0, 220, 0]) * 0.45).astype(np.uint8)
    base[uncertain] = (0.55 * base[uncertain] + np.array([255, 170, 0]) * 0.45).astype(np.uint8)
    output_path.parent.mkdir(parents=True, exist_ok=True)
    Image.fromarray(base).save(output_path)


def main():
    args = parse_args()
    out = Path(args.output_dir)
    out.mkdir(parents=True, exist_ok=True)
    predictions = pd.read_csv(args.predictions_csv)
    rows = []
    for sid, group in predictions.groupby("sample_id"):
        valid = group[group["valid_mask"].astype(bool)]
        masks = [(r["prompt"], read(r["pred_mask_path"])) for _, r in valid.iterrows()]
        pair = [(overlap(a, b), bbox_iou(bbox(a), bbox(b))) for (_, a), (_, b) in itertools.combinations(masks, 2)]
        areas = np.array([m.sum() for _, m in masks], dtype=float)
        centers = np.array([[(b[0] + b[2]) / 2, (b[1] + b[3]) / 2] for _, m in masks if (b := bbox(m)) is not None])
        selected = valid[valid["prompt"] == args.default_prompt]
        selected = selected.iloc[0] if len(selected) else valid.iloc[0]
        rows.append({
            "sample_id": sid,
            "selected_prompt": selected["prompt"],
            "selected_mask_path": selected["pred_mask_path"],
            "selected_mask_dice": selected["dice"],
            "mean_pairwise_dice": np.mean([x[0][0] for x in pair]) if pair else np.nan,
            "min_pairwise_dice": np.min([x[0][0] for x in pair]) if pair else np.nan,
            "mean_pairwise_iou": np.mean([x[0][1] for x in pair]) if pair else np.nan,
            "mask_area_mean": areas.mean() if len(areas) else np.nan,
            "mask_area_std": areas.std() if len(areas) else np.nan,
            "mask_area_cv": areas.std() / areas.mean() if len(areas) and areas.mean() else np.nan,
            "mask_center_std": centers.std(axis=0).mean() if len(centers) else np.nan,
            "mask_bbox_iou_mean": np.mean([x[1] for x in pair]) if pair else np.nan,
            "mask_bbox_iou_min": np.min([x[1] for x in pair]) if pair else np.nan,
            "valid_prompt_count": len(valid),
            "failure_prompt_count": len(group) - len(valid),
        })
    features = pd.DataFrame(rows)
    features["good_mask"] = features["selected_mask_dice"] >= 0.75
    features["bad_mask"] = features["selected_mask_dice"] < 0.50
    features["quality_label"] = np.where(features["good_mask"], "good", np.where(features["bad_mask"], "bad", "medium"))
    features.to_csv(out / "stability_features.csv", index=False)
    feature_cols = ["mean_pairwise_dice", "min_pairwise_dice", "mean_pairwise_iou", "mask_area_cv", "mask_center_std", "mask_bbox_iou_mean", "mask_bbox_iou_min"]
    summary = []
    for col in feature_cols:
        sign = -1 if col in {"mask_area_cv", "mask_center_std"} else 1
        summary.append({"feature": col, "dice_correlation": features[[col, "selected_mask_dice"]].corr().iloc[0, 1], "auc_good": auc(features["good_mask"], sign * features[col]), "auc_bad": auc(features["bad_mask"], -sign * features[col])})
    pd.DataFrame(summary).to_csv(out / "qc_metrics_summary.csv", index=False)
    auc_good = auc(features["good_mask"], features["mean_pairwise_dice"])
    auc_bad = auc(features["bad_mask"], -features["mean_pairwise_dice"])
    sweeps = []
    for threshold in [0.7, 0.8, 0.9]:
        passed = (features["mean_pairwise_dice"] >= threshold) & (features["mask_area_cv"] <= 0.3) & (features["failure_prompt_count"] == 0)
        p, f = features[passed], features[~passed]
        sweeps.append({"threshold": threshold, "pass_rate": passed.mean(), "pass_subset_mean_dice": p["selected_mask_dice"].mean(), "fail_subset_mean_dice": f["selected_mask_dice"].mean(), "bad_mask_rate_in_pass": p["bad_mask"].mean(), "good_mask_rate_in_pass": p["good_mask"].mean(), "auc_good": auc_good, "auc_bad": auc_bad, "bad_detection_sensitivity": ((~passed) & features["bad_mask"]).sum() / max(1, features["bad_mask"].sum()), "bad_detection_specificity": (passed & ~features["bad_mask"]).sum() / max(1, (~features["bad_mask"]).sum())})
    pd.DataFrame(sweeps).to_csv(out / "qc_rule_threshold_sweep.csv", index=False)
    features.sort_values("selected_mask_dice", ascending=False).head(30).to_csv(out / "pass_subset_examples.csv", index=False)
    features.sort_values("selected_mask_dice").head(30).to_csv(out / "fail_subset_examples.csv", index=False)
    features.sort_values("mean_pairwise_dice").head(30).to_csv(out / "worst_instability_examples.csv", index=False)
    examples = pd.concat([features.nlargest(10, "mean_pairwise_dice"), features.nsmallest(10, "mean_pairwise_dice")]).drop_duplicates("sample_id")
    for _, example in examples.iterrows():
        group = predictions[predictions["sample_id"] == example["sample_id"]]
        save_stability_overlay(group, out / "overlay_examples" / f"{example['sample_id']}.png")
    print(pd.DataFrame(sweeps).to_string(index=False))


if __name__ == "__main__":
    main()