candle-coreml 0.3.1

//! Model configuration system for candle-coreml
//!
//! This module provides a configuration system that replaces hardcoded model shapes
//! with discoverable, flexible configurations. It supports loading configurations from
//! JSON files generated by the shape discovery tool, as well as built-in configurations
//! for known models.

use anyhow::{Context, Result};
use candle_core::{Device, Error as CandleError, Tensor};
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::path::Path;
use tracing::{debug, trace};

/// Complete model configuration including shapes, components, and naming patterns
#[derive(Debug, Clone, Deserialize, Serialize)]
pub struct ModelConfig {
    pub model_info: ModelInfo,
    pub shapes: ShapeConfig,
    pub components: HashMap<String, ComponentConfig>,
    pub naming: NamingConfig,
    /// Execution mode for FFN: "unified" (single component/function) or "split" (separate prefill/infer components)
    #[serde(skip_serializing_if = "Option::is_none")]
    pub ffn_execution: Option<String>,
}

/// Model metadata and identification
#[derive(Debug, Clone, Deserialize, Serialize)]
pub struct ModelInfo {
    #[serde(default)]
    pub model_id: Option<String>,
    pub path: Option<String>,
    pub model_type: String,
    pub discovered_at: Option<String>,
}

/// Overall model shape parameters
#[derive(Debug, Clone, Deserialize, Serialize)]
pub struct ShapeConfig {
    pub batch_size: usize,
    pub context_length: usize,
    pub hidden_size: usize,
    pub vocab_size: usize,
}

/// Configuration for a single model component (embeddings, FFN, LM head, etc.)
#[derive(Debug, Clone, Deserialize, Serialize)]
pub struct ComponentConfig {
    pub file_path: Option<String>,
    pub inputs: HashMap<String, TensorConfig>,
    pub outputs: HashMap<String, TensorConfig>,
    pub functions: Vec<String>,
    /// Optional deterministic input order; if absent, caller must provide correct order.
    #[serde(default)]
    pub input_order: Option<Vec<String>>,
}

/// Configuration for a tensor (shape and data type)
#[derive(Debug, Clone, Deserialize, Serialize)]
pub struct TensorConfig {
    pub name: String,
    pub shape: Vec<usize>,
    pub data_type: String,
}

/// Model file naming patterns for component discovery
#[derive(Debug, Clone, Deserialize, Serialize)]
pub struct NamingConfig {
    // Deprecated: patterns removed; explicit file paths are required per component
    #[serde(skip_serializing_if = "Option::is_none")]
    pub embeddings_pattern: Option<String>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub ffn_prefill_pattern: Option<String>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub ffn_infer_pattern: Option<String>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub lm_head_pattern: Option<String>,
}

impl ModelConfig {
    /// Create a minimal default Qwen ModelConfig (no components). Useful for tests and fallbacks.
    pub fn default_qwen() -> Self {
        Self {
            model_info: ModelInfo {
                model_id: Some("default/qwen".to_string()),
                path: None,
                model_type: "qwen".to_string(),
                discovered_at: None,
            },
            shapes: ShapeConfig {
                batch_size: 1,
                context_length: 512,
                hidden_size: 1024,
                vocab_size: 151_936,
            },
            components: HashMap::new(),
            naming: NamingConfig {
                embeddings_pattern: None,
                ffn_prefill_pattern: None,
                ffn_infer_pattern: None,
                lm_head_pattern: None,
            },
            ffn_execution: None,
        }
    }
    /// Load configuration from a JSON file
    pub fn load_from_file<P: AsRef<Path>>(path: P) -> Result<Self> {
        let path = path.as_ref();
        let content = std::fs::read_to_string(path)
            .with_context(|| format!("Failed to read config file: {}", path.display()))?;

        let config: ModelConfig = serde_json::from_str(&content)
            .with_context(|| format!("Failed to parse config file: {}", path.display()))?;

        Ok(config)
    }

    /// Save configuration to a JSON file
    pub fn save_to_file<P: AsRef<Path>>(&self, path: P) -> Result<()> {
        let path = path.as_ref();
        let content =
            serde_json::to_string_pretty(self).context("Failed to serialize configuration")?;

        std::fs::write(path, content)
            .with_context(|| format!("Failed to write config file: {}", path.display()))?;

        Ok(())
    }

