"""
pre_process.py -- Part B, Sections 3.1 and 3.2
(Data preparation, annotation checks, and exploratory analysis)

1. Loads Kvasir-SEG images, masks and the bounding-box json, and converts the
   annotations into the xyxy + label format Faster R-CNN expects.
2. Annotation QC: images with no box, boxes outside the image, degenerate
   (zero/negative area) boxes, box-vs-mask agreement (IoU between the annotated
   box and the tight bounding box of the segmentation mask), and json/image
   size disagreements.
3. Frozen stratified split with SEED = 16 -> splits/kvasir_detection_seed16.csv
4. EDA: box count per image, box width/height, box-area fraction, aspect ratio,
   polyp centre position, image-size distribution, per-split distribution.
5. At least 8 seed-selected image + mask + box overlays.

Usage
-----
python pre_process.py
python pre_process.py --data-root ../Dataset/kvasir/Kvasir-SEG --n-examples 8
"""

from __future__ import annotations

import argparse
import json
from collections import Counter
from pathlib import Path

import numpy as np
import pandas as pd

from common import (DET_SPLIT_CSV, boxplot, METRICS_DIR, OVERLAY_DIR, PLOTS_DIR, SEED,
                    SEED_TAG, SPLIT_FRACTIONS, _mpl, box_iou_np, find_bbox_json,
                    find_kvasir_root, load_bboxes, progress, save_json,
                    seeded_rng, set_seed)

IMG_EXT = {".jpg", ".jpeg", ".png", ".bmp", ".tif"}


# ---------------------------------------------------------------------------
def mask_bbox(mask: np.ndarray) -> list | None:
    """Tight xyxy box of the (largest connected component of the) mask."""
    b = mask > 0.5
    if b.sum() == 0:
        return None
    ys, xs = np.where(b)
    return [float(xs.min()), float(ys.min()), float(xs.max() + 1), float(ys.max() + 1)]


def build_index(root: Path, bbox_json: Path) -> tuple[pd.DataFrame, dict]:
    from PIL import Image

    ann = load_bboxes(bbox_json)
    img_dir, mask_dir = root / "images", root / "masks"
    images = sorted(p for p in img_dir.iterdir() if p.suffix.lower() in IMG_EXT)

    rows, issues = [], {
        "images_without_annotation": [], "annotations_without_image": [],
        "images_with_zero_boxes": [], "degenerate_boxes": [],
        "boxes_out_of_bounds": [], "json_size_mismatch": [],
        "missing_masks": [], "empty_masks": [], "low_box_mask_iou": [],
    }

    ann_keys = set(ann.keys())
    for ip in progress(images, desc="index kvasir", total=len(images)):
        stem = ip.stem
        with Image.open(ip) as im:
            W, H = im.size
        mp = mask_dir / f"{stem}{ip.suffix}"
        if not mp.exists():
            for e in (".jpg", ".png", ".jpeg"):
                if (mask_dir / f"{stem}{e}").exists():
                    mp = mask_dir / f"{stem}{e}"
                    break
        has_mask = mp.exists()
        if not has_mask:
            issues["missing_masks"].append(stem)

        a = ann.get(stem)
        if a is None:
            issues["images_without_annotation"].append(stem)
            boxes = []
        else:
            ann_keys.discard(stem)
            boxes = a["boxes"]
            if a.get("width") and (int(a["width"]) != W or int(a["height"]) != H):
                issues["json_size_mismatch"].append(
                    {"image_id": stem, "json": [a.get("width"), a.get("height")],
                     "file": [W, H]})

        clean = []
        for b in boxes:
            x1, y1, x2, y2 = b
            x1, x2 = min(x1, x2), max(x1, x2)
            y1, y2 = min(y1, y2), max(y1, y2)
            if (x2 - x1) < 1 or (y2 - y1) < 1:
                issues["degenerate_boxes"].append({"image_id": stem, "box": b})
                continue
            if x1 < 0 or y1 < 0 or x2 > W or y2 > H:
                issues["boxes_out_of_bounds"].append(
                    {"image_id": stem, "box": [x1, y1, x2, y2], "image_size": [W, H]})
                x1, y1 = max(0.0, x1), max(0.0, y1)
                x2, y2 = min(float(W), x2), min(float(H), y2)
                if (x2 - x1) < 1 or (y2 - y1) < 1:
                    continue
            clean.append([x1, y1, x2, y2])
        if not clean:
            issues["images_with_zero_boxes"].append(stem)

        # --- box vs mask agreement -----------------------------------------
        box_mask_iou, mask_area_frac = np.nan, np.nan
        if has_mask:
            m = np.asarray(Image.open(mp).convert("L"), dtype=np.float32) / 255.0
            if m.shape[:2] != (H, W):
                from PIL import Image as _I
                m = np.asarray(_I.open(mp).convert("L").resize((W, H), _I.NEAREST),
                               dtype=np.float32) / 255.0
            mb = mask_bbox(m)
            mask_area_frac = float((m > 0.5).mean())
            if mb is None:
                issues["empty_masks"].append(stem)
            elif clean:
                box_mask_iou = float(box_iou_np(np.array(clean, np.float32),
                                                np.array([mb], np.float32)).max())
                if box_mask_iou < 0.5:
                    issues["low_box_mask_iou"].append(
                        {"image_id": stem, "iou": round(box_mask_iou, 3)})

        areas = [(b[2] - b[0]) * (b[3] - b[1]) for b in clean]
        rows.append({
            "image_id": stem,
            "image_path": str(ip.resolve()),
            "mask_path": str(mp.resolve()) if has_mask else "",
            "width": W, "height": H, "image_area": W * H,
            "n_boxes": len(clean),
            "boxes": json.dumps([[round(v, 2) for v in b] for b in clean]),
            "total_box_area_frac": float(sum(areas) / (W * H)) if areas else 0.0,
            "max_box_area_frac": float(max(areas) / (W * H)) if areas else 0.0,
            "mask_area_frac": mask_area_frac,
            "box_mask_iou": box_mask_iou,
        })

    issues["annotations_without_image"] = sorted(ann_keys)
    return pd.DataFrame(rows), issues


