"""Two-phase training for the EfficientNet-B3 brain tumor classifier.

v2 changes (over v1 which hit 84% test acc):
- Input resolution 300×300 (matches pretrain; v1 used 224).
- Simpler head (Dropout → Linear) — v1's wide+BN head was over-regularizing.
- Higher fine-tune LR (5e-5) so the backbone actually adapts to MRI features.
- Unfreeze the last 3 feature blocks (v1: 2).
- Lighter augmentation (no ColorJitter — MRI intensities are calibrated).
- Mixed-precision (autocast + GradScaler) so 300×300 fits on a 4 GB GPU.

Run:
    python model/train.py
"""

from __future__ import annotations

import json
import sys
import time
from pathlib import Path

import torch

# Force line-buffered stdout so progress shows up immediately even when piped
# (e.g. `python train.py | tee log`). Equivalent to running with `python -u`.
sys.stdout.reconfigure(line_buffering=True)
import torch.nn as nn
import torch.optim as optim
from torch.amp import GradScaler, autocast
from torch.utils.data import DataLoader, Subset
from torchvision import transforms
from torchvision.datasets import ImageFolder

from architecture import (
    CLASS_NAMES,
    IMG_SIZE,
    build_model,
    count_trainable,
    freeze_backbone,
    unfreeze_last_blocks,
)

# ─── Config ───────────────────────────────────────────────────────────────────
PROJECT_ROOT = Path(__file__).resolve().parent.parent
TRAIN_DIR = PROJECT_ROOT / "data" / "raw" / "Training"
TEST_DIR = PROJECT_ROOT / "data" / "raw" / "Testing"
SAVE_DIR = PROJECT_ROOT / "model" / "saved"
SAVE_DIR.mkdir(parents=True, exist_ok=True)
CHECKPOINT = SAVE_DIR / "brain_tumor_model.pth"
HISTORY = SAVE_DIR / "history.json"

BATCH_SIZE = 24            # 300×300 + last-3 unfrozen → tight on 4 GB
NUM_WORKERS = 2            # 8 workers (4 train + 4 val × persistent) was thrashing 16 GB RAM
EPOCHS_HEAD = 6            # head is now small, learns fast
EPOCHS_FINE = 25
LR_HEAD = 1e-3
LR_FINE = 5e-5             # was 1e-5 — too conservative, backbone barely moved
WEIGHT_DECAY = 1e-4
VAL_SPLIT = 0.15
EARLY_STOP_PATIENCE = 12
UNFREEZE_LAST_N_BLOCKS = 3
SEED = 42

DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
USE_AMP = DEVICE.type == "cuda"

torch.manual_seed(SEED)
if DEVICE.type == "cuda":
    torch.cuda.manual_seed_all(SEED)

IMAGENET_MEAN = [0.485, 0.456, 0.406]
IMAGENET_STD = [0.229, 0.224, 0.225]

# ─── Transforms ───────────────────────────────────────────────────────────────
# A bit larger resize so RandomCrop has room. No ColorJitter — MRI intensities
# are diagnostic, jittering them simulates a different scanner not a different patient.
train_transform = transforms.Compose(
    [
        transforms.Resize((IMG_SIZE + 24, IMG_SIZE + 24)),
        transforms.RandomCrop(IMG_SIZE),
        transforms.RandomHorizontalFlip(p=0.5),
        transforms.RandomRotation(degrees=12),
        transforms.RandomAffine(degrees=0, translate=(0.06, 0.06), scale=(0.95, 1.05)),
        transforms.ToTensor(),
        transforms.Normalize(IMAGENET_MEAN, IMAGENET_STD),
    ]
)

eval_transform = transforms.Compose(
    [
        transforms.Resize((IMG_SIZE + 24, IMG_SIZE + 24)),
        transforms.CenterCrop(IMG_SIZE),
        transforms.ToTensor(),
        transforms.Normalize(IMAGENET_MEAN, IMAGENET_STD),
    ]
)


