#!/usr/bin/env python3
"""
Infer disease clusters from RadIR image embeddings + disease labels.

Idea (matching the user request):
  Within each disease (and Normal), keep only images that carry that label,
  then cluster them in the same feature space as plot_tsne.py. Every cluster
  is disease-homogeneous by construction ("nearby images that are all pneumonia
  → label them Pneumonia").

Conflicts (an image in multiple disease clusters when using inclusive labels)
are resolved by exclusive single-label first, else largest cluster, else
highest within-disease density.

Usage (from RadIR/RadIR):
  # RadIR hyperbolic
  python scripts/infer_disease_clusters.py \\
    --load_feats outputs/tsne_feats.npz \\
    --labels_csv outputs/tsne_disease_labels.csv \\
    --out_csv outputs/tsne_cluster_disease.csv \\
    --out_plot outputs/tsne_cluster_disease.png

  # RadIR_compare Euclidean (same label CSV / jsonl order)
  python scripts/infer_disease_clusters.py \\
    --load_feats ../RadIR_compare/outputs/tsne_compare_feats.npz \\
    --labels_csv outputs/tsne_disease_labels.csv \\
    --jsonl ../RadIR_compare/dataset/valid.jsonl \\
    --feat_type euclidean \\
    --out_csv ../RadIR_compare/outputs/tsne_cluster_disease.csv \\
    --out_stats ../RadIR_compare/outputs/tsne_cluster_stats.json \\
    --out_plot ../RadIR_compare/outputs/tsne_cluster_disease.png
"""

from __future__ import annotations

import argparse
import json
import os
import sys
from collections import Counter, defaultdict

import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import torch
from sklearn.cluster import DBSCAN
from sklearn.manifold import TSNE
from sklearn.neighbors import NearestNeighbors
from sklearn.preprocessing import normalize

ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
RAD_IR_ROOT = os.path.join(ROOT, "Rad_IR")
for p in (ROOT, RAD_IR_ROOT):
    if p not in sys.path:
        sys.path.insert(0, p)

from radir import lorentz as L

DISEASE_NAMES = [
    "Pneumothorax",
    "Edema",
    "Pneumonia",
    "Effusion",
    "Cardiomegaly",
    "Atelectasis",
]
CLASS_ORDER = ["Normal"] + DISEASE_NAMES + ["Unassigned"]

CLASS_COLORS = {
    "Normal": "#00C853",
    "Pneumothorax": "#AA00FF",
    "Edema": "#00B8D4",
    "Pneumonia": "#D50000",
    "Effusion": "#2962FF",
    "Cardiomegaly": "#FF6D00",
    "Atelectasis": "#FFD600",
    "Unassigned": "#BDBDBD",
}


def dicom_from_img_path(img_path: str) -> str:
    return os.path.splitext(os.path.basename(img_path.replace("\\", "/")))[0]


def load_dicom_ids_from_jsonl(jsonl_path: str) -> list[str]:
    ids = []
    with open(jsonl_path, encoding="utf-8") as f:
        for line in f:
            item = json.loads(line.strip())
            ids.append(dicom_from_img_path(item["img_path"]))
    return ids


def prepare_image_feats(data, feat_type: str) -> np.ndarray:
    """
    Support both:
      - RadIR/RadIR tsne_feats.npz: eu_img / hyp_img / curv
      - RadIR_compare tsne_compare_feats.npz: img_feats (already L2 Euclidean)
    """
    files = set(data.files)

    # RadIR_compare cached format
    if "img_feats" in files:
        if feat_type == "hyperbolic":
            print(
                "WARNING: compare feats are Euclidean only; "
                "ignoring --feat_type hyperbolic"
            )
        return np.asarray(data["img_feats"])

    # RadIR/RadIR cached format
    eu_img = torch.from_numpy(data["eu_img"])
    hyp_img = torch.from_numpy(data["hyp_img"])
    curv = torch.tensor(float(data["curv"]))
    if feat_type == "hyperbolic":
        with torch.no_grad():
            return L.log_map0(hyp_img.float(), curv.detach().float()).cpu().numpy()
    if feat_type == "euclidean":
        return eu_img.detach().cpu().numpy()
    raise ValueError(feat_type)


