//! Activation Checkpointing for Memory Optimization
//!
//! This module implements gradient checkpointing (also known as activation
//! checkpointing or rematerialization) to reduce memory usage during training.
//! Instead of storing all activations for the backward pass, we selectively
//! store checkpoints and recompute activations when needed.
//!
//! # Memory Savings
//! - Without checkpointing: O(n) memory for n layers
//! - With sqrt checkpointing: O(sqrt(n)) memory
//! - With selective checkpointing: O(k) memory for k checkpoints
//!
//! # Example
//! ```rust,ignore
//! use rtx_distributed::activation_checkpointing::{
//!     CheckpointedFunction, CheckpointPolicy, CheckpointManager
//! };
//!
//! let policy = CheckpointPolicy::sqrt(); // Checkpoint every sqrt(n) layers
//! let manager = CheckpointManager::new(policy);
//!
//! // Wrap forward function
//! let output = manager.checkpoint(|| {
//!     model.forward(&input)
//! });
//! ```

use crate::error::{DistributedError, Result};
use parking_lot::{Mutex, RwLock};
use serde::{Deserialize, Serialize};
use std::any::Any;
use std::collections::HashMap;
use std::sync::Arc;

// =============================================================================
// Checkpoint Policy
// =============================================================================

/// Checkpointing policy determines which layers to checkpoint
#[derive(Debug, Clone, Serialize, Deserialize)]
pub enum CheckpointPolicy {
    /// No checkpointing (store all activations)
    None,
    /// Checkpoint every N layers
    Every { interval: usize },
    /// Checkpoint at sqrt(n) intervals (optimal for memory/compute tradeoff)
    Sqrt,
    /// Checkpoint specific layers by name or index
    Selective { layers: Vec<String> },
    /// Checkpoint all layers (maximum memory savings, highest compute cost)
    All,
    /// Adaptive based on available memory.
    ///
    /// Uses a tiered heuristic based on `target_memory_mb`:
    /// - `>4096 MB` (generous): checkpoint every `ceil(sqrt(n))` layers
    /// - `1024–4096 MB` (moderate): checkpoint every other layer
    /// - `<1024 MB` (tight): checkpoint every layer
    Adaptive { target_memory_mb: usize },
    /// Custom policy with a function
    Custom { name: String },
    /// Checkpoint only attention-family layers (identified by name matching).
    ///
    /// FFN/MLP layers are intentionally skipped — they are cheaper to recompute
    /// from a saved input than to store, yielding ~40% memory savings with
    /// ~60% less recompute overhead compared to `CheckpointPolicy::All`.
    AttentionSelective {
        /// Substrings that identify attention layers in their name.
        /// A layer is checkpointed when its name contains any of these patterns.
        attention_patterns: Vec<String>,
    },
}

impl CheckpointPolicy {
    /// Create policy with no checkpointing
    pub fn none() -> Self {
        CheckpointPolicy::None
    }

    /// Checkpoint every N layers
    pub fn every(n: usize) -> Self {
        CheckpointPolicy::Every { interval: n }
    }

    /// Optimal sqrt(n) checkpointing
    pub fn sqrt() -> Self {
        CheckpointPolicy::Sqrt
    }

    /// Checkpoint specific layers
    pub fn selective(layers: Vec<String>) -> Self {
        CheckpointPolicy::Selective { layers }
    }

    /// Checkpoint all layers (maximum memory savings)
    pub fn all() -> Self {
        CheckpointPolicy::All
    }

    /// Adaptive based on memory target
    pub fn adaptive(target_memory_mb: usize) -> Self {
        CheckpointPolicy::Adaptive { target_memory_mb }
    }

    /// Checkpoint only attention-family layers.
    ///
    /// Uses the standard set of attention-related name patterns:
    /// `"attn"`, `"attention"`, `"self_attn"`, `"cross_attn"`, `"mha"`.
    /// FFN/MLP blocks are not checkpointed, cutting recompute cost by ~60%
    /// relative to `CheckpointPolicy::All` while still saving ~40% of
    /// transformer activation memory.
    pub fn attention_selective() -> Self {
        CheckpointPolicy::AttentionSelective {
            attention_patterns: vec![
                "attn".to_string(),
                "attention".to_string(),
                "self_attn".to_string(),
                "cross_attn".to_string(),
                "mha".to_string(),
            ],
        }
    }

    /// Check if a layer should be checkpointed
    pub fn should_checkpoint(
        &self,
        layer_index: usize,
        layer_name: &str,
        total_layers: usize,
    ) -> bool {
        match self {
            CheckpointPolicy::None => false,
            CheckpointPolicy::Every { interval } => layer_index % interval == 0,
            CheckpointPolicy::Sqrt => {
                let sqrt_n = (total_layers as f64).sqrt().ceil() as usize;
                layer_index % sqrt_n == 0
            }
            CheckpointPolicy::Selective { layers } => {
                layers.contains(&layer_name.to_string())
                    || layers.contains(&layer_index.to_string())
            }
            CheckpointPolicy::All => true,
            CheckpointPolicy::Adaptive { target_memory_mb } => {
                // Tiered heuristic based on memory target.
                // We cannot query live GPU memory here without a device handle, so
                // we use the target as a proxy for how aggressively to checkpoint:
                //   - Generous target (>4096 MB): checkpoint every ceil(sqrt(n)) layers
                //   - Moderate target (1024–4096 MB): checkpoint every other layer
                //   - Tight target (<1024 MB): checkpoint every layer
                let tier = match *target_memory_mb {
                    t if t > 4096 => (total_layers as f64).sqrt().ceil() as usize,
                    t if t > 1024 => 2,
                    _ => 1,
                };
                // tier is always >= 1 (sqrt of at least 1 layer rounds up to 1)
                layer_index % tier == 0
            }
            CheckpointPolicy::Custom { .. } => false, // Custom logic handled elsewhere
            CheckpointPolicy::AttentionSelective { attention_patterns } => {
                // Checkpoint the layer if its name contains any attention pattern.
                // Layer index and total_layers are intentionally not used: the
                // decision is purely structural (attention vs. non-attention).
                attention_patterns
                    .iter()
                    .any(|pat| layer_name.contains(pat.as_str()))
            }
        }
    }