# ---------------------------------------------------------------------------
def _collapse_rare(strat: "pd.Series", min_count: int = 2) -> "pd.Series":
    """Merge any stratum below `min_count` members into a single 'rare' bucket.
    sklearn's stratified split requires at least 2 members per class, so this
    must be re-applied to whatever subset is about to be split -- a stratum
    that clears the threshold in the full dataset can still drop below it
    once a 70/30 (or 15/15) subset is drawn from it."""
    vc = strat.value_counts()
    return strat.where(strat.map(vc) >= min_count, "rare")


def make_splits(df: pd.DataFrame, seed: int = SEED, fracs=SPLIT_FRACTIONS) -> pd.DataFrame:
    """Stratified on a size bucket derived from the largest box, so that small,
    medium and large polyps are represented in every split (Kvasir-SEG has no
    class label to stratify on -- polyp size is the meaningful axis)."""
    from sklearn.model_selection import train_test_split

    d = df.copy().reset_index(drop=True)
    d["size_bucket"] = pd.cut(d["max_box_area_frac"],
                              bins=[-0.001, 0.01, 0.05, 0.20, 1.01],
                              labels=["tiny", "small", "medium", "large"]).astype(str)
    d.loc[d["n_boxes"] == 0, "size_bucket"] = "empty"
    d["strat"] = d["size_bucket"] + "_" + (d["n_boxes"] > 1).map({True: "multi", False: "single"})

    # Collapse rare strata for the first split (70 / 30).
    strat1 = _collapse_rare(d["strat"], min_count=6)
    try:
        tr, tmp = train_test_split(d, train_size=fracs[0], random_state=seed,
                                   stratify=strat1, shuffle=True)
    except ValueError:
        # even the collapsed bucketing has a class with < 2 members overall;
        # fall back to an unstratified split rather than crashing.
        print("[warn] falling back to an unstratified train split: a size "
              "bucket has fewer than 2 images even after collapsing rare "
              "buckets together")
        tr, tmp = train_test_split(d, train_size=fracs[0], random_state=seed, shuffle=True)

    # Re-collapse against the smaller `tmp` subset before the second split --
    # a stratum with 8 members overall can easily have only 1 in a 30% sample.
    rel_val = fracs[1] / (fracs[1] + fracs[2])
    strat2 = _collapse_rare(tmp["strat"], min_count=2)
    try:
        va, te = train_test_split(tmp, train_size=rel_val, random_state=seed,
                                  stratify=strat2, shuffle=True)
    except ValueError:
        print("[warn] falling back to an unstratified val/test split for the "
              "same reason")
        va, te = train_test_split(tmp, train_size=rel_val, random_state=seed, shuffle=True)

    d["split"] = ""
    d.loc[tr.index, "split"] = "train"
    d.loc[va.index, "split"] = "val"
    d.loc[te.index, "split"] = "test"
    return d


