backend/scripts/train_efficientnet.py

from __future__ import annotations

import json
import math
import random
from collections import Counter
from copy import deepcopy
from dataclasses import dataclass
from datetime import datetime, timezone
from pathlib import Path

import numpy as np
import torch
from PIL import Image
from sklearn.metrics import accuracy_score, f1_score, mean_absolute_error, precision_score, recall_score, roc_auc_score
from sklearn.model_selection import GroupShuffleSplit
from torch import nn
from torch.optim import AdamW
from torch.utils.data import DataLoader, Dataset, WeightedRandomSampler

from app.config import (
    BACKEND_ROOT,
    DEFAULT_EFFICIENTNET_MODEL_PATH,
    DEFAULT_EFFICIENTNET_REPORT_PATH,
    DEFAULT_TRAINING_REPORT_PATH,
)
from app.ml.archive_model import ANEMIA_HB_THRESHOLD, _first_path, _load_image_with_fallback, _parse_float, _parse_workbook
from app.ml.efficientnet_model import (
    EFFICIENTNET_VERSION,
    build_efficientnet_model,
    build_train_transform,
    build_val_transform,
)
from app.services.conjunctiva_roi import ConjunctivaRoiExtractor


DATA_ROOT = BACKEND_ROOT / "data"
ARCHIVE_ROOT = BACKEND_ROOT.parent / "archive" / "dataset anemia"
SEED = 42
BATCH_SIZE = 16
EPOCHS = 60
PATIENCE = 15
MAX_GRAD_NORM = 1.0
WARMUP_EPOCHS = 5
LABEL_SMOOTHING = 0.05
MIXUP_ALPHA = 0.3
# Hb spread loss: penalizes predictions that cluster near the mean
HB_SPREAD_WEIGHT = 0.15


@dataclass(frozen=True)
class ImageRecord:
    subject_id: str
    label: int
    hb: float
    image_path: Path
    source: str


class ConjunctivaDataset(Dataset):
    def __init__(self, records: list[ImageRecord], transform: object) -> None:
        self.records = records
        self.transform = transform
        self.roi_extractor = ConjunctivaRoiExtractor()

    def __len__(self) -> int:
        return len(self.records)

    def __getitem__(self, index: int) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
        record = self.records[index]
        image, label, hb = self._prepare_item(record)
        tensor = self.transform(image)
        return tensor, torch.tensor([label], dtype=torch.float32), torch.tensor([hb], dtype=torch.float32)

    def _prepare_item(self, record: ImageRecord) -> tuple[Image.Image, float, float]:
        image = _load_image_with_fallback(record.image_path)
        if record.source == "roi_original":
            image = self.roi_extractor.extract(image).image
        return image.convert("RGB"), float(record.label), float(record.hb)


class FocalLoss(nn.Module):
    def __init__(self, alpha: float = 0.25, gamma: float = 2.0, pos_weight: torch.Tensor | None = None) -> None:
        super().__init__()
        self.alpha = alpha
        self.gamma = gamma
        self.bce = nn.BCEWithLogitsLoss(pos_weight=pos_weight, reduction="none")

    def forward(self, inputs: torch.Tensor, targets: torch.Tensor) -> torch.Tensor:
        bce_loss = self.bce(inputs, targets)
        probabilities = torch.sigmoid(inputs)
        p_t = probabilities * targets + (1 - probabilities) * (1 - targets)
        loss = bce_loss * ((1 - p_t) ** self.gamma)
        if self.alpha >= 0:
            alpha_t = self.alpha * targets + (1 - self.alpha) * (1 - targets)
            loss = alpha_t * loss
        return loss.mean()


