"""REFUGE training/evaluation helper.

Usage examples (after activating .venv_refuge):

  python refuge_build.py --train-seg
  python refuge_build.py --train-clf
  python refuge_build.py --eval --with-ttt

The script expects the REFUGE folder and writes checkpoints under
models/refuge/segmentation and models/refuge/classifier.
"""

from __future__ import annotations

import argparse
import csv
from pathlib import Path
from typing import Dict, List, Optional, Sequence, Set, Tuple
import shutil
import sys

import torch
import numpy as np
from PIL import Image, ImageDraw
from torch.utils.data import DataLoader
from sklearn.metrics import roc_auc_score
from tqdm import tqdm
from torch import nn
from torchvision import models

REPO_ROOT = Path(__file__).resolve().parents[3]
if str(REPO_ROOT) not in sys.path:
    sys.path.insert(0, str(REPO_ROOT))

from classes.refuge_preprocessing import RefugePreprocessing, RefugeSample
from classes.refuge_segmentation import RefugeSegmentation
from classes.refuge_classification import (
    RefugeClassification,
    RefugeClassificationRecord,
    RefugeClassificationDataset,
    _default_image_transform,
    _geometry_from_mask,
    UNetGeometryProvider,
)
from classes.unet_segmenter import UNetSegmenter
from classes.papila_builders import build_papila_clinical

REFUGE_ROOT = Path("REFUGE")
SEG_CKPT = Path("models/refuge/segmentation/refuge_segmentation_best.pt")
CLF_DIR = Path("models/refuge/classifier")
UNET_WEIGHT_CANDIDATES = (
    Path("models/v2/refuge/segmentation/per_image/best.pt"),
    Path("models/v2/refuge/segmentation/best.pt"),
    Path("models/unet_segmenter/best.pt"),
)

CLASSIFIER_BACKBONES = {
    "resnet50": models.ResNet50_Weights.DEFAULT,
    "densenet121": models.DenseNet121_Weights.DEFAULT,
    "efficientnet_b0": models.EfficientNet_B0_Weights.DEFAULT,
    "efficientnet_b7": models.EfficientNet_B7_Weights.DEFAULT,
}


def build_classifier_backbone(name: str) -> nn.Module:
    name = name.lower()
    if name not in CLASSIFIER_BACKBONES:
        raise ValueError(f"Unsupported classifier backbone '{name}'")

    weights = CLASSIFIER_BACKBONES[name]

    if name == "resnet50":
        model = models.resnet50(weights=weights)
        feat_dim = model.fc.in_features
        model.fc = nn.Identity()
    elif name == "densenet121":
        model = models.densenet121(weights=weights)
        feat_dim = model.classifier.in_features
        model.classifier = nn.Identity()
    elif name == "efficientnet_b0":
        model = models.efficientnet_b0(weights=weights)
        feat_dim = model.classifier[-1].in_features  # type: ignore[index]
        model.classifier = nn.Identity()
    elif name == "efficientnet_b7":
        model = models.efficientnet_b7(weights=weights)
        feat_dim = model.classifier[-1].in_features  # type: ignore[index]
        model.classifier = nn.Identity()
    else:  # pragma: no cover
        raise ValueError(f"Unsupported classifier backbone '{name}'")

    setattr(model, "_feature_dim", int(feat_dim))
    return model


def classifier_checkpoint_dir(backbone_name: str) -> Path:
    return CLF_DIR / backbone_name


def classifier_checkpoint_path(backbone_name: str) -> Path:
    return classifier_checkpoint_dir(backbone_name) / "refuge_classifier_best.pt"


def resolve_unet_weights(explicit: Optional[Path]) -> Path:
    if explicit is not None:
        return explicit
    for cand in UNET_WEIGHT_CANDIDATES:
        if cand.exists():
            return cand
    return UNET_WEIGHT_CANDIDATES[0]


def ensure_preprocessing() -> RefugePreprocessing:
    if not REFUGE_ROOT.exists():
        raise FileNotFoundError(f"REFUGE directory not found at {REFUGE_ROOT}")
    return RefugePreprocessing(REFUGE_ROOT)


