//! LoRA-enabled Qwen3 model architecture.
//!
//! Implements Qwen3 with LoRA adapters on attention projections for efficient fine-tuning.
//! Key Qwen3-specific features:
//! - Q/K normalization before RoPE
//! - Higher default vocab size (151936)
//! - Higher default rope_theta (1_000_000)
//! - Metal FlashAttention integration for O(n) memory training

use std::collections::HashMap;
use std::rc::Rc;
use std::sync::Once;
use std::sync::atomic::{AtomicUsize, Ordering};

use mlx_rs::{
    Array,
    builder::Builder,
    error::Exception,
    module::{ModuleParamMut, ModuleParamRef, ModuleParameters, Param},
    nested::NestedValue,
    nn,
};

use pmetal_core::LoraConfig;
use pmetal_mlx::gradient_checkpoint::CheckpointConfig;
use pmetal_mlx::kernels::{
    AttentionMaskType, FusedAttentionConfig, differentiable_attention, fused_sdpa,
    get_training_context, rope::apply_rope,
};
use pmetal_mlx::kv_cache::{KVCache, KVCacheConfig};
use pmetal_models::ModelConfig;
use pmetal_models::architectures::qwen3::Qwen3Config;

use crate::{LoraError, LoraLinear};

/// Global counter for unique layer IDs (for FlashAttention caching).
static LAYER_ID_COUNTER: AtomicUsize = AtomicUsize::new(0);

/// Reset the global layer ID counter to 0.
///
/// Must be called at model initialization to ensure layer IDs are assigned
/// consistently from 0 for each new model instance.
pub fn reset_layer_ids() {
    LAYER_ID_COUNTER.store(0, Ordering::SeqCst);
}

/// LoRA-enabled attention layer for Qwen3.
///
/// Applies LoRA to q_proj, k_proj, v_proj, and o_proj.
/// Includes Qwen3-specific Q/K normalization before RoPE.
/// Supports Metal FlashAttention for O(n) memory training.
#[derive(Debug)]
pub struct Qwen3LoraAttention {
    /// Number of attention heads.
    pub n_heads: i32,
    /// Number of KV heads.
    pub n_kv_heads: i32,
    /// Head dimension.
    pub head_dim: i32,
    /// Attention scale factor.
    pub scale: f32,
    /// Unique layer ID for FlashAttention caching.
    pub layer_id: usize,

    /// Query projection with LoRA.
    pub q_proj: LoraLinear,
    /// Key projection with LoRA.
    pub k_proj: LoraLinear,
    /// Value projection with LoRA.
    pub v_proj: LoraLinear,
    /// Output projection with LoRA.
    pub o_proj: LoraLinear,
    /// Query normalization (Qwen3 specific).
    pub q_norm: nn::RmsNorm,
    /// Key normalization (Qwen3 specific).
    pub k_norm: nn::RmsNorm,
    /// RoPE layer.
    pub rope: nn::Rope,
}

impl Qwen3LoraAttention {
    /// Create a new LoRA attention layer.
    pub fn new(config: &Qwen3Config, lora_config: &LoraConfig) -> Result<Self, LoraError> {
        let n_heads = config.num_attention_heads;
        let n_kv_heads = config.num_kv_heads();
        let head_dim = config.get_head_dim();
        let scale = (head_dim as f32).sqrt().recip();

        // Assign unique layer ID for FlashAttention caching
        let layer_id = LAYER_ID_COUNTER.fetch_add(1, Ordering::SeqCst);

        let alpha = lora_config.alpha;
        let use_rslora = lora_config.use_rslora;
        let q_rank = crate::effective_rank(lora_config, "q_proj") as i32;
        let k_rank = crate::effective_rank(lora_config, "k_proj") as i32;
        let v_rank = crate::effective_rank(lora_config, "v_proj") as i32;
        let o_rank = crate::effective_rank(lora_config, "o_proj") as i32;

        // Create LoRA linear layers for projections
        let q_proj = LoraLinear::new(
            config.hidden_size,
            n_heads * head_dim,
            q_rank,
            alpha,
            use_rslora,
            false,
        )?;
        let k_proj = LoraLinear::new(
            config.hidden_size,
            n_kv_heads * head_dim,
            k_rank,
            alpha,
            use_rslora,
            false,
        )?;
        let v_proj = LoraLinear::new(
            config.hidden_size,
            n_kv_heads * head_dim,
            v_rank,
            alpha,
            use_rslora,
            false,
        )?;
        let o_proj = LoraLinear::new(
            n_heads * head_dim,
            config.hidden_size,
            o_rank,
            alpha,
            use_rslora,
            false,
        )?;

        // Qwen3-specific: Q and K normalization before RoPE
        let q_norm = nn::RmsNormBuilder::new(head_dim)
            .eps(config.rms_norm_eps)
            .build()
            .unwrap();
        let k_norm = nn::RmsNormBuilder::new(head_dim)
            .eps(config.rms_norm_eps)
            .build()
            .unwrap();

        // Initialize RoPE
        let rope = nn::RopeBuilder::new(head_dim)
            .base(config.rope_theta)
            .traditional(false)
            .build()
            .unwrap();

        Ok(Self {
            n_heads,
            n_kv_heads,
            head_dim,
            scale,
            layer_id,
            q_proj,
            k_proj,
            v_proj,
            o_proj,
            q_norm,
            k_norm,
            rope,
        })
    }