def main() -> None:
    _set_seed(SEED)
    dataset_root = DATA_ROOT if DATA_ROOT.exists() else ARCHIVE_ROOT
    if not dataset_root.exists():
        raise RuntimeError(f"No dataset directory found at {DATA_ROOT} or {ARCHIVE_ROOT}.")

    records = _build_records(dataset_root)
    if not records:
        raise RuntimeError(f"No training records found in {dataset_root}.")

    train_records, val_records = _balanced_group_split(records, test_size=0.2, n_splits=32)
    train_dataset = ConjunctivaDataset(train_records, build_train_transform())
    val_dataset = ConjunctivaDataset(val_records, build_val_transform())
    hb_mean, hb_std = _hb_normalization_stats(train_records)
    train_sampler = _build_weighted_sampler(train_records)

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model = build_efficientnet_model(pretrained=True).to(device)
    optimizer = AdamW(
        [
            {"params": list(model.classifier.parameters()), "lr": 1.5e-4}, # Slightly lower for stability
            {"params": [param for param in model.features.parameters() if param.requires_grad], "lr": 5e-6},
        ],
        weight_decay=4e-4, # Higher weight decay for better regularization
    )

    # Warmup then cosine annealing
    def warmup_cosine_lr(epoch: int) -> float:
        if epoch < WARMUP_EPOCHS:
            return float(epoch + 1) / WARMUP_EPOCHS
        progress = (epoch - WARMUP_EPOCHS) / max(EPOCHS - WARMUP_EPOCHS, 1)
        return 0.5 * (1.0 + math.cos(math.pi * progress))

    scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda=warmup_cosine_lr)

    # Use Focal Loss with positive weights
    pos_weight = torch.tensor([_positive_class_weight(train_records)], device=device)
    cls_loss_fn = FocalLoss(alpha=0.25, gamma=2.0, pos_weight=pos_weight)
    hb_loss_fn = nn.SmoothL1Loss(beta=0.5)

    train_loader = DataLoader(train_dataset, batch_size=BATCH_SIZE, sampler=train_sampler, num_workers=0)
    val_loader = DataLoader(val_dataset, batch_size=BATCH_SIZE, shuffle=False, num_workers=0)

    best_state: dict[str, torch.Tensor] | None = None
    best_metrics: dict[str, float] | None = None
    best_threshold = 0.5
    best_score = -1.0
    epochs_without_improvement = 0
    history: list[dict[str, float]] = []

    for epoch in range(1, EPOCHS + 1):
        model.train()
        train_loss_total = 0.0

        for images, labels, hbs in train_loader:
            images = images.to(device)
            labels = labels.to(device)
            hbs = hbs.to(device)
            normalized_hbs = (hbs - hb_mean) / hb_std

            # MixUp augmentation
            if MIXUP_ALPHA > 0 and np.random.random() < 0.5:
                lam = float(np.random.beta(MIXUP_ALPHA, MIXUP_ALPHA))
                idx = torch.randperm(images.size(0), device=device)
                images = lam * images + (1.0 - lam) * images[idx]
                labels_a, labels_b = labels, labels[idx]
                hbs_a, hbs_b = normalized_hbs, normalized_hbs[idx]

                optimizer.zero_grad(set_to_none=True)
                output = model(images)
                # Label smoothing applied to both MixUp targets
                smooth_a = labels_a * (1.0 - LABEL_SMOOTHING) + 0.5 * LABEL_SMOOTHING
                smooth_b = labels_b * (1.0 - LABEL_SMOOTHING) + 0.5 * LABEL_SMOOTHING
                cls_loss = lam * cls_loss_fn(output[:, 0:1], smooth_a) + (1.0 - lam) * cls_loss_fn(output[:, 0:1], smooth_b)
                hb_loss = lam * hb_loss_fn(output[:, 1:2], hbs_a) + (1.0 - lam) * hb_loss_fn(output[:, 1:2], hbs_b)
            else:
                optimizer.zero_grad(set_to_none=True)
                output = model(images)
                smooth_labels = labels * (1.0 - LABEL_SMOOTHING) + 0.5 * LABEL_SMOOTHING
                cls_loss = cls_loss_fn(output[:, 0:1], smooth_labels)
                hb_loss = hb_loss_fn(output[:, 1:2], normalized_hbs)

            # Hb spread loss: penalize predictions clustering near zero (normalized mean)
            # Encourages the model to predict a wider range of Hb values
            hb_pred_norm = output[:, 1:2]
            spread_loss = torch.clamp(0.5 - hb_pred_norm.std(), min=0.0)

            total_loss = (0.60 * cls_loss) + (0.30 * hb_loss) + (HB_SPREAD_WEIGHT * spread_loss)
            total_loss.backward()
            nn.utils.clip_grad_norm_(model.parameters(), MAX_GRAD_NORM)
            optimizer.step()
            train_loss_total += float(total_loss.item()) * images.size(0)

        scheduler.step()
        val_metrics = _evaluate_model(model, val_loader, device, hb_mean=hb_mean, hb_std=hb_std)
        history.append(
            {
                "epoch": float(epoch),
                "train_loss": round(train_loss_total / max(len(train_dataset), 1), 4),
                "val_f1": val_metrics["f1"],
                "val_auc": val_metrics["auc"],
                "val_hb_mae": val_metrics["hb_mae"],
            }
        )
        print(
            f"epoch={epoch:02d} train_loss={history[-1]['train_loss']:.4f} "
            f"val_f1={val_metrics['f1']:.4f} val_auc={val_metrics['auc']:.4f} "
            f"val_hb_mae={val_metrics['hb_mae']:.4f}"
        )

        # Use composite score: AUC weighted more heavily than F1 (more stable early on)
        composite_score = val_metrics["auc"] * 0.55 + val_metrics["f1"] * 0.35 + (1.0 - min(val_metrics["hb_mae"] / 4.0, 1.0)) * 0.10
        if composite_score > best_score:
            best_score = composite_score
            best_state = deepcopy(model.state_dict())
            best_metrics = val_metrics
            best_threshold = val_metrics["decision_threshold"]
            epochs_without_improvement = 0
        else:
            epochs_without_improvement += 1

        if epochs_without_improvement >= PATIENCE:
            print(f"Early stopping after {epoch} epochs.")
            break

    if best_state is None or best_metrics is None:
        raise RuntimeError("EfficientNet training did not produce a valid checkpoint.")

    checkpoint = {
        "version": EFFICIENTNET_VERSION,
        "created_at": datetime.now(timezone.utc).isoformat(),
        "state_dict": best_state,
        "decision_threshold": best_threshold,
        "hb_mean": hb_mean,
        "hb_std": hb_std,
        "hb_spread_factor": _compute_hb_spread_factor(val_records, hb_mean, hb_std),
        "val_metrics": best_metrics,
    }
    DEFAULT_EFFICIENTNET_MODEL_PATH.parent.mkdir(parents=True, exist_ok=True)
    torch.save(checkpoint, DEFAULT_EFFICIENTNET_MODEL_PATH)

    report = {
        "dataset_name": str(dataset_root.name),
        "record_count": len(records),
        "subject_count": len({record.subject_id for record in records}),
        "primary_model": EFFICIENTNET_VERSION,
        "selected_mode": "efficientnet_hybrid_dual",
        "source_counts": _source_counts(records),
        "metrics": {
            "accuracy": round(best_metrics["accuracy"], 4),
            "precision": round(best_metrics["precision"], 4),
            "recall": round(best_metrics["recall"], 4),
            "f1": round(best_metrics["f1"], 4),
            "auc": round(best_metrics["auc"], 4),
            "mae_hb": round(best_metrics["hb_mae"], 4),
            "validation_size": len(val_records),
            "split_strategy": "group-shuffle-balance-select",
            "sample_count": len(records),
            "subject_count": len({record.subject_id for record in records}),
            "decision_threshold": round(best_threshold, 4),
        },
        "training": {
            "epochs_requested": EPOCHS,
            "history": history,
            "batch_size": BATCH_SIZE,
            "patience": PATIENCE,
            "device": str(device),
            "hb_target_mean": round(hb_mean, 4),
            "hb_target_std": round(hb_std, 4),
            "class_positive_weight": round(_positive_class_weight(train_records), 4),
            "sampler": "weighted-random-balanced",
        },
    }
    DEFAULT_EFFICIENTNET_REPORT_PATH.write_text(json.dumps(report, indent=2), encoding="utf-8")
    DEFAULT_TRAINING_REPORT_PATH.write_text(json.dumps(report, indent=2), encoding="utf-8")

    print("\nBest validation metrics")
    for key in ("accuracy", "precision", "recall", "f1", "auc", "hb_mae", "decision_threshold"):
        print(f"{key}: {best_metrics[key]:.4f}")


