"""Export PET and SUVR embeddings from Stage 1 checkpoints for Figure 2."""
from __future__ import annotations

import argparse
from pathlib import Path

import numpy as np
import torch
from torch.utils.data import DataLoader

from pet_vlm_dataset import PETSUVRDataset, collate_pet_suvr
from train_pet_foundation import PETSUVRFoundationModel, build_encoder


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--checkpoint", type=Path, required=True)
    parser.add_argument("--split", type=Path, default=Path("data/metadata/splits/test.csv"))
    parser.add_argument("--backbone", default=None)
    parser.add_argument("--medicalnet-weights", type=Path, default=Path("pretrained/medicalnet/resnet_50_23dataset.pth"))
    parser.add_argument("--brainfm-weights", type=Path, default=Path("pretrained/brainfm/assets/brainfm_pretrained.pth"))
    parser.add_argument("--brainfm-code-root", type=Path, default=Path("pretrained/brainfm"))
    parser.add_argument("--batch-size", type=int, default=4)
    parser.add_argument("--num-workers", type=int, default=2)
    parser.add_argument("--output-size", type=int, nargs=3, default=None)
    parser.add_argument("--embed-dim", type=int, default=None)
    parser.add_argument("--freeze-encoder", type=bool, default=None)
    parser.add_argument("--out", type=Path, required=True)
    args = parser.parse_args()

    ckpt = torch.load(args.checkpoint, map_location="cpu", weights_only=False)
    saved_args = ckpt.get("args", {})
    for name in ("backbone", "embed_dim", "freeze_encoder"):
        if getattr(args, name, None) is None and name in saved_args:
            setattr(args, name, saved_args[name])
    if args.output_size is None:
        args.output_size = tuple(saved_args.get("output_size", (96, 96, 96)))

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    dataset = PETSUVRDataset(args.split, output_size=tuple(args.output_size))
    loader = DataLoader(dataset, batch_size=args.batch_size, shuffle=False,
                        num_workers=args.num_workers, collate_fn=collate_pet_suvr)
    n_regions = int(dataset[0]["suvr"].numel())

    encoder = build_encoder(args)
    model = PETSUVRFoundationModel(encoder, n_regions, args.embed_dim or 256, bool(args.freeze_encoder)).to(device)
    model.load_state_dict(ckpt["model"], strict=True)
    model.eval()

    pet_embeddings = []
    suvr_embeddings = []
    pred_suvr = []
    true_suvr = []

    with torch.no_grad():
        for batch in loader:
            image = batch["image"].to(device)
            suvr = batch["suvr"].to(device)
            outputs = model(image, suvr)

            pet_feat = model.pet_encoder(image)
            pet_z = torch.nn.functional.normalize(model.pet_projector(pet_feat), dim=-1)
            suvr_z = torch.nn.functional.normalize(model.suvr_encoder(suvr), dim=-1)

            pet_embeddings.append(pet_z.cpu().numpy())
            suvr_embeddings.append(suvr_z.cpu().numpy())
            pred_suvr.append(outputs["pred_suvr"].cpu().numpy())
            true_suvr.append(suvr.cpu().numpy())

    pet_z = np.concatenate(pet_embeddings, axis=0)
    suvr_z = np.concatenate(suvr_embeddings, axis=0)
    pred = np.concatenate(pred_suvr, axis=0)
    true = np.concatenate(true_suvr, axis=0)

    args.out.parent.mkdir(parents=True, exist_ok=True)
    np.savez(args.out, pet_z=pet_z, suvr_z=suvr_z, pred_suvr=pred, true_suvr=true)
    print(f"wrote {pet_z.shape[0]} samples, pet_z={pet_z.shape}, suvr_z={suvr_z.shape} to {args.out}")


if __name__ == "__main__":
    main()