def load_allowed_ids(
    csv_path: Optional[Path], dice_threshold: float
) -> Optional[Set[str]]:
    if csv_path is None or not csv_path.exists():
        return None
    allowed: Set[str] = set()
    with csv_path.open(newline="") as fh:
        reader = csv.DictReader(fh)
        for row in reader:
            sample_id = row.get("sample_id")
            if not sample_id or sample_id == "__mean__":
                continue
            try:
                disc = float(row.get("dice_disc", "nan"))
                cup = float(row.get("dice_cup", "nan"))
            except (TypeError, ValueError):
                continue
            if disc < dice_threshold and cup < dice_threshold:
                continue
            allowed.add(sample_id)
    return allowed


def build_papila_samples(
    image_dir: Path,
    clinical_dir: Path,
    label_col: str,
    positive_labels: Sequence[str],
    allowed_ids: Optional[Set[str]],
) -> List[RefugeSample]:
    clinical = build_papila_clinical(
        image_dir=str(image_dir),
        clinical_dir=str(clinical_dir),
        label_col=label_col,
        cat_cols=[],
    )
    positives = {lbl.lower() for lbl in positive_labels}
    samples: Dict[str, RefugeSample] = {}
    for _, row in clinical.df.iterrows():
        image_path = clinical.get_image_path(row)
        sample_id = f"papila_{Path(image_path).stem}"
        if allowed_ids is not None and sample_id not in allowed_ids:
            continue
        if sample_id in samples:
            continue
        value = row.get(label_col)
        if value is None or (isinstance(value, float) and np.isnan(value)):
            continue
        try:
            label_int = int(value)
            if label_int == 2:
                continue
            label = 1 if label_int > 0 else 0
        except (TypeError, ValueError):
            label = 1 if str(value).strip().lower() in positives else 0
        samples[sample_id] = RefugeSample(
            sample_id=sample_id,
            dataset="papila",
            split="holdout",
            image_path=Path(image_path),
            label=label,
            device=None,
            mask_path=None,
            fovea_coord=None,
        )
    return list(samples.values())


def load_contour(path: Path) -> np.ndarray:
    coords = np.loadtxt(path)
    if coords.ndim == 1:
        coords = coords.reshape(-1, 2)
    return coords


def contour_to_mask(coords: np.ndarray, size: Tuple[int, int]) -> np.ndarray:
    if coords is None or coords.size == 0:
        return np.zeros((size[1], size[0]), dtype=np.uint8)
    img = Image.new("L", size, 0)
    draw = ImageDraw.Draw(img)
    points = [tuple(map(float, pt)) for pt in coords]
    draw.polygon(points, outline=1, fill=1)
    return np.array(img, dtype=np.uint8)


class PapilaGTGeometryProvider:
    def __init__(self, contours_dir: Path) -> None:
        self.contours_dir = contours_dir

    def _pick(self, base: str, kind: str) -> Optional[Path]:
        for exp in ("exp2", "exp1"):
            cand = self.contours_dir / f"{base}_{kind}_{exp}.txt"
            if cand.exists():
                return cand
        return None

    def __call__(self, sample: RefugeSample, scale: float):
        base = Path(sample.image_path).stem
        disc_path = self._pick(base, "disc")
        cup_path = self._pick(base, "cup")
        if disc_path is None or cup_path is None:
            raise RuntimeError(f"Missing ground-truth contours for {sample.sample_id}")

        image = Image.open(sample.image_path).convert("RGB")
        disc_coords = load_contour(disc_path)
        cup_coords = load_contour(cup_path)
        disc_mask = contour_to_mask(disc_coords, image.size)
        cup_mask = contour_to_mask(cup_coords, image.size)
        cup_mask = ((cup_mask > 0) & (disc_mask > 0)).astype(np.uint8)
        geom = _geometry_from_mask(disc_mask, scale)
        return geom, disc_mask.astype(np.uint8), cup_mask.astype(np.uint8)