    /// Estimate memory savings factor
    pub fn memory_savings_factor(&self, total_layers: usize) -> f64 {
        match self {
            CheckpointPolicy::None => 1.0,
            CheckpointPolicy::Every { interval } => *interval as f64,
            CheckpointPolicy::Sqrt => (total_layers as f64).sqrt(),
            CheckpointPolicy::Selective { layers } => {
                total_layers as f64 / (layers.len().max(1) as f64)
            }
            CheckpointPolicy::All => total_layers as f64,
            CheckpointPolicy::Adaptive { .. } => 2.0, // Rough estimate
            CheckpointPolicy::Custom { .. } => 1.5,   // Rough estimate
            CheckpointPolicy::AttentionSelective { .. } => {
                // Attention blocks constitute ~40% of transformer activation memory.
                // Checkpointing only them therefore saves ~40% of total activations,
                // expressed here as the number of layers' worth of memory freed.
                total_layers as f64 * 0.4
            }
        }
    }
}

// =============================================================================
// Activation Storage
// =============================================================================

/// Stored activation data
#[derive(Debug)]
pub struct StoredActivation {
    /// Layer name or identifier
    pub layer_id: String,
    /// Activation tensor data (serialized)
    pub data: Vec<u8>,
    /// Original shape
    pub shape: Vec<usize>,
    /// Data type (e.g., "f32", "f16")
    pub dtype: String,
    /// Memory size in bytes
    pub size_bytes: usize,
    /// Whether this is a checkpoint (will be recomputed if false)
    pub is_checkpoint: bool,
    /// Timestamp when stored
    pub stored_at: std::time::Instant,
}

impl StoredActivation {
    /// Create new stored activation
    pub fn new(layer_id: String, data: Vec<u8>, shape: Vec<usize>, dtype: String) -> Self {
        let size_bytes = data.len();
        Self {
            layer_id,
            data,
            shape,
            dtype,
            size_bytes,
            is_checkpoint: false,
            stored_at: std::time::Instant::now(),
        }
    }

    /// Mark as checkpoint
    pub fn as_checkpoint(mut self) -> Self {
        self.is_checkpoint = true;
        self
    }
}

/// Activation storage manager
#[derive(Debug)]
pub struct ActivationStorage {
    /// Stored activations by layer ID
    activations: HashMap<String, StoredActivation>,
    /// Total memory used
    total_memory: usize,
    /// Memory limit
    memory_limit: usize,
    /// Statistics
    stats: ActivationStats,
}

/// Activation storage statistics
#[derive(Debug, Default, Clone)]
pub struct ActivationStats {
    /// Number of activations stored
    pub stored_count: usize,
    /// Number of activations recomputed
    pub recomputed_count: usize,
    /// Total bytes stored
    pub bytes_stored: usize,
    /// Bytes saved by checkpointing
    pub bytes_saved: usize,
    /// Peak memory usage
    pub peak_memory: usize,
}

impl ActivationStorage {
    /// Create new activation storage
    pub fn new(memory_limit: usize) -> Self {
        Self {
            activations: HashMap::new(),
            total_memory: 0,
            memory_limit,
            stats: ActivationStats::default(),
        }
    }

    /// Store an activation
    pub fn store(&mut self, activation: StoredActivation) -> Result<()> {
        let size = activation.size_bytes;

        // Check memory limit
        if self.total_memory + size > self.memory_limit {
            // Try to evict non-checkpoint activations
            self.evict_until_free(size)?;
        }

        let layer_id = activation.layer_id.clone();
        self.total_memory += size;
        self.stats.stored_count += 1;
        self.stats.bytes_stored += size;
        self.stats.peak_memory = self.stats.peak_memory.max(self.total_memory);

        self.activations.insert(layer_id, activation);
        Ok(())
    }

    /// Get an activation
    pub fn get(&self, layer_id: &str) -> Option<&StoredActivation> {
        self.activations.get(layer_id)
    }

    /// Remove an activation
    pub fn remove(&mut self, layer_id: &str) -> Option<StoredActivation> {
        if let Some(activation) = self.activations.remove(layer_id) {
            self.total_memory -= activation.size_bytes;
            Some(activation)
        } else {
            None
        }
    }

    /// Clear all activations
    pub fn clear(&mut self) {
        self.activations.clear();
        self.total_memory = 0;
    }