def make_loaders() -> tuple[DataLoader, DataLoader, DataLoader, list[str]]:
    # Two ImageFolders pointing at the same dir but with different transforms.
    # Split deterministically by index so the val subset gets eval_transform.
    train_full = ImageFolder(TRAIN_DIR, transform=train_transform)
    val_full = ImageFolder(TRAIN_DIR, transform=eval_transform)

    classes = train_full.classes
    if classes != CLASS_NAMES:
        raise RuntimeError(
            f"Class folder order {classes} != expected {CLASS_NAMES}. "
            "Rename folders so the alphabetical order matches."
        )

    n_total = len(train_full)
    n_val = int(n_total * VAL_SPLIT)
    g = torch.Generator().manual_seed(SEED)
    perm = torch.randperm(n_total, generator=g).tolist()
    val_idx = perm[:n_val]
    train_idx = perm[n_val:]

    train_subset = Subset(train_full, train_idx)
    val_subset = Subset(val_full, val_idx)

    test_dataset = ImageFolder(TEST_DIR, transform=eval_transform)

    common = dict(
        batch_size=BATCH_SIZE,
        num_workers=NUM_WORKERS,
        pin_memory=(DEVICE.type == "cuda"),
        persistent_workers=False,
    )
    train_loader = DataLoader(train_subset, shuffle=True, **common)
    val_loader = DataLoader(val_subset, shuffle=False, **common)
    test_loader = DataLoader(test_dataset, shuffle=False, **common)

    print(
        f"Train batches: {len(train_loader)} | "
        f"Val batches: {len(val_loader)} | "
        f"Test batches: {len(test_loader)} "
        f"| img {IMG_SIZE}×{IMG_SIZE}, bs {BATCH_SIZE}"
    )
    return train_loader, val_loader, test_loader, classes


def run_epoch(
    model: nn.Module,
    loader: DataLoader,
    criterion: nn.Module,
    optimizer: optim.Optimizer | None,
    scaler: GradScaler | None,
    train: bool,
) -> tuple[float, float]:
    model.train(train)
    total, correct, loss_sum = 0, 0, 0.0
    ctx = torch.enable_grad() if train else torch.no_grad()
    with ctx:
        for inputs, targets in loader:
            inputs = inputs.to(DEVICE, non_blocking=True)
            targets = targets.to(DEVICE, non_blocking=True)

            # AMP only during training. In eval, fp16 + label_smoothing can underflow
            # to NaN even when argmax accuracy is fine — keep eval in fp32.
            if USE_AMP and train:
                with autocast(device_type="cuda", dtype=torch.float16):
                    outputs = model(inputs)
                    loss = criterion(outputs, targets)
            else:
                outputs = model(inputs)
                loss = criterion(outputs, targets)

            if train:
                assert optimizer is not None
                optimizer.zero_grad(set_to_none=True)
                if USE_AMP and scaler is not None:
                    scaler.scale(loss).backward()
                    scaler.step(optimizer)
                    scaler.update()
                else:
                    loss.backward()
                    optimizer.step()

            loss_sum += loss.item() * inputs.size(0)
            correct += (outputs.argmax(1) == targets).sum().item()
            total += inputs.size(0)
    return loss_sum / total, correct / total


def train_phase(
    model: nn.Module,
    train_loader: DataLoader,
    val_loader: DataLoader,
    optimizer: optim.Optimizer,
    scheduler: optim.lr_scheduler._LRScheduler | None,
    criterion: nn.Module,
    scaler: GradScaler | None,
    epochs: int,
    phase_name: str,
    best_val_acc: float,
    history: list[dict],
) -> float:
    no_improve = 0
    for epoch in range(1, epochs + 1):
        t0 = time.time()
        train_loss, train_acc = run_epoch(
            model, train_loader, criterion, optimizer, scaler, train=True
        )
        val_loss, val_acc = run_epoch(
            model, val_loader, criterion, None, None, train=False
        )
        if scheduler is not None:
            scheduler.step()

        elapsed = time.time() - t0
        lr = optimizer.param_groups[0]["lr"]
        print(
            f"[{phase_name}] epoch {epoch:02d}/{epochs} "
            f"lr={lr:.2e} "
            f"train_loss={train_loss:.4f} train_acc={train_acc:.4f} "
            f"val_loss={val_loss:.4f} val_acc={val_acc:.4f} "
            f"({elapsed:.1f}s)",
            flush=True,
        )
        history.append(
            dict(
                phase=phase_name,
                epoch=epoch,
                lr=lr,
                train_loss=train_loss,
                train_acc=train_acc,
                val_loss=val_loss,
                val_acc=val_acc,
            )
        )

        if val_acc > best_val_acc:
            best_val_acc = val_acc
            torch.save(
                {
                    "model_state_dict": model.state_dict(),
                    "class_names": CLASS_NAMES,
                    "val_acc": val_acc,
                    "phase": phase_name,
                    "epoch": epoch,
                    "img_size": IMG_SIZE,
                },
                CHECKPOINT,
            )
            print(f"  ↳ new best val_acc={val_acc:.4f}, saved to {CHECKPOINT.name}", flush=True)
            no_improve = 0
        else:
            no_improve += 1
            if no_improve >= EARLY_STOP_PATIENCE:
                print(f"  ↳ early stop after {no_improve} epochs without improvement", flush=True)
                break

    return best_val_acc


