#!/usr/bin/env python3
"""figure_s2.py — PCA and t-SNE for all 5 CNN models, colored by class.
Usage: python scripts/visualizations/figure_s2.py [--tag TAG]"""

import os, sys, argparse
import numpy as np
import matplotlib; matplotlib.use("Agg")
import matplotlib.pyplot as plt
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")
SEED = 42
MODELS = ["VGG16", "MobileNetV2", "DenseNet121", "ResNet50", "EfficientNetB1"]
CLASS_COLORS = {"Bengin cases": "#2CA02C", "Malignant cases": "#D62728", "Normal cases": "#1F77B4"}

ap = argparse.ArgumentParser()
ap.add_argument("--tag", default="", help="Append tag to filename")
args = ap.parse_args()
tag = f"_{args.tag}" if args.tag else ""

fig, axes = plt.subplots(2, 5, figsize=(22, 9))
for mi, model_name in enumerate(MODELS):
    data = np.load(os.path.join(FEATURES_DIR, f"{model_name}_features.npz"), allow_pickle=True)
    X, Y = data["X"], data["Y"]
    print(f"{model_name}: {X.shape}")
    pca50 = PCA(n_components=50, random_state=SEED).fit_transform(X)
    pca2 = PCA(n_components=2, random_state=SEED).fit_transform(pca50)
    tsne = TSNE(n_components=2, random_state=SEED, perplexity=30, max_iter=800).fit_transform(pca50)
    for row, coords, title in [(0, pca2, f"{model_name} PCA"), (1, tsne, f"{model_name} t-SNE")]:
        ax = axes[row, mi]
        for cls in sorted(np.unique(Y)):
            mask = Y == cls; label = cls.replace("Bengin cases", "Benign")
            ax.scatter(coords[mask,0], coords[mask,1], c=CLASS_COLORS[cls], label=label, alpha=0.5, s=8, edgecolors="none")
        ax.set_title(title, fontsize=9)
        ax.set_xlabel("Component 1" if row==0 else "t-SNE 1")
        ax.set_ylabel("Component 2" if row==0 else "t-SNE 2")
        ax.grid(alpha=0.2)

handles = [plt.Line2D([0],[0], marker='o', color='w', markerfacecolor=c, markersize=8, label=l)
           for l,c in zip(["Benign","Malignant","Normal"], ["#2CA02C","#D62728","#1F77B4"])]
fig.legend(handles=handles, loc='lower center', ncol=3, fontsize=10, bbox_to_anchor=(0.5,-0.02))
fig.suptitle("Figure S2 — PCA (top) and t-SNE (bottom) per model", fontsize=13, y=1.01)
plt.tight_layout()
out = os.path.join(PLOTS_DIR, "figure_s2", f"figure_s2{tag}.png")
os.makedirs(os.path.dirname(out), exist_ok=True)
plt.savefig(out, dpi=150, bbox_inches="tight")
plt.close()
print(f"Saved → {out}")