    /// Forward pass through attention.
    ///
    /// # Arguments
    /// * `x` - Input tensor [batch, seq_len, hidden_size]
    /// * `mask` - Optional attention mask
    /// * `position_ids` - Optional position IDs for packed sequences
    pub fn forward(
        &mut self,
        x: &Array,
        mask: Option<&Array>,
        position_ids: Option<&Array>,
    ) -> Result<Array, LoraError> {
        let shape = x.shape();
        let batch = shape[0];
        let seq_len = shape[1];

        // Project to Q, K, V using LoRA layers
        let queries = self.q_proj.forward(x)?;
        let keys = self.k_proj.forward(x)?;
        let values = self.v_proj.forward(x)?;

        // Reshape for multi-head attention: [B, L, heads, head_dim]
        let queries = queries.reshape(&[batch, seq_len, self.n_heads, self.head_dim])?;
        let keys = keys.reshape(&[batch, seq_len, self.n_kv_heads, self.head_dim])?;
        let values = values.reshape(&[batch, seq_len, self.n_kv_heads, self.head_dim])?;

        // Qwen3-specific: Apply Q/K normalization before RoPE
        let queries = mlx_rs::module::Module::forward(&mut self.q_norm, &queries)?;
        let keys = mlx_rs::module::Module::forward(&mut self.k_norm, &keys)?;

        // Transpose for attention: [B, heads, L, head_dim]
        let queries = queries.transpose_axes(&[0, 2, 1, 3])?;
        let keys = keys.transpose_axes(&[0, 2, 1, 3])?;
        let values = values.transpose_axes(&[0, 2, 1, 3])?;

        // Apply RoPE - use position_ids if provided (for packed sequences)
        let (queries, keys) = if let Some(pos_ids) = position_ids {
            // Use custom RoPE with explicit position IDs
            use pmetal_mlx::kernels::rope::apply_rope_with_positions;
            let q = apply_rope_with_positions(
                &queries,
                pos_ids,
                self.head_dim,
                false,
                self.rope.base,
                1.0,
            )?;
            let k = apply_rope_with_positions(
                &keys,
                pos_ids,
                self.head_dim,
                false,
                self.rope.base,
                1.0,
            )?;
            (q, k)
        } else {
            // Use standard RoPE for sequential positions
            let q = mlx_rs::module::Module::forward(&mut self.rope, &queries)?;
            let k = mlx_rs::module::Module::forward(&mut self.rope, &keys)?;
            (q, k)
        };

        // Determine if Metal FlashAttention should be used
        // FlashAttention provides O(n) memory vs O(n²) for standard attention
        // However, it has overhead from f32<->f16 conversions and mutex locks
        // Only use it when sequence length is long enough to benefit (>= 2048)
        let is_training = get_training_context()
            .map(|ctx| ctx.lock().map(|c| c.is_training()).unwrap_or(false))
            .unwrap_or(false);

        // Threshold: Use Metal FA for seq_len >= 2048 during training
        // Below this, standard MLX SDPA is faster due to no conversion overhead
        const FLASH_ATTENTION_SEQ_THRESHOLD: i32 = 2048;
        let use_flash_attention = is_training && seq_len >= FLASH_ATTENTION_SEQ_THRESHOLD;

        let output = if use_flash_attention {
            // Metal FlashAttention path: O(n) memory, handles GQA internally
            // Pass non-expanded K, V - FlashAttention handles GQA via gqa_ratio
            let fa_config = FusedAttentionConfig {
                num_heads: self.n_heads,
                num_kv_heads: self.n_kv_heads,
                head_dim: self.head_dim,
                scale: self.scale,
                mask_type: AttentionMaskType::Causal, // Qwen3 uses causal attention
                logit_softcapping: None,
            };

            differentiable_attention(self.layer_id, &queries, &keys, &values, &fa_config)
                .map_err(|e| LoraError::Mlx(Exception::custom(e.to_string())))?
        } else {
            // Standard MLX attention path (inference or non-Metal training)
            // Expand KV heads for GQA if needed
            let keys = if self.n_kv_heads < self.n_heads {
                let repeats = self.n_heads / self.n_kv_heads;
                expand_kv_heads(&keys, repeats)?
            } else {
                keys
            };
            let values = if self.n_kv_heads < self.n_heads {
                let repeats = self.n_heads / self.n_kv_heads;
                expand_kv_heads(&values, repeats)?
            } else {
                values
            };

            // Scaled dot-product attention (O(n²) memory)
            let scores = queries.matmul(&keys.transpose_axes(&[0, 1, 3, 2])?)?;
            let scores = scores.multiply(Array::from_f32(self.scale))?;

            // Apply mask if provided
            let scores = if let Some(m) = mask {
                scores.add(m)?
            } else {
                scores
            };

            // Softmax
            let weights = mlx_rs::ops::softmax_axis(&scores, -1, None)?;

            // Attention output
            weights.matmul(&values)?
        };

        // Reshape back: [B, heads, L, head_dim] -> [B, L, hidden]
        let output = output
            .transpose_axes(&[0, 2, 1, 3])?
            .reshape(&[batch, seq_len, -1])?;

        // Output projection
        self.o_proj.forward(&output).map_err(LoraError::from)
    }

    /// Forward pass with KV cache for efficient inference.
    ///
    /// This method is optimized for autoregressive generation with O(n)
    /// complexity per token instead of O(n²).
    ///
    /// # Arguments
    /// * `x` - Input tensor [batch, seq_len, hidden_size]
    /// * `mask` - Optional attention mask
    /// * `cache` - Optional (KVCache, layer_idx) tuple for cached generation
    pub fn forward_with_cache(
        &mut self,
        x: &Array,
        mask: Option<&Array>,
        cache: Option<(&mut KVCache, usize)>,
    ) -> Result<Array, LoraError> {
        let shape = x.shape();
        let batch = shape[0];
        let seq_len = shape[1];

        // Project to Q, K, V using LoRA layers
        let queries = self.q_proj.forward(x)?;
        let keys = self.k_proj.forward(x)?;
        let values = self.v_proj.forward(x)?;

        // Reshape for multi-head attention: [B, L, heads, head_dim]
        let queries = queries.reshape(&[batch, seq_len, self.n_heads, self.head_dim])?;
        let keys = keys.reshape(&[batch, seq_len, self.n_kv_heads, self.head_dim])?;
        let values = values.reshape(&[batch, seq_len, self.n_kv_heads, self.head_dim])?;

        // Qwen3-specific: Apply Q/K normalization before RoPE
        let queries = mlx_rs::module::Module::forward(&mut self.q_norm, &queries)?;
        let keys = mlx_rs::module::Module::forward(&mut self.k_norm, &keys)?;

        // Transpose for attention: [B, heads, L, head_dim]
        let queries = queries.transpose_axes(&[0, 2, 1, 3])?;
        let keys = keys.transpose_axes(&[0, 2, 1, 3])?;
        let values = values.transpose_axes(&[0, 2, 1, 3])?;

        // Get RoPE offset and apply RoPE (after Q/K norm)
        let (queries, keys, values) = if let Some((ref cache_ref, _layer_idx)) = cache {
            let offset = cache_ref.rope_offset();
            let queries = apply_rope(&queries, self.head_dim, false, self.rope.base, 1.0, offset)?;
            let keys = apply_rope(&keys, self.head_dim, false, self.rope.base, 1.0, offset)?;
            (queries, keys, values)
        } else {
            let queries = mlx_rs::module::Module::forward(&mut self.rope, &queries)?;
            let keys = mlx_rs::module::Module::forward(&mut self.rope, &keys)?;
            (queries, keys, values)
        };

        // Handle KV cache update - keys/values are already in [B, heads, seq, head_dim] format
        let (keys, values) = if let Some((cache, layer_idx)) = cache {
            cache
                .update_and_fetch(layer_idx, &keys, &values)
                .map_err(|e| LoraError::Mlx(e))?
        } else {
            (keys, values)
        };

        // Use fused attention kernel for inference (more efficient than standard attention)
        let attn_config = FusedAttentionConfig::new(self.n_heads, self.n_kv_heads, self.head_dim)
            .with_scale(self.scale)
            .with_mask_type(AttentionMaskType::Causal);

        let output = fused_sdpa(&queries, &keys, &values, &attn_config, mask)
            .map_err(|e| LoraError::Mlx(e))?;

        // Reshape back: [B, heads, L, head_dim] -> [B, L, hidden]
        let output = output
            .transpose_axes(&[0, 2, 1, 3])?
            .reshape(&[batch, seq_len, -1])?;

        // Output projection
        self.o_proj.forward(&output).map_err(LoraError::from)
    }

    /// Get number of trainable parameters.
    pub fn num_trainable_params(&self) -> usize {
        self.q_proj.num_trainable_params()
            + self.k_proj.num_trainable_params()
            + self.v_proj.num_trainable_params()
            + self.o_proj.num_trainable_params()
    }
}

/// Expand KV heads for grouped query attention.
fn expand_kv_heads(x: &Array, repeats: i32) -> Result<Array, Exception> {
    let shape = x.shape();
    let batch = shape[0];
    let n_kv_heads = shape[1];
    let seq_len = shape[2];
    let head_dim = shape[3];

    let x = x.reshape(&[batch, n_kv_heads, 1, seq_len, head_dim])?;
    let x = mlx_rs::ops::broadcast_to(&x, &[batch, n_kv_heads, repeats, seq_len, head_dim])?;
    x.reshape(&[batch, n_kv_heads * repeats, seq_len, head_dim])
}

/// LoRA-enabled MLP layer for Qwen3.
#[derive(Debug)]
pub struct Qwen3LoraMLP {
    /// Gate projection with LoRA.
    pub gate_proj: LoraLinear,
    /// Up projection with LoRA.
    pub up_proj: LoraLinear,
    /// Down projection with LoRA.
    pub down_proj: LoraLinear,
}