    /// Evict non-checkpoint activations until enough space is free
    fn evict_until_free(&mut self, needed: usize) -> Result<()> {
        let mut to_remove = Vec::new();

        for (id, activation) in &self.activations {
            if !activation.is_checkpoint {
                to_remove.push(id.clone());
                self.stats.bytes_saved += activation.size_bytes;
            }
        }

        for id in to_remove {
            self.remove(&id);
            if self.memory_limit - self.total_memory >= needed {
                break;
            }
        }

        if self.memory_limit - self.total_memory < needed {
            return Err(DistributedError::configuration(format!(
                "Cannot free enough memory: need {} bytes, have {} available",
                needed,
                self.memory_limit - self.total_memory
            )));
        }

        Ok(())
    }

    /// Get current memory usage
    pub fn memory_usage(&self) -> usize {
        self.total_memory
    }

    /// Get statistics
    pub fn stats(&self) -> &ActivationStats {
        &self.stats
    }
}

// =============================================================================
// Checkpoint Context
// =============================================================================

/// Context for checkpointed computation
#[derive(Debug)]
pub struct CheckpointContext {
    /// Forward function inputs (saved for recomputation)
    pub inputs: Vec<Box<dyn Any + Send + Sync>>,
    /// Saved tensors for backward
    pub saved_tensors: Vec<Box<dyn Any + Send + Sync>>,
    /// Whether we need to recompute on backward
    pub needs_recompute: bool,
    /// Layer identifier
    pub layer_id: String,
}

impl CheckpointContext {
    /// Create new checkpoint context
    pub fn new(layer_id: String) -> Self {
        Self {
            inputs: Vec::new(),
            saved_tensors: Vec::new(),
            needs_recompute: false,
            layer_id,
        }
    }

    /// Save input for potential recomputation
    pub fn save_input<T: Any + Send + Sync + 'static>(&mut self, input: T) {
        self.inputs.push(Box::new(input));
    }

    /// Save tensor for backward pass
    pub fn save_for_backward<T: Any + Send + Sync + 'static>(&mut self, tensor: T) {
        self.saved_tensors.push(Box::new(tensor));
    }

    /// Mark that this context needs recomputation
    pub fn mark_recompute(&mut self) {
        self.needs_recompute = true;
    }
}

// =============================================================================
// Checkpointed Segment
// =============================================================================

/// A segment of the model that can be checkpointed
#[derive(Debug)]
pub struct CheckpointedSegment {
    /// Segment identifier
    pub id: String,
    /// Layers in this segment
    pub layers: Vec<String>,
    /// Input tensors (saved for recomputation)
    inputs: RwLock<Vec<Vec<u8>>>,
    /// Output tensors
    outputs: RwLock<Vec<Vec<u8>>>,
    /// Whether outputs are currently stored
    outputs_stored: RwLock<bool>,
}

impl CheckpointedSegment {
    /// Create new checkpointed segment
    pub fn new(id: String, layers: Vec<String>) -> Self {
        Self {
            id,
            layers,
            inputs: RwLock::new(Vec::new()),
            outputs: RwLock::new(Vec::new()),
            outputs_stored: RwLock::new(false),
        }
    }

    /// Save inputs for potential recomputation
    pub fn save_inputs(&self, inputs: Vec<Vec<u8>>) {
        *self.inputs.write() = inputs;
    }

    /// Save outputs (only done for checkpoints)
    pub fn save_outputs(&self, outputs: Vec<Vec<u8>>) {
        *self.outputs.write() = outputs;
        *self.outputs_stored.write() = true;
    }

    /// Get saved inputs
    pub fn get_inputs(&self) -> Vec<Vec<u8>> {
        self.inputs.read().clone()
    }

    /// Get saved outputs
    pub fn get_outputs(&self) -> Option<Vec<Vec<u8>>> {
        if *self.outputs_stored.read() {
            Some(self.outputs.read().clone())
        } else {
            None
        }
    }

    /// Clear saved outputs (after recomputation)
    pub fn clear_outputs(&self) {
        self.outputs.write().clear();
        *self.outputs_stored.write() = false;
    }

    /// Check if outputs are available
    pub fn has_outputs(&self) -> bool {
        *self.outputs_stored.read()
    }
}

// =============================================================================
// Checkpoint Manager
// =============================================================================

/// Manager for activation checkpointing
#[derive(Debug)]
pub struct ActivationCheckpointManager {
    /// Checkpointing policy
    policy: CheckpointPolicy,
    /// Activation storage
    storage: Mutex<ActivationStorage>,
    /// Registered segments
    segments: RwLock<HashMap<String, Arc<CheckpointedSegment>>>,
    /// Current layer index (for policy decisions)
    current_layer: RwLock<usize>,
    /// Total number of layers
    total_layers: RwLock<usize>,
    /// Recomputation count
    recompute_count: RwLock<usize>,
    /// Whether checkpointing is enabled
    enabled: RwLock<bool>,
}

impl ActivationCheckpointManager {
    /// Create new checkpoint manager
    pub fn new(policy: CheckpointPolicy) -> Self {
        Self {
            policy,
            storage: Mutex::new(ActivationStorage::new(8 * 1024 * 1024 * 1024)), // 8GB default
            segments: RwLock::new(HashMap::new()),
            current_layer: RwLock::new(0),
            total_layers: RwLock::new(0),
            recompute_count: RwLock::new(0),
            enabled: RwLock::new(true),
        }
    }

