#!/usr/bin/env python
"""
Trainer script for kavin-aravindhan/vit-oct-wamd, reproducing the training
procedure used in the paper.

This is a cleaned-up, path-parameterized version of the internal training
script. It reproduces the exact architecture, loss, hyperparameters, and
augmentation recipe used for the released checkpoint.

NOTE ON DATA: the training set (112 OCT images with clinical-finding
captions, TFRecord format) is clinical research data and is not bundled with
this repo. Contact the authors (see the model card) for access. Point
--tfrecord-path at your local copy to run this script.

Usage:
    pip install -r requirements-train.txt
    python train.py --tfrecord-path /path/to/VQA_v4.tfrecord --output-dir ./runs/my_run
"""
import argparse
import json
import os
import random
from datetime import datetime

import albumentations as A
import cv2
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from albumentations.pytorch import ToTensorV2
from tfrecord.torch.dataset import TFRecordDataset
from torch.utils.data import DataLoader, Dataset
from torch.utils.tensorboard import SummaryWriter
from transformers import SiglipVisionModel

from alignment import SigLIPLoss
from embedder import TextEmbedder  # noqa: F401  (used indirectly via SigLIPLoss)

# ---------------------------------------------------------------------------
# Hyperparameters -- selected via a 500-trial Optuna search (see
# training_config.json in this repo) and used for the released checkpoint.
# ---------------------------------------------------------------------------
BATCH_SIZE = 8
EPOCHS = 50
LEARNING_RATE = 1e-4
WEIGHT_DECAY = 1.6079555710533247e-06
ALPHA = 0.8840963962895334  # weight on classification loss vs. alignment loss
SCHEDULER_PATIENCE = 2
SCHEDULER_FACTOR = 0.7851246675328261
EARLY_STOPPING_PATIENCE = 20
DROPOUT_RATE = 0.057129660535791494
MAX_TEXT_LEN = 128
RANDOM_SAMPLES_PER_EPOCH = 1000

IMAGE_HEIGHT, IMAGE_WIDTH, IMAGE_CHANNELS = 703, 1055, 3
IMAGE_ENCODER = "google/siglip-so400m-patch14-384"
TEXT_MODEL = "google-t5/t5-base"

DESCRIPTION = {
    "input_ids": "int",
    "input_ids_shape": "int",
    "attn_mask": "int",
    "attn_mask_shape": "int",
    "class": "byte",
    "normalized_image": "byte",
}

AUG = A.Compose(
    [
        A.HorizontalFlip(p=0.9674733435973407),
        A.ShiftScaleRotate(
            shift_limit=0.028795483291628815,
            scale_limit=0.037814873748683406,
            rotate_limit=3,
            border_mode=cv2.BORDER_REPLICATE,
            p=0.028486520024779284,
        ),
        A.RandomBrightnessContrast(
            brightness_limit=0.09474000592838613,
            contrast_limit=0.03187889420730783,
            p=0.060697129192364106,
        ),
        A.GaussNoise(noise_limit=(0, 1e-4), p=0.342779329951288),
        A.MotionBlur(blur_limit=3, p=0.194405815151957),
        ToTensorV2(),
    ]
)


def parse_and_augment_image(img_bytes):
    img_array = np.frombuffer(img_bytes, dtype=np.float32)
    img = img_array.reshape(IMAGE_HEIGHT, IMAGE_WIDTH, IMAGE_CHANNELS).copy()
    img_uint8 = (img * 255).astype(np.uint8)
    augmented = AUG(image=img_uint8)
    return augmented["image"].float()


class RandomSampleDataset(Dataset):
    """Samples `samples_per_epoch` items per epoch, with replacement, from the
    (small) TFRecord. This matches the original training procedure -- the
    dataset has 112 unique images and is heavily oversampled with
    augmentation rather than trained on unique examples."""

    def __init__(self, tfrecord_path, description, samples_per_epoch):
        self.dataset = TFRecordDataset(tfrecord_path, None, description)
        self.items = list(self.dataset)
        self.samples_per_epoch = samples_per_epoch
        print(f"Loaded {len(self.items)} unique items from TFRecord")
        print(f"Will draw {samples_per_epoch} random (with-replacement) samples per epoch")

    def __len__(self):
        return self.samples_per_epoch

    def __getitem__(self, idx):
        random_idx = random.randint(0, len(self.items) - 1)
        return self.items[random_idx]


