unillm-runtime 0.1.0

//! LLaVA Model V2 - Clean implementation using solid abstractions
//!
//! This implements the LLaVA (Large Language and Vision Assistant) architecture including:
//! - LLaVA-1.5-7B, LLaVA-1.5-13B, LLaVA-1.6-7B, LLaVA-1.6-13B, LLaVA-1.6-34B
//!
//! Architecture components:
//! - Vision Tower: CLIP-style vision encoder with patch embeddings and bidirectional attention
//! - Multimodal Projector: Projects vision features to language model dimension (Linear or MLP)
//! - Language Model: LLaMA-style decoder with RoPE, GQA, and SwiGLU MLP

use crate::model_config;
use super::traits::*;
use anyhow::Result;
use serde::{Serialize, Deserialize};

// ============================================================================
// Configuration
// ============================================================================

model_config!(LLaVAConfig {
    // Language model config (based on Llama/Vicuna)
    vocab_size: usize = 32000,
    hidden_size: usize = 4096,
    intermediate_size: usize = 11008,
    num_hidden_layers: usize = 32,
    num_attention_heads: usize = 32,
    num_key_value_heads: usize = 32,
    hidden_act: String = "silu".to_string(),
    max_position_embeddings: usize = 4096,
    initializer_range: f32 = 0.02,
    rms_norm_eps: f32 = 1e-5,
    use_cache: bool = true,
    pad_token_id: i64 = 0,
    bos_token_id: i64 = 1,
    eos_token_id: i64 = 2,
    tie_word_embeddings: bool = false,
    rope_theta: f32 = 10000.0,
    attention_bias: bool = false,
    attention_dropout: f32 = 0.0,

    // Vision config
    vision_hidden_size: usize = 1024,
    vision_intermediate_size: usize = 4096,
    vision_num_hidden_layers: usize = 24,
    vision_num_attention_heads: usize = 16,
    vision_num_channels: usize = 3,
    vision_patch_size: usize = 14,
    vision_image_size: usize = 336,
    vision_layer_norm_eps: f32 = 1e-5,

    // Multimodal config
    mm_projector_type: String = "mlp2x_gelu".to_string(),
    mm_hidden_size: usize = 1024,
    mm_vision_select_layer: i32 = -2,
    mm_vision_select_feature: String = "patch".to_string(),
    image_token_len: usize = 576,
    im_patch_token: i64 = 32000,
    im_start_token: i64 = 32001,
    im_end_token: i64 = 32002,
});

impl LLaVAConfig {
    /// Create LLaVAConfig from GGUF model configuration
    pub fn from_gguf_config(gguf: &crate::weight_loader_core::GGUFModelConfig) -> Self {
        Self {
            vocab_size: gguf.vocab_size,
            hidden_size: gguf.hidden_size,
            intermediate_size: gguf.intermediate_size,
            num_hidden_layers: gguf.num_hidden_layers,
            num_attention_heads: gguf.num_attention_heads,
            num_key_value_heads: gguf.num_key_value_heads,
            rms_norm_eps: gguf.rms_norm_eps,
            rope_theta: gguf.rope_theta,
            max_position_embeddings: gguf.max_position_embeddings,
            ..Default::default()
        }
    }

    /// Get the number of vision patches
    pub fn num_patches(&self) -> usize {
        (self.vision_image_size / self.vision_patch_size).pow(2)
    }

    /// Get the number of vision positions (patches + CLS token)
    pub fn num_vision_positions(&self) -> usize {
        self.num_patches() + 1
    }
}

// ============================================================================
// Vision Tower (CLIP-style)
// ============================================================================

/// CLIP-style Vision Embeddings
pub struct LLaVAVisionEmbeddings {
    patch_embedding_weight: Tensor,  // Conv2d kernel: [hidden, channels, patch, patch]
    patch_embedding_bias: Tensor,    // Conv2d bias: [hidden]
    class_embedding: Tensor,         // [hidden]
    position_embedding: Tensor,      // [num_positions, hidden]
    config: LLaVAConfig,
}

impl LLaVAVisionEmbeddings {
    fn new(config: &LLaVAConfig, device: &Device) -> Result<Self> {
        let num_positions = config.num_vision_positions();
        let hidden = config.vision_hidden_size;
        let patch = config.vision_patch_size;
        let channels = config.vision_num_channels;

        Ok(Self {
            patch_embedding_weight: ops_fn::zeros(
                &[hidden, channels, patch, patch],
                DataType::Float32,
                device,
            )?,
            patch_embedding_bias: ops_fn::zeros(&[hidden], DataType::Float32, device)?,
            class_embedding: ops_fn::zeros(&[hidden], DataType::Float32, device)?,
            position_embedding: ops_fn::zeros(&[num_positions, hidden], DataType::Float32, device)?,
            config: config.clone(),
        })
    }

    fn forward(&self, pixel_values: &Tensor) -> Result<Tensor> {
        let shape = pixel_values.shape();
        let batch_size = shape[0];
        let hidden_size = self.config.vision_hidden_size;
        let num_patches = self.config.num_patches();
        let seq_len = num_patches + 1; // patches + CLS token

        // In a real implementation, we would:
        // 1. Apply Conv2d to extract patch embeddings
        // 2. Flatten patches to [batch, num_patches, hidden]
        // 3. Prepend class embedding
        // 4. Add position embeddings

        // For now, simulate the output shape
        // Real implementation would use conv2d operation
        let pixel_candle = pixel_values.to_candle()?;
        let device = pixel_candle.device();

        // Create patch embeddings by simulating conv2d output
        let patch_embeds = candle_core::Tensor::zeros(
            &[batch_size, num_patches, hidden_size],
            candle_core::DType::F32,
            device,
        )?;

        // Create class embedding for each batch
        let class_emb = self.class_embedding.to_candle()?;
        let class_emb = class_emb.unsqueeze(0)?.unsqueeze(0)?; // [1, 1, hidden]
        let class_emb = class_emb.broadcast_as(&[batch_size, 1, hidden_size])?;

        // Concatenate: [CLS, patches] -> [batch, seq_len, hidden]
        let embeddings = candle_core::Tensor::cat(&[&class_emb, &patch_embeds], 1)?;

        // Add position embeddings
        let pos_emb = self.position_embedding.to_candle()?;
        let pos_emb = pos_emb.unsqueeze(0)?; // [1, seq_len, hidden]
        let embeddings = embeddings.broadcast_add(&pos_emb)?;

        Ok(Tensor::from_candle(embeddings))
    }

    fn load_weights(&mut self, weights: &ModelWeights, prefix: &str) -> Result<()> {
        if let Some(w) = weights.get(&format!("{}.patch_embedding.weight", prefix)) {
            self.patch_embedding_weight = w.clone();
        }
        if let Some(w) = weights.get(&format!("{}.patch_embedding.bias", prefix)) {
            self.patch_embedding_bias = w.clone();
        }
        if let Some(w) = weights.get(&format!("{}.class_embedding", prefix)) {
            self.class_embedding = w.clone();
        }
        if let Some(w) = weights.get(&format!("{}.position_embedding.weight", prefix)) {
            self.position_embedding = w.clone();
        }
        Ok(())
    }

    fn to_device(&mut self, device: &Device) -> Result<()> {
        self.patch_embedding_weight = self.patch_embedding_weight.to_device(device)?;
        self.patch_embedding_bias = self.patch_embedding_bias.to_device(device)?;
        self.class_embedding = self.class_embedding.to_device(device)?;
        self.position_embedding = self.position_embedding.to_device(device)?;
        Ok(())
    }
}