impl Qwen3LoraMLP {
    /// Create a new LoRA MLP layer.
    pub fn new(config: &Qwen3Config, lora_config: &LoraConfig) -> Result<Self, LoraError> {
        let alpha = lora_config.alpha;
        let use_rslora = lora_config.use_rslora;
        let gate_rank = crate::effective_rank(lora_config, "gate_proj") as i32;
        let up_rank = crate::effective_rank(lora_config, "up_proj") as i32;
        let down_rank = crate::effective_rank(lora_config, "down_proj") as i32;

        let gate_proj = LoraLinear::new(
            config.hidden_size,
            config.intermediate_size,
            gate_rank,
            alpha,
            use_rslora,
            false,
        )?;
        let up_proj = LoraLinear::new(
            config.hidden_size,
            config.intermediate_size,
            up_rank,
            alpha,
            use_rslora,
            false,
        )?;
        let down_proj = LoraLinear::new(
            config.intermediate_size,
            config.hidden_size,
            down_rank,
            alpha,
            use_rslora,
            false,
        )?;

        Ok(Self {
            gate_proj,
            up_proj,
            down_proj,
        })
    }

    /// Forward pass (SwiGLU activation).
    ///
    /// **Autograd note**: This uses the standard unfused SwiGLU path where
    /// `gate_proj` and `up_proj` are computed separately. A fused
    /// `gate_up_proj` variant (single matmul projecting to `2 * intermediate`)
    /// would save one kernel launch but requires custom autograd support
    /// because MLX's automatic differentiation cannot split the fused output
    /// back into gate/up gradients. Until a custom VJP is implemented, the
    /// unfused path is required for correct gradient computation during training.
    pub fn forward(&mut self, x: &Array) -> Result<Array, LoraError> {
        let gate = self.gate_proj.forward(x)?;
        let gate = nn::silu(gate)?;
        let up = self.up_proj.forward(x)?;
        let hidden = gate.multiply(&up)?;
        self.down_proj.forward(&hidden)
    }

    /// Get number of trainable parameters.
    pub fn num_trainable_params(&self) -> usize {
        self.gate_proj.num_trainable_params()
            + self.up_proj.num_trainable_params()
            + self.down_proj.num_trainable_params()
    }
}

/// LoRA-enabled Qwen3 decoder layer.
#[derive(Debug)]
pub struct Qwen3LoraDecoderLayer {
    /// Self-attention layer with LoRA.
    pub self_attn: Qwen3LoraAttention,
    /// MLP layer with LoRA.
    pub mlp: Qwen3LoraMLP,
    /// Input layer norm.
    pub input_layernorm: nn::RmsNorm,
    /// Post-attention layer norm.
    pub post_attention_layernorm: nn::RmsNorm,
}

impl Qwen3LoraDecoderLayer {
    /// Create a new decoder layer with LoRA.
    pub fn new(config: &Qwen3Config, lora_config: &LoraConfig) -> Result<Self, LoraError> {
        let self_attn = Qwen3LoraAttention::new(config, lora_config)?;
        let mlp = Qwen3LoraMLP::new(config, lora_config)?;

        let input_layernorm = nn::RmsNormBuilder::new(config.hidden_size)
            .eps(config.rms_norm_eps)
            .build()
            .unwrap();
        let post_attention_layernorm = nn::RmsNormBuilder::new(config.hidden_size)
            .eps(config.rms_norm_eps)
            .build()
            .unwrap();

        Ok(Self {
            self_attn,
            mlp,
            input_layernorm,
            post_attention_layernorm,
        })
    }

    /// Forward pass.
    ///
    /// # Arguments
    /// * `x` - Input tensor
    /// * `mask` - Optional attention mask
    /// * `position_ids` - Optional position IDs for packed sequences
    pub fn forward(
        &mut self,
        x: &Array,
        mask: Option<&Array>,
        position_ids: Option<&Array>,
    ) -> Result<Array, LoraError> {
        // Pre-norm + attention + residual
        let normed = mlx_rs::module::Module::forward(&mut self.input_layernorm, x)?;
        let attn_out = self.self_attn.forward(&normed, mask, position_ids)?;
        let h = x.add(&attn_out)?;

        // Pre-norm + MLP + residual
        let normed = mlx_rs::module::Module::forward(&mut self.post_attention_layernorm, &h)?;
        let mlp_out = self.mlp.forward(&normed)?;
        Ok(h.add(&mlp_out)?)
    }

    /// Forward pass with KV cache for efficient inference.
    ///
    /// # Arguments
    /// * `x` - Input tensor
    /// * `mask` - Optional attention mask
    /// * `cache` - Optional (KVCache, layer_idx) tuple for cached generation
    pub fn forward_with_cache(
        &mut self,
        x: &Array,
        mask: Option<&Array>,
        cache: Option<(&mut KVCache, usize)>,
    ) -> Result<Array, LoraError> {
        // Pre-norm + attention + residual
        let normed = mlx_rs::module::Module::forward(&mut self.input_layernorm, x)?;
        let attn_out = self.self_attn.forward_with_cache(&normed, mask, cache)?;
        let h = x.add(&attn_out)?;

        // Pre-norm + MLP + residual
        let normed = mlx_rs::module::Module::forward(&mut self.post_attention_layernorm, &h)?;
        let mlp_out = self.mlp.forward(&normed)?;
        Ok(h.add(&mlp_out)?)
    }

    /// Get number of trainable parameters.
    pub fn num_trainable_params(&self) -> usize {
        self.self_attn.num_trainable_params() + self.mlp.num_trainable_params()
    }
}

static GRAD_CKPT_WARN: Once = Once::new();

/// LoRA-enabled Qwen3 model (without LM head).
#[derive(Debug)]
pub struct Qwen3LoraModel {
    /// Configuration.
    pub config: Qwen3Config,
    /// LoRA configuration.
    pub lora_config: LoraConfig,
    /// Token embeddings (frozen).
    pub embed_tokens: nn::Embedding,
    /// Transformer layers with LoRA.
    pub layers: Vec<Qwen3LoraDecoderLayer>,
    /// Final layer norm (frozen).
    pub norm: nn::RmsNorm,
}

impl Qwen3LoraModel {
    /// Create a new LoRA Qwen3 model.
    pub fn new(config: Qwen3Config, lora_config: LoraConfig) -> Result<Self, LoraError> {
        // Reset layer ID counter so each model instance assigns IDs starting from 0.
        reset_layer_ids();

        let embed_tokens = nn::Embedding::new(config.vocab_size, config.hidden_size)?;

        let layers = (0..config.num_hidden_layers)
            .map(|_| Qwen3LoraDecoderLayer::new(&config, &lora_config))
            .collect::<Result<Vec<_>, _>>()?;

        let norm = nn::RmsNormBuilder::new(config.hidden_size)
            .eps(config.rms_norm_eps)
            .build()
            .unwrap();

        Ok(Self {
            config,
            lora_config,
            embed_tokens,
            layers,
            norm,
        })
    }

    /// Forward pass.
    ///
    /// # Arguments
    /// * `input_ids` - Token IDs [batch, seq_len]
    /// * `mask` - Optional attention mask
    /// * `position_ids` - Optional position IDs for packed sequences
    pub fn forward(
        &mut self,
        input_ids: &Array,
        mask: Option<&Array>,
        position_ids: Option<&Array>,
    ) -> Result<Array, LoraError> {
        self.forward_with_checkpoint(input_ids, mask, position_ids, None)
    }