    /// Get the shape configuration for a specific component and tensor
    pub fn get_tensor_shape(
        &self,
        component: &str,
        tensor_name: &str,
        is_input: bool,
    ) -> Option<&Vec<usize>> {
        let component_config = self.components.get(component)?;

        let tensor_map = if is_input {
            &component_config.inputs
        } else {
            &component_config.outputs
        };

        tensor_map.get(tensor_name).map(|tensor| &tensor.shape)
    }

    /// Get the expected input shape for embeddings
    pub fn embeddings_input_shape(&self) -> Option<&Vec<usize>> {
        self.get_tensor_shape("embeddings", "input_ids", true)
    }

    /// Get the expected output shape for embeddings
    pub fn embeddings_output_shape(&self) -> Option<&Vec<usize>> {
        self.get_tensor_shape("embeddings", "hidden_states", false)
    }

    /// Get the expected input shape for FFN prefill
    pub fn ffn_prefill_input_shape(&self) -> Option<&Vec<usize>> {
        self.get_tensor_shape("ffn_prefill", "hidden_states", true)
    }

    /// Get the expected input shape for LM head
    pub fn lm_head_input_shape(&self) -> Option<&Vec<usize>> {
        self.get_tensor_shape("lm_head", "hidden_states", true)
    }

    /// Check if this model supports multi-part logits output
    pub fn has_multipart_logits(&self) -> bool {
        if let Some(lm_head) = self.components.get("lm_head") {
            // Check if there are multiple logits outputs (logits1, logits2, etc.)
            let logits_outputs: Vec<_> = lm_head
                .outputs
                .keys()
                .filter(|name| name.starts_with("logits") && name.len() > 6) // "logits" + number
                .collect();
            return logits_outputs.len() > 1;
        }
        false
    }

    /// Get the number of logits parts for multi-part output
    pub fn logits_part_count(&self) -> usize {
        if let Some(lm_head) = self.components.get("lm_head") {
            let logits_outputs: Vec<_> = lm_head
                .outputs
                .keys()
                .filter(|name| name.starts_with("logits"))
                .collect();
            if logits_outputs.is_empty() {
                1 // Single logits output
            } else {
                logits_outputs.len()
            }
        } else {
            1
        }
    }

    /// Select the primary logits output name for the LM head.
    /// Preference order: "logits1" (multipart), then "logits", otherwise the first available key.
    pub fn lm_head_primary_output_name(&self) -> Option<String> {
        let lm_head = self.components.get("lm_head")?;

        // Prefer explicit multipart naming starting with logits1
        if lm_head.outputs.contains_key("logits1") {
            return Some("logits1".to_string());
        }

        // Common single output name
        if lm_head.outputs.contains_key("logits") {
            return Some("logits".to_string());
        }

        // Fallback to the first key if any
        lm_head.outputs.keys().next().map(|k| k.to_string())
    }

    /// Validate the configuration for consistency
    pub fn validate(&self) -> Result<()> {
        // Check that required components exist
        let required_components = ["embeddings", "lm_head"];
        for component in required_components {
            if !self.components.contains_key(component) {
                return Err(anyhow::anyhow!("Missing required component: {}", component));
            }
        }

        // Check shape consistency
        if self.shapes.batch_size == 0 {
            return Err(anyhow::anyhow!("batch_size must be greater than 0"));
        }

        if self.shapes.context_length == 0 {
            return Err(anyhow::anyhow!("context_length must be greater than 0"));
        }

        if self.shapes.hidden_size == 0 {
            return Err(anyhow::anyhow!("hidden_size must be greater than 0"));
        }

        if self.shapes.vocab_size == 0 {
            return Err(anyhow::anyhow!("vocab_size must be greater than 0"));
        }

        // Validate tensor shapes make sense
        for (component_name, component) in &self.components {
            for (tensor_name, tensor) in &component.inputs {
                if tensor.shape.is_empty() {
                    return Err(anyhow::anyhow!(
                        "Empty shape for {}.inputs.{}",
                        component_name,
                        tensor_name
                    ));
                }
            }
            for (tensor_name, tensor) in &component.outputs {
                if tensor.shape.is_empty() {
                    return Err(anyhow::anyhow!(
                        "Empty shape for {}.outputs.{}",
                        component_name,
                        tensor_name
                    ));
                }
            }
        }

        Ok(())
    }