/// CLIP Vision Attention (bidirectional, no causal mask)
pub struct LLaVAVisionAttention {
    q_proj: Tensor,
    k_proj: Tensor,
    v_proj: Tensor,
    out_proj: Tensor,
    num_heads: usize,
    head_dim: usize,
    scale: f32,
}

impl LLaVAVisionAttention {
    fn new(config: &LLaVAConfig, device: &Device) -> Result<Self> {
        let hidden = config.vision_hidden_size;
        let num_heads = config.vision_num_attention_heads;
        let head_dim = hidden / num_heads;
        let scale = 1.0 / (head_dim as f32).sqrt();

        Ok(Self {
            q_proj: ops_fn::zeros(&[hidden, hidden], DataType::Float32, device)?,
            k_proj: ops_fn::zeros(&[hidden, hidden], DataType::Float32, device)?,
            v_proj: ops_fn::zeros(&[hidden, hidden], DataType::Float32, device)?,
            out_proj: ops_fn::zeros(&[hidden, hidden], DataType::Float32, device)?,
            num_heads,
            head_dim,
            scale,
        })
    }

    fn forward(&self, hidden_states: &Tensor) -> Result<Tensor> {
        let shape = hidden_states.shape();
        let (batch_size, seq_len, _) = (shape[0], shape[1], shape[2]);

        // Project to Q, K, V
        let query = ops_fn::matmul(hidden_states, &self.q_proj)?;
        let key = ops_fn::matmul(hidden_states, &self.k_proj)?;
        let value = ops_fn::matmul(hidden_states, &self.v_proj)?;

        // Reshape for multi-head attention
        let q_candle = query.to_candle()?;
        let k_candle = key.to_candle()?;
        let v_candle = value.to_candle()?;

        // [batch, seq, hidden] -> [batch, seq, heads, head_dim] -> [batch, heads, seq, head_dim]
        let q = q_candle
            .reshape(&[batch_size, seq_len, self.num_heads, self.head_dim])?
            .transpose(1, 2)?;
        let k = k_candle
            .reshape(&[batch_size, seq_len, self.num_heads, self.head_dim])?
            .transpose(1, 2)?;
        let v = v_candle
            .reshape(&[batch_size, seq_len, self.num_heads, self.head_dim])?
            .transpose(1, 2)?;

        // Attention scores: Q @ K^T
        let k_t = k.transpose(2, 3)?;
        let scores = q.contiguous()?.matmul(&k_t.contiguous()?)?;
        let scaled_scores = (scores * (self.scale as f64))?;

        // Softmax (bidirectional - no causal mask for vision)
        let attention_weights = candle_nn::ops::softmax_last_dim(&scaled_scores)?;

        // Apply attention to values
        let attn_output = attention_weights.matmul(&v.contiguous()?)?;

        // Reshape back: [batch, heads, seq, head_dim] -> [batch, seq, hidden]
        let attn_output = attn_output
            .transpose(1, 2)?
            .reshape(&[batch_size, seq_len, self.num_heads * self.head_dim])?;

        let attn_output = Tensor::from_candle(attn_output);

        // Output projection
        ops_fn::matmul(&attn_output, &self.out_proj)
    }

    fn load_weights(&mut self, weights: &ModelWeights, prefix: &str) -> Result<()> {
        // Load and transpose for matmul: [out, in] -> [in, out]
        if let Some(w) = weights.get(&format!("{}.q_proj.weight", prefix)) {
            self.q_proj = ops_fn::transpose(w)?;
        }
        if let Some(w) = weights.get(&format!("{}.k_proj.weight", prefix)) {
            self.k_proj = ops_fn::transpose(w)?;
        }
        if let Some(w) = weights.get(&format!("{}.v_proj.weight", prefix)) {
            self.v_proj = ops_fn::transpose(w)?;
        }
        if let Some(w) = weights.get(&format!("{}.out_proj.weight", prefix)) {
            self.out_proj = ops_fn::transpose(w)?;
        }
        Ok(())
    }

    fn to_device(&mut self, device: &Device) -> Result<()> {
        self.q_proj = self.q_proj.to_device(device)?;
        self.k_proj = self.k_proj.to_device(device)?;
        self.v_proj = self.v_proj.to_device(device)?;
        self.out_proj = self.out_proj.to_device(device)?;
        Ok(())
    }
}

/// CLIP Vision MLP (GELU activation)
pub struct LLaVAVisionMLP {
    fc1: Tensor,
    fc2: Tensor,
}

impl LLaVAVisionMLP {
    fn new(config: &LLaVAConfig, device: &Device) -> Result<Self> {
        let hidden = config.vision_hidden_size;
        let intermediate = config.vision_intermediate_size;

        Ok(Self {
            fc1: ops_fn::zeros(&[hidden, intermediate], DataType::Float32, device)?,
            fc2: ops_fn::zeros(&[intermediate, hidden], DataType::Float32, device)?,
        })
    }

    fn forward(&self, hidden_states: &Tensor) -> Result<Tensor> {
        let hidden = ops_fn::matmul(hidden_states, &self.fc1)?;
        let hidden = ops_fn::gelu(&hidden)?;
        ops_fn::matmul(&hidden, &self.fc2)
    }

    fn load_weights(&mut self, weights: &ModelWeights, prefix: &str) -> Result<()> {
        if let Some(w) = weights.get(&format!("{}.fc1.weight", prefix)) {
            self.fc1 = ops_fn::transpose(w)?;
        }
        if let Some(w) = weights.get(&format!("{}.fc2.weight", prefix)) {
            self.fc2 = ops_fn::transpose(w)?;
        }
        Ok(())
    }

    fn to_device(&mut self, device: &Device) -> Result<()> {
        self.fc1 = self.fc1.to_device(device)?;
        self.fc2 = self.fc2.to_device(device)?;
        Ok(())
    }
}

/// CLIP Vision Encoder Layer
pub struct LLaVAVisionLayer {
    self_attn: LLaVAVisionAttention,
    layer_norm1: Tensor,
    mlp: LLaVAVisionMLP,
    layer_norm2: Tensor,
    eps: f32,
}

impl LLaVAVisionLayer {
    fn new(config: &LLaVAConfig, device: &Device) -> Result<Self> {
        let hidden = config.vision_hidden_size;

        Ok(Self {
            self_attn: LLaVAVisionAttention::new(config, device)?,
            layer_norm1: ops_fn::zeros(&[hidden], DataType::Float32, device)?,
            mlp: LLaVAVisionMLP::new(config, device)?,
            layer_norm2: ops_fn::zeros(&[hidden], DataType::Float32, device)?,
            eps: config.vision_layer_norm_eps,
        })
    }

    fn forward(&self, hidden_states: &Tensor) -> Result<Tensor> {
        // Pre-norm attention with residual
        let residual = hidden_states.clone();
        let hidden_states = ops_fn::layer_norm(hidden_states, &self.layer_norm1, None, self.eps)?;
        let hidden_states = self.self_attn.forward(&hidden_states)?;
        let hidden_states = ops_fn::add(&residual, &hidden_states)?;

        // Pre-norm MLP with residual
        let residual = hidden_states.clone();
        let hidden_states = ops_fn::layer_norm(&hidden_states, &self.layer_norm2, None, self.eps)?;
        let hidden_states = self.mlp.forward(&hidden_states)?;
        ops_fn::add(&residual, &hidden_states)
    }

