from __future__ import annotations

import argparse
from pathlib import Path

import torch
from torch import nn
from torch.utils.data import DataLoader

from pet_vlm_dataset import PETSUVRDataset, collate_pet_suvr


class Small3DPETEncoder(nn.Module):
    def __init__(self, embed_dim: int = 256) -> None:
        super().__init__()
        self.net = nn.Sequential(
            nn.Conv3d(1, 16, 3, stride=2, padding=1),
            nn.BatchNorm3d(16),
            nn.GELU(),
            nn.Conv3d(16, 32, 3, stride=2, padding=1),
            nn.BatchNorm3d(32),
            nn.GELU(),
            nn.Conv3d(32, 64, 3, stride=2, padding=1),
            nn.BatchNorm3d(64),
            nn.GELU(),
            nn.Conv3d(64, 128, 3, stride=2, padding=1),
            nn.BatchNorm3d(128),
            nn.GELU(),
            nn.AdaptiveAvgPool3d(1),
        )
        self.proj = nn.Linear(128, embed_dim)

    def forward(self, image: torch.Tensor) -> torch.Tensor:
        x = self.net(image).flatten(1)
        return self.proj(x)


class RegionSUVREncoder(nn.Module):
    def __init__(self, n_regions: int = 120, embed_dim: int = 256) -> None:
        super().__init__()
        self.net = nn.Sequential(
            nn.LayerNorm(n_regions),
            nn.Linear(n_regions, 256),
            nn.GELU(),
            nn.Linear(256, embed_dim),
        )

    def forward(self, suvr: torch.Tensor) -> torch.Tensor:
        return self.net(suvr)


class PETSUVRAlignmentModel(nn.Module):
    def __init__(self, n_regions: int = 120, embed_dim: int = 256) -> None:
        super().__init__()
        self.pet_encoder = Small3DPETEncoder(embed_dim)
        self.suvr_encoder = RegionSUVREncoder(n_regions, embed_dim)
        self.suvr_head = nn.Linear(embed_dim, n_regions)
        self.temperature = nn.Parameter(torch.tensor(0.07))

    def forward(self, image: torch.Tensor, suvr: torch.Tensor) -> dict[str, torch.Tensor]:
        pet_z = nn.functional.normalize(self.pet_encoder(image), dim=-1)
        suvr_z = nn.functional.normalize(self.suvr_encoder(suvr), dim=-1)
        pred_suvr = self.suvr_head(pet_z)
        logits = pet_z @ suvr_z.T / self.temperature.clamp_min(0.01)
        return {"pet_z": pet_z, "suvr_z": suvr_z, "pred_suvr": pred_suvr, "logits": logits}


def alignment_loss(outputs: dict[str, torch.Tensor], suvr: torch.Tensor) -> torch.Tensor:
    labels = torch.arange(suvr.shape[0], device=suvr.device)
    loss_i = nn.functional.cross_entropy(outputs["logits"], labels)
    loss_t = nn.functional.cross_entropy(outputs["logits"].T, labels)
    loss_reg = nn.functional.mse_loss(outputs["pred_suvr"], suvr)
    return 0.5 * (loss_i + loss_t) + loss_reg


def main() -> None:
    parser = argparse.ArgumentParser(description="Train a minimal FDG-PET + SUVR alignment baseline.")
    parser.add_argument("--manifest", type=Path, default=Path("metadata/splits/train.csv"))
    parser.add_argument("--val-manifest", type=Path, default=Path("metadata/splits/val.csv"))
    parser.add_argument("--epochs", type=int, default=2)
    parser.add_argument("--batch-size", type=int, default=2)
    parser.add_argument("--lr", type=float, default=1e-4)
    parser.add_argument("--num-workers", type=int, default=0)
    parser.add_argument("--output-size", type=int, nargs=3, default=(96, 96, 96))
    parser.add_argument("--max-samples", type=int, default=0, help="Use the first N samples for a quick smoke test.")
    parser.add_argument("--log-every", type=int, default=10)
    parser.add_argument("--out", type=Path, default=Path("runs/pet_suvr_baseline.pt"))
    args = parser.parse_args()

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    dataset = PETSUVRDataset(args.manifest, output_size=tuple(args.output_size))
    val_dataset = PETSUVRDataset(args.val_manifest, output_size=tuple(args.output_size))
    if args.max_samples > 0:
        max_samples = min(args.max_samples, len(dataset))
        dataset = torch.utils.data.Subset(dataset, range(max_samples))
        val_max_samples = min(max(1, args.max_samples // 4), len(val_dataset))
        val_dataset = torch.utils.data.Subset(val_dataset, range(val_max_samples))
    train_loader = DataLoader(dataset, batch_size=args.batch_size, shuffle=True, num_workers=args.num_workers, collate_fn=collate_pet_suvr)
    val_loader = DataLoader(val_dataset, batch_size=args.batch_size, shuffle=False, num_workers=args.num_workers, collate_fn=collate_pet_suvr)

    sample = dataset[0]
    model = PETSUVRAlignmentModel(n_regions=int(sample["suvr"].numel())).to(device)
    optimizer = torch.optim.AdamW(model.parameters(), lr=args.lr, weight_decay=1e-4)

    for epoch in range(1, args.epochs + 1):
        model.train()
        train_loss = 0.0
        for step, batch in enumerate(train_loader, start=1):
            image = batch["image"].to(device)
            suvr = batch["suvr"].to(device)
            outputs = model(image, suvr)
            loss = alignment_loss(outputs, suvr)
            optimizer.zero_grad(set_to_none=True)
            loss.backward()
            optimizer.step()
            train_loss += float(loss.detach()) * image.shape[0]
            if args.log_every > 0 and step % args.log_every == 0:
                print(
                    f"epoch={epoch} step={step}/{len(train_loader)} loss={float(loss.detach()):.4f}",
                    flush=True,
                )
        train_loss /= len(dataset)

        model.eval()
        val_loss = 0.0
        with torch.no_grad():
            for batch in val_loader:
                image = batch["image"].to(device)
                suvr = batch["suvr"].to(device)
                outputs = model(image, suvr)
                val_loss += float(alignment_loss(outputs, suvr)) * image.shape[0]
        val_loss = val_loss / max(1, len(val_dataset))
        print(f"epoch={epoch} train_loss={train_loss:.4f} val_loss={val_loss:.4f}")

    args.out.parent.mkdir(parents=True, exist_ok=True)
    torch.save({"model": model.state_dict(), "args": vars(args)}, args.out)
    print(f"saved {args.out}", flush=True)


if __name__ == "__main__":
    main()