    /// Validate internal wiring between components for basic shape compatibility.
    /// Examples:
    ///  - embeddings.outputs.hidden_states == ffn_prefill.inputs.hidden_states
    ///  - ffn_infer.outputs.output_hidden_states == lm_head.inputs.hidden_states (when ffn_infer exists)
    pub fn validate_internal_wiring(&self) -> Result<()> {
        // Embeddings -> FFN prefill hidden_states flow
        if let (Some(emb_out), Some(ffn_in_hidden)) = (
            self.get_tensor_shape("embeddings", "hidden_states", false),
            self.get_tensor_shape("ffn_prefill", "hidden_states", true),
        ) {
            if emb_out != ffn_in_hidden {
                return Err(anyhow::anyhow!(
                    "Shape mismatch: embeddings.hidden_states {:?} != ffn_prefill.hidden_states {:?}",
                    emb_out, ffn_in_hidden
                ));
            }
        }

        // FFN infer -> LM head hidden_states flow
        if self.components.contains_key("ffn_infer") {
            if let (Some(ffn_out), Some(lm_in)) = (
                self.get_tensor_shape("ffn_infer", "output_hidden_states", false),
                self.get_tensor_shape("lm_head", "hidden_states", true),
            ) {
                if ffn_out != lm_in {
                    return Err(anyhow::anyhow!(
                        "Shape mismatch: ffn_infer.output_hidden_states {:?} != lm_head.hidden_states {:?}",
                        ffn_out, lm_in
                    ));
                }
            }
        } else {
            // If there's no separate ffn_infer, check ffn_prefill output shape matches lm_head (single-token path)
            if let (Some(ffn_out), Some(lm_in)) = (
                self.get_tensor_shape("ffn_prefill", "output_hidden_states", false),
                self.get_tensor_shape("lm_head", "hidden_states", true),
            ) {
                if ffn_out != lm_in {
                    return Err(anyhow::anyhow!(
                        "Shape mismatch: ffn_prefill.output_hidden_states {:?} != lm_head.hidden_states {:?}",
                        ffn_out, lm_in
                    ));
                }
            }
        }

        Ok(())
    }

    /// Determine if FFN execution should be treated as split (separate infer component)
    pub fn ffn_is_split(&self) -> bool {
        if let Some(mode) = self.ffn_execution.as_deref() {
            return mode == "split";
        }
        if let (Some(prefill), Some(infer)) = (
            self.components.get("ffn_prefill"),
            self.components.get("ffn_infer"),
        ) {
            match (&prefill.file_path, &infer.file_path) {
                (Some(p), Some(i)) => p != i, // different files => split
                _ => false,
            }
        } else {
            false
        }
    }

    /// Detect if prefill should run in single-token sequential mode based on configured shapes
    pub fn prefill_is_single_token(&self) -> bool {
        if let Some(prefill) = self.components.get("ffn_prefill") {
            if let Some(hs) = prefill.inputs.get("hidden_states") {
                let is_single = hs.shape.len() == 3 && hs.shape.get(1) == Some(&1);
                debug!(
                    "🔍 prefill_is_single_token: shape={:?}, len={}, dim[1]={:?}, result={}",
                    hs.shape,
                    hs.shape.len(),
                    hs.shape.get(1),
                    is_single
                );
                return is_single;
            }
        }
        debug!(
            "🔍 prefill_is_single_token: no ffn_prefill or hidden_states found, returning false"
        );
        false
    }