    fn load_weights(&mut self, weights: &ModelWeights, prefix: &str) -> Result<()> {
        if let Some(w) = weights.get(&format!("{}.layer_norm1.weight", prefix)) {
            self.layer_norm1 = w.clone();
        }
        if let Some(w) = weights.get(&format!("{}.layer_norm2.weight", prefix)) {
            self.layer_norm2 = w.clone();
        }
        self.self_attn.load_weights(weights, &format!("{}.self_attn", prefix))?;
        self.mlp.load_weights(weights, &format!("{}.mlp", prefix))?;
        Ok(())
    }

    fn to_device(&mut self, device: &Device) -> Result<()> {
        self.layer_norm1 = self.layer_norm1.to_device(device)?;
        self.layer_norm2 = self.layer_norm2.to_device(device)?;
        self.self_attn.to_device(device)?;
        self.mlp.to_device(device)?;
        Ok(())
    }
}

/// CLIP Vision Encoder
pub struct LLaVAVisionEncoder {
    layers: Vec<LLaVAVisionLayer>,
}

impl LLaVAVisionEncoder {
    fn new(config: &LLaVAConfig, device: &Device) -> Result<Self> {
        let mut layers = Vec::with_capacity(config.vision_num_hidden_layers);
        for _ in 0..config.vision_num_hidden_layers {
            layers.push(LLaVAVisionLayer::new(config, device)?);
        }
        Ok(Self { layers })
    }

    fn forward(&self, hidden_states: &Tensor, select_layer: i32) -> Result<Tensor> {
        let num_layers = self.layers.len() as i32;
        let target_layer = if select_layer < 0 {
            (num_layers + select_layer) as usize
        } else {
            select_layer as usize
        };

        let mut hidden_states = hidden_states.clone();
        for (i, layer) in self.layers.iter().enumerate() {
            hidden_states = layer.forward(&hidden_states)?;
            // Return early if we reached the target layer
            if i == target_layer {
                return Ok(hidden_states);
            }
        }
        Ok(hidden_states)
    }

    fn load_weights(&mut self, weights: &ModelWeights, prefix: &str) -> Result<()> {
        for (i, layer) in self.layers.iter_mut().enumerate() {
            layer.load_weights(weights, &format!("{}.layers.{}", prefix, i))?;
        }
        Ok(())
    }

    fn to_device(&mut self, device: &Device) -> Result<()> {
        for layer in &mut self.layers {
            layer.to_device(device)?;
        }
        Ok(())
    }
}

/// Complete Vision Tower (CLIP ViT)
pub struct LLaVAVisionTower {
    embeddings: LLaVAVisionEmbeddings,
    encoder: LLaVAVisionEncoder,
    post_layernorm: Tensor,
    config: LLaVAConfig,
}

impl LLaVAVisionTower {
    fn new(config: &LLaVAConfig, device: &Device) -> Result<Self> {
        Ok(Self {
            embeddings: LLaVAVisionEmbeddings::new(config, device)?,
            encoder: LLaVAVisionEncoder::new(config, device)?,
            post_layernorm: ops_fn::zeros(
                &[config.vision_hidden_size],
                DataType::Float32,
                device,
            )?,
            config: config.clone(),
        })
    }

    fn forward(&self, pixel_values: &Tensor) -> Result<Tensor> {
        // Embed patches
        let hidden_states = self.embeddings.forward(pixel_values)?;

        // Encode through transformer layers (select layer for multimodal)
        let hidden_states = self.encoder.forward(
            &hidden_states,
            self.config.mm_vision_select_layer,
        )?;

        // Post layer norm
        let hidden_states = ops_fn::layer_norm(
            &hidden_states,
            &self.post_layernorm,
            None,
            self.config.vision_layer_norm_eps,
        )?;

        // Select features based on config
        if self.config.mm_vision_select_feature == "patch" {
            // Remove CLS token, keep only patch features
            let candle_tensor = hidden_states.to_candle()?;
            let shape = candle_tensor.dims();
            let patch_features = candle_tensor.narrow(1, 1, shape[1] - 1)?;
            Ok(Tensor::from_candle(patch_features))
        } else {
            // Return all features including CLS
            Ok(hidden_states)
        }
    }

    fn load_weights(&mut self, weights: &ModelWeights) -> Result<()> {
        let prefix = "vision_tower.vision_model";
        self.embeddings.load_weights(weights, &format!("{}.embeddings", prefix))?;
        self.encoder.load_weights(weights, &format!("{}.encoder", prefix))?;
        if let Some(w) = weights.get(&format!("{}.post_layernorm.weight", prefix)) {
            self.post_layernorm = w.clone();
        }
        Ok(())
    }

    fn to_device(&mut self, device: &Device) -> Result<()> {
        self.embeddings.to_device(device)?;
        self.encoder.to_device(device)?;
        self.post_layernorm = self.post_layernorm.to_device(device)?;
        Ok(())
    }
}

// ============================================================================
// Multimodal Projector
// ============================================================================

/// Projects vision features to language model dimension
pub struct LLaVAMultiModalProjector {
    projector_type: String,
    // For "linear" type
    linear: Option<Tensor>,
    linear_bias: Option<Tensor>,
    // For "mlp2x_gelu" type (2-layer MLP with GELU)
    mlp_fc1: Option<Tensor>,
    mlp_fc1_bias: Option<Tensor>,
    mlp_fc2: Option<Tensor>,
    mlp_fc2_bias: Option<Tensor>,
}

impl LLaVAMultiModalProjector {
    fn new(config: &LLaVAConfig, device: &Device) -> Result<Self> {
        let vision_hidden = config.vision_hidden_size;
        let text_hidden = config.hidden_size;

        match config.mm_projector_type.as_str() {
            "linear" => Ok(Self {
                projector_type: "linear".to_string(),
                linear: Some(ops_fn::zeros(
                    &[vision_hidden, text_hidden],
                    DataType::Float32,
                    device,
                )?),
                linear_bias: Some(ops_fn::zeros(&[text_hidden], DataType::Float32, device)?),
                mlp_fc1: None,
                mlp_fc1_bias: None,
                mlp_fc2: None,
                mlp_fc2_bias: None,
            }),
            "mlp2x_gelu" => Ok(Self {
                projector_type: "mlp2x_gelu".to_string(),
                linear: None,
                linear_bias: None,
                mlp_fc1: Some(ops_fn::zeros(
                    &[vision_hidden, text_hidden],
                    DataType::Float32,
                    device,
                )?),
                mlp_fc1_bias: Some(ops_fn::zeros(&[text_hidden], DataType::Float32, device)?),
                mlp_fc2: Some(ops_fn::zeros(
                    &[text_hidden, text_hidden],
                    DataType::Float32,
                    device,
                )?),
                mlp_fc2_bias: Some(ops_fn::zeros(&[text_hidden], DataType::Float32, device)?),
            }),
            _ => Err(anyhow::anyhow!(
                "Unsupported projector type: {}",
                config.mm_projector_type
            )),
        }
    }