    /// Create with memory limit
    pub fn with_memory_limit(policy: CheckpointPolicy, memory_limit: usize) -> Self {
        Self {
            policy,
            storage: Mutex::new(ActivationStorage::new(memory_limit)),
            segments: RwLock::new(HashMap::new()),
            current_layer: RwLock::new(0),
            total_layers: RwLock::new(0),
            recompute_count: RwLock::new(0),
            enabled: RwLock::new(true),
        }
    }

    /// Set total number of layers
    pub fn set_total_layers(&self, n: usize) {
        *self.total_layers.write() = n;
    }

    /// Enable or disable checkpointing
    pub fn set_enabled(&self, enabled: bool) {
        *self.enabled.write() = enabled;
    }

    /// Check if checkpointing is enabled
    pub fn is_enabled(&self) -> bool {
        *self.enabled.read()
    }

    /// Register a segment for checkpointing
    pub fn register_segment(&self, segment: CheckpointedSegment) -> Arc<CheckpointedSegment> {
        let segment = Arc::new(segment);
        self.segments
            .write()
            .insert(segment.id.clone(), segment.clone());
        segment
    }

    /// Get a registered segment
    pub fn get_segment(&self, id: &str) -> Option<Arc<CheckpointedSegment>> {
        self.segments.read().get(id).cloned()
    }

    /// Check if current layer should be checkpointed
    pub fn should_checkpoint_current(&self, layer_name: &str) -> bool {
        if !*self.enabled.read() {
            return false;
        }

        let layer_index = *self.current_layer.read();
        let total = *self.total_layers.read();

        self.policy
            .should_checkpoint(layer_index, layer_name, total)
    }

    /// Advance to next layer
    pub fn next_layer(&self) {
        *self.current_layer.write() += 1;
    }

    /// Reset layer counter
    pub fn reset_layers(&self) {
        *self.current_layer.write() = 0;
    }

    /// Store activation
    pub fn store_activation(
        &self,
        layer_id: String,
        data: Vec<u8>,
        shape: Vec<usize>,
        dtype: String,
    ) -> Result<()> {
        let activation = StoredActivation::new(layer_id, data, shape, dtype);
        self.storage.lock().store(activation)
    }

    /// Store checkpoint activation
    pub fn store_checkpoint(
        &self,
        layer_id: String,
        data: Vec<u8>,
        shape: Vec<usize>,
        dtype: String,
    ) -> Result<()> {
        let activation = StoredActivation::new(layer_id, data, shape, dtype).as_checkpoint();
        self.storage.lock().store(activation)
    }

    /// Get stored activation
    pub fn get_activation(&self, layer_id: &str) -> Option<(Vec<u8>, Vec<usize>, String)> {
        self.storage
            .lock()
            .get(layer_id)
            .map(|a| (a.data.clone(), a.shape.clone(), a.dtype.clone()))
    }

    /// Record a recomputation
    pub fn record_recompute(&self) {
        *self.recompute_count.write() += 1;
        self.storage.lock().stats.recomputed_count += 1;
    }

    /// Clear all stored activations
    pub fn clear(&self) {
        self.storage.lock().clear();
        self.reset_layers();
    }

    /// Get memory usage
    pub fn memory_usage(&self) -> usize {
        self.storage.lock().memory_usage()
    }

    /// Get statistics
    pub fn stats(&self) -> CheckpointStats {
        let storage_stats = self.storage.lock().stats().clone();
        CheckpointStats {
            activations_stored: storage_stats.stored_count,
            activations_recomputed: storage_stats.recomputed_count,
            bytes_stored: storage_stats.bytes_stored,
            bytes_saved: storage_stats.bytes_saved,
            peak_memory: storage_stats.peak_memory,
            recompute_count: *self.recompute_count.read(),
            memory_savings_factor: self.policy.memory_savings_factor(*self.total_layers.read()),
        }
    }

    /// Get estimated memory savings
    pub fn estimated_memory_savings(&self) -> f64 {
        let total_layers = *self.total_layers.read();
        if total_layers == 0 {
            return 1.0;
        }
        self.policy.memory_savings_factor(total_layers)
    }
}

/// Checkpoint statistics
#[derive(Debug, Clone)]
pub struct CheckpointStats {
    /// Number of activations stored
    pub activations_stored: usize,
    /// Number of activations recomputed
    pub activations_recomputed: usize,
    /// Bytes stored
    pub bytes_stored: usize,
    /// Bytes saved by checkpointing
    pub bytes_saved: usize,
    /// Peak memory usage
    pub peak_memory: usize,
    /// Number of recomputations
    pub recompute_count: usize,
    /// Memory savings factor
    pub memory_savings_factor: f64,
}

// =============================================================================
// Checkpoint Function Wrapper
// =============================================================================

/// Trait for functions that can be checkpointed
pub trait CheckpointedFunction {
    /// Output type
    type Output;

    /// Run the function with checkpointing
    fn run(&self, ctx: &mut CheckpointContext) -> Self::Output;

    /// Recompute the function (for backward pass)
    fn recompute(&self, ctx: &CheckpointContext) -> Self::Output;
}

/// Wrapper to checkpoint a closure
pub struct CheckpointedClosure<F, O>
where
    F: Fn() -> O,
{
    /// The function to checkpoint
    func: F,
    /// Layer ID for this checkpoint
    layer_id: String,
}

