#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
Real vs. synthetic disease-mask comparison for AURAD.
Nature / NPG (ggsci) style.

Two improvements over the original two-panel t-SNE:

  1) SINGLE shared embedding, plotted in ONE panel:
        color  = disease
        marker = source   (real = filled dot, syn = open ring)
     so you can directly see whether each disease's real & syn points
     co-localize.

  2) Better disease separation, via (pick with EMBED_METHOD):
        "tsne"            - PCA -> t-SNE                  (unsupervised)
        "umap"            - PCA -> UMAP                   (unsupervised, tighter clusters)
        "umap_supervised" - PCA -> UMAP supervised on DISEASE labels only
                            (source/real-syn is NEVER given to the embedding,
                             so real<->syn overlap stays an emergent signal)
     Plus an optional Gaussian blur on the masks (BLUR_SIGMA) so that
     "nearby location" becomes a smooth similarity instead of 0/1 overlap.

  3) Quantitative co-clustering evidence (COMPUTE_METRICS), per disease:
        - C2ST AUC          : classifier two-sample test, ~0.5 = real/syn indistinguishable
        - MMD^2 (RBF)       : distribution discrepancy, lower = closer
        - norm. centroid d  : Mahalanobis-normalized real<->syn centroid distance
        - disease silhouette: how separable the diseases are (combined real+syn)
     Saved to <prefix>_metrics.csv and printed.

JSON record example (list of these, or one JSON per line):
  {
    "file_name": "chestdet/dataset/test/39796/39796.png",
    "mask":      "chestdet/dataset/test/39796/mask_Consolidation_39796.png",
    "attn_list": [["Consolidation", "chestdet/dataset/test/39796/39796_Consolidation.png"]],
    ...
  }