    fn forward(&self, vision_features: &Tensor) -> Result<Tensor> {
        match self.projector_type.as_str() {
            "linear" => {
                let linear = self.linear.as_ref().ok_or_else(|| {
                    anyhow::anyhow!("Linear projector not initialized")
                })?;
                let output = ops_fn::matmul(vision_features, linear)?;
                if let Some(ref bias) = self.linear_bias {
                    ops_fn::add(&output, bias)
                } else {
                    Ok(output)
                }
            }
            "mlp2x_gelu" => {
                let fc1 = self.mlp_fc1.as_ref().ok_or_else(|| {
                    anyhow::anyhow!("MLP fc1 not initialized")
                })?;
                let fc2 = self.mlp_fc2.as_ref().ok_or_else(|| {
                    anyhow::anyhow!("MLP fc2 not initialized")
                })?;

                // First layer with GELU
                let mut hidden = ops_fn::matmul(vision_features, fc1)?;
                if let Some(ref bias) = self.mlp_fc1_bias {
                    hidden = ops_fn::add(&hidden, bias)?;
                }
                hidden = ops_fn::gelu(&hidden)?;

                // Second layer
                let mut output = ops_fn::matmul(&hidden, fc2)?;
                if let Some(ref bias) = self.mlp_fc2_bias {
                    output = ops_fn::add(&output, bias)?;
                }
                Ok(output)
            }
            _ => Err(anyhow::anyhow!(
                "Unsupported projector type: {}",
                self.projector_type
            )),
        }
    }

    fn load_weights(&mut self, weights: &ModelWeights) -> Result<()> {
        match self.projector_type.as_str() {
            "linear" => {
                if let Some(w) = weights.get("mm_projector.weight") {
                    self.linear = Some(ops_fn::transpose(w)?);
                }
                if let Some(w) = weights.get("mm_projector.bias") {
                    self.linear_bias = Some(w.clone());
                }
            }
            "mlp2x_gelu" => {
                // MLP projector has 2 layers: 0 and 2 (with GELU at index 1)
                if let Some(w) = weights.get("mm_projector.0.weight") {
                    self.mlp_fc1 = Some(ops_fn::transpose(w)?);
                }
                if let Some(w) = weights.get("mm_projector.0.bias") {
                    self.mlp_fc1_bias = Some(w.clone());
                }
                if let Some(w) = weights.get("mm_projector.2.weight") {
                    self.mlp_fc2 = Some(ops_fn::transpose(w)?);
                }
                if let Some(w) = weights.get("mm_projector.2.bias") {
                    self.mlp_fc2_bias = Some(w.clone());
                }
            }
            _ => {}
        }
        Ok(())
    }

    fn to_device(&mut self, device: &Device) -> Result<()> {
        if let Some(ref mut w) = self.linear {
            *w = w.to_device(device)?;
        }
        if let Some(ref mut w) = self.linear_bias {
            *w = w.to_device(device)?;
        }
        if let Some(ref mut w) = self.mlp_fc1 {
            *w = w.to_device(device)?;
        }
        if let Some(ref mut w) = self.mlp_fc1_bias {
            *w = w.to_device(device)?;
        }
        if let Some(ref mut w) = self.mlp_fc2 {
            *w = w.to_device(device)?;
        }
        if let Some(ref mut w) = self.mlp_fc2_bias {
            *w = w.to_device(device)?;
        }
        Ok(())
    }
}

// ============================================================================
// Language Model (LLaMA-style)
// ============================================================================

/// LLaMA-style Attention with RoPE and GQA
pub struct LLaVALanguageAttention {
    q_proj: Tensor,
    k_proj: Tensor,
    v_proj: Tensor,
    o_proj: Tensor,
    num_heads: usize,
    num_key_value_heads: usize,
    head_dim: usize,
    scale: f32,
}

impl LLaVALanguageAttention {
    fn new(config: &LLaVAConfig, device: &Device) -> Result<Self> {
        let num_heads = config.num_attention_heads;
        let num_kv_heads = config.num_key_value_heads;
        let head_dim = config.hidden_size / num_heads;
        let scale = 1.0 / (head_dim as f32).sqrt();

        Ok(Self {
            q_proj: ops_fn::zeros(
                &[config.hidden_size, num_heads * head_dim],
                DataType::Float32,
                device,
            )?,
            k_proj: ops_fn::zeros(
                &[config.hidden_size, num_kv_heads * head_dim],
                DataType::Float32,
                device,
            )?,
            v_proj: ops_fn::zeros(
                &[config.hidden_size, num_kv_heads * head_dim],
                DataType::Float32,
                device,
            )?,
            o_proj: ops_fn::zeros(
                &[num_heads * head_dim, config.hidden_size],
                DataType::Float32,
                device,
            )?,
            num_heads,
            num_key_value_heads: num_kv_heads,
            head_dim,
            scale,
        })
    }

    fn forward(&self, hidden_states: &Tensor, rope_theta: f32) -> Result<Tensor> {
        let shape = hidden_states.shape();
        let (batch_size, seq_len, _) = if shape.len() == 3 {
            (shape[0], shape[1], shape[2])
        } else if shape.len() == 2 {
            (1, shape[0], shape[1])
        } else {
            return Err(anyhow::anyhow!("Invalid hidden_states shape: {:?}", shape));
        };

        // Project to Q, K, V
        let query_states = ops_fn::matmul(hidden_states, &self.q_proj)?;
        let key_states = ops_fn::matmul(hidden_states, &self.k_proj)?;
        let value_states = ops_fn::matmul(hidden_states, &self.v_proj)?;

        // Reshape for multi-head attention
        let q_candle = query_states.to_candle()?;
        let k_candle = key_states.to_candle()?;
        let v_candle = value_states.to_candle()?;

        let q_reshaped = q_candle
            .reshape(&[batch_size, seq_len, self.num_heads, self.head_dim])?
            .transpose(1, 2)?;
        let k_reshaped = k_candle
            .reshape(&[batch_size, seq_len, self.num_key_value_heads, self.head_dim])?
            .transpose(1, 2)?;
        let v_reshaped = v_candle
            .reshape(&[batch_size, seq_len, self.num_key_value_heads, self.head_dim])?
            .transpose(1, 2)?;

        // Apply RoPE to Q and K
        let (q_with_rope, k_with_rope) = apply_rope(
            &q_reshaped,
            &k_reshaped,
            seq_len,
            self.head_dim,
            rope_theta,
        )?;

        // Handle GQA - repeat K/V heads to match Q heads
        let num_groups = self.num_heads / self.num_key_value_heads;
        let (k_expanded, v_expanded) = if num_groups > 1 {
            let k_rep = k_with_rope
                .unsqueeze(2)?
                .broadcast_as(&[
                    batch_size,
                    self.num_key_value_heads,
                    num_groups,
                    seq_len,
                    self.head_dim,
                ])?
                .reshape(&[batch_size, self.num_heads, seq_len, self.head_dim])?;
            let v_rep = v_reshaped
                .unsqueeze(2)?
                .broadcast_as(&[
                    batch_size,
                    self.num_key_value_heads,
                    num_groups,
                    seq_len,
                    self.head_dim,
                ])?
                .reshape(&[batch_size, self.num_heads, seq_len, self.head_dim])?;
            (k_rep, v_rep)
        } else {
            (k_with_rope, v_reshaped)
        };

        // Attention scores with causal mask
        let k_t = k_expanded.transpose(2, 3)?;
        let scores = q_with_rope.contiguous()?.matmul(&k_t.contiguous()?)?;
        let scaled_scores = (scores * (self.scale as f64))?;

        // Apply causal mask
        let device = scaled_scores.device();
        let causal_mask = {
            let mut mask_data = vec![0.0f32; seq_len * seq_len];
            for i in 0..seq_len {
                for j in 0..seq_len {
                    if j > i {
                        mask_data[i * seq_len + j] = f32::NEG_INFINITY;
                    }
                }
            }
            candle_core::Tensor::from_vec(mask_data, &[1, 1, seq_len, seq_len], device)?
        };
        let masked_scores = scaled_scores.broadcast_add(&causal_mask)?;

        // Softmax and apply to values
        let attention_weights = candle_nn::ops::softmax_last_dim(&masked_scores)?;
        let attn_output = attention_weights.matmul(&v_expanded.contiguous()?)?;

        // Reshape back
        let attn_output = attn_output
            .transpose(1, 2)?
            .reshape(&[batch_size, seq_len, self.num_heads * self.head_dim])?;

        let attn_output = Tensor::from_candle(attn_output);

        // Output projection
        ops_fn::matmul(&attn_output, &self.o_proj)
    }

