#!/usr/bin/env python3
"""
t-SNE visualization for RadIR/RadIR image-text embeddings.

Extracts global features aligned with forward(is_condition=False), then plots
image vs text points (optionally with pair connectors).

Usage (from RadIR/RadIR project root):
  python scripts/plot_tsne.py \\
    --radir_ckpt outputs/radir_final.pt \\
    --out outputs/tsne_radir.png

  # Reuse saved features:
  python scripts/plot_tsne.py --load_feats outputs/tsne_feats.npz --out outputs/tsne_radir.png
"""

from __future__ import annotations

import argparse
import os
import sys

import matplotlib.pyplot as plt
import numpy as np
import torch
from sklearn.manifold import TSNE
from torch.utils.data import DataLoader
from transformers import AutoModel, AutoTokenizer

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
from radir.RadIR import RADIR, l2norm
from radir.data.hmclip_mimic_box import HyCoClipMimicBoxDataset_JPG, collate_hmclip_box
from transformer_maskgit.transformer_maskgit.ctvit import CTViT


def preprocess_image_like_forward(image: torch.Tensor) -> torch.Tensor:
    if image.ndim == 3:
        image = image.unsqueeze(0)
    if image.ndim == 4:
        if image.shape[1] != 1:
            image = image.mean(dim=1, keepdim=True)
    elif image.ndim == 5:
        if image.shape[1] != 1:
            image = image.mean(dim=1, keepdim=True)
    else:
        raise ValueError(
            f"Unexpected image ndim={image.ndim}, shape={tuple(image.shape)}"
        )
    return image


@torch.no_grad()
def extract_features(model, dataloader, device, modal_embedding: bool = False):
    """
    Aligned with RadIR.forward(is_condition=False) global path:
      visual_transformer(..., is_condition=False) -> l2norm -> project_img
      enc_text[:, 0, :] -> project_txt
    """
    all_eu_img, all_eu_txt = [], []
    all_hyp_img, all_hyp_txt = [], []

    model.eval()
    for batch in dataloader:
        imgs = batch["imgs"].to(device, non_blocking=True)
        imgs = preprocess_image_like_forward(imgs)

        enc_image_global = model.visual_transformer(
            imgs,
            return_encoded_tokens=True,
            modal_embedding=modal_embedding,
            modal_indexs=None,
            is_condition=False,
        )
        if enc_image_global.dim() == 2:
            image_global = l2norm(enc_image_global)
        elif enc_image_global.dim() == 3:
            image_global = l2norm(enc_image_global[:, 0, :])
        else:
            raise ValueError(
                f"Unexpected global visual feature shape: {enc_image_global.shape}"
            )

        input_ids = batch["caption_ids"].to(device, non_blocking=True)
        token_type_ids = batch["token_type_ids"].to(device, non_blocking=True)
        attention_mask = batch["attention_mask"].to(device, non_blocking=True)

        text_embeddings = model.text_transformer(
            input_ids=input_ids,
            token_type_ids=token_type_ids,
            attention_mask=attention_mask,
        )
        enc_text = text_embeddings[0]
        text_latents_global = enc_text[:, 0, :]
        text_latents = l2norm(text_latents_global)

        all_eu_img.append(image_global.cpu())
        all_eu_txt.append(text_latents.cpu())
        all_hyp_img.append(model.project_img(image_global).cpu())
        all_hyp_txt.append(model.project_txt(text_latents_global).cpu())

    return (
        torch.cat(all_eu_img, dim=0),
        torch.cat(all_eu_txt, dim=0),
        torch.cat(all_hyp_img, dim=0),
        torch.cat(all_hyp_txt, dim=0),
    )


def prepare_tsne_features(
    eu_img: torch.Tensor,
    eu_txt: torch.Tensor,
    hyp_img: torch.Tensor,
    hyp_txt: torch.Tensor,
    feat_type: str,
    curv: torch.Tensor,
) -> tuple[np.ndarray, np.ndarray]:
    if feat_type == "hyperbolic":
        with torch.no_grad():
            curv_val = curv.detach().float()
            img = L.log_map0(hyp_img.float(), curv_val).cpu().numpy()
            txt = L.log_map0(hyp_txt.float(), curv_val).cpu().numpy()
    elif feat_type == "euclidean":
        img = eu_img.detach().cpu().numpy()
        txt = eu_txt.detach().cpu().numpy()
    else:
        raise ValueError(f"Unknown feat_type={feat_type!r}, use 'hyperbolic' or 'euclidean'")
    return img, txt