    /// Check if the model expects full sequence prefill (as opposed to single-token processing)
    /// This is typically true for CoreML models with fixed-shape inputs like [1, 128, 1024]
    pub fn expects_full_sequence_prefill(&self) -> bool {
        if let Some(prefill) = self.components.get("ffn_prefill") {
            if let Some(hs) = prefill.inputs.get("hidden_states") {
                // If the model expects a fixed sequence length > 1, it needs full-sequence prefill
                let expects_full =
                    hs.shape.len() == 3 && hs.shape.get(1).is_some_and(|&seq_len| seq_len > 1);
                trace!(
                    "🔍 expects_full_sequence_prefill: shape={:?}, len={}, dim[1]={:?}, result={}",
                    hs.shape,
                    hs.shape.len(),
                    hs.shape.get(1),
                    expects_full
                );
                return expects_full;
            }
        }
        trace!("🔍 expects_full_sequence_prefill: no ffn_prefill or hidden_states found, returning false");
        false
    }

    // Tensor Creation Methods (moved from QwenConfig for consolidation)

    /// Create embeddings input tensor with proper shape from configuration
    pub fn create_embeddings_input_tensor(
        &self,
        tokens: &[i64],
        device: &Device,
    ) -> Result<Tensor, CandleError> {
        let expected_shape = self
            .embeddings_input_shape()
            .ok_or_else(|| CandleError::Msg("No embeddings input shape found".to_string()))?;
        let expected_len = expected_shape[1]; // [batch, seq_len] -> seq_len

        // Pad or truncate tokens to match expected length
        let mut padded_tokens = tokens.to_vec();
        padded_tokens.resize(expected_len, 0); // Pad with 0s

        Tensor::from_vec(
            padded_tokens,
            (expected_shape[0], expected_shape[1]),
            device,
        )
    }

    /// Create position IDs tensor for FFN prefill with proper shape
    pub fn create_ffn_position_ids_tensor(
        &self,
        positions: &[i64],
        device: &Device,
    ) -> Result<Tensor, CandleError> {
        let expected_shape = self
            .get_tensor_shape("ffn_prefill", "position_ids", true)
            .ok_or_else(|| {
                CandleError::Msg("No FFN prefill position_ids shape found".to_string())
            })?;

        // Heuristic: some manifests report position_ids length as [1] even for prefill.
        // When that happens, derive the true sequence length from other known shapes.
        let mut expected_len = expected_shape[0];
        if expected_len == 1 {
            // Prefer prefill hidden_states seq_len if available and > 1
            if let Some(hs_shape) = self.get_tensor_shape("ffn_prefill", "hidden_states", true) {
                if hs_shape.len() == 3 && hs_shape[1] > 1 {
                    expected_len = hs_shape[1];
                }
            }
            // Or embeddings input seq_len if available and > 1
            if expected_len == 1 {
                if let Some(emb) = self.embeddings_input_shape() {
                    if emb.len() == 2 && emb[1] > 1 {
                        expected_len = emb[1];
                    }
                }
            }
            // As a final fallback, keep 1 (some exotic models may actually expect [1])
        }

        // Create position sequence up to expected length
        let mut position_ids = Vec::with_capacity(expected_len);
        for i in 0..expected_len {
            if i < positions.len() {
                position_ids.push(positions[i]);
            } else {
                position_ids.push(0); // Pad with 0s
            }
        }

        Tensor::from_vec(position_ids, (expected_len,), device)
    }

    /// Create causal mask tensor for FFN with proper shape
    pub fn create_ffn_causal_mask_tensor(
        &self,
        _batch_size: usize,
        _context_length: usize,
        device: &Device,
    ) -> Result<Tensor, CandleError> {
        // Prefer explicit shape from config; otherwise synthesize a reasonable default
        let expected_shape_vec =
            if let Some(shape) = self.get_tensor_shape("ffn_prefill", "causal_mask", true) {
                shape.clone()
            } else {
                // Derive a default square mask [1,1,seq_len,seq_len]
                let mut seq_len = 0usize;
                if let Some(hs) = self.get_tensor_shape("ffn_prefill", "hidden_states", true) {
                    if hs.len() == 3 && hs[1] > 0 {
                        seq_len = hs[1];
                    }
                }
                if seq_len == 0 {
                    if let Some(emb) = self.embeddings_input_shape() {
                        if emb.len() == 2 && emb[1] > 0 {
                            seq_len = emb[1];
                        }
                    }
                }
                if seq_len == 0 {
                    seq_len = self.shapes.context_length;
                }
                vec![1, 1, seq_len, seq_len]
            };

        let mask_rows = expected_shape_vec[2];
        let mask_context_length = expected_shape_vec[3];

        // Create causal mask data
        let mut mask_data = vec![f32::NEG_INFINITY; mask_rows * mask_context_length];
        for i in 0..mask_rows {
            for j in 0..=i.min(mask_context_length - 1) {
                mask_data[i * mask_context_length + j] = 0.0;
            }
        }

        Tensor::from_vec(
            mask_data,
            (
                expected_shape_vec[0],
                expected_shape_vec[1],
                expected_shape_vec[2],
                expected_shape_vec[3],
            ),
            device,
        )
    }