    fn load_weights(&mut self, weights: &ModelWeights, layer_idx: usize) -> Result<()> {
        let prefix = format!("language_model.model.layers.{}.self_attn", layer_idx);
        if let Some(w) = weights.get(&format!("{}.q_proj.weight", prefix)) {
            self.q_proj = ops_fn::transpose(w)?;
        }
        if let Some(w) = weights.get(&format!("{}.k_proj.weight", prefix)) {
            self.k_proj = ops_fn::transpose(w)?;
        }
        if let Some(w) = weights.get(&format!("{}.v_proj.weight", prefix)) {
            self.v_proj = ops_fn::transpose(w)?;
        }
        if let Some(w) = weights.get(&format!("{}.o_proj.weight", prefix)) {
            self.o_proj = ops_fn::transpose(w)?;
        }
        Ok(())
    }

    fn to_device(&mut self, device: &Device) -> Result<()> {
        self.q_proj = self.q_proj.to_device(device)?;
        self.k_proj = self.k_proj.to_device(device)?;
        self.v_proj = self.v_proj.to_device(device)?;
        self.o_proj = self.o_proj.to_device(device)?;
        Ok(())
    }
}

/// LLaMA-style MLP with SwiGLU
pub struct LLaVALanguageMLP {
    gate_proj: Tensor,
    up_proj: Tensor,
    down_proj: Tensor,
}

impl LLaVALanguageMLP {
    fn new(config: &LLaVAConfig, device: &Device) -> Result<Self> {
        Ok(Self {
            gate_proj: ops_fn::zeros(
                &[config.hidden_size, config.intermediate_size],
                DataType::Float32,
                device,
            )?,
            up_proj: ops_fn::zeros(
                &[config.hidden_size, config.intermediate_size],
                DataType::Float32,
                device,
            )?,
            down_proj: ops_fn::zeros(
                &[config.intermediate_size, config.hidden_size],
                DataType::Float32,
                device,
            )?,
        })
    }

    fn forward(&self, hidden_states: &Tensor) -> Result<Tensor> {
        // SwiGLU: down(silu(gate(x)) * up(x))
        let gate_output = ops_fn::matmul(hidden_states, &self.gate_proj)?;
        let up_output = ops_fn::matmul(hidden_states, &self.up_proj)?;
        let gate_activated = ops_fn::silu(&gate_output)?;
        let gated = ops_fn::mul(&gate_activated, &up_output)?;
        ops_fn::matmul(&gated, &self.down_proj)
    }

    fn load_weights(&mut self, weights: &ModelWeights, layer_idx: usize) -> Result<()> {
        let prefix = format!("language_model.model.layers.{}.mlp", layer_idx);
        if let Some(w) = weights.get(&format!("{}.gate_proj.weight", prefix)) {
            self.gate_proj = ops_fn::transpose(w)?;
        }
        if let Some(w) = weights.get(&format!("{}.up_proj.weight", prefix)) {
            self.up_proj = ops_fn::transpose(w)?;
        }
        if let Some(w) = weights.get(&format!("{}.down_proj.weight", prefix)) {
            self.down_proj = ops_fn::transpose(w)?;
        }
        Ok(())
    }

    fn to_device(&mut self, device: &Device) -> Result<()> {
        self.gate_proj = self.gate_proj.to_device(device)?;
        self.up_proj = self.up_proj.to_device(device)?;
        self.down_proj = self.down_proj.to_device(device)?;
        Ok(())
    }
}

/// LLaMA-style Transformer Layer
pub struct LLaVALanguageLayer {
    self_attn: LLaVALanguageAttention,
    mlp: LLaVALanguageMLP,
    input_layernorm: Tensor,
    post_attention_layernorm: Tensor,
    rms_norm_eps: f32,
}

impl LLaVALanguageLayer {
    fn new(config: &LLaVAConfig, device: &Device) -> Result<Self> {
        Ok(Self {
            self_attn: LLaVALanguageAttention::new(config, device)?,
            mlp: LLaVALanguageMLP::new(config, device)?,
            input_layernorm: ops_fn::zeros(&[config.hidden_size], DataType::Float32, device)?,
            post_attention_layernorm: ops_fn::zeros(
                &[config.hidden_size],
                DataType::Float32,
                device,
            )?,
            rms_norm_eps: config.rms_norm_eps,
        })
    }

    fn forward(&self, hidden_states: &Tensor, rope_theta: f32) -> Result<Tensor> {
        // Pre-norm attention with residual
        let normed = ops_fn::layer_norm(
            hidden_states,
            &self.input_layernorm,
            None,
            self.rms_norm_eps,
        )?;
        let attn_output = self.self_attn.forward(&normed, rope_theta)?;
        let hidden_states = ops_fn::add(hidden_states, &attn_output)?;

        // Pre-norm MLP with residual
        let normed = ops_fn::layer_norm(
            &hidden_states,
            &self.post_attention_layernorm,
            None,
            self.rms_norm_eps,
        )?;
        let mlp_output = self.mlp.forward(&normed)?;
        ops_fn::add(&hidden_states, &mlp_output)
    }

    fn load_weights(&mut self, weights: &ModelWeights, layer_idx: usize) -> Result<()> {
        let prefix = format!("language_model.model.layers.{}", layer_idx);
        if let Some(w) = weights.get(&format!("{}.input_layernorm.weight", prefix)) {
            self.input_layernorm = w.clone();
        }
        if let Some(w) = weights.get(&format!("{}.post_attention_layernorm.weight", prefix)) {
            self.post_attention_layernorm = w.clone();
        }
        self.self_attn.load_weights(weights, layer_idx)?;
        self.mlp.load_weights(weights, layer_idx)?;
        Ok(())
    }

    fn to_device(&mut self, device: &Device) -> Result<()> {
        self.input_layernorm = self.input_layernorm.to_device(device)?;
        self.post_attention_layernorm = self.post_attention_layernorm.to_device(device)?;
        self.self_attn.to_device(device)?;
        self.mlp.to_device(device)?;
        Ok(())
    }
}

/// Complete Language Model
pub struct LLaVALanguageModel {
    embed_tokens: Tensor,
    layers: Vec<LLaVALanguageLayer>,
    norm: Tensor,
    lm_head: Tensor,
    config: LLaVAConfig,
}

