#!/usr/bin/env python3
"""manifest_tsne.py — visualize any patient-grouping manifest on the VGG16 t-SNE.

Plots the same VGG16 feature t-SNE used elsewhere, but colored by the groups in
a given manifest (siamese / pca50 / thumbnail). Because every manifest is drawn
on the *identical* layout (same features, PCA, seed, perplexity), the resulting
per-class figures are directly comparable across methods.

For the siamese grouping this is diagnostic: if a siamese "patient" is a coherent
patient it forms a tight island; if it's an over-merged chain, its color is
smeared across feature space (VGG16 sees images the siamese wrongly linked).

Usage:
    conda activate fundus_imaging
    python scripts/visualizations/manifest_tsne.py \
        --manifest results/siamese_manifest.csv --name siamese
    python scripts/visualizations/manifest_tsne.py \
        --manifest results/simple_patient_manifest.csv --name pca50
"""

import os
import sys
import re
import csv
import argparse
from collections import defaultdict

import numpy as np
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
from matplotlib.lines import Line2D
from sklearn.decomposition import PCA
from sklearn.manifold import TSNE

sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.dirname(
    os.path.abspath(__file__)))))
ROOT = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
FEATURES_DIR = os.path.join(ROOT, "features")
PLOTS_DIR = os.path.join(ROOT, "plots")

RANDOM_STATE = 42
CLASS_OF = {"B": "Benign", "M": "Malignant", "N": "Normal"}


def f2n(fname):
    """VGG16 feature filename -> manifest short name (e.g. 'B_001')."""
    m = re.search(r"\((\d+)\)", fname)
    num = int(m.group(1)) if m else None
    for cls_key, prefix in [("Bengin cases", "B"), ("Malignant cases", "M"),
                            ("Normal cases", "N")]:
        if fname.startswith(cls_key.rstrip("s")):
            return f"{prefix}_{num:03d}" if num else fname
    return fname


def load_manifest(path):
    """Return {image_short_name: group_id}."""
    mapping = {}
    with open(path, newline="") as f:
        for row in csv.DictReader(f):
            for img in row["images"].split(";"):
                if img:
                    mapping[img] = row["patient_id"]
    return mapping


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--manifest", required=True, help="Path to a grouping manifest CSV.")
    ap.add_argument("--name", required=True, help="Short method name for titles/filenames.")
    ap.add_argument("--features", default="VGG16", help="CNN feature set for the layout.")
    ap.add_argument("--no-centroids", dest="centroids", action="store_false",
                    help="Do not draw per-group centroid diamonds.")
    ap.add_argument("--tag", default="")
    args = ap.parse_args()
    tag = f"_{args.tag}" if args.tag else ""

    print(f"Loading {args.features} features ...")
    data = np.load(os.path.join(FEATURES_DIR, f"{args.features}_features.npz"),
                   allow_pickle=True)
    X, filenames = data["X"], data["filenames"]
    img_ids = np.array([f2n(f) for f in filenames])
    img_class = np.array([CLASS_OF.get(i.split("_")[0], "?") for i in img_ids])

    print("Computing PCA-50 + t-SNE (shared layout) ...")
    n_pca = min(50, X.shape[0] - 1, X.shape[1])
    X_pca = PCA(n_components=n_pca, random_state=RANDOM_STATE).fit_transform(X)
    X_tsne = TSNE(n_components=2, perplexity=35, learning_rate="auto",
                  init="pca", random_state=RANDOM_STATE).fit_transform(X_pca)

    groups = load_manifest(args.manifest)
    img_group = np.array([groups.get(i, "unassigned") for i in img_ids])

    for cls in ["Benign", "Malignant", "Normal"]:
        fig, ax = plt.subplots(figsize=(14, 10))
        ax.scatter(X_tsne[:, 0], X_tsne[:, 1], c="lightgray", s=3, alpha=0.15)

        mask = img_class == cls
        class_groups = sorted(set(img_group[mask]))
        n = len(class_groups)
        # Order groups by size so the biggest (most likely over-merged) is obvious.
        sizes = {g: int((img_group[mask] == g).sum()) for g in class_groups}
        class_groups = sorted(class_groups, key=lambda g: -sizes[g])
        cmap = plt.cm.tab20 if n <= 20 else plt.cm.gist_ncar

        for gi, g in enumerate(class_groups):
            color = cmap(gi % 20) if n <= 20 else cmap(gi / max(n - 1, 1))
            gm = mask & (img_group == g)
            ax.scatter(X_tsne[gm, 0], X_tsne[gm, 1], c=[color], s=18, alpha=0.8)
            if args.centroids:
                cx, cy = X_tsne[gm, 0].mean(), X_tsne[gm, 1].mean()
                ax.scatter(cx, cy, c=[color], s=60, marker="D", edgecolors="black",
                           linewidths=0.6, zorder=5)

        biggest = class_groups[0]
        legend_handles = [
            Line2D([0], [0], marker="o", color="w", markerfacecolor="gray",
                   markersize=8, label="Group images (dots)"),
        ]
        if args.centroids:
            legend_handles.append(
                Line2D([0], [0], marker="D", color="w", markerfacecolor="gray",
                       markersize=8, label="Group centroids (diamonds)"))
        ax.legend(handles=legend_handles, loc="lower right")
        ax.set_title(f"{cls} — {args.name} groups on {args.features} t-SNE "
                     f"({n} groups; largest={sizes[biggest]} imgs)")
        ax.set_xlabel("t-SNE dim 1")
        ax.set_ylabel("t-SNE dim 2")
        plt.tight_layout()
        out = os.path.join(PLOTS_DIR, "tsne", f"tsne_{args.name}_{cls}{tag}.png")
        os.makedirs(os.path.dirname(out), exist_ok=True)
        plt.savefig(out, dpi=150)
        plt.close()
        print(f"  Saved {out}")
    print("DONE")


if __name__ == "__main__":
    main()