def build_papila_records(
    args: argparse.Namespace,
    pre: RefugePreprocessing,
    checkpoint_path: Path,
) -> Tuple[List[RefugeClassificationRecord], Optional[RefugeClassification]]:
    allowed = load_allowed_ids(
        getattr(args, "papila_metrics", None),
        getattr(args, "papila_dice_threshold", 0.01),
    )
    samples = build_papila_samples(
        args.papila_image_dir,
        args.papila_clinical_dir,
        args.papila_label_col,
        args.papila_positive_labels,
        allowed,
    )
    if not samples:
        return [], None

    cache_dir = args.clf_cache_dir
    if cache_dir is not None and getattr(args, "papila_use_gt", False):
        cache_dir = cache_dir / "gt"

    if getattr(args, "papila_use_gt", False):
        geometry_fn = PapilaGTGeometryProvider(args.papila_contours_dir)
        provider = geometry_fn
    else:
        seg_manifest = getattr(args, "seg_manifest", None)
        seg_weights = resolve_unet_weights(getattr(args, "seg_weights", None))
        if seg_manifest is None or seg_weights is None:
            raise SystemExit(
                "Papila evaluation without GT masks requires --seg-manifest and --seg-weights"
            )
        segmenter = UNetSegmenter(
            manifest_path=seg_manifest,
            device=args.device,
            normalize=args.seg_normalize,
        )
        seg_state = torch.load(seg_weights, map_location=args.device)
        seg_state_dict = seg_state.get("model", seg_state)
        segmenter.model.load_state_dict(seg_state_dict)
        segmenter.model.to(args.device)
        provider = UNetGeometryProvider(
            segmenter=segmenter,
            threshold=args.segmenter_threshold,
            tta=args.segmenter_tta,
        )
        geometry_fn = provider

    papila_seg = RefugeSegmentation(pre)
    backbone = build_classifier_backbone(args.clf_backbone)
    papila_clf = RefugeClassification(
        pre,
        papila_seg,
        backbone=backbone,
        geometry_fn=provider,
        cache_dir=cache_dir,
    )
    papila_clf.crop_scale = args.crop_scale
    papila_clf.crop_size = args.crop_size
    papila_clf.eval_transform = _default_image_transform(args.crop_size)
    papila_clf.ttt_transform = papila_clf.eval_transform
    papila_state = torch.load(checkpoint_path, map_location=args.device)
    papila_clf.backbone.load_state_dict(papila_state["backbone"])
    papila_clf.classifier_head.load_state_dict(papila_state["classifier"])
    papila_clf.rotation_head.load_state_dict(papila_state["rotation"])
    papila_clf.backbone.to(args.device)
    papila_clf.classifier_head.to(args.device)
    papila_clf.rotation_head.to(args.device)

    records = papila_clf.build_records_for_samples(
        samples, crop_scale=args.crop_scale, progress_prefix="papila"
    )
    print(f"[eval] Prepared {len(records)} PAPILA records")
    return records, papila_clf


def train_segmentation(args: argparse.Namespace) -> None:
    pre = ensure_preprocessing()
    seg = RefugeSegmentation(pre)
    seg.build_datasets(
        image_size=args.seg_image_size,
        batch_size=args.seg_batch_size,
        num_workers=args.num_workers,
    )
    history = seg.train(
        epochs=args.seg_epochs,
        lr=args.seg_lr,
        weight_decay=args.seg_weight_decay,
        checkpoint_dir=SEG_CKPT.parent,
        device=args.device,
    )
    print("Segmentation training complete. Best Dice:", history.get("best_dice"))


def train_unet_segmenter(args: argparse.Namespace) -> None:
    manifest_path = args.seg_manifest or Path("manifest.csv")
    mask_cache_dir = None if args.in_memory_cache else args.mask_cache_dir
    image_cache_dir = None if args.in_memory_cache else args.image_cache_dir
    if args.in_memory_cache and (args.mask_cache_dir or args.image_cache_dir):
        print("[unet-seg] in_memory_cache enabled: disk caches disabled for this run.")

    segmenter = UNetSegmenter(
        manifest_path=manifest_path,
        device=args.device,
        target_size=args.seg_image_size,
        normalize=args.seg_normalize,
        use_stronger_aug=args.seg_strong_aug,
        train_datasets=args.seg_train_datasets,
        val_datasets=args.seg_val_datasets,
        holdout_datasets=args.seg_holdout_datasets,
        mask_cache_dir=mask_cache_dir,
        image_cache_dir=image_cache_dir,
        in_memory_cache=args.in_memory_cache,
        loader_workers=args.loader_workers,
    )
    if mask_cache_dir:
        print(f"[unet-seg] mask_cache_dir={mask_cache_dir}")
    if image_cache_dir:
        print(f"[unet-seg] image_cache_dir={image_cache_dir}")
    if args.in_memory_cache:
        print("[unet-seg] prebuilding in-memory cache")
        segmenter.prebuild_in_memory_cache(
            cache_workers=max(0, int(args.cache_workers)),
            include_train=True,
            include_val=True,
            include_holdout=False,
        )

    segmenter.train(
        epochs=args.seg_epochs,
        batch_size=args.seg_batch_size,
        lr=args.seg_lr,
        weight_decay=args.seg_weight_decay,
        checkpoint_dir=args.seg_checkpoint_dir,
    )
    print(
        "[unet-seg] Training complete. Best checkpoint stored at",
        (args.seg_checkpoint_dir / "best.pt").resolve(),
    )


