src/utils.py

"""
Utility functions for ACL-LKNet.

Includes: reproducibility seeding, EMA model, checkpoint save/load,
logging helpers, and Google Drive integration.
"""

import os
import copy
import random
import logging
from typing import Dict, Any, Optional

import numpy as np
import torch
import torch.nn as nn


# ── Reproducibility ─────────────────────────────────────────────────

def set_seed(seed: int = 42):
    """Set all random seeds for reproducibility."""
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)
    torch.backends.cudnn.deterministic = True
    torch.backends.cudnn.benchmark = False
    os.environ["PYTHONHASHSEED"] = str(seed)


def get_rng_states() -> Dict[str, Any]:
    """Capture all RNG states for exact checkpoint reproducibility."""
    states = {
        "python": random.getstate(),
        "numpy": np.random.get_state(),
        "torch": torch.get_rng_state(),
    }
    if torch.cuda.is_available():
        states["cuda"] = torch.cuda.get_rng_state_all()
    return states


def set_rng_states(states: Dict[str, Any]):
    """Restore RNG states from checkpoint."""
    random.setstate(states["python"])
    np.random.set_state(states["numpy"])
    torch.set_rng_state(states["torch"])
    if "cuda" in states and torch.cuda.is_available():
        torch.cuda.set_rng_state_all(states["cuda"])


# ── Exponential Moving Average ──────────────────────────────────────

class EMAModel:
    """
    Exponential Moving Average of model parameters.
    
    Maintains a shadow copy of model weights that is updated as:
        shadow = decay * shadow + (1 - decay) * current
    
    Use the EMA model for evaluation — it typically generalizes better.
    """

    def __init__(self, model: nn.Module, decay: float = 0.999):
        self.decay = decay
        self.shadow = copy.deepcopy(model)
        self.shadow.eval()
        for p in self.shadow.parameters():
            p.requires_grad_(False)

    @torch.no_grad()
    def update(self, model: nn.Module):
        """Update shadow weights with current model weights and buffers."""
        for s_param, m_param in zip(self.shadow.parameters(), model.parameters()):
            s_param.data.mul_(self.decay).add_(m_param.data, alpha=1.0 - self.decay)
        # Sync BatchNorm running statistics (buffers are not EMA-averaged,
        # they should directly mirror the training model's batch statistics)
        for s_buf, m_buf in zip(self.shadow.buffers(), model.buffers()):
            s_buf.data.copy_(m_buf.data)

    def state_dict(self):
        return self.shadow.state_dict()

    def load_state_dict(self, state_dict):
        self.shadow.load_state_dict(state_dict)

    def eval_model(self) -> nn.Module:
        """Return the shadow model for evaluation."""
        return self.shadow


# ── Checkpoint Management ───────────────────────────────────────────

def save_checkpoint(
    path: str,
    epoch: int,
    phase: str,
    model: nn.Module,
    optimizer: torch.optim.Optimizer,
    scheduler: Any,
    scaler: Optional[torch.amp.GradScaler],
    ema: Optional[EMAModel],
    best_metric: float,
    best_epoch: int,
    train_history: list,
    val_history: list,
    patience_counter: int,
    config: Any,
):
    """
    Save a full training checkpoint to Google Drive.
    
    Captures everything needed to resume training exactly:
    model, optimizer, scheduler, AMP scaler, EMA, RNG states, histories.
    """
    os.makedirs(os.path.dirname(path), exist_ok=True)

    checkpoint = {
        "epoch": epoch,
        "phase": phase,
        "model_state_dict": model.state_dict(),
        "optimizer_state_dict": optimizer.state_dict(),
        "scheduler_state_dict": scheduler.state_dict() if scheduler else None,
        "scaler_state_dict": scaler.state_dict() if scaler else None,
        "ema_state_dict": ema.state_dict() if ema else None,
        "best_metric": best_metric,
        "best_epoch": best_epoch,
        "train_history": train_history,
        "val_history": val_history,
        "patience_counter": patience_counter,
        "rng_states": get_rng_states(),
        "config": config.to_dict() if hasattr(config, "to_dict") else str(config),
    }
    torch.save(checkpoint, path)
    logging.info(f"Checkpoint saved: {path}")


