"""Extract backbone features on test set, project to 2D with t-SNE."""
import sys, os
sys.path.insert(0, "/mnt/d/SpiceNet" if os.path.exists("/mnt/d/SpiceNet") else "D:/SpiceNet")

import numpy as np
import torch
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
from sklearn.manifold import TSNE

import config
from src.dataset import get_dataloaders
from src.model import load_checkpoint


@torch.no_grad()
def extract_features(model, loader, device):
    model.eval()
    feats, labels = [], []
    for imgs, tex, col, lbl in loader:
        imgs = imgs.to(device)
        f = model.backbone(imgs)         # (B, 1792) pre-classifier features
        feats.append(f.cpu().numpy())
        labels.extend(lbl.tolist())
    return np.concatenate(feats), np.array(labels)


def main():
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model, *_ = load_checkpoint(str(config.CHECKPOINT_DIR / "best.pth"), device)

    _, _, test_loader, _, _ = get_dataloaders(multimodal=True)
    print(f"Extracting features on {len(test_loader.dataset)} test samples...")
    feats, labels = extract_features(model, test_loader, device)
    print(f"Features: {feats.shape}, labels: {labels.shape}")

    print("Running t-SNE (perplexity=30, ~30s)...")
    proj = TSNE(n_components=2, perplexity=30, init="pca",
                learning_rate="auto", random_state=config.RANDOM_SEED).fit_transform(feats)

    fig, ax = plt.subplots(figsize=(11, 9))
    cmap = plt.get_cmap("tab20")
    for i, cls in enumerate(config.CLASSES):
        mask = labels == i
        ax.scatter(proj[mask, 0], proj[mask, 1], s=8, c=[cmap(i)], label=cls, alpha=0.7)
    ax.set_xlabel("t-SNE 1"); ax.set_ylabel("t-SNE 2")
    ax.set_title("EfficientNet-B4 backbone features (test set) — strong-aug model")
    ax.legend(loc="best", fontsize=9, markerscale=1.8, framealpha=0.85)
    ax.grid(alpha=0.2)
    out_path = config.OUTPUT_DIR / "feature_tsne.png"
    plt.tight_layout()
    plt.savefig(out_path, dpi=150)
    plt.close()
    print(f"Saved -> {out_path}")

    # Per-class centroid distances (just for the analysis table)
    centroids = np.stack([proj[labels == i].mean(0) for i in range(len(config.CLASSES))])
    np.savez(config.OUTPUT_DIR / "feature_tsne_data.npz",
             proj=proj, labels=labels, centroids=centroids)


if __name__ == "__main__":
    main()