    /// Forward pass with optional gradient checkpointing.
    ///
    /// Gradient checkpointing trades compute for memory by breaking the computation
    /// graph at layer boundaries. This allows training with larger batch sizes
    /// at the cost of recomputing activations during the backward pass.
    ///
    /// # Arguments
    /// * `input_ids` - Token IDs [batch, seq_len]
    /// * `mask` - Optional attention mask
    /// * `position_ids` - Optional position IDs for packed sequences
    /// * `checkpoint_config` - Optional checkpointing configuration
    ///
    /// # Memory Savings
    /// With checkpointing enabled (layers_per_block=4 for 28-layer model):
    /// - Without: Store all 28 layer activations
    /// - With: Store 7 checkpoint boundaries + recompute
    /// - Typical savings: 50-75% activation memory
    pub fn forward_with_checkpoint(
        &mut self,
        input_ids: &Array,
        mask: Option<&Array>,
        position_ids: Option<&Array>,
        checkpoint_config: Option<&CheckpointConfig>,
    ) -> Result<Array, LoraError> {
        // Get embeddings
        let mut hidden_states = mlx_rs::module::Module::forward(&mut self.embed_tokens, input_ids)?;

        // Create causal mask if not provided
        let mask = if mask.is_none() {
            let seq_len = input_ids.dim(1);
            Some(create_causal_mask(seq_len)?)
        } else {
            mask.cloned()
        };

        // Pass through transformer layers with optional checkpointing
        let layers_per_block = checkpoint_config
            .map(|c| c.layers_per_block)
            .unwrap_or(usize::MAX);
        let checkpointing_enabled = checkpoint_config.map(|c| c.enabled).unwrap_or(false);

        for (idx, layer) in self.layers.iter_mut().enumerate() {
            hidden_states = layer.forward(&hidden_states, mask.as_ref(), position_ids)?;

            // Checkpoint boundary marker (NO-OP — see note below).
            //
            // Gradient checkpointing is not yet implemented for Qwen3 LoRA.
            // True checkpointing requires a custom VJP that saves activations
            // at block boundaries and recomputes intermediate activations
            // during the backward pass. MLX does not provide a built-in
            // `checkpoint()` primitive, so implementing this requires manually
            // wrapping the per-block forward in a closure registered with
            // `mlx_rs::transforms::vjp`. Until that is done, this code path
            // only logs a warning — no memory savings are applied.
            //
            // Calling `eval()` here would materialize intermediates and break
            // the lazy computation graph needed for automatic differentiation.
            if checkpointing_enabled && (idx + 1) % layers_per_block == 0 {
                GRAD_CKPT_WARN.call_once(|| {
                    tracing::warn!(
                        "Gradient checkpointing not yet implemented for Qwen3 - \
                         no memory savings applied. Use --no-gradient-checkpointing to suppress."
                    );
                });
            }
        }

        // Final norm
        Ok(mlx_rs::module::Module::forward(
            &mut self.norm,
            &hidden_states,
        )?)
    }

    /// Forward pass with KV cache for efficient inference.
    ///
    /// # Arguments
    /// * `input_ids` - Token IDs [batch, seq_len]
    /// * `mask` - Optional attention mask
    /// * `cache` - Optional mutable reference to KV cache
    pub fn forward_with_cache(
        &mut self,
        input_ids: &Array,
        mask: Option<&Array>,
        cache: Option<&mut KVCache>,
    ) -> Result<Array, LoraError> {
        // Get embeddings
        let mut hidden_states = mlx_rs::module::Module::forward(&mut self.embed_tokens, input_ids)?;

        // Don't create explicit causal mask - fused SDPA handles it internally
        // with proper dtype handling. Only pass through user-provided masks.

        // Pass through transformer layers
        match cache {
            Some(cache) => {
                for (layer_idx, layer) in self.layers.iter_mut().enumerate() {
                    hidden_states =
                        layer.forward_with_cache(&hidden_states, mask, Some((cache, layer_idx)))?;
                }
            }
            None => {
                for layer in &mut self.layers {
                    hidden_states = layer.forward(&hidden_states, mask, None)?;
                }
            }
        }

        // Final norm
        Ok(mlx_rs::module::Module::forward(
            &mut self.norm,
            &hidden_states,
        )?)
    }

    /// Get number of trainable parameters.
    pub fn num_trainable_params(&self) -> usize {
        self.layers.iter().map(|l| l.num_trainable_params()).sum()
    }
}

/// LoRA-enabled Qwen3 model with LM head.
#[derive(Debug)]
pub struct Qwen3LoraForCausalLM {
    /// Base model with LoRA.
    pub model: Qwen3LoraModel,
    /// LM head (frozen, optional for tied weights).
    pub lm_head: Option<nn::Linear>,
    /// Gradient checkpointing configuration.
    /// When enabled, breaks computation graph at layer boundaries to save memory.
    pub checkpoint_config: Option<CheckpointConfig>,
}

impl Qwen3LoraForCausalLM {
    /// Create a new LoRA Qwen3 model with LM head.
    pub fn new(config: Qwen3Config, lora_config: LoraConfig) -> Result<Self, LoraError> {
        let tie_weights = config.tie_word_embeddings;
        let model = Qwen3LoraModel::new(config.clone(), lora_config)?;

        let lm_head = if !tie_weights {
            let head = nn::LinearBuilder::new(config.hidden_size, config.vocab_size)
                .bias(false)
                .build()
                .unwrap();
            Some(head)
        } else {
            None
        };

        Ok(Self {
            model,
            lm_head,
            checkpoint_config: None,
        })
    }

    /// Enable gradient checkpointing for memory-efficient training.
    ///
    /// # Arguments
    /// * `layers_per_block` - Number of layers per checkpoint block (typically 2-4)
    ///
    /// # Example
    /// ```ignore
    /// model.enable_gradient_checkpointing(4); // Checkpoint every 4 layers
    /// ```
    pub fn enable_gradient_checkpointing(&mut self, layers_per_block: usize) {
        self.checkpoint_config = Some(CheckpointConfig {
            enabled: true,
            layers_per_block,
            eval_at_boundaries: true,
        });
    }

    /// Disable gradient checkpointing.
    pub fn disable_gradient_checkpointing(&mut self) {
        self.checkpoint_config = None;
    }

    /// Forward pass producing logits.
    ///
    /// Uses the model's stored checkpoint_config if set via `enable_gradient_checkpointing()`.
    ///
    /// # Arguments
    /// * `input_ids` - Token IDs [batch, seq_len]
    /// * `mask` - Optional attention mask
    /// * `position_ids` - Optional position IDs for packed sequences
    pub fn forward(
        &mut self,
        input_ids: &Array,
        mask: Option<&Array>,
        position_ids: Option<&Array>,
    ) -> Result<Array, LoraError> {
        // Clone checkpoint config to avoid borrow conflicts
        let checkpoint_config = self.checkpoint_config.clone();
        self.forward_with_checkpoint(input_ids, mask, position_ids, checkpoint_config.as_ref())
    }

    /// Forward pass with optional gradient checkpointing.
    ///
    /// # Arguments
    /// * `input_ids` - Token IDs [batch, seq_len]
    /// * `mask` - Optional attention mask
    /// * `position_ids` - Optional position IDs for packed sequences
    /// * `checkpoint_config` - Optional checkpointing configuration
    pub fn forward_with_checkpoint(
        &mut self,
        input_ids: &Array,
        mask: Option<&Array>,
        position_ids: Option<&Array>,
        checkpoint_config: Option<&CheckpointConfig>,
    ) -> Result<Array, LoraError> {
        let hidden_states =
            self.model
                .forward_with_checkpoint(input_ids, mask, position_ids, checkpoint_config)?;

        // Get logits from LM head or shared embeddings
        if let Some(ref mut lm_head) = self.lm_head {
            Ok(mlx_rs::module::Module::forward(lm_head, &hidden_states)?)
        } else {
            // Tie weights: use embedding weight transposed
            Ok(self.model.embed_tokens.as_linear(&hidden_states)?)
        }
    }