def load_label_matrix(labels_csv: str, dicom_ids: list[str]):
    df = pd.read_csv(labels_csv)
    df = df.drop_duplicates(subset=["dicom_id"], keep="first").set_index("dicom_id")
    n = len(dicom_ids)
    normal_only = np.zeros(n, dtype=bool)
    disease_mat = np.zeros((n, len(DISEASE_NAMES)), dtype=bool)
    missing = 0
    for i, did in enumerate(dicom_ids):
        if did not in df.index:
            missing += 1
            continue
        row = df.loc[did]
        normal_only[i] = bool(int(row.get("normal_only", 0)))
        for j, name in enumerate(DISEASE_NAMES):
            disease_mat[i, j] = bool(int(row.get(f"{name}_pos", 0)))
    if missing:
        print(f"WARNING: {missing} dicom_ids missing from labels CSV")
    return normal_only, disease_mat


def exclusive_labels(normal_only: np.ndarray, disease_mat: np.ndarray) -> np.ndarray:
    """Single exclusive label or None."""
    n = len(normal_only)
    out = np.array([None] * n, dtype=object)
    for i in range(n):
        hits = []
        if normal_only[i]:
            hits.append("Normal")
        for j, name in enumerate(DISEASE_NAMES):
            if disease_mat[i, j]:
                hits.append(name)
        if len(hits) == 1:
            out[i] = hits[0]
    return out


def auto_eps(feats_sub: np.ndarray, min_samples: int, quantile: float) -> float:
    if feats_sub.shape[0] < min_samples + 1:
        return 0.1
    x = normalize(feats_sub.astype(np.float64), norm="l2", axis=1)
    k = min(min_samples, feats_sub.shape[0] - 1)
    nn = NearestNeighbors(n_neighbors=k + 1, metric="cosine").fit(x)
    dists, _ = nn.kneighbors(x)
    return float(np.quantile(dists[:, -1], quantile))


def cluster_subset(
    feats_all: np.ndarray,
    member_idx: np.ndarray,
    eps: float | None,
    min_samples: int,
    eps_quantile: float,
    min_cluster_size: int,
) -> list[dict]:
    """
    DBSCAN on a disease subset. Returns list of
      {global_indices, size, eps_used}
    for clusters with size >= min_cluster_size.
    """
    if member_idx.size < min_cluster_size:
        return []

    sub = feats_all[member_idx]
    eps_used = eps
    if eps_used is None:
        eps_used = auto_eps(sub, min_samples, eps_quantile)

    x = normalize(sub.astype(np.float64), norm="l2", axis=1)
    labels = DBSCAN(
        eps=eps_used, min_samples=min_samples, metric="cosine"
    ).fit_predict(x)

    clusters = []
    for cid in sorted(set(labels.tolist()) - {-1}):
        local = np.where(labels == cid)[0]
        if local.size < min_cluster_size:
            continue
        global_idx = member_idx[local]
        clusters.append(
            {
                "global_indices": global_idx.astype(int),
                "size": int(local.size),
                "eps_used": float(eps_used),
                "local_cluster_id": int(cid),
            }
        )
    return clusters