def _load_segmentation(
    pre: RefugePreprocessing, args: argparse.Namespace
) -> RefugeSegmentation:
    seg = RefugeSegmentation(pre)
    seg.build_datasets(
        image_size=args.seg_image_size,
        batch_size=args.seg_batch_size,
        num_workers=args.num_workers,
    )
    if not SEG_CKPT.exists():
        raise FileNotFoundError(f"Segmentation checkpoint missing: {SEG_CKPT}")
    state = torch.load(SEG_CKPT, map_location=args.device)
    seg.model.load_state_dict(state)
    seg.model.to(args.device)
    return seg


def train_classifier(args: argparse.Namespace) -> None:
    pre = ensure_preprocessing()
    seg = _load_segmentation(pre, args)
    backbone = build_classifier_backbone(args.clf_backbone)

    print(f"[classifier] Using backbone: {args.clf_backbone}")

    clf = RefugeClassification(
        pre,
        seg,
        backbone=backbone,
        cache_dir=args.clf_cache_dir,
        use_all_labeled=args.clf_use_all,
        auto_val_ratio=args.clf_auto_val_ratio,
    )
    clf.build_datasets(
        crop_scale=args.crop_scale,
        crop_size=args.crop_size,
        batch_size=args.clf_batch_size,
        num_workers=args.num_workers,
    )

    default_ckpt_path = classifier_checkpoint_path(args.clf_backbone)
    ckpt_path = args.clf_checkpoint_path or default_ckpt_path
    ckpt_dir = ckpt_path.parent
    history = clf.train(
        epochs=args.clf_epochs,
        lr=args.clf_lr,
        weight_decay=args.clf_weight_decay,
        rotation_weight=args.rotation_weight,
        checkpoint_dir=ckpt_dir,
        device=args.device,
    )
    print("Classifier training complete. Best AUC:", history.get("best_auc"))
    print(f"Checkpoint directory: {ckpt_dir}")
    saved_path = ckpt_dir / "refuge_classifier_best.pt"
    if ckpt_path != saved_path:
        ckpt_path.parent.mkdir(parents=True, exist_ok=True)
        shutil.copy2(saved_path, ckpt_path)
        print(f"Checkpoint copied to: {ckpt_path}")


def _load_classifier(
    pre: RefugePreprocessing, seg: RefugeSegmentation, args: argparse.Namespace
) -> Tuple[RefugeClassification, Path]:
    backbone = build_classifier_backbone(args.clf_backbone)
    clf = RefugeClassification(
        pre,
        seg,
        backbone=backbone,
        cache_dir=args.clf_cache_dir,
        use_all_labeled=args.clf_use_all,
        auto_val_ratio=args.clf_auto_val_ratio,
    )
    clf.build_datasets(
        crop_scale=args.crop_scale,
        crop_size=args.crop_size,
        batch_size=args.clf_batch_size,
        num_workers=args.num_workers,
    )
    ckpt_path = args.clf_checkpoint_path or classifier_checkpoint_path(
        args.clf_backbone
    )
    if not ckpt_path.exists():
        raise FileNotFoundError(f"Classifier checkpoint missing: {ckpt_path}")
    print(f"[classifier] Loading checkpoint: {ckpt_path}")
    state = torch.load(ckpt_path, map_location=args.device)
    clf.backbone.load_state_dict(state["backbone"])
    clf.classifier_head.load_state_dict(state["classifier"])
    clf.rotation_head.load_state_dict(state["rotation"])
    clf.backbone.to(args.device)
    clf.classifier_head.to(args.device)
    clf.rotation_head.to(args.device)
    return clf, ckpt_path