def collate_fn(batch):
    images, labels, input_ids_list, attention_masks = [], [], [], []
    for item in batch:
        img_tensor = parse_and_augment_image(item["normalized_image"])
        if img_tensor.shape[1] != 384 or img_tensor.shape[2] != 384:
            img_tensor = F.interpolate(
                img_tensor.unsqueeze(0), size=(384, 384), mode="bilinear", align_corners=False
            ).squeeze(0)
        img_tensor = (img_tensor - 0.5) / 0.5
        images.append(img_tensor)

        labels.append(0 if item["class"].decode("utf-8") == "n" else 1)

        input_ids_array = np.array(item["input_ids"]).reshape(tuple(item["input_ids_shape"]))
        attn_mask_array = np.array(item["attn_mask"]).reshape(tuple(item["attn_mask_shape"]))
        selected_input_ids = input_ids_array[0]
        selected_attn_mask = attn_mask_array[0]

        input_ids_list.append(torch.tensor(selected_input_ids, dtype=torch.long))
        attention_masks.append(torch.tensor(selected_attn_mask, dtype=torch.bool))

    return (
        torch.stack(images),
        torch.tensor(labels, dtype=torch.long),
        torch.stack(input_ids_list),
        torch.stack(attention_masks),
    )


class SigLIPModel(nn.Module):
    """Same architecture as modeling.py, plus the auxiliary alignment loss
    used only during training."""

    def __init__(self, dropout_rate=DROPOUT_RATE):
        super().__init__()
        self.image_encoder = SiglipVisionModel.from_pretrained(IMAGE_ENCODER)
        self.dropout = nn.Dropout(dropout_rate)
        self.cls_head = nn.Linear(1152, 2)
        self.siglip_loss = SigLIPLoss(
            latent_dim=1152, text_model=TEXT_MODEL, max_txt_len=MAX_TEXT_LEN, pool="mean", dtype=torch.float32
        )

    def forward(self, images, input_ids, attention_mask):
        img_features = self.image_encoder(pixel_values=images).last_hidden_state
        cls_features = self.dropout(img_features[:, 0])
        cls_logits = self.cls_head(cls_features)
        align_loss, _, _ = self.siglip_loss(img_features, input_ids, attention_mask)
        return cls_logits, align_loss


def train_epoch(model, dataloader, optimizer, device, epoch):
    model.train()
    total_loss = total_cls = total_align = 0.0
    num_batches = 0
    for batch_idx, (images, labels, input_ids, attention_mask) in enumerate(dataloader):
        images, labels = images.to(device), labels.to(device)
        input_ids, attention_mask = input_ids.to(device), attention_mask.to(device)

        cls_logits, align_loss = model(images, input_ids, attention_mask)
        cls_loss = F.cross_entropy(cls_logits, labels)
        loss = ALPHA * cls_loss + (1 - ALPHA) * align_loss

        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        total_loss += loss.item()
        total_cls += cls_loss.item()
        total_align += align_loss.item()
        num_batches += 1

        if batch_idx % 100 == 0:
            print(f"Epoch {epoch}, Batch {batch_idx}: Total={loss.item():.4f}, Cls={cls_loss.item():.4f}, Align={align_loss.item():.4f}")

    if num_batches == 0:
        return 0.0, 0.0, 0.0
    return total_loss / num_batches, total_cls / num_batches, total_align / num_batches


def save_checkpoint(model, optimizer, scheduler, epoch, loss, path):
    torch.save(
        {
            "epoch": epoch,
            "model_state_dict": model.state_dict(),
            "optimizer_state_dict": optimizer.state_dict(),
            "scheduler_state_dict": scheduler.state_dict(),
            "loss": loss,
            "alpha": ALPHA,
        },
        path,
    )
    print(f"Saved: {path}")