    /// Create single token hidden states tensor for LM head
    pub fn create_single_token_hidden_states(
        &self,
        _tokens: &[i64],
        device: &Device,
    ) -> Result<Tensor, CandleError> {
        let expected_shape = self
            .get_tensor_shape("lm_head", "hidden_states", true)
            .ok_or_else(|| CandleError::Msg("No LM head hidden_states shape found".to_string()))?;

        // Create dummy tensor with correct shape (would be filled by actual embeddings)
        let tensor_data = vec![0.0f32; expected_shape.iter().product()];
        let shape = (expected_shape[0], expected_shape[1], expected_shape[2]);

        Tensor::from_vec(tensor_data, shape, device)
    }

    /// Create position IDs tensor for inference (single position)
    pub fn create_infer_position_ids_tensor(
        &self,
        position: i64,
        device: &Device,
    ) -> Result<Tensor, CandleError> {
        // Check if we have a dedicated ffn_infer component with specific shape
        if let Some(infer_shape) = self.get_tensor_shape("ffn_infer", "position_ids", true) {
            // Use the infer-specific shape
            if infer_shape.len() == 1 {
                Tensor::from_vec(vec![position], (infer_shape[0],), device)
            } else {
                let size = infer_shape.iter().product();
                let mut data = vec![0i64; size];
                data[0] = position;
                Tensor::from_vec(data, infer_shape.as_slice(), device)
            }
        } else {
            // No dedicated infer component - use single position for inference (original QwenConfig behavior)
            Tensor::from_vec(vec![position], (1,), device)
        }
    }