def _build_records(dataset_root: Path) -> list[ImageRecord]:
    records: list[ImageRecord] = []
    for country in ("India", "Italy"):
        workbook_path = dataset_root / country / f"{country}.xlsx"
        if not workbook_path.exists():
            continue
        metadata = _parse_workbook(workbook_path)
        for subject_number, row in metadata.items():
            hb = _parse_float(row.get("Hgb"))
            if hb is None:
                continue

            subject_dir = dataset_root / country / subject_number
            if not subject_dir.exists():
                continue

            subject_id = f"{country}-{subject_number}"
            label = int(hb < ANEMIA_HB_THRESHOLD)
            original_path = _first_path(subject_dir.glob("*.jpg"))
            palpebral_path = _first_path(
                path
                for path in subject_dir.glob("*_palpebral.png")
                if "forniceal_palpebral" not in path.name.lower()
            )

            if original_path is not None:
                records.append(
                    ImageRecord(
                        subject_id=subject_id,
                        label=label,
                        hb=float(hb),
                        image_path=original_path,
                        source="roi_original",
                    )
                )
            if palpebral_path is not None:
                records.append(
                    ImageRecord(
                        subject_id=subject_id,
                        label=label,
                        hb=float(hb),
                        image_path=palpebral_path,
                        source="palpebral",
                    )
                )
    return records