impl<F, O> CheckpointedClosure<F, O>
where
    F: Fn() -> O,
{
    /// Create new checkpointed closure
    pub fn new(layer_id: String, func: F) -> Self {
        Self { func, layer_id }
    }

    /// Execute the closure
    pub fn execute(&self) -> O {
        (self.func)()
    }

    /// Get layer ID
    pub fn layer_id(&self) -> &str {
        &self.layer_id
    }
}

// =============================================================================
// Shared Checkpoint Manager
// =============================================================================

/// Thread-safe shared checkpoint manager
pub type SharedActivationCheckpointManager = Arc<ActivationCheckpointManager>;

/// Create a shared checkpoint manager
pub fn shared_activation_checkpoint_manager(
    policy: CheckpointPolicy,
) -> SharedActivationCheckpointManager {
    Arc::new(ActivationCheckpointManager::new(policy))
}

/// Create a shared checkpoint manager with memory limit
pub fn shared_activation_checkpoint_manager_with_limit(
    policy: CheckpointPolicy,
    memory_limit: usize,
) -> SharedActivationCheckpointManager {
    Arc::new(ActivationCheckpointManager::with_memory_limit(
        policy,
        memory_limit,
    ))
}

// =============================================================================
// Utility Functions
// =============================================================================

/// Calculate optimal checkpoint interval for given memory budget
pub fn optimal_checkpoint_interval(
    num_layers: usize,
    activation_size_per_layer: usize,
    memory_budget: usize,
) -> usize {
    // With checkpointing every k layers:
    // Memory = k * activation_size (for storing between checkpoints)
    // Recomputation cost = O(k) per backward step

    // Optimal is sqrt(n) checkpoints for balanced memory/compute
    let sqrt_n = (num_layers as f64).sqrt().ceil() as usize;

    // Adjust based on memory budget
    let max_interval = memory_budget / activation_size_per_layer.max(1);

    sqrt_n.min(max_interval).max(1)
}

/// Estimate memory savings for a given configuration
pub fn estimate_memory_savings(
    num_layers: usize,
    activation_size_per_layer: usize,
    checkpoint_interval: usize,
) -> (usize, usize) {
    // Without checkpointing
    let without = num_layers * activation_size_per_layer;

    // With checkpointing: store only checkpoint interval activations
    let num_checkpoints = num_layers / checkpoint_interval.max(1);
    let with = checkpoint_interval * activation_size_per_layer
        + num_checkpoints * activation_size_per_layer;

    (without, with)
}

// =============================================================================
// Memory-Aware Checkpointer
// =============================================================================

/// Applies adaptive checkpointing decisions based on live memory tracking.
///
/// `MemoryAwareCheckpointer` wraps a [`CheckpointPolicy`] and augments it with
/// a running estimate of activation memory consumed during a forward pass.
/// When that estimate exceeds `target_bytes`, every layer is checkpointed
/// regardless of the underlying policy, preventing OOM conditions.
///
/// The internal counter uses `AtomicUsize` so the checkpointer can be shared
/// across threads without locking overhead.
///
/// # Example
/// ```rust,ignore
/// use rtx_distributed::activation_checkpointing::{
///     CheckpointPolicy, MemoryAwareCheckpointer
/// };
///
/// let checkpointer = MemoryAwareCheckpointer::new(
///     CheckpointPolicy::attention_selective(),
///     24,     // 24 transformer layers
///     4096,   // 4 GB target
/// );
///
/// // During forward pass:
/// checkpointer.record_activation(64 * 1024 * 1024); // 64 MB for this layer
/// let do_ckpt = checkpointer.should_checkpoint(3, "self_attn_0");
/// ```
pub struct MemoryAwareCheckpointer {
    /// The base policy used when memory is below the target.
    pub policy: CheckpointPolicy,
    /// Total number of layers in the model (used by the base policy).
    pub total_layers: usize,
    /// Estimated bytes of activation memory currently live.
    /// Incremented by [`record_activation`] and decremented by
    /// [`free_activation`].  Uses `Relaxed` ordering — exact byte counts
    /// are not required for correctness, only approximate pressure detection.
    estimated_memory_bytes: std::sync::atomic::AtomicUsize,
    /// Memory target in bytes. When `estimated_memory_bytes` exceeds this,
    /// the checkpointer forces checkpointing on every layer.
    pub target_bytes: usize,
}

impl std::fmt::Debug for MemoryAwareCheckpointer {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("MemoryAwareCheckpointer")
            .field("policy", &self.policy)
            .field("total_layers", &self.total_layers)
            .field(
                "estimated_memory_bytes",
                &self
                    .estimated_memory_bytes
                    .load(std::sync::atomic::Ordering::Relaxed),
            )
            .field("target_bytes", &self.target_bytes)
            .finish()
    }
}

impl MemoryAwareCheckpointer {
    /// Create a new `MemoryAwareCheckpointer`.
    ///
    /// # Parameters
    /// - `policy`: Base checkpointing policy applied when memory is within budget.
    /// - `total_layers`: Number of layers in the model.
    /// - `target_mb`: Memory budget in mebibytes (MiB = 1024 * 1024 bytes).
    pub fn new(policy: CheckpointPolicy, total_layers: usize, target_mb: usize) -> Self {
        Self {
            policy,
            total_layers,
            estimated_memory_bytes: std::sync::atomic::AtomicUsize::new(0),
            target_bytes: target_mb * 1024 * 1024,
        }
    }