def _collect_records(
    pre: RefugePreprocessing,
    seg: RefugeSegmentation,
    clf: RefugeClassification,
    dataset_name: str,
    split: str,
    scale: float,
) -> List[RefugeClassificationRecord]:
    manifest = pre.build_manifest()
    samples = [
        sample
        for sample in manifest
        if sample.dataset == dataset_name
        and sample.split == split
        and sample.label is not None
    ]
    if not samples:
        return []
    print(f"[eval] Preparing {len(samples)} samples for {dataset_name.upper()} {split}")
    return clf.build_records_for_samples(
        samples, crop_scale=scale, progress_prefix=f"{dataset_name}_{split}"
    )


def _auc_for_records(
    clf: RefugeClassification,
    records: List[RefugeClassificationRecord],
    device: str,
) -> float:
    if not records:
        return float("nan")
    dataset = RefugeClassificationDataset(
        records,
        transform=clf.eval_transform,
        polar_transform=clf.polar_transform,
        size=clf.crop_size,
    )
    loader = DataLoader(dataset, batch_size=64, shuffle=False, num_workers=0)
    clf.backbone.to(device).eval()
    clf.classifier_head.to(device).eval()
    preds: List[float] = []
    targets: List[int] = []
    with torch.no_grad():
        for batch in tqdm(loader, desc="Eval", leave=False, unit="batch"):
            images = batch["image"].to(device)
            polars = batch["polar"].to(device)
            extra_feats = batch["features"].to(device)
            labels = batch["label"].cpu().numpy().tolist()
            feats_img = clf.backbone(images)
            feats = feats_img
            if getattr(clf, "use_polar", False):
                feats_polar = clf.backbone(polars)
                feats = torch.cat([feats, feats_polar], dim=1)
            if getattr(clf, "extra_feature_dim", 0) > 0:
                feats = torch.cat([feats, extra_feats], dim=1)
            logits = clf.classifier_head(feats)
            probs = torch.softmax(logits, dim=1)[:, 1].cpu().numpy().tolist()
            preds.extend(probs)
            targets.extend(labels)
    if len(set(targets)) < 2:
        return float("nan")
    return float(roc_auc_score(targets, preds))


def evaluate(args: argparse.Namespace) -> None:
    pre = ensure_preprocessing()
    seg = _load_segmentation(pre, args)
    clf, clf_ckpt = _load_classifier(pre, seg, args)

    def evaluate_subset(
        clf_obj: RefugeClassification,
        records: List[RefugeClassificationRecord],
        label: str,
    ) -> None:
        if not records:
            print(f"[eval] No samples found for {label}; skipping.")
            return

        base_state = {
            "backbone": clf_obj.backbone.state_dict(),
            "rotation": clf_obj.rotation_head.state_dict(),
        }

        auc_no_ttt = _auc_for_records(clf_obj, records, device=args.device)

        auc_ttt = float("nan")
        if args.with_ttt:
            ttt_loader = DataLoader(
                RefugeClassificationDataset(
                    records,
                    transform=clf_obj.ttt_transform,
                    polar_transform=clf_obj.polar_transform,
                    size=clf_obj.crop_size,
                ),
                batch_size=16,
                shuffle=False,
                num_workers=0,
            )
            ttt_iter = tqdm(range(args.ttt_steps), desc="TTT", unit="step")
            for _ in ttt_iter:
                clf_obj.apply_ttt(ttt_loader, device=args.device, steps=1)
            auc_ttt = _auc_for_records(clf_obj, records, device=args.device)
            clf_obj.backbone.load_state_dict(base_state["backbone"])
            clf_obj.rotation_head.load_state_dict(base_state["rotation"])

        print(
            f"{label}: AUC (no TTT) = {auc_no_ttt:.4f}"
            + (f", AUC (TTT) = {auc_ttt:.4f}" if args.with_ttt else "")
        )

    if args.eval_datasets:
        for dataset_name in dict.fromkeys(args.eval_datasets):
            if dataset_name.lower() == "papila":
                papila_records, papila_clf = build_papila_records(args, pre, clf_ckpt)
                if papila_clf is None:
                    print("[eval] Papila evaluation aborted; no samples built.")
                else:
                    evaluate_subset(papila_clf, papila_records, "PAPILA holdout")
            else:
                records = _collect_records(
                    pre,
                    seg,
                    clf,
                    dataset_name,
                    "holdout",
                    scale=args.crop_scale,
                )
                evaluate_subset(clf, records, f"{dataset_name.upper()} holdout")
        return

    # Do not mix splits: report per dataset + split
    subsets = [
        ("refuge1", "val"),
        ("refuge2", "val"),
        ("refuge2", "test"),
    ]

    for dataset_name, split in subsets:
        records = _collect_records(
            pre, seg, clf, dataset_name, split, scale=args.crop_scale
        )
        evaluate_subset(clf, records, f"{dataset_name.upper()} {split}")

        if args.dump_masks and dataset_name == "refuge1" and split == "val":
            out_dir = Path(args.dump_masks)
            out_dir.mkdir(parents=True, exist_ok=True)
            for rec in records:
                sample = rec.sample
                if sample is None:
                    continue
                pred = seg.predict_mask(sample, device=args.device).numpy()
                Image.fromarray((pred * 255).astype(np.uint8)).save(
                    out_dir / f"{sample.sample_id}_pred.png"
                )
                if sample.mask_path and sample.mask_path.exists():
                    Image.open(sample.mask_path).convert("L").save(
                        out_dir / f"{sample.sample_id}_gt.png"
                    )