    /// Forward pass with KV cache for efficient inference.
    ///
    /// KV caching provides O(n) complexity per token instead of O(n²),
    /// enabling fast autoregressive generation.
    ///
    /// # Arguments
    /// * `input_ids` - Token IDs [batch, seq_len]
    /// * `mask` - Optional attention mask
    /// * `cache` - Optional mutable reference to KV cache
    pub fn forward_with_cache(
        &mut self,
        input_ids: &Array,
        mask: Option<&Array>,
        cache: Option<&mut KVCache>,
    ) -> Result<Array, LoraError> {
        let hidden_states = self.model.forward_with_cache(input_ids, mask, cache)?;

        // Get logits from LM head or shared embeddings
        if let Some(ref mut lm_head) = self.lm_head {
            Ok(mlx_rs::module::Module::forward(lm_head, &hidden_states)?)
        } else {
            // Tie weights: use embedding weight transposed
            Ok(self.model.embed_tokens.as_linear(&hidden_states)?)
        }
    }

    /// Create a KV cache for this model.
    ///
    /// # Arguments
    /// * `max_seq_len` - Maximum sequence length to cache
    pub fn create_cache(&self, max_seq_len: usize) -> KVCache {
        let config = KVCacheConfig::new(
            self.model.config.num_hidden_layers as usize,
            max_seq_len,
            self.model.config.num_kv_heads() as usize,
            self.model.config.get_head_dim() as usize,
        );
        KVCache::new(config)
    }

    /// Merge LoRA weights into base weights.
    pub fn merge_lora(&mut self) -> Result<(), LoraError> {
        for layer in &mut self.model.layers {
            layer.self_attn.q_proj.merge()?;
            layer.self_attn.k_proj.merge()?;
            layer.self_attn.v_proj.merge()?;
            layer.self_attn.o_proj.merge()?;
            layer.mlp.gate_proj.merge()?;
            layer.mlp.up_proj.merge()?;
            layer.mlp.down_proj.merge()?;
        }
        Ok(())
    }

    /// Unmerge is not supported.
    pub fn unmerge_lora(&mut self) -> Result<(), LoraError> {
        Err(LoraError::InvalidState(
            "unmerge_lora is not supported: reload base model weights to undo a merge".to_string(),
        ))
    }

    /// Get all trainable LoRA parameters as a flat HashMap.
    pub fn lora_parameters(&self) -> HashMap<Rc<str>, Array> {
        let mut params = HashMap::new();

        for (i, layer) in self.model.layers.iter().enumerate() {
            let prefix = format!("layers.{}", i);

            // Attention LoRA params
            params.insert(
                Rc::from(format!("{}.self_attn.q_proj.lora_a", prefix)),
                layer.self_attn.q_proj.lora_a.clone(),
            );
            params.insert(
                Rc::from(format!("{}.self_attn.q_proj.lora_b", prefix)),
                layer.self_attn.q_proj.lora_b.clone(),
            );
            params.insert(
                Rc::from(format!("{}.self_attn.k_proj.lora_a", prefix)),
                layer.self_attn.k_proj.lora_a.clone(),
            );
            params.insert(
                Rc::from(format!("{}.self_attn.k_proj.lora_b", prefix)),
                layer.self_attn.k_proj.lora_b.clone(),
            );
            params.insert(
                Rc::from(format!("{}.self_attn.v_proj.lora_a", prefix)),
                layer.self_attn.v_proj.lora_a.clone(),
            );
            params.insert(
                Rc::from(format!("{}.self_attn.v_proj.lora_b", prefix)),
                layer.self_attn.v_proj.lora_b.clone(),
            );
            params.insert(
                Rc::from(format!("{}.self_attn.o_proj.lora_a", prefix)),
                layer.self_attn.o_proj.lora_a.clone(),
            );
            params.insert(
                Rc::from(format!("{}.self_attn.o_proj.lora_b", prefix)),
                layer.self_attn.o_proj.lora_b.clone(),
            );

            // MLP LoRA params
            params.insert(
                Rc::from(format!("{}.mlp.gate_proj.lora_a", prefix)),
                layer.mlp.gate_proj.lora_a.clone(),
            );
            params.insert(
                Rc::from(format!("{}.mlp.gate_proj.lora_b", prefix)),
                layer.mlp.gate_proj.lora_b.clone(),
            );
            params.insert(
                Rc::from(format!("{}.mlp.up_proj.lora_a", prefix)),
                layer.mlp.up_proj.lora_a.clone(),
            );
            params.insert(
                Rc::from(format!("{}.mlp.up_proj.lora_b", prefix)),
                layer.mlp.up_proj.lora_b.clone(),
            );
            params.insert(
                Rc::from(format!("{}.mlp.down_proj.lora_a", prefix)),
                layer.mlp.down_proj.lora_a.clone(),
            );
            params.insert(
                Rc::from(format!("{}.mlp.down_proj.lora_b", prefix)),
                layer.mlp.down_proj.lora_b.clone(),
            );
        }

        params
    }

    /// Set LoRA parameters from a HashMap.
    pub fn set_lora_parameters(&mut self, params: &HashMap<Rc<str>, Array>) {
        for (i, layer) in self.model.layers.iter_mut().enumerate() {
            let prefix = format!("layers.{}", i);

            macro_rules! set_param {
                ($param:expr, $key:expr) => {
                    if let Some(value) = params.get(&Rc::from($key)) {
                        $param = value.clone();
                    }
                };
            }

            // Attention LoRA params
            set_param!(
                layer.self_attn.q_proj.lora_a,
                format!("{}.self_attn.q_proj.lora_a", prefix)
            );
            set_param!(
                layer.self_attn.q_proj.lora_b,
                format!("{}.self_attn.q_proj.lora_b", prefix)
            );
            set_param!(
                layer.self_attn.k_proj.lora_a,
                format!("{}.self_attn.k_proj.lora_a", prefix)
            );
            set_param!(
                layer.self_attn.k_proj.lora_b,
                format!("{}.self_attn.k_proj.lora_b", prefix)
            );
            set_param!(
                layer.self_attn.v_proj.lora_a,
                format!("{}.self_attn.v_proj.lora_a", prefix)
            );
            set_param!(
                layer.self_attn.v_proj.lora_b,
                format!("{}.self_attn.v_proj.lora_b", prefix)
            );
            set_param!(
                layer.self_attn.o_proj.lora_a,
                format!("{}.self_attn.o_proj.lora_a", prefix)
            );
            set_param!(
                layer.self_attn.o_proj.lora_b,
                format!("{}.self_attn.o_proj.lora_b", prefix)
            );

            // MLP LoRA params
            set_param!(
                layer.mlp.gate_proj.lora_a,
                format!("{}.mlp.gate_proj.lora_a", prefix)
            );
            set_param!(
                layer.mlp.gate_proj.lora_b,
                format!("{}.mlp.gate_proj.lora_b", prefix)
            );
            set_param!(
                layer.mlp.up_proj.lora_a,
                format!("{}.mlp.up_proj.lora_a", prefix)
            );
            set_param!(
                layer.mlp.up_proj.lora_b,
                format!("{}.mlp.up_proj.lora_b", prefix)
            );
            set_param!(
                layer.mlp.down_proj.lora_a,
                format!("{}.mlp.down_proj.lora_a", prefix)
            );
            set_param!(
                layer.mlp.down_proj.lora_b,
                format!("{}.mlp.down_proj.lora_b", prefix)
            );
        }
    }

    /// Get number of trainable parameters.
    pub fn num_trainable_params(&self) -> usize {
        self.model.num_trainable_params()
    }