impl LLaVALanguageModel {
    fn new(config: &LLaVAConfig, device: &Device) -> Result<Self> {
        let embed_tokens = ops_fn::zeros(
            &[config.vocab_size, config.hidden_size],
            DataType::Float32,
            device,
        )?;
        let norm = ops_fn::zeros(&[config.hidden_size], DataType::Float32, device)?;
        let lm_head = if config.tie_word_embeddings {
            embed_tokens.clone()
        } else {
            ops_fn::zeros(
                &[config.hidden_size, config.vocab_size],
                DataType::Float32,
                device,
            )?
        };

        let mut layers = Vec::with_capacity(config.num_hidden_layers);
        for _ in 0..config.num_hidden_layers {
            layers.push(LLaVALanguageLayer::new(config, device)?);
        }

        Ok(Self {
            embed_tokens,
            layers,
            norm,
            lm_head,
            config: config.clone(),
        })
    }

    fn forward(&self, hidden_states: &Tensor) -> Result<Tensor> {
        let mut hidden_states = hidden_states.clone();

        // Apply transformer layers
        for layer in &self.layers {
            hidden_states = layer.forward(&hidden_states, self.config.rope_theta)?;
        }

        // Final norm
        hidden_states = ops_fn::layer_norm(
            &hidden_states,
            &self.norm,
            None,
            self.config.rms_norm_eps,
        )?;

        // LM head
        ops_fn::matmul(&hidden_states, &self.lm_head)
    }

    fn forward_from_ids(&self, input_ids: &Tensor) -> Result<Tensor> {
        let hidden_states = ops_fn::embedding(input_ids, &self.embed_tokens)?;
        self.forward(&hidden_states)
    }

    fn embed(&self, input_ids: &Tensor) -> Result<Tensor> {
        ops_fn::embedding(input_ids, &self.embed_tokens)
    }

    fn load_weights(&mut self, weights: &ModelWeights) -> Result<()> {
        if let Some(w) = weights.get("language_model.model.embed_tokens.weight") {
            self.embed_tokens = w.clone();
        }
        if let Some(w) = weights.get("language_model.model.norm.weight") {
            self.norm = w.clone();
        }
        if let Some(w) = weights.get("language_model.lm_head.weight") {
            self.lm_head = ops_fn::transpose(w)?;
        }
        for (i, layer) in self.layers.iter_mut().enumerate() {
            layer.load_weights(weights, i)?;
        }
        Ok(())
    }

    fn to_device(&mut self, device: &Device) -> Result<()> {
        self.embed_tokens = self.embed_tokens.to_device(device)?;
        self.norm = self.norm.to_device(device)?;
        self.lm_head = self.lm_head.to_device(device)?;
        for layer in &mut self.layers {
            layer.to_device(device)?;
        }
        Ok(())
    }
}

// ============================================================================
// RoPE Implementation
// ============================================================================

/// Apply Rotary Position Embedding (RoPE) to Q and K tensors
fn apply_rope(
    q: &candle_core::Tensor,
    k: &candle_core::Tensor,
    seq_len: usize,
    head_dim: usize,
    rope_theta: f32,
) -> Result<(candle_core::Tensor, candle_core::Tensor)> {
    let device = q.device();
    let half_dim = head_dim / 2;

    // Compute inverse frequencies
    let inv_freq: Vec<f32> = (0..half_dim)
        .map(|i| 1.0 / rope_theta.powf((2 * i) as f32 / head_dim as f32))
        .collect();

    // Create position indices
    let positions: Vec<f32> = (0..seq_len).map(|p| p as f32).collect();

    // Compute angles: pos * inv_freq -> [seq_len, half_dim]
    let mut angles = Vec::with_capacity(seq_len * half_dim);
    for pos in &positions {
        for freq in &inv_freq {
            angles.push(pos * freq);
        }
    }

    let angles_tensor = candle_core::Tensor::from_vec(angles, &[seq_len, half_dim], device)?;
    let cos = angles_tensor.cos()?;
    let sin = angles_tensor.sin()?;

    // Reshape for broadcasting: [1, 1, seq_len, half_dim]
    let cos = cos.unsqueeze(0)?.unsqueeze(0)?;
    let sin = sin.unsqueeze(0)?.unsqueeze(0)?;

    // Apply rotation
    let q_half1 = q.narrow(3, 0, half_dim)?;
    let q_half2 = q.narrow(3, half_dim, half_dim)?;
    let k_half1 = k.narrow(3, 0, half_dim)?;
    let k_half2 = k.narrow(3, half_dim, half_dim)?;

    let q_rot1 = (q_half1.broadcast_mul(&cos)? - q_half2.broadcast_mul(&sin)?)?;
    let q_rot2 = (q_half1.broadcast_mul(&sin)? + q_half2.broadcast_mul(&cos)?)?;
    let k_rot1 = (k_half1.broadcast_mul(&cos)? - k_half2.broadcast_mul(&sin)?)?;
    let k_rot2 = (k_half1.broadcast_mul(&sin)? + k_half2.broadcast_mul(&cos)?)?;

    let q_rotated = candle_core::Tensor::cat(&[&q_rot1, &q_rot2], 3)?;
    let k_rotated = candle_core::Tensor::cat(&[&k_rot1, &k_rot2], 3)?;

    Ok((q_rotated, k_rotated))
}

// ============================================================================
// Main LLaVA Model
// ============================================================================

/// Complete LLaVA Model
pub struct LLaVAModelV2 {
    config: LLaVAConfig,
    device: Device,
    vision_tower: LLaVAVisionTower,
    mm_projector: LLaVAMultiModalProjector,
    language_model: LLaVALanguageModel,
}

impl Model for LLaVAModelV2 {
    type Config = LLaVAConfig;

    fn new(config: LLaVAConfig) -> Result<Self> {
        let device = Device::CPU;
        let vision_tower = LLaVAVisionTower::new(&config, &device)?;
        let mm_projector = LLaVAMultiModalProjector::new(&config, &device)?;
        let language_model = LLaVALanguageModel::new(&config, &device)?;

        Ok(Self {
            config,
            device,
            vision_tower,
            mm_projector,
            language_model,
        })
    }

    fn from_weights(config: LLaVAConfig, weights: ModelWeights) -> Result<Self> {
        let mut model = Self::new(config)?;
        model.vision_tower.load_weights(&weights)?;
        model.mm_projector.load_weights(&weights)?;
        model.language_model.load_weights(&weights)?;
        Ok(model)
    }

    fn forward(&self, inputs: &ModelInputs) -> Result<ModelOutputs> {
        match inputs {
            ModelInputs::Multimodal {
                input_ids,
                pixel_values,
                ..
            } => {
                if let Some(pixel_values) = pixel_values {
                    // Process image through vision tower
                    let image_features = self.vision_tower.forward(pixel_values)?;

                    // Project to language model dimension
                    let image_features = self.mm_projector.forward(&image_features)?;

                    // Get text embeddings
                    let text_embeds = self.language_model.embed(input_ids)?;

                    // Merge multimodal inputs
                    let merged_embeds = self.merge_multimodal_inputs(
                        input_ids,
                        &text_embeds,
                        &image_features,
                    )?;

                    // Forward through language model
                    let logits = self.language_model.forward(&merged_embeds)?;

                    Ok(ModelOutputs::Logits {
                        logits,
                        hidden_states: Some(merged_embeds),
                    })
                } else {
                    // No image, just process text
                    let logits = self.language_model.forward_from_ids(input_ids)?;
                    Ok(ModelOutputs::Logits {
                        logits,
                        hidden_states: None,
                    })
                }
            }
            ModelInputs::Text { input_ids, .. } => {
                let logits = self.language_model.forward_from_ids(input_ids)?;
                Ok(ModelOutputs::Logits {
                    logits,
                    hidden_states: None,
                })
            }
            ModelInputs::Image { pixel_values, .. } => {
                // Image-only: return image embeddings
                let image_features = self.vision_tower.forward(pixel_values)?;
                let image_features = self.mm_projector.forward(&image_features)?;
                Ok(ModelOutputs::Embeddings {
                    embeddings: image_features,
                    pooled: None,
                })
            }
            _ => Err(anyhow::anyhow!("LLaVA expects text, image, or multimodal input")),
        }
    }