    /// Create current position tensor for FFN
    pub fn create_current_pos_tensor(
        &self,
        position: i64,
        device: &Device,
    ) -> Result<Tensor, CandleError> {
        // Most models expect [1] shape for current_pos
        Tensor::from_vec(vec![position], (1,), device)
    }
}

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

    fn create_test_config() -> ModelConfig {
        let mut components = HashMap::new();

        // Embeddings component
        let mut embeddings_inputs = HashMap::new();
        embeddings_inputs.insert(
            "input_ids".to_string(),
            TensorConfig {
                name: "input_ids".to_string(),
                shape: vec![1, 64],
                data_type: "INT32".to_string(),
            },
        );

        let mut embeddings_outputs = HashMap::new();
        embeddings_outputs.insert(
            "hidden_states".to_string(),
            TensorConfig {
                name: "hidden_states".to_string(),
                shape: vec![1, 64, 1024],
                data_type: "FLOAT16".to_string(),
            },
        );

        components.insert(
            "embeddings".to_string(),
            ComponentConfig {
                file_path: None,
                inputs: embeddings_inputs,
                outputs: embeddings_outputs,
                functions: vec![],
                input_order: None,
            },
        );

        // LM Head component
        let mut lm_head_inputs = HashMap::new();
        lm_head_inputs.insert(
            "hidden_states".to_string(),
            TensorConfig {
                name: "hidden_states".to_string(),
                shape: vec![1, 1, 1024],
                data_type: "FLOAT16".to_string(),
            },
        );

        let mut lm_head_outputs = HashMap::new();
        lm_head_outputs.insert(
            "logits".to_string(),
            TensorConfig {
                name: "logits".to_string(),
                shape: vec![1, 1, 151936],
                data_type: "FLOAT32".to_string(),
            },
        );

        components.insert(
            "lm_head".to_string(),
            ComponentConfig {
                file_path: None,
                inputs: lm_head_inputs,
                outputs: lm_head_outputs,
                functions: vec![],
                input_order: None,
            },
        );

        ModelConfig {
            model_info: ModelInfo {
                model_id: Some("test/model".to_string()),
                path: Some("/test/path".to_string()),
                model_type: "qwen".to_string(),
                discovered_at: Some("2025-08-07T00:00:00".to_string()),
            },
            shapes: ShapeConfig {
                batch_size: 1,
                context_length: 512,
                hidden_size: 1024,
                vocab_size: 151936,
            },
            components,
            naming: NamingConfig {
                embeddings_pattern: None,
                ffn_prefill_pattern: None,
                ffn_infer_pattern: None,
                lm_head_pattern: None,
            },
            ffn_execution: Some("unified".to_string()),
        }
    }

    #[test]
    fn test_config_serialization() {
        let config = create_test_config();

        // Test JSON serialization
        let json = serde_json::to_string_pretty(&config).unwrap();
        assert!(json.contains("test/model"));
        assert!(json.contains("batch_size"));
        assert!(json.contains("embeddings"));

        // Test deserialization
        let parsed: ModelConfig = serde_json::from_str(&json).unwrap();
        assert_eq!(parsed.model_info.model_id, config.model_info.model_id);
        assert_eq!(parsed.shapes.batch_size, config.shapes.batch_size);
        assert_eq!(parsed.components.len(), config.components.len());
    }

    #[test]
    fn test_config_file_io() {
        let config = create_test_config();
        let temp_file = NamedTempFile::new().unwrap();

        // Save configuration
        config.save_to_file(temp_file.path()).unwrap();

        // Load configuration
        let loaded = ModelConfig::load_from_file(temp_file.path()).unwrap();
        assert_eq!(loaded.model_info.model_id, config.model_info.model_id);
        assert_eq!(loaded.shapes.hidden_size, config.shapes.hidden_size);
    }

    #[test]
    fn test_shape_accessors() {
        let config = create_test_config();

        // Test embeddings shapes
        let embeddings_input = config.embeddings_input_shape().unwrap();
        assert_eq!(embeddings_input, &vec![1, 64]);

        let embeddings_output = config.embeddings_output_shape().unwrap();
        assert_eq!(embeddings_output, &vec![1, 64, 1024]);

        let lm_head_input = config.lm_head_input_shape().unwrap();
        assert_eq!(lm_head_input, &vec![1, 1, 1024]);
    }

    #[test]
    fn test_multipart_logits_detection() {
        let config = create_test_config();
        assert!(!config.has_multipart_logits()); // Single logits output

        // Create config with multipart logits
        let mut config_multipart = config;
        let lm_head = config_multipart.components.get_mut("lm_head").unwrap();
        lm_head.outputs.clear();
        lm_head.outputs.insert(
            "logits1".to_string(),
            TensorConfig {
                name: "logits1".to_string(),
                shape: vec![1, 1, 9480],
                data_type: "FLOAT32".to_string(),
            },
        );
        lm_head.outputs.insert(
            "logits2".to_string(),
            TensorConfig {
                name: "logits2".to_string(),
                shape: vec![1, 1, 9479],
                data_type: "FLOAT32".to_string(),
            },
        );

        assert!(config_multipart.has_multipart_logits());
        assert_eq!(config_multipart.logits_part_count(), 2);
    }

    #[test]
    fn test_config_validation() {
        let config = create_test_config();
        assert!(config.validate().is_ok());

        // Internal wiring should be consistent in this synthetic setup
        assert!(config.validate_internal_wiring().is_ok());

        // Test missing component
        let mut invalid_config = config.clone();
        invalid_config.components.remove("embeddings");
        assert!(invalid_config.validate().is_err());

        // Test invalid shapes
        let mut invalid_shapes = config;
        invalid_shapes.shapes.batch_size = 0;
        assert!(invalid_shapes.validate().is_err());
    }
}