    /// Record that a layer consumed `bytes` of activation memory.
    ///
    /// Call this after each layer's forward pass to keep the pressure
    /// estimate current.
    pub fn record_activation(&self, bytes: usize) {
        self.estimated_memory_bytes
            .fetch_add(bytes, std::sync::atomic::Ordering::Relaxed);
    }

    /// Free tracked memory when activations are released (e.g. after a
    /// segment's backward pass completes).
    ///
    /// Saturates at zero; cannot underflow.
    pub fn free_activation(&self, bytes: usize) {
        // Clamp the subtraction so we never wrap on underflow.
        let current = self
            .estimated_memory_bytes
            .load(std::sync::atomic::Ordering::Relaxed);
        let to_sub = bytes.min(current);
        self.estimated_memory_bytes
            .fetch_sub(to_sub, std::sync::atomic::Ordering::Relaxed);
    }

    /// Decide whether to checkpoint this layer.
    ///
    /// If `estimated_memory_bytes > target_bytes`, returns `true` unconditionally
    /// (memory-pressure override). Otherwise delegates to `self.policy`.
    pub fn should_checkpoint(&self, layer_index: usize, layer_name: &str) -> bool {
        let current = self
            .estimated_memory_bytes
            .load(std::sync::atomic::Ordering::Relaxed);
        if current > self.target_bytes {
            // Memory pressure: checkpoint everything to stay under target.
            return true;
        }
        self.policy
            .should_checkpoint(layer_index, layer_name, self.total_layers)
    }

    /// Current memory pressure as a fraction of the target (0.0 = empty, 1.0 = at target).
    ///
    /// Values above 1.0 indicate the memory budget has been exceeded.
    pub fn memory_pressure(&self) -> f32 {
        let current = self
            .estimated_memory_bytes
            .load(std::sync::atomic::Ordering::Relaxed);
        current as f32 / self.target_bytes as f32
    }
}

