# file: memory_utils.py
"""Memory management utilities for NVIDIA RTX 4050 (6GB VRAM)."""

from __future__ import annotations

import gc
import logging
from typing import Any, Callable, Dict, Tuple

import torch
import torch.nn as nn
from torch.optim import Optimizer

logger = logging.getLogger(__name__)


# ===================================================================
# MemoryManager
# ===================================================================


class MemoryManager:
    """CUDA memory management for constrained VRAM (6GB)."""

    # ------------------------------------------------------------------
    # cache clearing
    # ------------------------------------------------------------------

    @staticmethod
    def clear_memory() -> None:
        """Clear CUDA cache and run garbage collection."""
        if torch.cuda.is_available():
            torch.cuda.empty_cache()
            torch.cuda.synchronize()
        gc.collect()

    # ------------------------------------------------------------------
    # diagnostics
    # ------------------------------------------------------------------

    @staticmethod
    def get_memory_usage() -> Dict[str, Any]:
        """Return CUDA memory stats."""
        if torch.cuda.is_available():
            allocated = torch.cuda.memory_allocated() / 1024**2
            reserved = torch.cuda.memory_reserved() / 1024**2
            max_allocated = torch.cuda.max_memory_allocated() / 1024**2
            return {
                "allocated_mb": round(allocated, 1),
                "reserved_mb": round(reserved, 1),
                "max_allocated_mb": round(max_allocated, 1),
                "device": torch.cuda.get_device_name(0),
            }
        return {"status": "cuda_unavailable"}

    # ------------------------------------------------------------------
    # gradient checkpointing
    # ------------------------------------------------------------------

    @staticmethod
    def enable_gradient_checkpointing(model: nn.Module) -> nn.Module:
        """Model already uses checkpoint in forward — no-op here."""
        return model

    # ------------------------------------------------------------------
    # mixed precision
    # ------------------------------------------------------------------

    @staticmethod
    def enable_mixed_precision() -> torch.amp.autocast:
        """Return bfloat16 autocast for RTX 40-series (native bfloat16 support)."""
        return torch.autocast(device_type='cuda', dtype=torch.bfloat16)

    # ------------------------------------------------------------------
    # model memory estimation
    # ------------------------------------------------------------------

    @staticmethod
    def estimate_model_memory(model: nn.Module) -> float:
        """Estimate model memory in MiB (params only, bfloat16)."""
        total = sum(p.numel() * 2 for p in model.parameters())  # bfloat16 = 2 bytes
        total += sum(p.numel() * 2 for p in model.buffers())
        return total / 1024**2


# ===================================================================
# Model optimisation helper
# ===================================================================


def optimize_for_rtx4050(model: nn.Module, _config: Any = None) -> nn.Module:
    """Apply RTX 40-series optimizations.

    - torch.compile for Ada Lovelace speedup
    - Log model size and VRAM estimate
    - Verify forward pass fits in 6GB

    Note: torch.compile won't compile on first call (requires CUDA warmup).
    If it fails, falls back to eager mode gracefully.
    """
    param_count = sum(p.numel() for p in model.parameters())
    logger.info("Model params: %.2fM", param_count / 1e6)
    mem_mb = MemoryManager.estimate_model_memory(model)
    logger.info("Model memory (bf16): %.1f MiB", mem_mb)

    if torch.cuda.is_available():
        free, total = torch.cuda.mem_get_info()
        logger.info("VRAM: %.1f GiB free / %.1f GiB total", free / 1024**3, total / 1024**3)
        if mem_mb > free / 1024**2 * 0.8:
            logger.warning("Model may exceed 80%% of free VRAM — consider reducing batch size")

    # torch.compile handled by model's create_small_model factory
    return model


def verify_cuda_compatibility(
    model: nn.Module,
    batch_size: int = 2,
    seq_len: int = 128,
    vocab_size: int = 5000,
) -> bool:
    """Run dummy forward+backward to verify CUDA compatibility and VRAM fit.

    Returns True if the model fits and gradients flow correctly.
    """
    if not torch.cuda.is_available():
        logger.warning("CUDA not available — skipping verification")
        return True

    try:
        model = model.cuda()
        model.train()
        x = torch.randint(0, vocab_size, (batch_size, seq_len), device='cuda')
        with torch.autocast(device_type='cuda', dtype=torch.bfloat16):
            logits = model(x)
            loss = logits.mean()
        loss.backward()
        MemoryManager.clear_memory()
        logger.info("CUDA verification passed — forward+backward OK")
        return True
    except torch.cuda.OutOfMemoryError:
        logger.error("OOM during verification — reduce batch size or model size")
        return False
    except Exception as e:
        logger.error("CUDA verification failed: %s", e)
        return False


# ===================================================================
# Logging helper
# ===================================================================


def log_memory_usage(logger_instance: logging.Logger) -> None:
    """Log current CUDA memory usage."""
    if torch.cuda.is_available():
        stats = MemoryManager.get_memory_usage()
        logger_instance.debug(
            "VRAM: allocated=%(allocated_mb).1fMB reserved=%(reserved_mb).1fMB max=%(max_allocated_mb).1fMB",
            stats,
        )


# ===================================================================
# Gradient accumulation
# ===================================================================


def create_gradient_accumulator(
    accumulation_steps: int,
) -> Tuple[Callable[[Optimizer], None], Callable[[int], bool]]:
    """Return (step_fn, should_step_fn) for gradient accumulation.

    Usage::

        step_fn, should_step = create_gradient_accumulator(4)
        for i, (x, y) in enumerate(dataloader):
            loss = model(x, y)
            loss.backward()
            if should_step(i):
                step_fn(optimizer)   # optimizer.step() + zero_grad()
    """

    def step_fn(optimizer: Optimizer) -> None:
        optimizer.step()
        optimizer.zero_grad()
        MemoryManager.clear_memory()

    def should_step_fn(batch_idx: int) -> bool:
        return (batch_idx + 1) % accumulation_steps == 0

    return step_fn, should_step_fn


# ===================================================================
# Backward compat aliases
# ===================================================================

optimize_for_m2 = optimize_for_rtx4050