    /// Get configuration.
    pub fn config(&self) -> &Qwen3Config {
        &self.model.config
    }

    /// Get LoRA configuration.
    pub fn lora_config(&self) -> &LoraConfig {
        &self.model.lora_config
    }

    /// Save LoRA weights to safetensors.
    pub fn save_lora_weights(&self, path: impl AsRef<std::path::Path>) -> Result<(), LoraError> {
        let params = self.lora_parameters();
        Array::save_safetensors(params, None, path)?;
        Ok(())
    }

    /// Load LoRA weights from safetensors.
    ///
    /// # Arguments
    /// * `path` - Path to either:
    ///   - A directory containing `lora_weights.safetensors`
    ///   - A direct path to a `.safetensors` file
    pub fn load_lora_weights(
        &mut self,
        path: impl AsRef<std::path::Path>,
    ) -> Result<(), LoraError> {
        let path = path.as_ref();
        let file_path = if path.is_dir() {
            path.join("lora_weights.safetensors")
        } else {
            path.to_path_buf()
        };
        let loaded = Array::load_safetensors(&file_path)?;

        for (i, layer) in self.model.layers.iter_mut().enumerate() {
            let prefix = format!("layers.{}", i);

            macro_rules! load_param {
                ($param:expr, $key:expr) => {
                    if let Some(value) = loaded.get(&Rc::from($key) as &str) {
                        $param = value.clone();
                    }
                };
            }

            // Attention LoRA params
            load_param!(
                layer.self_attn.q_proj.lora_a,
                format!("{}.self_attn.q_proj.lora_a", prefix)
            );
            load_param!(
                layer.self_attn.q_proj.lora_b,
                format!("{}.self_attn.q_proj.lora_b", prefix)
            );
            load_param!(
                layer.self_attn.k_proj.lora_a,
                format!("{}.self_attn.k_proj.lora_a", prefix)
            );
            load_param!(
                layer.self_attn.k_proj.lora_b,
                format!("{}.self_attn.k_proj.lora_b", prefix)
            );
            load_param!(
                layer.self_attn.v_proj.lora_a,
                format!("{}.self_attn.v_proj.lora_a", prefix)
            );
            load_param!(
                layer.self_attn.v_proj.lora_b,
                format!("{}.self_attn.v_proj.lora_b", prefix)
            );
            load_param!(
                layer.self_attn.o_proj.lora_a,
                format!("{}.self_attn.o_proj.lora_a", prefix)
            );
            load_param!(
                layer.self_attn.o_proj.lora_b,
                format!("{}.self_attn.o_proj.lora_b", prefix)
            );

            // MLP LoRA params
            load_param!(
                layer.mlp.gate_proj.lora_a,
                format!("{}.mlp.gate_proj.lora_a", prefix)
            );
            load_param!(
                layer.mlp.gate_proj.lora_b,
                format!("{}.mlp.gate_proj.lora_b", prefix)
            );
            load_param!(
                layer.mlp.up_proj.lora_a,
                format!("{}.mlp.up_proj.lora_a", prefix)
            );
            load_param!(
                layer.mlp.up_proj.lora_b,
                format!("{}.mlp.up_proj.lora_b", prefix)
            );
            load_param!(
                layer.mlp.down_proj.lora_a,
                format!("{}.mlp.down_proj.lora_a", prefix)
            );
            load_param!(
                layer.mlp.down_proj.lora_b,
                format!("{}.mlp.down_proj.lora_b", prefix)
            );
        }

        Ok(())
    }

    /// Load base model weights from a HashMap of weight tensors.
    pub fn load_base_weights(
        &mut self,
        weights: &std::collections::HashMap<String, Array>,
    ) -> Result<(), LoraError> {
        // Load embed_tokens
        if let Some(w) = weights.get("model.embed_tokens.weight") {
            self.model.embed_tokens.weight = Param::new(w.clone());
        }

        // Load transformer layers
        for (i, layer) in self.model.layers.iter_mut().enumerate() {
            let prefix = format!("model.layers.{}", i);

            // Self-attention projections
            if let Some(w) = weights.get(&format!("{}.self_attn.q_proj.weight", prefix)) {
                layer.self_attn.q_proj.weight = w.clone();
            }
            if let Some(w) = weights.get(&format!("{}.self_attn.k_proj.weight", prefix)) {
                layer.self_attn.k_proj.weight = w.clone();
            }
            if let Some(w) = weights.get(&format!("{}.self_attn.v_proj.weight", prefix)) {
                layer.self_attn.v_proj.weight = w.clone();
            }
            if let Some(w) = weights.get(&format!("{}.self_attn.o_proj.weight", prefix)) {
                layer.self_attn.o_proj.weight = w.clone();
            }

            // Qwen3-specific: Q/K norms
            if let Some(w) = weights.get(&format!("{}.self_attn.q_norm.weight", prefix)) {
                layer.self_attn.q_norm.weight = Param::new(w.clone());
            }
            if let Some(w) = weights.get(&format!("{}.self_attn.k_norm.weight", prefix)) {
                layer.self_attn.k_norm.weight = Param::new(w.clone());
            }

            // MLP projections
            if let Some(w) = weights.get(&format!("{}.mlp.gate_proj.weight", prefix)) {
                layer.mlp.gate_proj.weight = w.clone();
            }
            if let Some(w) = weights.get(&format!("{}.mlp.up_proj.weight", prefix)) {
                layer.mlp.up_proj.weight = w.clone();
            }
            if let Some(w) = weights.get(&format!("{}.mlp.down_proj.weight", prefix)) {
                layer.mlp.down_proj.weight = w.clone();
            }

            // Layer norms
            if let Some(w) = weights.get(&format!("{}.input_layernorm.weight", prefix)) {
                layer.input_layernorm.weight = Param::new(w.clone());
            }
            if let Some(w) = weights.get(&format!("{}.post_attention_layernorm.weight", prefix)) {
                layer.post_attention_layernorm.weight = Param::new(w.clone());
            }
        }

        // Load final norm
        if let Some(w) = weights.get("model.norm.weight") {
            self.model.norm.weight = Param::new(w.clone());
        }

        // Load lm_head if present and not tied
        if let Some(ref mut lm_head) = self.lm_head {
            if let Some(w) = weights.get("lm_head.weight") {
                lm_head.weight = Param::new(w.clone());
            }
        }

        Ok(())
    }

    /// Load base model weights from safetensor files in a directory.
    pub fn load_base_weights_from_dir(
        &mut self,
        model_dir: impl AsRef<std::path::Path>,
    ) -> Result<(), LoraError> {
        let model_dir = model_dir.as_ref();

        // Check for single file model
        let single_file = model_dir.join("model.safetensors");
        if single_file.exists() {
            let weights = Array::load_safetensors(&single_file)?;
            return self.load_base_weights(&weights);
        }

        // Load sharded model
        let index_path = model_dir.join("model.safetensors.index.json");
        if !index_path.exists() {
            return Err(LoraError::Mlx(Exception::custom(
                "No model.safetensors or model.safetensors.index.json found".to_string(),
            )));
        }

        let index_content = std::fs::read_to_string(&index_path)
            .map_err(|e| LoraError::Mlx(Exception::custom(e.to_string())))?;

        #[derive(serde::Deserialize)]
        struct WeightIndex {
            weight_map: std::collections::HashMap<String, String>,
        }

        let index: WeightIndex = serde_json::from_str(&index_content)
            .map_err(|e| LoraError::Mlx(Exception::custom(e.to_string())))?;

        // Get unique shard files
        let shard_files: std::collections::HashSet<&String> = index.weight_map.values().collect();

        // Load each shard and combine weights
        let mut all_weights = std::collections::HashMap::new();
        for shard_file in shard_files {
            let shard_path = model_dir.join(shard_file);
            let shard_weights = Array::load_safetensors(&shard_path)?;
            all_weights.extend(shard_weights);
        }

        self.load_base_weights(&all_weights)
    }