// =============================================================================
// Tests
// =============================================================================

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_checkpoint_policy_none() {
        let policy = CheckpointPolicy::none();
        assert!(!policy.should_checkpoint(0, "layer0", 10));
        assert!(!policy.should_checkpoint(5, "layer5", 10));
    }

    #[test]
    fn test_checkpoint_policy_every() {
        let policy = CheckpointPolicy::every(3);
        assert!(policy.should_checkpoint(0, "layer0", 10));
        assert!(!policy.should_checkpoint(1, "layer1", 10));
        assert!(!policy.should_checkpoint(2, "layer2", 10));
        assert!(policy.should_checkpoint(3, "layer3", 10));
        assert!(policy.should_checkpoint(6, "layer6", 10));
    }

    #[test]
    fn test_checkpoint_policy_sqrt() {
        let policy = CheckpointPolicy::sqrt();
        // For 16 layers, sqrt is 4, so checkpoint every 4 layers
        assert!(policy.should_checkpoint(0, "layer0", 16));
        assert!(!policy.should_checkpoint(1, "layer1", 16));
        assert!(policy.should_checkpoint(4, "layer4", 16));
        assert!(policy.should_checkpoint(8, "layer8", 16));
    }

    #[test]
    fn test_checkpoint_policy_selective() {
        let policy = CheckpointPolicy::selective(vec!["attention".to_string(), "5".to_string()]);
        assert!(policy.should_checkpoint(0, "attention", 10));
        assert!(policy.should_checkpoint(5, "mlp", 10));
        assert!(!policy.should_checkpoint(3, "norm", 10));
    }

    #[test]
    fn test_checkpoint_policy_all() {
        let policy = CheckpointPolicy::all();
        for i in 0..10 {
            assert!(policy.should_checkpoint(i, &format!("layer{}", i), 10));
        }
    }

    #[test]
    fn test_memory_savings_factor() {
        let none = CheckpointPolicy::none();
        assert_eq!(none.memory_savings_factor(10), 1.0);

        let every_2 = CheckpointPolicy::every(2);
        assert_eq!(every_2.memory_savings_factor(10), 2.0);

        let sqrt = CheckpointPolicy::sqrt();
        let factor = sqrt.memory_savings_factor(16);
        assert!((factor - 4.0).abs() < 0.01); // sqrt(16) = 4
    }

    #[test]
    fn test_activation_storage() {
        let mut storage = ActivationStorage::new(1024 * 1024); // 1MB limit

        let activation = StoredActivation::new(
            "layer1".to_string(),
            vec![0u8; 1024],
            vec![32, 32],
            "f32".to_string(),
        );

        assert!(storage.store(activation).is_ok());
        assert_eq!(storage.memory_usage(), 1024);

        let retrieved = storage.get("layer1");
        assert!(retrieved.is_some());
        assert_eq!(retrieved.unwrap().shape, vec![32, 32]);
    }

    #[test]
    fn test_activation_storage_eviction() {
        let mut storage = ActivationStorage::new(2048); // 2KB limit

        // Store non-checkpoint activation
        let act1 = StoredActivation::new(
            "layer1".to_string(),
            vec![0u8; 1024],
            vec![32, 32],
            "f32".to_string(),
        );
        storage.store(act1).unwrap();

        // Store checkpoint activation
        let act2 = StoredActivation::new(
            "layer2".to_string(),
            vec![0u8; 512],
            vec![16, 32],
            "f32".to_string(),
        )
        .as_checkpoint();
        storage.store(act2).unwrap();

        // Try to store another large activation - should evict non-checkpoint
        let act3 = StoredActivation::new(
            "layer3".to_string(),
            vec![0u8; 1024],
            vec![32, 32],
            "f32".to_string(),
        );
        storage.store(act3).unwrap();

        // layer1 should be evicted, layer2 (checkpoint) should remain
        assert!(storage.get("layer1").is_none());
        assert!(storage.get("layer2").is_some());
        assert!(storage.get("layer3").is_some());
    }

    #[test]
    fn test_checkpoint_manager() {
        let manager = ActivationCheckpointManager::new(CheckpointPolicy::every(2));
        manager.set_total_layers(10);

        assert!(manager.is_enabled());
        assert!(manager.should_checkpoint_current("layer0")); // Index 0

        manager.next_layer();
        assert!(!manager.should_checkpoint_current("layer1")); // Index 1

        manager.next_layer();
        assert!(manager.should_checkpoint_current("layer2")); // Index 2
    }

    #[test]
    fn test_checkpoint_manager_disable() {
        let manager = ActivationCheckpointManager::new(CheckpointPolicy::all());
        manager.set_total_layers(10);

        assert!(manager.should_checkpoint_current("layer0"));

        manager.set_enabled(false);
        assert!(!manager.should_checkpoint_current("layer0"));

        manager.set_enabled(true);
        assert!(manager.should_checkpoint_current("layer0"));
    }

    #[test]
    fn test_checkpointed_segment() {
        let segment = CheckpointedSegment::new(
            "segment1".to_string(),
            vec!["layer1".to_string(), "layer2".to_string()],
        );

        assert!(!segment.has_outputs());

        segment.save_inputs(vec![vec![1, 2, 3]]);
        segment.save_outputs(vec![vec![4, 5, 6]]);

        assert!(segment.has_outputs());
        assert_eq!(segment.get_inputs(), vec![vec![1, 2, 3]]);
        assert_eq!(segment.get_outputs(), Some(vec![vec![4, 5, 6]]));

        segment.clear_outputs();
        assert!(!segment.has_outputs());
    }

    #[test]
    fn test_optimal_checkpoint_interval() {
        // For 16 layers, optimal is sqrt(16) = 4
        let interval = optimal_checkpoint_interval(16, 1024 * 1024, 100 * 1024 * 1024);
        assert!(interval >= 4);

        // With very limited memory, should reduce interval
        let interval = optimal_checkpoint_interval(16, 1024 * 1024, 2 * 1024 * 1024);
        assert!(interval <= 2);
    }

    #[test]
    fn test_estimate_memory_savings() {
        let (without, with) = estimate_memory_savings(16, 1024 * 1024, 4);

        // Without: 16 * 1MB = 16MB
        assert_eq!(without, 16 * 1024 * 1024);

        // With: 4 * 1MB (interval) + 4 * 1MB (checkpoints) = 8MB
        // Actually: interval activations + checkpoint activations
        assert!(with < without);
    }

    #[test]
    fn test_checkpoint_context() {
        let mut ctx = CheckpointContext::new("layer1".to_string());

        ctx.save_input(vec![1.0f32, 2.0, 3.0]);
        ctx.save_for_backward(vec![4.0f32, 5.0, 6.0]);
        ctx.mark_recompute();

        assert!(ctx.needs_recompute);
        assert_eq!(ctx.inputs.len(), 1);
        assert_eq!(ctx.saved_tensors.len(), 1);
    }

    #[test]
    fn test_checkpoint_stats() {
        let manager = ActivationCheckpointManager::new(CheckpointPolicy::sqrt());
        manager.set_total_layers(16);

        manager
            .store_activation(
                "layer0".to_string(),
                vec![0u8; 1024],
                vec![32, 32],
                "f32".to_string(),
            )
            .unwrap();

        manager.record_recompute();
        manager.record_recompute();

        let stats = manager.stats();
        assert_eq!(stats.activations_stored, 1);
        assert_eq!(stats.recompute_count, 2);
        assert!((stats.memory_savings_factor - 4.0).abs() < 0.01);
    }

    // -------------------------------------------------------------------------
    // AttentionSelective tests
    // -------------------------------------------------------------------------

    #[test]
    fn test_attention_selective_matches_attn_layers() {
        let policy = CheckpointPolicy::attention_selective();
        // Standard attention name patterns must be recognised.
        assert!(policy.should_checkpoint(0, "self_attn", 12));
        assert!(policy.should_checkpoint(0, "cross_attention_layer", 12));
        assert!(policy.should_checkpoint(0, "mha_block", 12));
        // FFN / MLP layers must not be checkpointed.
        assert!(!policy.should_checkpoint(0, "ffn_layer", 12));
        assert!(!policy.should_checkpoint(0, "mlp_dense", 12));
        assert!(!policy.should_checkpoint(0, "layer_norm", 12));
    }

    #[test]
    fn test_attention_selective_memory_savings_factor() {
        let policy = CheckpointPolicy::attention_selective();
        let factor = policy.memory_savings_factor(12);
        // 12 * 0.4 = 4.8
        assert!((factor - 4.8).abs() < 1e-6, "expected 4.8, got {factor}");
    }

    #[test]
    fn test_custom_attention_patterns() {
        let policy = CheckpointPolicy::AttentionSelective {
            attention_patterns: vec!["transformer_block".to_string()],
        };
        assert!(policy.should_checkpoint(0, "transformer_block_0", 12));
        assert!(!policy.should_checkpoint(0, "linear_layer", 12));
    }

    // -------------------------------------------------------------------------
    // Adaptive (fixed) tests
    // -------------------------------------------------------------------------

    #[test]
    fn test_adaptive_tight_memory_checkpoints_all() {
        let policy = CheckpointPolicy::Adaptive {
            target_memory_mb: 512,
        };
        // tier == 1 → every layer is checkpointed
        for i in 0..8 {
            assert!(
                policy.should_checkpoint(i, "any_layer", 8),
                "layer {i} should be checkpointed under tight target"
            );
        }
    }

    #[test]
    fn test_adaptive_moderate_memory_checkpoints_every_other() {
        let policy = CheckpointPolicy::Adaptive {
            target_memory_mb: 2048,
        };
        // tier == 2 → even indices only
        for i in 0..8usize {
            let expected = i % 2 == 0;
            assert_eq!(
                policy.should_checkpoint(i, "layer", 8),
                expected,
                "layer {i}: expected={expected}"
            );
        }
    }

    #[test]
    fn test_adaptive_generous_memory_checkpoints_sqrt() {
        let policy = CheckpointPolicy::Adaptive {
            target_memory_mb: 8192,
        };
        let total = 16usize;
        // ceil(sqrt(16)) = 4 → checkpoint at indices 0, 4, 8, 12
        assert!(policy.should_checkpoint(0, "layer", total));
        assert!(!policy.should_checkpoint(1, "layer", total));
        assert!(!policy.should_checkpoint(2, "layer", total));
        assert!(!policy.should_checkpoint(3, "layer", total));
        assert!(policy.should_checkpoint(4, "layer", total));
        assert!(!policy.should_checkpoint(5, "layer", total));
        assert!(policy.should_checkpoint(8, "layer", total));
        assert!(policy.should_checkpoint(12, "layer", total));
        assert!(!policy.should_checkpoint(13, "layer", total));
    }

    // -------------------------------------------------------------------------
    // MemoryAwareCheckpointer tests
    // -------------------------------------------------------------------------

    #[test]
    fn test_memory_aware_checkpointer_pressure() {
        let checkpointer = MemoryAwareCheckpointer::new(
            CheckpointPolicy::attention_selective(),
            12,
            1024, // 1 GiB target
        );
        // Freshly created — no memory recorded yet.
        assert!(
            checkpointer.memory_pressure() < 0.01,
            "initial pressure should be near zero"
        );

        checkpointer.record_activation(512 * 1024 * 1024); // 512 MiB
        let pressure = checkpointer.memory_pressure();
        assert!(
            pressure > 0.49 && pressure < 0.51,
            "pressure after 512 MiB / 1024 MiB should be ~0.5, got {pressure}"
        );
    }

    #[test]
    fn test_memory_aware_checkpointer_fallback_under_pressure() {
        let checkpointer = MemoryAwareCheckpointer::new(
            CheckpointPolicy::None, // normally checkpoints nothing
            12,
            256, // 256 MiB target
        );
        checkpointer.record_activation(300 * 1024 * 1024); // exceeds target
        // Memory pressure override must force checkpointing even though policy is None.
        assert!(
            checkpointer.should_checkpoint(0, "ffn_layer"),
            "should force-checkpoint under memory pressure regardless of policy"
        );
    }

    #[test]
    fn test_memory_aware_free_reduces_pressure() {
        let checkpointer = MemoryAwareCheckpointer::new(CheckpointPolicy::None, 12, 1024);
        checkpointer.record_activation(200 * 1024 * 1024);
        checkpointer.free_activation(200 * 1024 * 1024);
        assert!(
            checkpointer.memory_pressure() < 0.01,
            "pressure should return to near zero after freeing all recorded memory"
        );
    }

    #[test]
    fn test_memory_aware_free_does_not_underflow() {
        // Freeing more than was recorded must saturate at zero, not wrap.
        let checkpointer = MemoryAwareCheckpointer::new(CheckpointPolicy::None, 4, 512);
        checkpointer.record_activation(10 * 1024 * 1024);
        checkpointer.free_activation(50 * 1024 * 1024); // larger than recorded
        assert_eq!(
            checkpointer
                .estimated_memory_bytes
                .load(std::sync::atomic::Ordering::Relaxed),
            0,
            "counter must not underflow to usize::MAX"
        );
    }

    #[test]
    fn test_memory_aware_checkpointer_delegates_to_policy() {
        // Below pressure threshold the policy should be respected.
        let checkpointer = MemoryAwareCheckpointer::new(
            CheckpointPolicy::attention_selective(),
            12,
            4096, // large target — won't be exceeded
        );
        checkpointer.record_activation(1024); // trivial amount
        // Attention layer — policy says checkpoint.
        assert!(checkpointer.should_checkpoint(0, "self_attn"));
        // FFN layer — policy says don't checkpoint.
        assert!(!checkpointer.should_checkpoint(0, "ffn_dense"));
    }

    #[test]
    fn test_memory_aware_checkpointer_debug_impl() {
        // Ensure the manual Debug impl doesn't panic.
        let c = MemoryAwareCheckpointer::new(CheckpointPolicy::None, 8, 512);
        let _ = format!("{c:?}");
    }
}