def main():
    parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    parser.add_argument("--tfrecord-path", required=True, help="Path to VQA_v4.tfrecord (see NOTE above).")
    parser.add_argument("--output-dir", default="./runs", help="Where to save checkpoints/config/history.")
    parser.add_argument("--epochs", type=int, default=EPOCHS)
    parser.add_argument("--seed", type=int, default=None, help="Unset by default, matching the original run.")
    args = parser.parse_args()

    if args.seed is not None:
        random.seed(args.seed)
        np.random.seed(args.seed)
        torch.manual_seed(args.seed)

    device = "cuda" if torch.cuda.is_available() else ("mps" if torch.backends.mps.is_available() else "cpu")
    print(f"Using device: {device}")

    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    run_dir = os.path.join(args.output_dir, f"run_{timestamp}")
    os.makedirs(run_dir, exist_ok=True)
    writer = SummaryWriter(os.path.join(run_dir, "tensorboard"))

    config = {
        "timestamp": timestamp,
        "tfrecord_path": args.tfrecord_path,
        "random_samples_per_epoch": RANDOM_SAMPLES_PER_EPOCH,
        "batch_size": BATCH_SIZE,
        "learning_rate": LEARNING_RATE,
        "weight_decay": WEIGHT_DECAY,
        "alpha": ALPHA,
        "dropout_rate": DROPOUT_RATE,
        "scheduler_patience": SCHEDULER_PATIENCE,
        "scheduler_factor": SCHEDULER_FACTOR,
        "early_stopping_patience": EARLY_STOPPING_PATIENCE,
        "image_encoder": IMAGE_ENCODER,
        "text_model": TEXT_MODEL,
    }
    with open(os.path.join(run_dir, "config.json"), "w") as f:
        json.dump(config, f, indent=2)

    dataset = RandomSampleDataset(args.tfrecord_path, DESCRIPTION, RANDOM_SAMPLES_PER_EPOCH)
    dataloader = DataLoader(dataset, batch_size=BATCH_SIZE, shuffle=True, collate_fn=collate_fn)

    model = SigLIPModel().to(device)
    optimizer = torch.optim.AdamW(model.parameters(), lr=LEARNING_RATE, weight_decay=WEIGHT_DECAY)
    scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(
        optimizer, mode="min", factor=SCHEDULER_FACTOR, patience=SCHEDULER_PATIENCE
    )

    print(f"Starting training for {args.epochs} epochs...")
    best_loss = float("inf")
    patience_counter = 0
    history = {"total_loss": [], "cls_loss": [], "align_loss": []}
    epoch = 0

    for epoch in range(1, args.epochs + 1):
        avg_total, avg_cls, avg_align = train_epoch(model, dataloader, optimizer, device, epoch)
        history["total_loss"].append(avg_total)
        history["cls_loss"].append(avg_cls)
        history["align_loss"].append(avg_align)

        print(f"Epoch {epoch}: Total={avg_total:.4f}, Cls={avg_cls:.4f}, Align={avg_align:.4f}")
        writer.add_scalar("Loss/Total", avg_total, epoch)
        writer.add_scalar("Loss/Classification", avg_cls, epoch)
        writer.add_scalar("Loss/Alignment", avg_align, epoch)
        writer.add_scalar("Learning_Rate", optimizer.param_groups[0]["lr"], epoch)
        scheduler.step(avg_total)

        if avg_total < best_loss and avg_total > 0:
            best_loss = avg_total
            save_checkpoint(model, optimizer, scheduler, epoch, avg_total, os.path.join(run_dir, "best_model.pt"))
            patience_counter = 0
        else:
            patience_counter += 1

        if epoch % 5 == 0:
            save_checkpoint(model, optimizer, scheduler, epoch, avg_total, os.path.join(run_dir, f"checkpoint_epoch_{epoch}.pt"))

        if patience_counter >= EARLY_STOPPING_PATIENCE:
            print(f"Early stopping triggered at epoch {epoch}")
            break

    save_checkpoint(model, optimizer, scheduler, epoch, avg_total, os.path.join(run_dir, "final_model.pt"))
    with open(os.path.join(run_dir, "training_history.json"), "w") as f:
        json.dump(history, f, indent=2)

    print(f"Training complete. Best loss: {best_loss:.4f}. Checkpoints in: {run_dir}")
    writer.close()


if __name__ == "__main__":
    main()