def main() -> None:
    print(f"Device: {DEVICE} | AMP: {USE_AMP}")
    if DEVICE.type == "cuda":
        print(f"GPU: {torch.cuda.get_device_name(0)}")

    train_loader, val_loader, test_loader, classes = make_loaders()
    print(f"Classes: {classes}")

    model = build_model(pretrained=True).to(DEVICE)
    criterion = nn.CrossEntropyLoss(label_smoothing=0.05)
    scaler = GradScaler() if USE_AMP else None
    history: list[dict] = []
    best_val_acc = 0.0

    # ── Phase 1: head only ────────────────────────────────────────────────
    freeze_backbone(model)
    print(f"\nPhase 1 (head-only): trainable params = {count_trainable(model):,}")
    optimizer = optim.AdamW(
        filter(lambda p: p.requires_grad, model.parameters()),
        lr=LR_HEAD,
        weight_decay=WEIGHT_DECAY,
    )
    best_val_acc = train_phase(
        model,
        train_loader,
        val_loader,
        optimizer,
        scheduler=None,
        criterion=criterion,
        scaler=scaler,
        epochs=EPOCHS_HEAD,
        phase_name="head",
        best_val_acc=best_val_acc,
        history=history,
    )

    # ── Phase 2: fine-tune last blocks ────────────────────────────────────
    unfreeze_last_blocks(model, num_blocks=UNFREEZE_LAST_N_BLOCKS)
    print(f"\nPhase 2 (fine-tune): trainable params = {count_trainable(model):,}")
    optimizer = optim.AdamW(
        filter(lambda p: p.requires_grad, model.parameters()),
        lr=LR_FINE,
        weight_decay=WEIGHT_DECAY,
    )
    scheduler = optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=EPOCHS_FINE)
    best_val_acc = train_phase(
        model,
        train_loader,
        val_loader,
        optimizer,
        scheduler=scheduler,
        criterion=criterion,
        scaler=scaler,
        epochs=EPOCHS_FINE,
        phase_name="finetune",
        best_val_acc=best_val_acc,
        history=history,
    )

    # ── Final test evaluation with best checkpoint ────────────────────────
    print("\nLoading best checkpoint and evaluating on test set...")
    ckpt = torch.load(CHECKPOINT, map_location=DEVICE)
    model.load_state_dict(ckpt["model_state_dict"])
    test_loss, test_acc = run_epoch(model, test_loader, criterion, None, None, train=False)
    print(f"Test loss: {test_loss:.4f} | Test accuracy: {test_acc*100:.2f}%")

    HISTORY.write_text(
        json.dumps(
            dict(
                history=history,
                best_val_acc=best_val_acc,
                test_acc=test_acc,
                test_loss=test_loss,
                class_names=CLASS_NAMES,
                img_size=IMG_SIZE,
                batch_size=BATCH_SIZE,
            ),
            indent=2,
        )
    )
    print(f"History written to {HISTORY}")
    print("\nDone. Next: run `python model/evaluate.py` for detailed metrics.")


if __name__ == "__main__":
    main()