    /// Evaluate all model parameters (force computation).
    pub fn eval_all(&self) -> Result<(), LoraError> {
        // Eval embeddings
        self.model.embed_tokens.weight.value.as_ref().eval()?;

        // Eval layers
        for layer in &self.model.layers {
            // Base weights
            layer.self_attn.q_proj.weight.eval()?;
            layer.self_attn.k_proj.weight.eval()?;
            layer.self_attn.v_proj.weight.eval()?;
            layer.self_attn.o_proj.weight.eval()?;
            layer.mlp.gate_proj.weight.eval()?;
            layer.mlp.up_proj.weight.eval()?;
            layer.mlp.down_proj.weight.eval()?;

            // Q/K norms
            layer.self_attn.q_norm.weight.value.as_ref().eval()?;
            layer.self_attn.k_norm.weight.value.as_ref().eval()?;

            // LoRA weights
            layer.self_attn.q_proj.lora_a.eval()?;
            layer.self_attn.q_proj.lora_b.eval()?;
            layer.self_attn.k_proj.lora_a.eval()?;
            layer.self_attn.k_proj.lora_b.eval()?;
            layer.self_attn.v_proj.lora_a.eval()?;
            layer.self_attn.v_proj.lora_b.eval()?;
            layer.self_attn.o_proj.lora_a.eval()?;
            layer.self_attn.o_proj.lora_b.eval()?;
            layer.mlp.gate_proj.lora_a.eval()?;
            layer.mlp.gate_proj.lora_b.eval()?;
            layer.mlp.up_proj.lora_a.eval()?;
            layer.mlp.up_proj.lora_b.eval()?;
            layer.mlp.down_proj.lora_a.eval()?;
            layer.mlp.down_proj.lora_b.eval()?;

            // Layer norms
            layer.input_layernorm.weight.value.as_ref().eval()?;
            layer
                .post_attention_layernorm
                .weight
                .value
                .as_ref()
                .eval()?;
        }

        // Final norm
        self.model.norm.weight.value.as_ref().eval()?;

        // LM head if present
        if let Some(ref lm_head) = self.lm_head {
            lm_head.weight.value.as_ref().eval()?;
        }

        Ok(())
    }
}

/// Implement ModuleParameters for Qwen3LoraForCausalLM.
impl ModuleParameters for Qwen3LoraForCausalLM {
    /// Returns the number of trainable (LoRA) parameters.
    fn num_parameters(&self) -> usize {
        self.model.num_trainable_params()
    }

    fn parameters(&self) -> ModuleParamRef<'_> {
        let mut params = ModuleParamRef::new();

        for (i, layer) in self.model.layers.iter().enumerate() {
            let prefix: Rc<str> = Rc::from(format!("layers.{}", i));
            let mut layer_params = HashMap::new();

            // Attention LoRA params
            let mut attn_params = HashMap::new();
            let mut q_params = HashMap::new();
            q_params.insert(
                Rc::from("lora_a"),
                NestedValue::Value(&layer.self_attn.q_proj.lora_a),
            );
            q_params.insert(
                Rc::from("lora_b"),
                NestedValue::Value(&layer.self_attn.q_proj.lora_b),
            );
            attn_params.insert(Rc::from("q_proj"), NestedValue::Map(q_params));

            let mut k_params = HashMap::new();
            k_params.insert(
                Rc::from("lora_a"),
                NestedValue::Value(&layer.self_attn.k_proj.lora_a),
            );
            k_params.insert(
                Rc::from("lora_b"),
                NestedValue::Value(&layer.self_attn.k_proj.lora_b),
            );
            attn_params.insert(Rc::from("k_proj"), NestedValue::Map(k_params));

            let mut v_params = HashMap::new();
            v_params.insert(
                Rc::from("lora_a"),
                NestedValue::Value(&layer.self_attn.v_proj.lora_a),
            );
            v_params.insert(
                Rc::from("lora_b"),
                NestedValue::Value(&layer.self_attn.v_proj.lora_b),
            );
            attn_params.insert(Rc::from("v_proj"), NestedValue::Map(v_params));

            let mut o_params = HashMap::new();
            o_params.insert(
                Rc::from("lora_a"),
                NestedValue::Value(&layer.self_attn.o_proj.lora_a),
            );
            o_params.insert(
                Rc::from("lora_b"),
                NestedValue::Value(&layer.self_attn.o_proj.lora_b),
            );
            attn_params.insert(Rc::from("o_proj"), NestedValue::Map(o_params));

            layer_params.insert(Rc::from("self_attn"), NestedValue::Map(attn_params));

            // MLP LoRA params
            let mut mlp_params = HashMap::new();
            let mut gate_params = HashMap::new();
            gate_params.insert(
                Rc::from("lora_a"),
                NestedValue::Value(&layer.mlp.gate_proj.lora_a),
            );
            gate_params.insert(
                Rc::from("lora_b"),
                NestedValue::Value(&layer.mlp.gate_proj.lora_b),
            );
            mlp_params.insert(Rc::from("gate_proj"), NestedValue::Map(gate_params));

            let mut up_params = HashMap::new();
            up_params.insert(
                Rc::from("lora_a"),
                NestedValue::Value(&layer.mlp.up_proj.lora_a),
            );
            up_params.insert(
                Rc::from("lora_b"),
                NestedValue::Value(&layer.mlp.up_proj.lora_b),
            );
            mlp_params.insert(Rc::from("up_proj"), NestedValue::Map(up_params));

            let mut down_params = HashMap::new();
            down_params.insert(
                Rc::from("lora_a"),
                NestedValue::Value(&layer.mlp.down_proj.lora_a),
            );
            down_params.insert(
                Rc::from("lora_b"),
                NestedValue::Value(&layer.mlp.down_proj.lora_b),
            );
            mlp_params.insert(Rc::from("down_proj"), NestedValue::Map(down_params));

            layer_params.insert(Rc::from("mlp"), NestedValue::Map(mlp_params));

            params.insert(prefix, NestedValue::Map(layer_params));
        }