def subsample_features(
    img_feats: np.ndarray,
    txt_feats: np.ndarray,
    max_samples: int | None,
    seed: int,
) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
    n = img_feats.shape[0]
    if max_samples is None or max_samples >= n:
        return img_feats, txt_feats, np.arange(n)
    rng = np.random.default_rng(seed)
    idx = np.sort(rng.choice(n, size=max_samples, replace=False))
    return img_feats[idx], txt_feats[idx], idx


def run_tsne(
    img_feats: np.ndarray,
    txt_feats: np.ndarray,
    perplexity: float,
    seed: int,
) -> tuple[np.ndarray, np.ndarray]:
    n = img_feats.shape[0]
    all_feats = np.concatenate([img_feats, txt_feats], axis=0)

    # sklearn requires perplexity < n_samples (here n_samples = 2 * n)
    max_perp = max(2.0, (n * 2 - 1) / 3.0)
    eff_perp = min(perplexity, max_perp)

    tsne = TSNE(
        n_components=2,
        perplexity=eff_perp,
        random_state=seed,
        init="pca",
        learning_rate="auto",
    )
    emb = tsne.fit_transform(all_feats)
    emb_img = emb[:n]
    emb_txt = emb[n:]
    return emb_img, emb_txt


def plot_tsne(
    emb_img: np.ndarray,
    emb_txt: np.ndarray,
    out_path: str,
    title: str,
    max_pair_lines: int,
    dpi: int,
) -> None:
    n = emb_img.shape[0]
    fig, ax = plt.subplots(figsize=(10, 8))

    ax.scatter(
        emb_img[:, 0],
        emb_img[:, 1],
        c="#1f77b4",
        alpha=0.55,
        s=18,
        label="Image",
        edgecolors="none",
    )
    ax.scatter(
        emb_txt[:, 0],
        emb_txt[:, 1],
        c="#d62728",
        alpha=0.55,
        s=18,
        label="Text",
        edgecolors="none",
    )

    n_lines = min(max_pair_lines, n)
    for i in range(n_lines):
        ax.plot(
            [emb_img[i, 0], emb_txt[i, 0]],
            [emb_img[i, 1], emb_txt[i, 1]],
            color="gray",
            alpha=0.15,
            linewidth=0.6,
            zorder=0,
        )

    ax.set_title(title)
    ax.set_xlabel("t-SNE 1")
    ax.set_ylabel("t-SNE 2")
    ax.legend(loc="best", 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 t-SNE plot to {out_path}")


def build_model(args, device: str) -> RADIR:
    tokenizer = AutoTokenizer.from_pretrained(
        args.model_path,
        trust_remote_code=True,
        local_files_only=True,
    )
    text_model = AutoModel.from_pretrained(
        args.model_path,
        trust_remote_code=True,
        local_files_only=True,
    )
    image_encoder = CTViT(
        dim=768,
        codebook_size=8192,
        image_size=args.imsize,
        patch_size=16,
        temporal_patch_size=10,
        spatial_depth=8,
        temporal_depth=6,
        cls_depth=4,
        dim_head=32,
        heads=8,
        channels=1,
    )
    model = RADIR(
        tokenizer=tokenizer,
        image_encoder=image_encoder,
        text_encoder=text_model,
        dim_text=768,
        dim_image=512,
        dim_latent=512,
        use_mlm=False,
        use_all_token_embeds=False,
    )
    print(f"Loading checkpoint: {args.radir_ckpt}")
    model.load(args.radir_ckpt)
    model.to(device)
    return model


def parse_args():
    parser = argparse.ArgumentParser(description="t-SNE plot for RadIR/RadIR")
    parser.add_argument("--device", default="cuda", type=str)
    parser.add_argument(
        "--model_path",
        default=os.path.join(
            ROOT, "hf_models/microsoft/BiomedVLP-CXR-BERT-specialized"
        ),
    )
    parser.add_argument(
        "--radir_ckpt",
        default=os.path.join(ROOT, "outputs/radir_final.pt"),
    )
    parser.add_argument(
        "--dataset_root",
        default=os.path.join(ROOT, "dataset"),
    )
    parser.add_argument("--caption_file", default="valid.jsonl")
    parser.add_argument(
        "--cig_matched_dir",
        default=os.path.join(ROOT, "dataset/processed/matched"),
    )
    parser.add_argument("--split", default="valid", type=str)
    parser.add_argument("--batch_size", default=16, type=int)
    parser.add_argument("--num_workers", default=8, type=int)
    parser.add_argument("--imsize", default=224, type=int)
    parser.add_argument("--max_words", default=128, type=int)
    parser.add_argument(
        "--modal_embedding",
        action="store_true",
        help="Set if training used --modal_embedding",
    )
    parser.add_argument(
        "--feat_type",
        choices=["hyperbolic", "euclidean"],
        default="hyperbolic",
        help="hyperbolic: log_map0 on project_* feats (default, matches retrieval)",
    )
    parser.add_argument(
        "--max_samples",
        type=int,
        default=0,
        help="Subsample for t-SNE speed; set 0 to use all samples",
    )
    parser.add_argument("--perplexity", type=float, default=30.0)
    parser.add_argument("--seed", type=int, default=42)
    parser.add_argument(
        "--max_pair_lines",
        type=int,
        default=50,
        help="Number of image-text pair connectors to draw",
    )
    parser.add_argument(
        "--out",
        default=os.path.join(ROOT, "outputs/tsne_radir.png"),
    )
    parser.add_argument(
        "--save_feats",
        default=None,
        help="Optional .npz path to cache extracted features",
    )
    parser.add_argument(
        "--load_feats",
        default=None,
        help="Skip extraction; load features from .npz",
    )
    parser.add_argument("--dpi", type=int, default=150)
    return parser.parse_args()


def main():
    args = parse_args()
    device = args.device
    max_samples = None if args.max_samples <= 0 else args.max_samples

    if args.load_feats:
        print(f"Loading features from {args.load_feats}")
        data = np.load(args.load_feats)
        eu_img = torch.from_numpy(data["eu_img"])
        eu_txt = torch.from_numpy(data["eu_txt"])
        hyp_img = torch.from_numpy(data["hyp_img"])
        hyp_txt = torch.from_numpy(data["hyp_txt"])
        curv = torch.tensor(float(data["curv"]))
        feat_type = str(data.get("feat_type", args.feat_type))
    else:
        model = build_model(args, device)
        curv = model.get_curv().cpu()
        print(f"Curvature = {curv.item():.6f}")

        tokenizer = model.tokenizer
        dataset = HyCoClipMimicBoxDataset_JPG(
            split=args.split,
            tokenizer=tokenizer,
            dataset_root=args.dataset_root,
            cig_matched_dir=args.cig_matched_dir,
            caption_file=args.caption_file,
            imsize=args.imsize,
            max_words=args.max_words,
            transform=None,
        )
        dataloader = DataLoader(
            dataset,
            batch_size=args.batch_size,
            shuffle=False,
            num_workers=args.num_workers,
            pin_memory=True,
            collate_fn=collate_hmclip_box,
        )
        print(f"Extracting features from {len(dataset)} samples ...")
        eu_img, eu_txt, hyp_img, hyp_txt = extract_features(
            model,
            dataloader,
            device,
            modal_embedding=args.modal_embedding,
        )
        feat_type = args.feat_type

        if args.save_feats:
            os.makedirs(os.path.dirname(os.path.abspath(args.save_feats)) or ".", exist_ok=True)
            np.savez(
                args.save_feats,
                eu_img=eu_img.numpy(),
                eu_txt=eu_txt.numpy(),
                hyp_img=hyp_img.numpy(),
                hyp_txt=hyp_txt.numpy(),
                curv=curv.item(),
                feat_type=feat_type,
            )
            print(f"Saved features to {args.save_feats}")

    img_feats, txt_feats = prepare_tsne_features(
        eu_img, eu_txt, hyp_img, hyp_txt, feat_type, curv
    )
    img_feats, txt_feats, _ = subsample_features(
        img_feats, txt_feats, max_samples, args.seed
    )
    print(f"t-SNE on {img_feats.shape[0]} image-text pairs (feat_type={feat_type})")

    emb_img, emb_txt = run_tsne(img_feats, txt_feats, args.perplexity, args.seed)

    pair_dist = np.linalg.norm(emb_img - emb_txt, axis=1).mean()
    print(f"Mean 2D pair distance (image-text): {pair_dist:.4f}")

    ckpt_name = os.path.basename(args.radir_ckpt) if not args.load_feats else "cached"
    title = f"RadIR t-SNE ({feat_type}) | n={img_feats.shape[0]} | {ckpt_name}"
    plot_tsne(
        emb_img,
        emb_txt,
        args.out,
        title=title,
        max_pair_lines=args.max_pair_lines,
        dpi=args.dpi,
    )


if __name__ == "__main__":
    main()