    fn generate(&self, prompt: &str, config: &GenerationConfig) -> Result<String> {
        use crate::tokenizer::Tokenizer;
        use rand::Rng;

        // Tokenize prompt
        let tokenizer = Tokenizer::new();
        let mut tokens: Vec<u32> = tokenizer.encode(prompt);

        // Generation loop
        for _ in 0..config.max_new_tokens {
            let tokens_i64: Vec<i64> = tokens.iter().map(|&t| t as i64).collect();
            let input_tensor = Tensor::from_i64_slice(&tokens_i64, &[1, tokens.len()], &self.device)?;

            let inputs = ModelInputs::Text {
                input_ids: input_tensor,
                attention_mask: None,
                position_ids: None,
            };

            let outputs = self.forward(&inputs)?;
            let logits = match outputs {
                ModelOutputs::Logits { logits, .. } => logits,
                _ => return Err(anyhow::anyhow!("Expected logits output")),
            };

            // Get last token logits
            let logits_candle = logits.to_candle()?;
            let shape = logits_candle.dims();
            let last_logits = if shape.len() == 3 {
                let seq_len = shape[1];
                logits_candle
                    .narrow(1, seq_len - 1, 1)?
                    .squeeze(1)?
                    .squeeze(0)?
            } else {
                let seq_len = shape[0];
                logits_candle.narrow(0, seq_len - 1, 1)?.squeeze(0)?
            };

            let logits_vec: Vec<f32> = last_logits.to_vec1()?;

            let next_token = if config.do_sample && config.temperature > 0.0 {
                // Temperature sampling
                let scaled: Vec<f32> = logits_vec
                    .iter()
                    .map(|&x| x / config.temperature)
                    .collect();
                let max_val = scaled.iter().cloned().fold(f32::NEG_INFINITY, f32::max);
                let exp_sum: f32 = scaled.iter().map(|&x| (x - max_val).exp()).sum();
                let probs: Vec<f32> = scaled
                    .iter()
                    .map(|&x| (x - max_val).exp() / exp_sum)
                    .collect();

                let mut rng = rand::thread_rng();
                let random_val: f32 = rng.gen();
                let mut cumulative = 0.0;
                let mut sampled = 0u32;
                for (idx, &prob) in probs.iter().enumerate() {
                    cumulative += prob;
                    if random_val <= cumulative {
                        sampled = idx as u32;
                        break;
                    }
                }
                sampled
            } else {
                // Greedy sampling
                let mut max_idx = 0;
                let mut max_val = logits_vec[0];
                for (idx, &val) in logits_vec.iter().enumerate() {
                    if val > max_val {
                        max_val = val;
                        max_idx = idx;
                    }
                }
                max_idx as u32
            };

            if next_token == config.eos_token_id {
                break;
            }

            tokens.push(next_token);
        }

        Ok(tokenizer.decode(&tokens))
    }

    fn config(&self) -> &Self::Config {
        &self.config
    }

    fn memory_requirements(&self) -> MemoryRequirements {
        let language_params = self.config.vocab_size * self.config.hidden_size
            + self.config.num_hidden_layers
                * (4 * self.config.hidden_size * self.config.hidden_size
                    + 3 * self.config.hidden_size * self.config.intermediate_size);

        let vision_params = self.config.vision_num_hidden_layers
            * (4 * self.config.vision_hidden_size * self.config.vision_hidden_size
                + 2 * self.config.vision_hidden_size * self.config.vision_intermediate_size);

        let projector_params = self.config.vision_hidden_size * self.config.hidden_size * 2;

        let total_params = language_params + vision_params + projector_params;
        let param_bytes = total_params * 4; // float32

        let kv_cache_bytes = 2
            * self.config.num_hidden_layers
            * self.config.max_position_embeddings
            * self.config.hidden_size
            * 4;

        MemoryRequirements {
            gpu_memory: param_bytes,
            cpu_memory: param_bytes / 4,
            kv_cache_memory: kv_cache_bytes,
            peak_memory: param_bytes + kv_cache_bytes,
        }
    }

    fn to_device(&mut self, device: &Device) -> Result<()> {
        self.vision_tower.to_device(device)?;
        self.mm_projector.to_device(device)?;
        self.language_model.to_device(device)?;
        self.device = device.clone();
        Ok(())
    }
}

impl LLaVAModelV2 {
    /// Merge vision features into text sequence at image token positions
    ///
    /// This replaces <image> tokens in the text sequence with actual image features.
    /// The resulting sequence is: [text_before_image, image_features, text_after_image]
    fn merge_multimodal_inputs(
        &self,
        input_ids: &Tensor,
        text_embeds: &Tensor,
        image_features: &Tensor,
    ) -> Result<Tensor> {
        let input_candle = input_ids.to_candle()?;
        let text_candle = text_embeds.to_candle()?;
        let image_candle = image_features.to_candle()?;

        let batch_size = text_candle.dims()[0];
        let text_seq_len = text_candle.dims()[1];
        let hidden_size = text_candle.dims()[2];
        let image_seq_len = image_candle.dims()[1];

        // Find image token positions
        let input_flat = input_candle.flatten_all()?;
        let input_vec: Vec<i64> = input_flat.to_vec1()?;

        let image_token_id = self.config.im_patch_token;

        // Find first image token position (if any)
        let image_pos = input_vec.iter().position(|&id| id == image_token_id);

        let merged = if let Some(pos) = image_pos {
            // Split text at image position
            let pos = pos % text_seq_len; // Handle batch dimension

            if pos == 0 {
                // Image at start: [image, text]
                let text_after = text_candle.narrow(1, 1, text_seq_len - 1)?;
                candle_core::Tensor::cat(&[&image_candle, &text_after], 1)?
            } else if pos >= text_seq_len - 1 {
                // Image at end: [text, image]
                let text_before = text_candle.narrow(1, 0, text_seq_len - 1)?;
                candle_core::Tensor::cat(&[&text_before, &image_candle], 1)?
            } else {
                // Image in middle: [text_before, image, text_after]
                let text_before = text_candle.narrow(1, 0, pos)?;
                let text_after = text_candle.narrow(1, pos + 1, text_seq_len - pos - 1)?;
                candle_core::Tensor::cat(&[&text_before, &image_candle, &text_after], 1)?
            }
        } else {
            // No image token found, prepend image features
            candle_core::Tensor::cat(&[&image_candle, &text_candle], 1)?
        };

        Ok(Tensor::from_candle(merged))
    }