Real mask  = REAL_ROOT / record["mask"]              (REAL_MASK_SOURCE = "mask")
Synth mask = SYN_ROOT  / attn_list[i][1]             (SYN_MASK_SOURCE  = "attn")
"""

import os
import json
import argparse

import numpy as np
from PIL import Image
from scipy.ndimage import gaussian_filter
from scipy.spatial.distance import cdist

from sklearn.decomposition import PCA
from sklearn.manifold import TSNE
from sklearn.metrics import silhouette_score
from sklearn.linear_model import LogisticRegression
from sklearn.preprocessing import StandardScaler
from sklearn.pipeline import make_pipeline
from sklearn.model_selection import cross_val_score, StratifiedKFold

import matplotlib as mpl
import matplotlib.pyplot as plt
from matplotlib.lines import Line2D

# ----------------------------- CONFIG -----------------------------------
JSON_PATH = "data.json"
REAL_ROOT = "real_masks"
SYN_ROOT  = "syn_masks"

# How to build the relative mask path (joined with the corresponding root):
#   "file_name" -> record["file_name"]   (the raw CXR, NOT a mask)
#   "mask"      -> record["mask"]         (real disease mask)            <-- real
#   "attn"      -> attn_list[i][1]        (syn disease mask, per disease) <-- syn
REAL_MASK_SOURCE = "mask"
SYN_MASK_SOURCE  = "attn"

TARGET_DISEASES = [
    "Cardiomegaly",
    "Pneumothorax",
    "Consolidation",
    "Effusion",
    "Atelectasis",
]
ATTN_KEY = "attn_list"

# ---- feature extraction ----
MASK_SIZE  = 64       # downsample masks to MASK_SIZE x MASK_SIZE
BLUR_SIGMA = 2.0      # Gaussian blur sigma on the mask before flatten; 0 = off

# ---- embedding ----
EMBED_METHOD    = "umap"   # "tsne" | "umap" | "umap_supervised"
PCA_DIM         = 50
TSNE_PERPLEXITY = 30
UMAP_NEIGHBORS  = 30
UMAP_MIN_DIST   = 0.40    # larger -> clusters spread into visible clouds (less dot-in-whitespace)
UMAP_SPREAD     = 1.30    # scales overall spacing; raise together with min_dist
RANDOM_STATE    = 42

# ---- plotting ----
SYN_MARKER = "x"      # how to draw synthetic points: "x" | "plus" | "ring"
                      #   "x"/"plus" read clearly on top of filled real dots (best when dense)
                      #   "ring"     = same-color open circle (cleaner when sparse)
DRAW_KDE   = False    # overlay per-disease density contours (real solid / syn dashed)
ALLOW_MULTI   = False # if a record lists >1 target disease: False=first only, True=duplicate
MAX_PER_GROUP = None  # cap points per (disease, source); None = all

# ---- metrics ----
COMPUTE_METRICS = True

OUT_PREFIX = "real_vs_syn"
# ------------------------------------------------------------------------

# NPG categorical palette — first 5 colors (matches the order of TARGET_DISEASES).
NPG_PALETTE = ["#E64B35", "#4DBBD5", "#00A087", "#3C5488", "#F39B7F"]


def set_nature_style():
    mpl.rcParams.update({
        "font.family": "sans-serif",
        "font.sans-serif": ["Arial", "Helvetica", "DejaVu Sans"],
        "font.size": 8,
        "axes.linewidth": 0.8,
        "axes.spines.top": False,
        "axes.spines.right": False,
        "axes.titlesize": 9,
        "axes.titleweight": "bold",
        "axes.labelsize": 8,
        "xtick.major.width": 0.8,
        "ytick.major.width": 0.8,
        "xtick.major.size": 3,
        "ytick.major.size": 3,
        "legend.frameon": False,
        "legend.fontsize": 7.5,
        "svg.fonttype": "none",
        "pdf.fonttype": 42,
    })


# --------------------------- data loading -------------------------------
def load_mask_vector(path, size, blur_sigma=0.0):
    """Load mask -> grayscale -> resize -> (optional blur) -> /255 -> flatten."""
    with Image.open(path) as im:
        im = im.convert("L").resize((size, size), Image.BILINEAR)
        arr = np.asarray(im, dtype=np.float32) / 255.0
    if blur_sigma and blur_sigma > 0:
        arr = gaussian_filter(arr, sigma=blur_sigma)
    return arr.reshape(-1)


def parse_attn(attn_list):
    if not isinstance(attn_list, (list, tuple)):
        return
    for entry in attn_list:
        if isinstance(entry, (list, tuple)) and len(entry) >= 1:
            name = entry[0]
            path = entry[1] if len(entry) >= 2 else None
        else:
            name, path = entry, None
        yield name, path


def record_items(rec, target_set, allow_multi, mask_source):
    """[(relative_mask_path, disease), ...] for one record per mask_source."""
    diseases = [(n, p) for n, p in parse_attn(rec.get(ATTN_KEY, []))
                if n in target_set]
    if not diseases:
        return []

    if mask_source == "attn":
        return [(p, n) for n, p in diseases if p]

    rel = rec.get(mask_source)
    if not rel:
        return []
    labels = [n for n, _ in diseases]
    labels = labels if allow_multi else labels[:1]
    return [(rel, n) for n in labels]


def collect(records, root, source_tag, target_set, size, allow_multi,
            mask_source, blur_sigma):
    feats, diseases, missing, no_label = [], [], 0, 0
    for rec in records:
        items = record_items(rec, target_set, allow_multi, mask_source)
        if not items:
            no_label += 1
            continue
        for rel, dis in items:
            mpath = os.path.join(root, rel)
            if not os.path.isfile(mpath):
                missing += 1
                continue
            try:
                vec = load_mask_vector(mpath, size, blur_sigma)
            except Exception as e:
                print(f"  [warn] failed to read {mpath}: {e}")
                missing += 1
                continue
            feats.append(vec)
            diseases.append(dis)
    print(f"  [{source_tag}] loaded={len(feats)}  missing_files={missing}  "
          f"no_target_label={no_label}")
    return feats, diseases


def subsample(feats, diseases, sources, max_per_group, seed):
    if not max_per_group:
        return feats, diseases, sources
    rng = np.random.default_rng(seed)
    keep, groups = [], {}
    for i, (d, s) in enumerate(zip(diseases, sources)):
        groups.setdefault((d, s), []).append(i)
    for idxs in groups.values():
        if len(idxs) > max_per_group:
            idxs = list(rng.choice(idxs, size=max_per_group, replace=False))
        keep.extend(idxs)
    keep.sort()
    return ([feats[i] for i in keep],
            [diseases[i] for i in keep],
            [sources[i] for i in keep])


def load_records(path):
    """Robust loader: whole-file JSON (array or dict) OR JSONL (one record/line)."""
    with open(path, "r") as f:
        text = f.read()
    # Try the whole file as a single JSON value first.
    try:
        data = json.loads(text)
        if isinstance(data, dict):
            return list(data.values())
        if isinstance(data, list):
            return data
        return [data]
    except json.JSONDecodeError:
        pass
    # Fall back to line-delimited JSON (JSONL).
    records = []
    for line in text.splitlines():
        line = line.strip()
        if line:
            records.append(json.loads(line))
    return records


# ----------------------------- embedding --------------------------------
def embed(X, diseases, method, seed,
          n_neighbors=UMAP_NEIGHBORS, min_dist=UMAP_MIN_DIST, spread=UMAP_SPREAD):
    """PCA -> {t-SNE | UMAP | supervised UMAP}. Returns (N,2) array."""
    n_comp = min(PCA_DIM, X.shape[0] - 1, X.shape[1])
    X_red = PCA(n_components=n_comp, random_state=seed).fit_transform(X)

    if method == "tsne":
        perp = min(TSNE_PERPLEXITY, max(5, (X.shape[0] - 1) // 3))
        print(f"  t-SNE (perplexity={perp})")
        emb = TSNE(n_components=2, perplexity=perp, init="pca",
                   learning_rate="auto", random_state=seed).fit_transform(X_red)
        return emb, X_red

    try:
        import umap
    except ImportError:
        raise SystemExit("UMAP not installed. Run:  pip install umap-learn\n"
                         "or set EMBED_METHOD = 'tsne'.")

    if method == "umap":
        print(f"  UMAP (n_neighbors={n_neighbors}, min_dist={min_dist}, spread={spread})")
        reducer = umap.UMAP(n_neighbors=n_neighbors, min_dist=min_dist,
                            spread=spread, random_state=seed)
        emb = reducer.fit_transform(X_red)
    elif method == "umap_supervised":
        print(f"  UMAP supervised on DISEASE only "
              f"(n_neighbors={n_neighbors}, min_dist={min_dist}, spread={spread})")
        y = np.array([TARGET_DISEASES.index(d) for d in diseases])
        reducer = umap.UMAP(n_neighbors=n_neighbors, min_dist=min_dist,
                            spread=spread, random_state=seed)
        emb = reducer.fit_transform(X_red, y=y)   # source NOT passed in
    else:
        raise SystemExit(f"Unknown EMBED_METHOD: {method}")
    return emb, X_red


# ------------------------------ metrics ---------------------------------
def rbf_mmd2(A, B):
    """Biased RBF-MMD^2 with the median-distance heuristic bandwidth."""
    if len(A) < 2 or len(B) < 2:
        return float("nan")
    XX = cdist(A, A, "sqeuclidean")
    YY = cdist(B, B, "sqeuclidean")
    XY = cdist(A, B, "sqeuclidean")
    pooled = np.concatenate([XX[np.triu_indices_from(XX, 1)],
                             YY[np.triu_indices_from(YY, 1)],
                             XY.ravel()])
    med = np.median(pooled)
    gamma = 1.0 / (med + 1e-12)
    k = lambda D: np.exp(-gamma * D)
    mmd2 = k(XX).mean() + k(YY).mean() - 2.0 * k(XY).mean()
    return float(max(mmd2, 0.0))


def c2st_auc(A, B, seed):
    """Classifier two-sample test: 5-fold CV AUC for real-vs-syn. ~0.5 = good."""
    n = min(len(A), len(B))
    if n < 5:
        return float("nan")
    Xc = np.vstack([A, B])
    yc = np.concatenate([np.zeros(len(A)), np.ones(len(B))]).astype(int)
    folds = min(5, np.bincount(yc).min())
    if folds < 2:
        return float("nan")
    clf = make_pipeline(StandardScaler(),
                        LogisticRegression(max_iter=2000))
    cv = StratifiedKFold(n_splits=folds, shuffle=True, random_state=seed)
    return float(cross_val_score(clf, Xc, yc, cv=cv, scoring="roc_auc").mean())


def norm_centroid_dist(A, B):
    """||mu_A - mu_B|| / sqrt(0.5*(tr cov_A + tr cov_B)). Scale-free."""
    if len(A) < 2 or len(B) < 2:
        return float("nan")
    d = np.linalg.norm(A.mean(0) - B.mean(0))
    spread = np.sqrt(0.5 * (A.var(0).sum() + B.var(0).sum())) + 1e-12
    return float(d / spread)


def compute_metrics(X_red, diseases, sources, seed):
    rows = []
    for d in TARGET_DISEASES:
        m = diseases == d
        A = X_red[m & (sources == "real")]
        B = X_red[m & (sources == "syn")]
        rows.append({
            "disease": d,
            "n_real": int(len(A)),
            "n_syn": int(len(B)),
            "c2st_auc": c2st_auc(A, B, seed),
            "mmd2": rbf_mmd2(A, B),
            "norm_centroid_dist": norm_centroid_dist(A, B),
        })
    # disease separability (combined real+syn) in feature space
    dis_int = np.array([TARGET_DISEASES.index(d) for d in diseases])
    sil = float(silhouette_score(X_red, dis_int)) if len(set(dis_int)) > 1 else float("nan")
    return rows, sil


def print_metrics(rows, sil):
    print("\n" + "=" * 64)
    print(f"{'disease':14s} {'n_real':>6s} {'n_syn':>6s} "
          f"{'C2ST':>6s} {'MMD^2':>8s} {'cdist':>7s}")
    print("-" * 64)
    aucs = []
    for r in rows:
        auc = r["c2st_auc"]
        if not np.isnan(auc):
            aucs.append(auc)
        print(f"{r['disease']:14s} {r['n_real']:6d} {r['n_syn']:6d} "
              f"{auc:6.3f} {r['mmd2']:8.4f} {r['norm_centroid_dist']:7.3f}")
    print("-" * 64)
    if aucs:
        print(f"mean C2ST AUC = {np.mean(aucs):.3f}  (0.5 = real/syn indistinguishable)")
    print(f"disease silhouette (real+syn) = {sil:.3f}")
    print("=" * 64 + "\n")


def save_metrics_csv(rows, sil, path):
    import csv
    with open(path, "w", newline="") as f:
        w = csv.DictWriter(f, fieldnames=list(rows[0].keys()))
        w.writeheader()
        for r in rows:
            w.writerow(r)
        f.write(f"# disease_silhouette_real_plus_syn,{sil:.6f}\n")
    print(f"Saved metrics: {path}")


# ------------------------------ plotting --------------------------------
def add_kde_contours(ax, emb, diseases, sources, color_map):
    """Per-disease density contour: real solid, syn dashed (single level)."""
    from scipy.stats import gaussian_kde
    xmin, xmax = emb[:, 0].min(), emb[:, 0].max()
    ymin, ymax = emb[:, 1].min(), emb[:, 1].max()
    xx, yy = np.mgrid[xmin:xmax:120j, ymin:ymax:120j]
    grid = np.vstack([xx.ravel(), yy.ravel()])
    for d in TARGET_DISEASES:
        for src, ls in [("real", "-"), ("syn", "--")]:
            m = (diseases == d) & (sources == src)
            if m.sum() < 8:
                continue
            try:
                kde = gaussian_kde(emb[m].T)
                z = kde(grid).reshape(xx.shape)
                ax.contour(xx, yy, z, levels=[z.max() * 0.30],
                           colors=[color_map[d]], linewidths=0.7,
                           linestyles=ls, alpha=0.55)
            except Exception:
                pass


def plot_overlay(emb, diseases, sources, color_map, axis_label, subtitle,
                 out_prefix, draw_kde=False, syn_marker="x"):
    fig, ax = plt.subplots(figsize=(5.2, 5.4), dpi=150)

    for d in TARGET_DISEASES:
        c = color_map[d]
        mr = (diseases == d) & (sources == "real")
        ms = (diseases == d) & (sources == "syn")
        # real = filled dot
        ax.scatter(emb[mr, 0], emb[mr, 1], s=15, c=c, marker="o",
                   linewidths=0, alpha=0.55, rasterized=True, zorder=2)
        # syn = high-contrast marker drawn on top
        if syn_marker == "ring":
            ax.scatter(emb[ms, 0], emb[ms, 1], s=30, facecolors="none",
                       edgecolors=c, marker="o", linewidths=0.9,
                       alpha=0.9, rasterized=True, zorder=3)
        else:
            mk = "x" if syn_marker == "x" else "+"
            ax.scatter(emb[ms, 0], emb[ms, 1], s=20, c=c, marker=mk,
                       linewidths=0.9, alpha=0.85, rasterized=True, zorder=3)

    if draw_kde:
        add_kde_contours(ax, emb, diseases, sources, color_map)

    # ---- tight, balanced canvas: equal aspect + small symmetric padding ----
    ax.set_aspect("equal", adjustable="box")
    xr = emb[:, 0].max() - emb[:, 0].min()
    yr = emb[:, 1].max() - emb[:, 1].min()
    pad = 0.06 * max(xr, yr)
    ax.set_xlim(emb[:, 0].min() - pad, emb[:, 0].max() + pad)
    ax.set_ylim(emb[:, 1].min() - pad, emb[:, 1].max() + pad)

    ax.set_xticks([]); ax.set_yticks([])
    ax.set_xlabel(f"{axis_label} 1")
    ax.set_ylabel(f"{axis_label} 2")
    if subtitle:
        ax.set_title(subtitle, fontsize=8, fontweight="normal", color="0.35")

    # ---- legends below the plot (no overlap with data) ----
    dis_handles = [Line2D([0], [0], marker="o", linestyle="", markersize=5,
                          markerfacecolor=color_map[d], markeredgecolor="none",
                          label=d) for d in TARGET_DISEASES]
    if syn_marker == "ring":
        syn_h = Line2D([0], [0], marker="o", linestyle="", markersize=6,
                       markerfacecolor="none", markeredgecolor="0.3", label="Synthetic")
    else:
        syn_h = Line2D([0], [0], marker=("x" if syn_marker == "x" else "+"),
                       linestyle="", markersize=6, color="0.3", label="Synthetic")
    src_handles = [
        Line2D([0], [0], marker="o", linestyle="", markersize=5,
               markerfacecolor="0.3", markeredgecolor="none", label="Real"),
        syn_h,
    ]
    leg1 = fig.legend(handles=dis_handles, loc="lower center",
                      bbox_to_anchor=(0.5, 0.04), ncol=5,
                      columnspacing=1.0, handletextpad=0.3)
    fig.add_artist(leg1)
    fig.legend(handles=src_handles, loc="lower center",
               bbox_to_anchor=(0.5, -0.01), ncol=2,
               columnspacing=1.5, handletextpad=0.3)

    fig.tight_layout(rect=[0, 0.10, 1, 1])
    png, svg = f"{out_prefix}.png", f"{out_prefix}.svg"
    fig.savefig(png, dpi=300, bbox_inches="tight")
    fig.savefig(svg, bbox_inches="tight")
    print(f"Saved figure: {png}  /  {svg}")


# -------------------------------- main ----------------------------------
def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--json", default=JSON_PATH)
    ap.add_argument("--real_root", default=REAL_ROOT)
    ap.add_argument("--syn_root", default=SYN_ROOT)
    ap.add_argument("--out", default=OUT_PREFIX)
    ap.add_argument("--method", default=EMBED_METHOD,
                    choices=["tsne", "umap", "umap_supervised"])
    ap.add_argument("--blur", type=float, default=BLUR_SIGMA)
    ap.add_argument("--min_dist", type=float, default=UMAP_MIN_DIST)
    ap.add_argument("--spread", type=float, default=UMAP_SPREAD)
    ap.add_argument("--n_neighbors", type=int, default=UMAP_NEIGHBORS)
    ap.add_argument("--syn_marker", default=SYN_MARKER,
                    choices=["x", "plus", "ring"])
    ap.add_argument("--kde", action="store_true", default=DRAW_KDE)
    ap.add_argument("--no_metrics", action="store_true")
    args = ap.parse_args()

    set_nature_style()
    target_set = set(TARGET_DISEASES)

    records = load_records(args.json)
    print(f"Loaded {len(records)} records from {args.json}  "
          f"(real={REAL_MASK_SOURCE}, syn={SYN_MASK_SOURCE}, "
          f"method={args.method}, blur={args.blur})")

    print("Reading masks...")
    real_f, real_d = collect(records, args.real_root, "real", target_set,
                             MASK_SIZE, ALLOW_MULTI, REAL_MASK_SOURCE, args.blur)
    syn_f,  syn_d  = collect(records, args.syn_root,  "syn",  target_set,
                             MASK_SIZE, ALLOW_MULTI, SYN_MASK_SOURCE, args.blur)

    feats    = real_f + syn_f
    diseases = real_d + syn_d
    sources  = ["real"] * len(real_f) + ["syn"] * len(syn_f)
    if len(feats) < 10:
        raise SystemExit("Too few masks loaded — check paths / MASK_SOURCE / disease names.")

    feats, diseases, sources = subsample(feats, diseases, sources,
                                         MAX_PER_GROUP, RANDOM_STATE)

    X = np.stack(feats).astype(np.float32)
    diseases = np.array(diseases)
    sources  = np.array(sources)
    print(f"Feature matrix: {X.shape}")

    print("Embedding...")
    emb, X_red = embed(X, diseases, args.method, RANDOM_STATE,
                       n_neighbors=args.n_neighbors, min_dist=args.min_dist,
                       spread=args.spread)
    color_map = {d: NPG_PALETTE[i] for i, d in enumerate(TARGET_DISEASES)}
    axis_label = "t-SNE" if args.method == "tsne" else "UMAP"

    subtitle = None
    if COMPUTE_METRICS and not args.no_metrics:
        print("Computing metrics...")
        rows, sil = compute_metrics(X_red, diseases, sources, RANDOM_STATE)
        print_metrics(rows, sil)
        save_metrics_csv(rows, sil, f"{args.out}_metrics.csv")
        aucs = [r["c2st_auc"] for r in rows if not np.isnan(r["c2st_auc"])]
        if aucs:
            subtitle = f"mean C2ST AUC = {np.mean(aucs):.3f}   silhouette = {sil:.3f}"

    plot_overlay(emb, diseases, sources, color_map, axis_label, subtitle,
                 args.out, draw_kde=args.kde, syn_marker=args.syn_marker)


if __name__ == "__main__":
    main()

"""
# recommended: UMAP, spread-out clusters, x-markers for syn, metrics
python draw_cluster.py \
    --json /home/jovyan/AURAD_infer/mask/det-test-cp9000/test_prompt_text2layout_single.json \
    --syn_root  /home/jovyan/AURAD_infer/mask/det-test-cp9000 \
    --real_root /home/jovyan/AURAD_dataset \
    --method umap --syn_marker x --min_dist 0.25 --spread 1.0 

# if clusters still look too spread / too tight, tune these:
python draw_cluster.py ... --min_dist 0.6 --spread 1.5      # looser, fills canvas more
python draw_cluster.py ... --min_dist 0.25 --spread 1.0     # tighter clusters

# other looks
python draw_cluster.py ... --syn_marker ring                # open-circle syn (sparse data)
python draw_cluster.py ... --method umap_supervised --kde   # diseases pushed apart + contours
python draw_cluster.py ... --method tsne --blur 0           # original t-SNE style
"""