# ---------------------------------------------------------------------------
def eda(d: pd.DataFrame, out_dir: Path = PLOTS_DIR) -> dict:
    plt = _mpl()
    stats: dict = {}

    boxes, wfrac, hfrac, afrac, ar, cx, cy, splits = [], [], [], [], [], [], [], []
    for _, r in d.iterrows():
        for b in json.loads(r["boxes"]):
            x1, y1, x2, y2 = b
            w, h = x2 - x1, y2 - y1
            boxes.append([w, h])
            wfrac.append(w / r["width"]); hfrac.append(h / r["height"])
            afrac.append(w * h / r["image_area"]); ar.append(w / max(1e-6, h))
            cx.append((x1 + x2) / 2 / r["width"]); cy.append((y1 + y2) / 2 / r["height"])
            splits.append(r["split"])
    boxes = np.array(boxes) if boxes else np.zeros((0, 2))
    afrac = np.array(afrac); ar = np.array(ar)

    # ---- figure 1: counts and sizes ----
    fig, axes = plt.subplots(2, 3, figsize=(15, 8))
    cnt = d["n_boxes"].value_counts().sort_index()
    axes[0, 0].bar(cnt.index.astype(str), cnt.values, color="#4C72B0")
    axes[0, 0].set_title("Bounding boxes per image"); axes[0, 0].set_xlabel("n boxes")
    for i, v in enumerate(cnt.values):
        axes[0, 0].text(i, v, str(v), ha="center", va="bottom", fontsize=8)

    sz = d.groupby(["width", "height"]).size().sort_values(ascending=False).head(10)
    axes[0, 1].bar([f"{w}x{h}" for w, h in sz.index], sz.values, color="#55A868")
    axes[0, 1].set_title("Top-10 image resolutions"); axes[0, 1].set_yscale("log")
    plt.setp(axes[0, 1].get_xticklabels(), rotation=40, ha="right", fontsize=7)

    axes[0, 2].hist(afrac, bins=50, color="#C44E52")
    axes[0, 2].set_title("Box area fraction of image"); axes[0, 2].set_xlabel("area fraction")
    axes[0, 2].axvline(0.01, ls="--", c="k", lw=1)
    axes[0, 2].axvline(0.05, ls="--", c="k", lw=1)
    axes[0, 2].text(0.011, axes[0, 2].get_ylim()[1] * 0.8, " small | medium | large", fontsize=7)

    if len(boxes):
        axes[1, 0].scatter(boxes[:, 0], boxes[:, 1], s=6, alpha=0.35, color="#8172B2")
        axes[1, 0].set_xlabel("box width (px)"); axes[1, 0].set_ylabel("box height (px)")
        axes[1, 0].set_title("Box width vs height")
    axes[1, 1].hist(np.clip(ar, 0, 4), bins=50, color="#937860")
    axes[1, 1].set_title("Aspect ratio (w/h)"); axes[1, 1].axvline(1.0, ls="--", c="k", lw=1)
    hb = axes[1, 2].hist2d(cx, cy, bins=20, range=[[0, 1], [0, 1]], cmap="magma")
    axes[1, 2].invert_yaxis()
    axes[1, 2].set_title("Polyp centre position (normalised)")
    fig.colorbar(hb[3], ax=axes[1, 2], fraction=0.046)
    fig.suptitle(f"Kvasir-SEG exploratory analysis (seed={SEED})", fontsize=13)
    fig.tight_layout()
    fig.savefig(out_dir / f"B1_kvasir_eda_{SEED_TAG}.png", dpi=155); plt.close(fig)

    # ---- figure 2: split distribution ----
    fig, axes = plt.subplots(1, 2, figsize=(11, 4.2))
    piv = (d.pivot_table(index="size_bucket", columns="split", values="image_id",
                         aggfunc="count", fill_value=0)
             .reindex(columns=["train", "val", "test"], fill_value=0))
    piv.plot(kind="bar", stacked=True, ax=axes[0], colormap="viridis")
    axes[0].set_title("Polyp size bucket per split"); axes[0].set_ylabel("images")
    plt.setp(axes[0].get_xticklabels(), rotation=0)
    sdf = pd.DataFrame({"split": splits, "area_frac": afrac})
    boxplot(axes[1], [sdf[sdf.split == s]["area_frac"].values for s in ["train", "val", "test"]],
            ["train", "val", "test"], showfliers=False)
    axes[1].set_title("Box area fraction per split")
    fig.tight_layout()
    fig.savefig(out_dir / f"B2_kvasir_split_distribution_{SEED_TAG}.png", dpi=155); plt.close(fig)

    stats["n_images"] = int(len(d))
    stats["n_boxes_total"] = int(d["n_boxes"].sum())
    stats["boxes_per_image"] = {str(k): int(v) for k, v in cnt.items()}
    stats["mean_boxes_per_image"] = float(d["n_boxes"].mean())
    stats["n_images_with_zero_boxes"] = int((d["n_boxes"] == 0).sum())
    stats["n_images_multi_box"] = int((d["n_boxes"] > 1).sum())
    stats["image_resolutions_top10"] = {f"{w}x{h}": int(v) for (w, h), v in sz.items()}
    stats["n_distinct_resolutions"] = int(d.groupby(["width", "height"]).ngroups)
    stats["box_width_px"] = _desc(boxes[:, 0] if len(boxes) else np.zeros(0))
    stats["box_height_px"] = _desc(boxes[:, 1] if len(boxes) else np.zeros(0))
    stats["box_area_fraction"] = _desc(afrac)
    stats["aspect_ratio"] = _desc(ar)
    stats["centre_x"] = _desc(np.array(cx)); stats["centre_y"] = _desc(np.array(cy))
    stats["size_buckets"] = {
        "small_lt_1pct": int((afrac < 0.01).sum()),
        "medium_1_to_5pct": int(((afrac >= 0.01) & (afrac < 0.05)).sum()),
        "large_ge_5pct": int((afrac >= 0.05).sum())}
    stats["split_counts"] = d["split"].value_counts().to_dict()
    stats["split_box_counts"] = d.groupby("split")["n_boxes"].sum().to_dict()
    stats["size_bucket_by_split"] = piv.to_dict()
    stats["foreground_pixel_fraction_mean"] = float(d["mask_area_frac"].dropna().mean())
    stats["class_imbalance_note"] = (
        "Kvasir-SEG is single-class, so 'imbalance' in detection is the "
        "foreground/background imbalance: on average only "
        f"{100 * float(np.mean(afrac)):.1f}% of an image is inside a polyp box, and the "
        "detector scores thousands of mostly-background anchors/proposals per image. "
        "This is precisely the imbalance the focal-loss ablation targets.")
    return stats