    /// Generate response for multimodal input (text + image)
    pub fn generate_multimodal(
        &self,
        prompt: &str,
        image: &Tensor,
        config: &GenerationConfig,
    ) -> Result<String> {
        use crate::tokenizer::Tokenizer;
        use rand::Rng;

        let tokenizer = Tokenizer::new();
        let mut tokens: Vec<u32> = tokenizer.encode(prompt);

        // Process image once
        let image_features = self.vision_tower.forward(image)?;
        let image_features = self.mm_projector.forward(&image_features)?;

        for _ in 0..config.max_new_tokens {
            let tokens_i64: Vec<i64> = tokens.iter().map(|&t| t as i64).collect();
            let input_tensor = Tensor::from_i64_slice(&tokens_i64, &[1, tokens.len()], &self.device)?;

            // Get text embeddings
            let text_embeds = self.language_model.embed(&input_tensor)?;

            // Merge with image features
            let merged_embeds = self.merge_multimodal_inputs(
                &input_tensor,
                &text_embeds,
                &image_features,
            )?;

            // Forward through language model
            let logits = self.language_model.forward(&merged_embeds)?;

            // Sample next token
            let logits_candle = logits.to_candle()?;
            let shape = logits_candle.dims();
            let last_logits = if shape.len() == 3 {
                let seq_len = shape[1];
                logits_candle
                    .narrow(1, seq_len - 1, 1)?
                    .squeeze(1)?
                    .squeeze(0)?
            } else {
                let seq_len = shape[0];
                logits_candle.narrow(0, seq_len - 1, 1)?.squeeze(0)?
            };

            let logits_vec: Vec<f32> = last_logits.to_vec1()?;

            let next_token = if config.do_sample && config.temperature > 0.0 {
                let scaled: Vec<f32> = logits_vec
                    .iter()
                    .map(|&x| x / config.temperature)
                    .collect();
                let max_val = scaled.iter().cloned().fold(f32::NEG_INFINITY, f32::max);
                let exp_sum: f32 = scaled.iter().map(|&x| (x - max_val).exp()).sum();
                let probs: Vec<f32> = scaled
                    .iter()
                    .map(|&x| (x - max_val).exp() / exp_sum)
                    .collect();

                let mut rng = rand::thread_rng();
                let random_val: f32 = rng.gen();
                let mut cumulative = 0.0;
                let mut sampled = 0u32;
                for (idx, &prob) in probs.iter().enumerate() {
                    cumulative += prob;
                    if random_val <= cumulative {
                        sampled = idx as u32;
                        break;
                    }
                }
                sampled
            } else {
                let mut max_idx = 0;
                let mut max_val = logits_vec[0];
                for (idx, &val) in logits_vec.iter().enumerate() {
                    if val > max_val {
                        max_val = val;
                        max_idx = idx;
                    }
                }
                max_idx as u32
            };

            if next_token == config.eos_token_id {
                break;
            }

            tokens.push(next_token);
        }

        Ok(tokenizer.decode(&tokens))
    }
}

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

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

    #[test]
    fn test_llava_config_defaults() {
        let config = LLaVAConfig::default();
        assert_eq!(config.vocab_size, 32000);
        assert_eq!(config.hidden_size, 4096);
        assert_eq!(config.vision_hidden_size, 1024);
        assert_eq!(config.num_patches(), 576); // (336/14)^2
    }

    #[test]
    fn test_llava_model_creation() {
        let config = LLaVAConfig {
            vocab_size: 1000,
            hidden_size: 128,
            intermediate_size: 512,
            num_hidden_layers: 2,
            num_attention_heads: 4,
            num_key_value_heads: 4,
            vision_hidden_size: 64,
            vision_intermediate_size: 256,
            vision_num_hidden_layers: 2,
            vision_num_attention_heads: 4,
            vision_patch_size: 14,
            vision_image_size: 224,
            ..Default::default()
        };

        let model = LLaVAModelV2::new(config).unwrap();
        assert_eq!(model.config().vocab_size(), 1000);
        assert_eq!(model.config().hidden_size(), 128);
    }

    #[test]
    fn test_llava_text_forward() {
        let config = LLaVAConfig {
            vocab_size: 100,
            hidden_size: 64,
            intermediate_size: 256,
            num_hidden_layers: 1,
            num_attention_heads: 4,
            num_key_value_heads: 4,
            vision_hidden_size: 32,
            vision_intermediate_size: 128,
            vision_num_hidden_layers: 1,
            vision_num_attention_heads: 4,
            vision_patch_size: 14,
            vision_image_size: 56,
            ..Default::default()
        };

        let model = LLaVAModelV2::new(config).unwrap();
        let input_ids = ops_fn::zeros(&[1, 8], DataType::Int64, &Device::CPU).unwrap();
        let inputs = ModelInputs::text(input_ids);

        let outputs = model.forward(&inputs).unwrap();
        match outputs {
            ModelOutputs::Logits { logits, .. } => {
                assert_eq!(logits.shape()[0], 1);
                assert_eq!(logits.shape()[1], 8);
                assert_eq!(logits.shape()[2], 100);
            }
            _ => panic!("Expected logits output"),
        }
    }

    #[test]
    fn test_llava_multimodal_forward() {
        let config = LLaVAConfig {
            vocab_size: 100,
            hidden_size: 64,
            intermediate_size: 256,
            num_hidden_layers: 1,
            num_attention_heads: 4,
            num_key_value_heads: 4,
            vision_hidden_size: 32,
            vision_intermediate_size: 128,
            vision_num_hidden_layers: 1,
            vision_num_attention_heads: 4,
            vision_patch_size: 14,
            vision_image_size: 56,
            ..Default::default()
        };

        let model = LLaVAModelV2::new(config.clone()).unwrap();

        // Create multimodal inputs
        let input_ids = ops_fn::zeros(&[1, 8], DataType::Int64, &Device::CPU).unwrap();
        let pixel_values = ops_fn::zeros(
            &[1, 3, config.vision_image_size, config.vision_image_size],
            DataType::Float32,
            &Device::CPU,
        )
        .unwrap();

        let inputs = ModelInputs::Multimodal {
            input_ids,
            pixel_values: Some(pixel_values),
            attention_mask: None,
            image_mask: None,
        };

        let outputs = model.forward(&inputs).unwrap();
        match outputs {
            ModelOutputs::Logits { logits, .. } => {
                // Output should have combined sequence length
                assert_eq!(logits.shape()[0], 1);
                // Sequence length = text_len + image_patches - 1 (one text token replaced)
                // But our simple merge prepends image if no token found
                assert!(logits.shape()[1] > 0);
                assert_eq!(logits.shape()[2], 100);
            }
            _ => panic!("Expected logits output"),
        }
    }

    #[test]
    fn test_llava_generation() {
        let config = LLaVAConfig {
            vocab_size: 256,
            hidden_size: 64,
            intermediate_size: 256,
            num_hidden_layers: 1,
            num_attention_heads: 4,
            num_key_value_heads: 4,
            vision_hidden_size: 32,
            vision_intermediate_size: 128,
            vision_num_hidden_layers: 1,
            vision_num_attention_heads: 4,
            vision_patch_size: 14,
            vision_image_size: 56,
            ..Default::default()
        };

        let model = LLaVAModelV2::new(config).unwrap();
        let gen_config = GenerationConfig {
            max_new_tokens: 5,
            ..Default::default()
        };

        let output = model.generate("Hello", &gen_config).unwrap();
        assert!(!output.is_empty());
    }
}