#!/usr/bin/env python3
"""figure5.py — PCA and t-SNE of VGG16 features, colored by class.

Usage:
    python scripts/visualizations/figure5.py
    python scripts/visualizations/figure5.py --tag v2
"""

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
CLASS_COLORS = {"Bengin cases": "#2CA02C", "Malignant cases": "#D62728",
                "Normal cases": "#1F77B4"}


def main():
    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 ""

    data = np.load(os.path.join(FEATURES_DIR, "VGG16_features.npz"),
                   allow_pickle=True)
    X, Y = data["X"], data["Y"]
    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=1000).fit_transform(pca50)

    fig, axes = plt.subplots(1, 2, figsize=(12, 5))
    for ax, coords, title in [(axes[0], pca2, "(a) PCA"),
                               (axes[1], tsne, "(b) t-SNE")]:
        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.6, s=15, edgecolors="none")
        ax.set_title(title, fontsize=12)
        ax.set_xlabel("Component 1")
        ax.set_ylabel("Component 2")
        ax.legend(markerscale=2, fontsize=9)
        ax.grid(alpha=0.2)

    fig.suptitle("Figure 5 — PCA and t-SNE of VGG16 Features (by class)",
                 fontsize=13)
    plt.tight_layout()
    out = os.path.join(PLOTS_DIR, "figure5", f"figure5{tag}.png")
    os.makedirs(os.path.dirname(out), exist_ok=True)
    plt.savefig(out, dpi=150)
    plt.close()
    print(f"Saved → {out}")


if __name__ == "__main__":
    main()