"""Extract PET/SUVR embeddings and per-subject predictions for paper figures."""
import sys, os, argparse
sys.path.insert(0, os.path.dirname(__file__))
import numpy as np
import torch
from pathlib import Path
from torch.utils.data import DataLoader
from pet_vlm_dataset import PETSUVRDataset, collate_pet_suvr
from train_pet_foundation import PETSUVRFoundationModel, build_encoder

def make_args(backbone):
    args = argparse.Namespace()
    args.backbone = backbone
    args.medicalnet_weights = Path("pretrained/medicalnet/resnet_50_23dataset.pth")
    args.brainiac_weights = Path("pretrained/brainiac/backbone.safetensors")
    args.brainfm_weights = Path("pretrained/brainfm/assets/brainfm_pretrained.pth")
    args.brainfm_code_root = Path("pretrained/brainfm")
    args.swinunetr_weights = Path("pretrained/swinunetr/model_swinvit.pt")
    args.sam_med3d_weights = Path("pretrained/sam-med3d/sam_med3d_turbo.pth")
    args.output_size = [96, 96, 96]
    return args

def extract(checkpoint_path, backbone, output_prefix):
    device = torch.device("cuda:0")
    args = make_args(backbone)
    output_size = tuple(args.output_size)

    encoder = build_encoder(args)
    n_regions = 120
    model = PETSUVRFoundationModel(encoder, n_regions=n_regions).to(device)

    ckpt = torch.load(checkpoint_path, map_location=device, weights_only=False)
    model.load_state_dict(ckpt["model"], strict=False)
    model.eval()

    ds = PETSUVRDataset(Path("metadata/splits/test.csv"), output_size=output_size)
    loader = DataLoader(ds, batch_size=4, shuffle=False, num_workers=2, collate_fn=collate_pet_suvr)

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

            pet_feat = model.pet_encoder(img)
            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_zs.append(pet_z.cpu().numpy())
            suvr_zs.append(suvr_z.cpu().numpy())
            preds.append(outputs["pred_suvr"].cpu().numpy())
            targets.append(suvr.cpu().numpy())

    pet_zs = np.concatenate(pet_zs)
    suvr_zs = np.concatenate(suvr_zs)
    preds = np.concatenate(preds)
    targets = np.concatenate(targets)

    out_dir = Path("runs/paper_figures")
    out_dir.mkdir(exist_ok=True)
    np.savez(out_dir / f"{output_prefix}_embeddings.npz",
             pet_z=pet_zs, suvr_z=suvr_zs, pred_suvr=preds, true_suvr=targets)
    print(f"Saved {output_prefix}: pet_z={pet_zs.shape}, suvr_z={suvr_zs.shape}, pred={preds.shape}")

if __name__ == "__main__":
    os.chdir("/data/Albus/Brain")
    extract("runs/foundation/medicalnet_layer4_regalign_best.pt", "medicalnet", "remap_pet")
    extract("runs/foundation/medicalnet_frozen_mlp.pt", "medicalnet", "medicalnet_frozen")
    print("Done!")