def _balanced_group_split(
    records: list[ImageRecord],
    *,
    test_size: float,
    n_splits: int,
) -> tuple[list[ImageRecord], list[ImageRecord]]:
    labels = np.asarray([record.label for record in records], dtype=np.int32)
    groups = np.asarray([record.subject_id for record in records], dtype=object)
    target_ratio = float(labels.mean())
    splitter = GroupShuffleSplit(n_splits=n_splits, test_size=test_size, random_state=SEED)

    best: tuple[np.ndarray, np.ndarray] | None = None
    best_score = float("inf")
    for train_index, val_index in splitter.split(np.zeros(len(records)), labels, groups):
        train_labels = labels[train_index]
        val_labels = labels[val_index]
        if len(np.unique(train_labels)) < 2 or len(np.unique(val_labels)) < 2:
            continue
        score = abs(float(train_labels.mean()) - target_ratio) + abs(float(val_labels.mean()) - target_ratio)
        if score < best_score:
            best_score = score
            best = (train_index, val_index)

    if best is None:
        raise RuntimeError("Unable to create a grouped train/validation split.")

    train_index, val_index = best
    return [records[i] for i in train_index], [records[i] for i in val_index]


def _evaluate_model(
    model: nn.Module,
    loader: DataLoader,
    device: torch.device,
    *,
    hb_mean: float,
    hb_std: float,
) -> dict[str, float]:
    model.eval()
    probabilities: list[float] = []
    labels: list[int] = []
    hb_predictions: list[float] = []
    hb_targets: list[float] = []

    with torch.no_grad():
        for images, batch_labels, batch_hbs in loader:
            images = images.to(device)
            output = model(images)
            probabilities.extend(torch.sigmoid(output[:, 0]).cpu().tolist())
            hb_predictions.extend(((output[:, 1].cpu() * hb_std) + hb_mean).tolist())
            labels.extend(batch_labels.squeeze(1).cpu().int().tolist())
            hb_targets.extend(batch_hbs.squeeze(1).cpu().tolist())

    threshold = _best_threshold(np.asarray(labels), np.asarray(probabilities))
    predicted_labels = [1 if probability >= threshold else 0 for probability in probabilities]
    auc = roc_auc_score(labels, probabilities) if len(set(labels)) > 1 else 0.5

    return {
        "accuracy": float(accuracy_score(labels, predicted_labels)),
        "precision": float(precision_score(labels, predicted_labels, zero_division=0)),
        "recall": float(recall_score(labels, predicted_labels, zero_division=0)),
        "f1": float(f1_score(labels, predicted_labels, zero_division=0)),
        "auc": float(auc),
        "hb_mae": float(mean_absolute_error(hb_targets, hb_predictions)),
        "decision_threshold": float(threshold),
    }


def _best_threshold(labels: np.ndarray, probabilities: np.ndarray) -> float:
    best_threshold = 0.5
    best_score = -1.0
    for threshold in np.linspace(0.25, 0.75, 51):
        predictions = (probabilities >= threshold).astype(np.int32)
        score = f1_score(labels, predictions, zero_division=0)
        if score > best_score:
            best_score = float(score)
            best_threshold = float(threshold)
    return best_threshold


def _source_counts(records: list[ImageRecord]) -> dict[str, int]:
    counts: dict[str, int] = {}
    for record in records:
        counts[record.source] = counts.get(record.source, 0) + 1
    return counts


def _positive_class_weight(records: list[ImageRecord]) -> float:
    counts = Counter(record.label for record in records)
    positive = max(counts.get(1, 0), 1)
    negative = max(counts.get(0, 0), 1)
    return float(negative / positive)


def _build_weighted_sampler(records: list[ImageRecord]) -> WeightedRandomSampler:
    counts = Counter(record.label for record in records)
    total = sum(counts.values())
    weights = [
        float(total / max(counts[record.label], 1))
        for record in records
    ]
    return WeightedRandomSampler(
        torch.as_tensor(weights, dtype=torch.double),
        num_samples=len(weights),
        replacement=True,
    )


def _hb_normalization_stats(records: list[ImageRecord]) -> tuple[float, float]:
    values = np.asarray([record.hb for record in records], dtype=np.float32)
    mean = float(values.mean())
    std = float(values.std())
    return mean, max(std, 1e-3)


def _set_seed(seed: int) -> None:
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    if torch.cuda.is_available():
        torch.cuda.manual_seed_all(seed)


def _compute_hb_spread_factor(records: list[ImageRecord], hb_mean: float, hb_std: float) -> float:
    """
    Estimate the spread amplification factor needed to correct regression-to-mean.
    Uses the ratio of true Hb std to the expected model output std (hb_std * 0.75 heuristic).
    """
    true_std = float(np.std([r.hb for r in records]))
    # Models typically predict ~75% of true std due to averaging
    predicted_std_estimate = max(hb_std * 0.75, 0.5)
    factor = float(np.clip(true_std / predicted_std_estimate, 1.0, 2.0))
    return round(factor, 3)


if __name__ == "__main__":
    main()