        params
    }

    fn parameters_mut(&mut self) -> ModuleParamMut<'_> {
        let mut params = ModuleParamMut::new();

        for (i, layer) in self.model.layers.iter_mut().enumerate() {
            let prefix: Rc<str> = Rc::from(format!("layers.{}", i));
            let mut layer_params = HashMap::new();

            // Attention LoRA params
            let mut attn_params = HashMap::new();
            let mut q_params = HashMap::new();
            q_params.insert(
                Rc::from("lora_a"),
                NestedValue::Value(&mut layer.self_attn.q_proj.lora_a),
            );
            q_params.insert(
                Rc::from("lora_b"),
                NestedValue::Value(&mut layer.self_attn.q_proj.lora_b),
            );
            attn_params.insert(Rc::from("q_proj"), NestedValue::Map(q_params));

            let mut k_params = HashMap::new();
            k_params.insert(
                Rc::from("lora_a"),
                NestedValue::Value(&mut layer.self_attn.k_proj.lora_a),
            );
            k_params.insert(
                Rc::from("lora_b"),
                NestedValue::Value(&mut layer.self_attn.k_proj.lora_b),
            );
            attn_params.insert(Rc::from("k_proj"), NestedValue::Map(k_params));

            let mut v_params = HashMap::new();
            v_params.insert(
                Rc::from("lora_a"),
                NestedValue::Value(&mut layer.self_attn.v_proj.lora_a),
            );
            v_params.insert(
                Rc::from("lora_b"),
                NestedValue::Value(&mut layer.self_attn.v_proj.lora_b),
            );
            attn_params.insert(Rc::from("v_proj"), NestedValue::Map(v_params));

            let mut o_params = HashMap::new();
            o_params.insert(
                Rc::from("lora_a"),
                NestedValue::Value(&mut layer.self_attn.o_proj.lora_a),
            );
            o_params.insert(
                Rc::from("lora_b"),
                NestedValue::Value(&mut layer.self_attn.o_proj.lora_b),
            );
            attn_params.insert(Rc::from("o_proj"), NestedValue::Map(o_params));

            layer_params.insert(Rc::from("self_attn"), NestedValue::Map(attn_params));

            // MLP LoRA params
            let mut mlp_params = HashMap::new();
            let mut gate_params = HashMap::new();
            gate_params.insert(
                Rc::from("lora_a"),
                NestedValue::Value(&mut layer.mlp.gate_proj.lora_a),
            );
            gate_params.insert(
                Rc::from("lora_b"),
                NestedValue::Value(&mut layer.mlp.gate_proj.lora_b),
            );
            mlp_params.insert(Rc::from("gate_proj"), NestedValue::Map(gate_params));

            let mut up_params = HashMap::new();
            up_params.insert(
                Rc::from("lora_a"),
                NestedValue::Value(&mut layer.mlp.up_proj.lora_a),
            );
            up_params.insert(
                Rc::from("lora_b"),
                NestedValue::Value(&mut layer.mlp.up_proj.lora_b),
            );
            mlp_params.insert(Rc::from("up_proj"), NestedValue::Map(up_params));

            let mut down_params = HashMap::new();
            down_params.insert(
                Rc::from("lora_a"),
                NestedValue::Value(&mut layer.mlp.down_proj.lora_a),
            );
            down_params.insert(
                Rc::from("lora_b"),
                NestedValue::Value(&mut layer.mlp.down_proj.lora_b),
            );
            mlp_params.insert(Rc::from("down_proj"), NestedValue::Map(down_params));

            layer_params.insert(Rc::from("mlp"), NestedValue::Map(mlp_params));

            params.insert(prefix, NestedValue::Map(layer_params));
        }

        params
    }

    fn trainable_parameters(&self) -> ModuleParamRef<'_> {
        self.parameters()
    }

    fn freeze_parameters(&mut self, _recursive: bool) {}
    fn unfreeze_parameters(&mut self, _recursive: bool) {}
    fn all_frozen(&self) -> Option<bool> {
        Some(false)
    }
    fn any_frozen(&self) -> Option<bool> {
        Some(false)
    }
}

/// Implement TrainableModel for Qwen3LoraForCausalLM.
impl crate::TrainableModel for Qwen3LoraForCausalLM {
    fn forward(&mut self, input_ids: &Array, mask: Option<&Array>) -> Result<Array, LoraError> {
        Qwen3LoraForCausalLM::forward(self, input_ids, mask, None)
    }

    fn forward_with_positions(
        &mut self,
        input_ids: &Array,
        mask: Option<&Array>,
        position_ids: &Array,
    ) -> Result<Array, LoraError> {
        Qwen3LoraForCausalLM::forward(self, input_ids, mask, Some(position_ids))
    }

    fn num_trainable_params(&self) -> usize {
        Qwen3LoraForCausalLM::num_trainable_params(self)
    }

    fn lora_parameters(&self) -> HashMap<Rc<str>, Array> {
        Qwen3LoraForCausalLM::lora_parameters(self)
    }

    fn set_lora_parameters(&mut self, params: &HashMap<Rc<str>, Array>) {
        Qwen3LoraForCausalLM::set_lora_parameters(self, params)
    }

    fn save_lora_weights(&self, path: impl AsRef<std::path::Path>) -> Result<(), LoraError> {
        Qwen3LoraForCausalLM::save_lora_weights(self, path)
    }

    fn load_lora_weights(&mut self, path: impl AsRef<std::path::Path>) -> Result<(), LoraError> {
        Qwen3LoraForCausalLM::load_lora_weights(self, path)
    }

    fn enable_gradient_checkpointing(&mut self, layers_per_block: usize) {
        Qwen3LoraForCausalLM::enable_gradient_checkpointing(self, layers_per_block)
    }

    fn disable_gradient_checkpointing(&mut self) {
        Qwen3LoraForCausalLM::disable_gradient_checkpointing(self)
    }

    fn supports_gradient_checkpointing(&self) -> bool {
        true
    }

    fn forward_with_cache(
        &mut self,
        input_ids: &Array,
        mask: Option<&Array>,
        cache: Option<&mut KVCache>,
    ) -> Result<Array, LoraError> {
        Qwen3LoraForCausalLM::forward_with_cache(self, input_ids, mask, cache)
    }

    fn create_cache(&self, max_seq_len: usize) -> Option<KVCache> {
        Some(Qwen3LoraForCausalLM::create_cache(self, max_seq_len))
    }

    fn supports_kv_cache(&self) -> bool {
        true
    }
}

/// Create a causal attention mask.
fn create_causal_mask(seq_len: i32) -> Result<Array, Exception> {
    let mask = mlx_rs::ops::tri::<f32>(seq_len, None, None)?;
    let neg_inf = Array::from_f32(f32::NEG_INFINITY);
    let zero = Array::from_f32(0.0);
    mlx_rs::ops::r#where(&mask.eq(&zero)?, &neg_inf, &zero)
}

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

    fn small_config() -> Qwen3Config {
        Qwen3Config {
            vocab_size: 1000,
            hidden_size: 64,
            intermediate_size: 128,
            num_hidden_layers: 2,
            num_attention_heads: 4,
            num_key_value_heads: Some(2),
            head_dim: 16,
            max_position_embeddings: 512,
            rms_norm_eps: 1e-6,
            rope_theta: 1_000_000.0,
            ..Default::default()
        }
    }

    fn small_lora_config() -> LoraConfig {
        LoraConfig {
            r: 8,
            alpha: 16.0,
            dropout: 0.0,
            use_rslora: false,
            target_modules: vec![
                "q_proj".to_string(),
                "k_proj".to_string(),
                "v_proj".to_string(),
                "o_proj".to_string(),
            ],
            bias: pmetal_core::LoraBias::None,
            init_lora_weights: true,
            use_dora: false,
        }
    }

    #[test]
    fn test_qwen3_lora_attention() {
        let config = small_config();
        let lora_config = small_lora_config();
        let mut attn = Qwen3LoraAttention::new(&config, &lora_config).unwrap();

        let x = mlx_rs::random::normal::<f32>(&[1, 4, 64], None, None, None).unwrap();
        let output = attn.forward(&x, None, None).unwrap();

        assert_eq!(output.shape(), &[1, 4, 64]);
    }

    #[test]
    fn test_qwen3_lora_model() {
        let config = small_config();
        let lora_config = small_lora_config();
        let mut model = Qwen3LoraForCausalLM::new(config, lora_config).unwrap();

        let input_ids = mlx_rs::Array::from_slice(&[1_i32, 2, 3, 4], &[1, 4]);
        let logits = model.forward(&input_ids, None, None).unwrap();

        assert_eq!(logits.shape(), &[1, 4, 1000]);
    }

    #[test]
    fn test_lora_param_count() {
        let config = small_config();
        let lora_config = small_lora_config();
        let model = Qwen3LoraForCausalLM::new(config, lora_config).unwrap();

        assert!(model.num_trainable_params() > 0);

        let params = model.lora_parameters();
        assert!(!params.is_empty());
    }
}