def evaluate_segmentation(args: argparse.Namespace) -> None:
    manifest_path = args.seg_manifest or Path("manifest.csv")
    mask_cache_dir = None if args.in_memory_cache else args.mask_cache_dir
    image_cache_dir = None if args.in_memory_cache else args.image_cache_dir

    segmenter = UNetSegmenter(
        manifest_path=manifest_path,
        normalize=args.seg_normalize,
        device=args.device,
        mask_cache_dir=mask_cache_dir,
        image_cache_dir=image_cache_dir,
        in_memory_cache=args.in_memory_cache,
        loader_workers=args.loader_workers,
    )
    if args.in_memory_cache:
        segmenter.prebuild_in_memory_cache(
            cache_workers=max(0, int(args.cache_workers)),
            include_train=False,
            include_val=bool(args.eval_seg_splits is None or "val" in args.eval_seg_splits),
            include_holdout=bool(args.eval_seg_splits is None or "holdout" in args.eval_seg_splits),
        )

    ckpt = resolve_unet_weights(args.seg_weights)
    if ckpt.exists():
        state = torch.load(ckpt, map_location=segmenter.device)
        state_dict = state.get("model", state)
        segmenter.model.load_state_dict(state_dict, strict=False)
        print(f"[seg-eval] Loaded weights from {ckpt}")
    else:
        raise FileNotFoundError(f"Segmentation weights not found at {ckpt}")

    dataset_filter = args.eval_seg_datasets
    split_filter = args.eval_seg_splits
    output_dir = args.eval_seg_output or Path("analysis_data/segmenter_eval")
    metrics_path = args.eval_seg_metrics_path

    segmenter.evaluate_dataset(
        dataset_filter=dataset_filter,
        split_filter=split_filter,
        output_dir=output_dir,
        save_overlays=not args.eval_seg_no_overlays,
        metrics_path=metrics_path,
        threshold=args.eval_seg_threshold,
        tta=args.eval_seg_tta,
    )


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="REFUGE pipeline helper")
    parser.add_argument(
        "--train-seg", action="store_true", help="Train the segmentation model"
    )
    parser.add_argument(
        "--train-unet-seg",
        action="store_true",
        help="Train the UNet segmenter (replacement for scripts/run_unet_segmenter.py)",
    )
    parser.add_argument(
        "--train-clf", action="store_true", help="Train the classification model"
    )
    parser.add_argument(
        "--eval", action="store_true", help="Run evaluation on stored checkpoints"
    )
    parser.add_argument(
        "--with-ttt",
        action="store_true",
        help="Apply test-time training during evaluation",
    )
    parser.add_argument(
        "--ttt-steps", type=int, default=1, help="TTT epochs over evaluation loader"
    )
    parser.add_argument(
        "--export-backbone",
        type=Path,
        default=None,
        help="Optional path to export the trained backbone weights",
    )
    parser.add_argument(
        "--dump-masks",
        type=Path,
        default=None,
        help="Optional directory to dump predicted/GT masks during eval",
    )
    parser.add_argument(
        "--device", default="cuda" if torch.cuda.is_available() else "cpu"
    )
    parser.add_argument("--num-workers", type=int, default=4)
    # Segmentation hyperparameters
    parser.add_argument("--seg-epochs", type=int, default=40)
    parser.add_argument("--seg-lr", type=float, default=1e-3)
    parser.add_argument("--seg-weight-decay", type=float, default=1e-5)
    parser.add_argument("--seg-image-size", type=int, default=512)
    parser.add_argument("--seg-batch-size", type=int, default=4)
    parser.add_argument(
        "--seg-manifest",
        type=Path,
        default=Path("manifest.csv"),
        help="Manifest CSV for the UNet segmenter (default: manifest.csv)",
    )
    parser.add_argument(
        "--seg-weights",
        type=Path,
        default=None,
        help="Path to UNet segmenter weights (default: models/unet_segmenter/best.pt)",
    )
    parser.add_argument(
        "--seg-normalize",
        choices=["none", "imagenet", "per_image"],
        default="none",
        help="Normalization mode used when running the UNet segmenter",
    )
    parser.add_argument(
        "--seg-strong-aug",
        action="store_true",
        help="Enable stronger geometric augmentations when training the UNet segmenter",
    )
    parser.add_argument(
        "--seg-train-datasets",
        nargs="+",
        default=["refuge"],
        help="Datasets to use for UNet segmenter training (default: refuge)",
    )
    parser.add_argument(
        "--seg-val-datasets",
        nargs="+",
        default=["refuge"],
        help="Datasets eligible for validation sampling (default: refuge)",
    )
    parser.add_argument(
        "--seg-holdout-datasets",
        nargs="+",
        default=["refuge"],
        help="Datasets reserved for holdout set during UNet segmenter training (default: refuge)",
    )
    parser.add_argument(
        "--seg-checkpoint-dir",
        type=Path,
        default=Path("models/v2/refuge/segmentation/per_image"),
        help="Directory to store UNet segmenter checkpoints",
    )
    parser.add_argument(
        "--loader-workers",
        type=int,
        default=0,
        help="DataLoader workers for UNet segmenter train/eval.",
    )
    parser.add_argument(
        "--mask-cache-dir",
        type=Path,
        default=None,
        help="Optional cache dir for parsed/resized disc+cup masks.",
    )
    parser.add_argument(
        "--image-cache-dir",
        type=Path,
        default=None,
        help="Optional cache dir for resized RGB images before augmentation.",
    )
    parser.add_argument(
        "--in-memory-cache",
        action="store_true",
        help="Cache preprocessed images and masks in RAM (per DataLoader worker process).",
    )
    parser.add_argument(
        "--cache-workers",
        type=int,
        default=0,
        help="Worker threads for prebuilding in-memory cache before training/eval.",
    )
    # Classification hyperparameters
    parser.add_argument("--clf-epochs", type=int, default=30)
    parser.add_argument("--clf-lr", type=float, default=1e-4)
    parser.add_argument("--clf-weight-decay", type=float, default=1e-4)
    parser.add_argument("--clf-batch-size", type=int, default=16)
    parser.add_argument(
        "--clf-backbone",
        choices=sorted(CLASSIFIER_BACKBONES.keys()),
        default="resnet50",
        help="Backbone architecture for the REFUGE classifier",
    )
    parser.add_argument("--rotation-weight", type=float, default=0.5)
    parser.add_argument("--crop-scale", type=float, default=2.5)
    parser.add_argument("--crop-size", type=int, default=224)
    parser.add_argument(
        "--clf-cache-dir",
        type=Path,
        default=Path("analysis_data/classifier_cache"),
        help="Directory to cache classifier preprocessing artifacts",
    )
    parser.add_argument(
        "--clf-use-all",
        action="store_true",
        help="Use all labelled samples (train+val) when building classifier dataset",
    )
    parser.add_argument(
        "--clf-auto-val-ratio",
        type=float,
        default=0.1,
        help="Fraction for automatic validation split when no explicit val set is used",
    )
    parser.add_argument(
        "--clf-checkpoint-path",
        type=Path,
        default=None,
        help="Optional explicit path for the classifier checkpoint (defaults to models/refuge/classifier/<backbone>/refuge_classifier_best.pt)",
    )
    parser.add_argument(
        "--eval-datasets",
        nargs="+",
        help="Datasets to evaluate during --eval (e.g. papila). Defaults to REFUGE splits.",
    )
    # Segmentation evaluation parameters
    parser.add_argument(
        "--eval-seg",
        action="store_true",
        help="Evaluate the segmentation model on specified datasets/splits",
    )
    parser.add_argument(
        "--eval-seg-datasets",
        nargs="+",
        default=["refuge"],
        help="Segmentation datasets to evaluate (default: refuge)",
    )
    parser.add_argument(
        "--eval-seg-splits",
        nargs="+",
        choices=["train", "val", "holdout"],
        help="Segmentation splits to evaluate (default: val)",
    )
    parser.add_argument(
        "--eval-seg-output",
        type=Path,
        default=Path("analysis_data/segmenter_eval"),
        help="Directory to store segmentation metrics CSVs",
    )
    parser.add_argument(
        "--eval-seg-threshold",
        type=float,
        default=0.5,
        help="Threshold for binarising predicted masks during segmentation eval",
    )
    parser.add_argument(
        "--eval-seg-metrics-path",
        type=Path,
        default=None,
        help="Optional explicit CSV path for segmentation metrics output",
    )
    parser.add_argument(
        "--eval-seg-no-overlays",
        action="store_true",
        help="Skip saving GT/pred overlay images during segmentation evaluation",
    )
    parser.add_argument(
        "--eval-seg-tta",
        action="store_true",
        help="Enable horizontal/vertical flip TTA during segmentation evaluation",
    )
    parser.add_argument(
        "--papila-metrics",
        type=Path,
        default=None,
        help="Optional CSV of Papila Dice metrics used to filter samples",
    )
    parser.add_argument(
        "--papila-dice-threshold",
        type=float,
        default=0.01,
        help="Minimum Dice required (disc or cup) when filtering Papila metrics",
    )
    parser.add_argument(
        "--papila-positive-labels",
        nargs="+",
        default=["glaucoma", "glaucoma suspect", "suspect"],
        help="Papila label values treated as positive when labels are non-numeric",
    )
    parser.add_argument(
        "--papila-image-dir",
        type=Path,
        default=Path("Papila/FundusImages"),
        help="Path to Papila fundus images",
    )
    parser.add_argument(
        "--papila-clinical-dir",
        type=Path,
        default=Path("Papila/ClinicalData"),
        help="Path to Papila clinical CSVs",
    )
    parser.add_argument(
        "--papila-label-col",
        type=str,
        default="Diagnosis",
        help="Column name containing Papila labels",
    )
    parser.add_argument(
        "--papila-use-gt",
        action="store_true",
        help="Use Papila ground-truth contours when evaluating classifiers",
    )
    parser.add_argument(
        "--papila-contours-dir",
        type=Path,
        default=Path("Papila/ExpertsSegmentations/Contours"),
        help="Directory containing Papila contour text files",
    )
    return parser.parse_args()


def main() -> None:
    args = parse_args()

    if not any(
        [
            args.train_seg,
            args.train_unet_seg,
            args.train_clf,
            args.eval,
            args.eval_seg,
            args.export_backbone,
        ]
    ):
        raise SystemExit(
            "Specify at least one action: --train-seg, --train-unet-seg, --train-clf, --eval, --eval-seg, or --export-backbone"
        )

    if args.train_seg:
        train_segmentation(args)

    if args.train_unet_seg:
        train_unet_segmenter(args)

    if args.train_clf:
        train_classifier(args)

    if args.eval:
        evaluate(args)

    if args.eval_seg:
        evaluate_segmentation(args)

    if args.export_backbone:
        pre = ensure_preprocessing()
        seg = _load_segmentation(pre, args)
        clf = _load_classifier(pre, seg, args)
        out_path = args.export_backbone
        out_path.parent.mkdir(parents=True, exist_ok=True)
        torch.save(clf.extract_backbone().state_dict(), out_path)
        print(f"Backbone weights exported to {out_path}")


if __name__ == "__main__":
    main()