"""
Advanced Gradient Checkpointing Techniques

Sophisticated checkpointing strategies for memory optimization:
- Selective gradient checkpointing
- Adaptive checkpointing based on memory pressure
- Dynamic activation offloading
"""

from collections.abc import Callable

import torch
import torch.nn as nn


class SelectiveGradientCheckpointing:
    """
    Selective checkpointing based on layer importance and memory usage
    """

    def __init__(self, importance_threshold: float = 0.5):
        self.importance_threshold = importance_threshold
        self.layer_importance = {}

    def should_checkpoint(self, layer_name: str, memory_pressure: float) -> bool:
        """Decide whether to checkpoint a layer"""
        importance = self.layer_importance.get(layer_name, 1.0)
        return importance < self.importance_threshold or memory_pressure > 0.8

    def update_importance(self, layer_name: str, importance: float):
        """Update layer importance score"""
        self.layer_importance[layer_name] = importance


class AdaptiveCheckpointing:
    """Adaptive checkpointing based on runtime conditions"""

    def __init__(self):
        self.enabled = True

    def forward(self, func: Callable, *args, **kwargs):
        """Apply adaptive checkpointing"""
        if self.enabled and torch.cuda.is_available():
            return torch.utils.checkpoint.checkpoint(func, *args, use_reentrant=False, **kwargs)
        else:
            return func(*args, **kwargs)


class MemoryEfficientBackprop:
    """Memory-efficient backpropagation"""

    def __init__(self):
        self.enabled = True

    def apply(self, module: nn.Module):
        """Apply memory-efficient backprop to module"""
        if hasattr(module, 'gradient_checkpointing_enable'):
            module.gradient_checkpointing_enable()


class DynamicActivationOffloading:
    """Dynamic activation offloading"""

    def __init__(self, offload_device: str = "cpu"):
        self.offload_device = torch.device(offload_device)

    def offload_activations(self, activations: torch.Tensor) -> torch.Tensor:
        """Offload activations to specified device"""
        return activations.to(self.offload_device)

    def reload_activations(self, activations: torch.Tensor, target_device: torch.device) -> torch.Tensor:
        """Reload activations to target device"""
        return activations.to(target_device)