def infer_disease_clusters(
    feats: np.ndarray,
    normal_only: np.ndarray,
    disease_mat: np.ndarray,
    seed_mode: str,
    eps: float | None,
    min_samples: int,
    eps_quantile: float,
    min_cluster_size: int,
) -> tuple[np.ndarray, np.ndarray, list[dict]]:
    """
    Returns:
      assigned: [N] disease name or 'Unassigned'
      cluster_uid: [N] int, -1 if unassigned
      stats: list of cluster summaries
    """
    n = feats.shape[0]
    excl = exclusive_labels(normal_only, disease_mat)

    # candidates: list of (disease, global_indices, size, eps, seed_mode)
    candidates: list[dict] = []
    class_names = ["Normal"] + DISEASE_NAMES

    for name in class_names:
        if seed_mode == "exclusive":
            member_idx = np.where(excl == name)[0]
        elif seed_mode == "inclusive":
            if name == "Normal":
                member_idx = np.where(normal_only)[0]
            else:
                j = DISEASE_NAMES.index(name)
                member_idx = np.where(disease_mat[:, j])[0]
        else:
            raise ValueError(seed_mode)

        print(
            f"  Seed {name:14s} ({seed_mode}): {member_idx.size} images",
            end="",
        )
        clusters = cluster_subset(
            feats,
            member_idx,
            eps=eps,
            min_samples=min_samples,
            eps_quantile=eps_quantile,
            min_cluster_size=min_cluster_size,
        )
        print(f" → {len(clusters)} clusters ≥{min_cluster_size}")
        for c in clusters:
            candidates.append(
                {
                    "disease": name,
                    "global_indices": c["global_indices"],
                    "size": c["size"],
                    "eps_used": c["eps_used"],
                    "seed_mode": seed_mode,
                }
            )

    # Resolve conflicts: each sample may be claimed by multiple disease clusters.
    # Score = (is_exclusive_match, cluster_size, -mean_dist_to_centroid)
    claims: dict[int, list[tuple]] = defaultdict(list)
    for cid, c in enumerate(candidates):
        disease = c["disease"]
        idxs = c["global_indices"]
        sub = normalize(feats[idxs].astype(np.float64), norm="l2", axis=1)
        centroid = normalize(sub.mean(axis=0, keepdims=True), norm="l2")[0]
        # cosine distance to centroid
        dist = 1.0 - (sub @ centroid)
        for k, i in enumerate(idxs):
            is_excl = 1 if excl[i] == disease else 0
            claims[int(i)].append((is_excl, c["size"], -float(dist[k]), cid, disease))

    assigned = np.array(["Unassigned"] * n, dtype=object)
    cluster_uid = np.full(n, -1, dtype=int)
    for i, opts in claims.items():
        opts.sort(reverse=True)  # prefer exclusive match, then larger cluster
        _, _, _, cid, disease = opts[0]
        assigned[i] = disease
        cluster_uid[i] = cid

    stats = []
    for cid, c in enumerate(candidates):
        idxs = c["global_indices"]
        kept = idxs[assigned[idxs] == c["disease"]]
        stats.append(
            {
                "cluster_uid": cid,
                "disease": c["disease"],
                "seed_mode": c["seed_mode"],
                "seed_size": c["size"],
                "assigned_size": int(kept.size),
                "eps_used": c["eps_used"],
                "exclusive_in_seed": int(
                    sum(1 for i in idxs if excl[i] == c["disease"])
                ),
            }
        )
    return assigned, cluster_uid, stats


def run_tsne_2d(feats: np.ndarray, perplexity: float, seed: int) -> np.ndarray:
    n = feats.shape[0]
    max_perp = max(2.0, (n - 1) / 3.0)
    eff = min(perplexity, max_perp)
    return TSNE(
        n_components=2,
        perplexity=eff,
        random_state=seed,
        init="pca",
        learning_rate="auto",
    ).fit_transform(feats)


def plot_assigned(
    emb: np.ndarray,
    assigned: np.ndarray,
    out_path: str,
    title: str,
    dpi: int,
) -> None:
    fig, ax = plt.subplots(figsize=(11, 9))
    # draw Unassigned first (background)
    for cls in ["Unassigned"] + [c for c in CLASS_ORDER if c != "Unassigned"]:
        mask = assigned == cls
        if not mask.any():
            continue
        ax.scatter(
            emb[mask, 0],
            emb[mask, 1],
            c=CLASS_COLORS.get(cls, "#333"),
            marker="o",
            alpha=0.35 if cls == "Unassigned" else 0.9,
            s=18 if cls == "Unassigned" else 40,
            edgecolors="none" if cls == "Unassigned" else "black",
            linewidths=0.0 if cls == "Unassigned" else 0.35,
            label=f"{cls} (n={mask.sum()})",
            zorder=1 if cls == "Unassigned" else 2,
        )
    ax.set_title(title)
    ax.set_xlabel("t-SNE 1")
    ax.set_ylabel("t-SNE 2")
    ax.legend(loc="best", fontsize=8, frameon=True)
    ax.grid(True, alpha=0.2)
    fig.tight_layout()
    os.makedirs(os.path.dirname(os.path.abspath(out_path)) or ".", exist_ok=True)
    fig.savefig(out_path, dpi=dpi, bbox_inches="tight")
    plt.close(fig)
    print(f"Saved plot to {out_path}")