def load_checkpoint(
    path: str,
    model: nn.Module,
    optimizer: Optional[torch.optim.Optimizer] = None,
    scheduler: Any = None,
    scaler: Optional[torch.amp.GradScaler] = None,
    ema: Optional[EMAModel] = None,
) -> Dict[str, Any]:
    """
    Load a checkpoint and restore all training state.
    
    Returns the checkpoint dict for extracting histories, epoch, etc.
    """
    checkpoint = torch.load(path, map_location="cpu", weights_only=False)

    model.load_state_dict(checkpoint["model_state_dict"])
    if optimizer and "optimizer_state_dict" in checkpoint:
        optimizer.load_state_dict(checkpoint["optimizer_state_dict"])
    if scheduler and checkpoint.get("scheduler_state_dict"):
        scheduler.load_state_dict(checkpoint["scheduler_state_dict"])
    if scaler and checkpoint.get("scaler_state_dict"):
        scaler.load_state_dict(checkpoint["scaler_state_dict"])
    if ema and checkpoint.get("ema_state_dict"):
        ema.load_state_dict(checkpoint["ema_state_dict"])
    if "rng_states" in checkpoint:
        set_rng_states(checkpoint["rng_states"])

    logging.info(f"Checkpoint loaded: {path} (epoch {checkpoint['epoch']})")
    return checkpoint


def find_latest_checkpoint(checkpoint_dir: str, phase: str = "finetune") -> Optional[str]:
    """Find the latest checkpoint file in the checkpoint directory."""
    if not os.path.exists(checkpoint_dir):
        return None
    
    checkpoints = [
        f for f in os.listdir(checkpoint_dir)
        if f.startswith(f"{phase}_") and f.endswith(".pt")
    ]
    if not checkpoints:
        return None

    # Sort by epoch number
    def extract_epoch(fname):
        try:
            parts = fname.replace(".pt", "").split("_epoch")
            return int(parts[-1])
        except (ValueError, IndexError):
            return -1

    checkpoints.sort(key=extract_epoch)
    return os.path.join(checkpoint_dir, checkpoints[-1])


# ── Logging ─────────────────────────────────────────────────────────

def setup_logging(log_dir: Optional[str] = None, level=logging.INFO):
    """Configure logging to console and optionally to file."""
    handlers = [logging.StreamHandler()]
    if log_dir:
        os.makedirs(log_dir, exist_ok=True)
        handlers.append(logging.FileHandler(os.path.join(log_dir, "training.log")))

    logging.basicConfig(
        level=level,
        format="%(asctime)s [%(levelname)s] %(message)s",
        datefmt="%Y-%m-%d %H:%M:%S",
        handlers=handlers,
        force=True,
    )


# ── Metrics Formatting ─────────────────────────────────────────────

def format_metrics(metrics: Dict[str, float]) -> str:
    """Format a metrics dict into a readable string."""
    parts = []
    for k, v in metrics.items():
        if isinstance(v, float):
            parts.append(f"{k}: {v:.4f}")
        else:
            parts.append(f"{k}: {v}")
    return " | ".join(parts)


# ── Memory Utils ────────────────────────────────────────────────────

def get_gpu_memory_info() -> Dict[str, float]:
    """Get GPU memory usage in MB."""
    if not torch.cuda.is_available():
        return {"allocated_mb": 0, "reserved_mb": 0, "total_mb": 0}
    try:
        props = torch.cuda.get_device_properties(0)
        total = getattr(props, "total_memory", getattr(props, "total_mem", 0)) / 1024**2
        return {
            "allocated_mb": torch.cuda.memory_allocated() / 1024**2,
            "reserved_mb": torch.cuda.memory_reserved() / 1024**2,
            "total_mb": total,
        }
    except Exception:
        return {"allocated_mb": 0, "reserved_mb": 0, "total_mb": 0}


def clear_gpu_memory():
    """Force GPU memory cleanup."""
    if torch.cuda.is_available():
        torch.cuda.empty_cache()
        torch.cuda.synchronize()