def _desc(a: np.ndarray) -> dict:
    if a.size == 0:
        return {}
    return {"mean": float(a.mean()), "std": float(a.std()), "min": float(a.min()),
            "p05": float(np.percentile(a, 5)), "median": float(np.median(a)),
            "p95": float(np.percentile(a, 95)), "max": float(a.max())}


# ---------------------------------------------------------------------------
def qualitative(d: pd.DataFrame, n: int = 8, out_dir: Path = OVERLAY_DIR):
    """Seed-selected image / mask / box overlays spanning the size buckets."""
    plt = _mpl()
    from PIL import Image

    rng = seeded_rng(offset=5)
    picks = []
    for bucket in ["tiny", "small", "medium", "large"]:
        sub = d[d["size_bucket"] == bucket]
        if len(sub) == 0:
            continue
        k = min(max(1, n // 4), len(sub))
        idx = rng.choice(len(sub), size=k, replace=False)
        picks.append(sub.iloc[idx])
    sel = pd.concat(picks).head(n) if picks else d.head(n)
    if len(sel) < n:
        extra = d[~d["image_id"].isin(sel["image_id"])].head(n - len(sel))
        sel = pd.concat([sel, extra])

    ncol = 4
    nrow = int(np.ceil(len(sel) / ncol))
    fig, axes = plt.subplots(nrow, ncol, figsize=(4.0 * ncol, 4.3 * nrow))
    axes = np.atleast_1d(axes).ravel()
    for k, (_, r) in enumerate(sel.iterrows()):
        img = np.asarray(Image.open(r["image_path"]).convert("RGB"), np.float32) / 255.0
        if r["mask_path"]:
            m = np.asarray(Image.open(r["mask_path"]).convert("L")
                           .resize((r["width"], r["height"]), Image.NEAREST), np.float32) / 255.0
            img[..., 1] = np.where(m > 0.5, 0.55 * img[..., 1] + 0.45, img[..., 1])
        axes[k].imshow(np.clip(img, 0, 1))
        for b in json.loads(r["boxes"]):
            x1, y1, x2, y2 = b
            axes[k].add_patch(plt.Rectangle((x1, y1), x2 - x1, y2 - y1, fill=False,
                                            edgecolor="red", lw=2.0))
        axes[k].set_title(f"{r['image_id'][:16]}...\n{r['width']}x{r['height']} | "
                          f"{r['n_boxes']} box | area={r['max_box_area_frac']:.3f} | "
                          f"{r['split']}", fontsize=7.5)
        axes[k].axis("off")
    for k in range(len(sel), len(axes)):
        axes[k].axis("off")
    fig.suptitle(f"Kvasir-SEG: image + mask (green) + box (red), seed={SEED}", fontsize=12)
    fig.tight_layout()
    out = out_dir / f"B3_kvasir_annotations_{SEED_TAG}.png"
    fig.savefig(out, dpi=150); plt.close(fig)
    print(f"[saved] {out}")
    return sel[["image_id", "split", "width", "height", "n_boxes", "max_box_area_frac"]]


# ---------------------------------------------------------------------------
def main():
    ap = argparse.ArgumentParser(description="Part B data prep + EDA (seed 16)")
    ap.add_argument("--data-root", default=None)
    ap.add_argument("--bbox-json", default=None)
    ap.add_argument("--n-examples", type=int, default=8)
    args = ap.parse_args()

    set_seed(SEED)
    root = find_kvasir_root(args.data_root)
    bj = Path(args.bbox_json) if args.bbox_json else find_bbox_json(root)
    print(f"[info] kvasir root : {root}")
    print(f"[info] bbox json   : {bj}")
    print(f"[info] seed        : {SEED}")

    df, issues = build_index(root, bj)
    d = make_splits(df, seed=SEED)
    d.to_csv(DET_SPLIT_CSV, index=False)
    print(f"[saved] {DET_SPLIT_CSV}")

    print("\n=== Split counts (seed=%d) ===" % SEED)
    tab = d.groupby("split").agg(images=("image_id", "count"), boxes=("n_boxes", "sum"))
    tab = tab.reindex(["train", "val", "test"])
    tab["mean_boxes_per_image"] = (tab["boxes"] / tab["images"]).round(3)
    print(tab.to_string())

    stats = eda(d)
    audit = {"seed": SEED, "dataset_root": str(root), "bbox_json": str(bj),
             "n_images": int(len(d)),
             "n_images_without_annotation": len(issues["images_without_annotation"]),
             "n_annotations_without_image": len(issues["annotations_without_image"]),
             "n_images_with_zero_boxes": len(issues["images_with_zero_boxes"]),
             "n_degenerate_boxes": len(issues["degenerate_boxes"]),
             "n_boxes_out_of_bounds": len(issues["boxes_out_of_bounds"]),
             "n_json_size_mismatch": len(issues["json_size_mismatch"]),
             "n_missing_masks": len(issues["missing_masks"]),
             "n_empty_masks": len(issues["empty_masks"]),
             "n_low_box_mask_iou_lt_0.5": len(issues["low_box_mask_iou"]),
             "median_box_mask_iou": float(d["box_mask_iou"].median(skipna=True)),
             "split_table": tab.to_dict(),
             "eda": stats,
             "details": {k: (v[:100] if isinstance(v, list) else v) for k, v in issues.items()}}
    save_json(audit, METRICS_DIR / f"B_data_audit_{SEED_TAG}.json")

    sel = qualitative(d, n=args.n_examples)
    sel.to_csv(METRICS_DIR / f"B_seed_selected_examples_{SEED_TAG}.csv", index=False)

    print("\n=== ANNOTATION AUDIT ===")
    for k, v in audit.items():
        if k.startswith("n_") or k.startswith("median_"):
            print(f"  {k:34s}: {v}")
    print(f"\n  box area fraction: mean={stats['box_area_fraction']['mean']:.4f} "
          f"median={stats['box_area_fraction']['median']:.4f}")
    print(f"  size buckets     : {stats['size_buckets']}")
    print("\n" + stats["class_imbalance_note"])


if __name__ == "__main__":
    main()