def parse_args():
    p = argparse.ArgumentParser(
        description="Disease-seeded cluster inference on RadIR image embeddings"
    )
    p.add_argument(
        "--load_feats",
        default=os.path.join(ROOT, "outputs/tsne_feats.npz"),
    )
    p.add_argument(
        "--labels_csv",
        default=os.path.join(ROOT, "outputs/tsne_disease_labels.csv"),
    )
    p.add_argument(
        "--jsonl",
        default=os.path.join(ROOT, "dataset/valid.jsonl"),
    )
    p.add_argument(
        "--feat_type",
        choices=["hyperbolic", "euclidean"],
        default="hyperbolic",
    )
    p.add_argument(
        "--seed_mode",
        choices=["exclusive", "inclusive"],
        default="exclusive",
        help="exclusive: only single-label images as seeds (cleaner); "
        "inclusive: all positives of each disease",
    )
    p.add_argument(
        "--eps",
        type=float,
        default=None,
        help="DBSCAN cosine eps within each disease; default=auto per disease",
    )
    p.add_argument("--eps_quantile", type=float, default=0.6)
    p.add_argument("--min_samples", type=int, default=5)
    p.add_argument("--min_cluster_size", type=int, default=8)
    p.add_argument("--perplexity", type=float, default=30.0)
    p.add_argument("--seed", type=int, default=42)
    p.add_argument(
        "--out_csv",
        default=os.path.join(ROOT, "outputs/tsne_cluster_disease.csv"),
    )
    p.add_argument(
        "--out_stats",
        default=os.path.join(ROOT, "outputs/tsne_cluster_stats.json"),
    )
    p.add_argument(
        "--out_plot",
        default=os.path.join(ROOT, "outputs/tsne_cluster_disease.png"),
    )
    p.add_argument("--dpi", type=int, default=150)
    return p.parse_args()


def main():
    args = parse_args()

    print(f"Loading features: {args.load_feats}")
    data = np.load(args.load_feats, allow_pickle=True)
    img_feats = prepare_image_feats(data, args.feat_type)
    n = img_feats.shape[0]
    print(f"  image feats: {img_feats.shape}, feat_type={args.feat_type}")

    dicom_ids = load_dicom_ids_from_jsonl(args.jsonl)
    if len(dicom_ids) != n:
        raise RuntimeError(f"jsonl ({len(dicom_ids)}) != feats ({n})")

    normal_only, disease_mat = load_label_matrix(args.labels_csv, dicom_ids)
    excl = exclusive_labels(normal_only, disease_mat)
    print(
        f"Exclusive single-label: {sum(e is not None for e in excl)} / {n}"
    )
    print(f"Seed mode: {args.seed_mode}")

    print("\n===== Within-disease clustering =====")
    assigned, cluster_uid, stats = infer_disease_clusters(
        img_feats,
        normal_only,
        disease_mat,
        seed_mode=args.seed_mode,
        eps=args.eps,
        min_samples=args.min_samples,
        eps_quantile=args.eps_quantile,
        min_cluster_size=args.min_cluster_size,
    )

    print("\n===== Clusters kept =====")
    stats_sorted = sorted(stats, key=lambda s: (-s["assigned_size"], s["disease"]))
    for s in stats_sorted:
        if s["assigned_size"] <= 0:
            continue
        print(
            f"  [{s['cluster_uid']:>3d}] {s['disease']:14s} "
            f"assigned={s['assigned_size']:<4d} "
            f"(seed={s['seed_size']}, excl_in_seed={s['exclusive_in_seed']}, "
            f"eps={s['eps_used']:.3f})"
        )

    print("\n===== Assigned counts =====")
    for cls in CLASS_ORDER:
        c = int((assigned == cls).sum())
        if c:
            print(f"  {cls:14s} {c}")

    os.makedirs(os.path.dirname(os.path.abspath(args.out_csv)) or ".", exist_ok=True)
    rows = {
        "dicom_id": dicom_ids,
        "cluster_uid": cluster_uid,
        "cluster_disease": assigned,
        "exclusive_label": [
            e if e is not None else "" for e in excl.tolist()
        ],
        "normal_only": normal_only.astype(int),
    }
    for j, name in enumerate(DISEASE_NAMES):
        rows[f"{name}_pos"] = disease_mat[:, j].astype(int)
    pd.DataFrame(rows).to_csv(args.out_csv, index=False)
    print(f"\nSaved {args.out_csv}")

    with open(args.out_stats, "w", encoding="utf-8") as f:
        json.dump(
            {
                "seed_mode": args.seed_mode,
                "feat_type": args.feat_type,
                "eps": args.eps,
                "eps_quantile": args.eps_quantile,
                "min_samples": args.min_samples,
                "min_cluster_size": args.min_cluster_size,
                "n_samples": n,
                "assigned_counts": {
                    cls: int((assigned == cls).sum()) for cls in CLASS_ORDER
                },
                "clusters": stats_sorted,
            },
            f,
            indent=2,
        )
    print(f"Saved {args.out_stats}")

    print("Running t-SNE for visualization ...")
    emb = run_tsne_2d(img_feats, args.perplexity, args.seed)
    title = (
        f"Disease-seeded image clusters "
        f"({args.seed_mode}, {args.feat_type}, min_size={args.min_cluster_size})"
    )
    plot_assigned(emb, assigned, args.out_plot, title, args.dpi)


if __name__ == "__main__":
    main()