rustytorch/crates/training/rtx-model-merging/src/config.rs

//! Configuration system for model merging operations

use crate::error::{MergeError, Result};
use serde::{Deserialize, Serialize};
use std::collections::HashMap;

/// Main configuration for model merging operations
#[derive(Debug, Clone, Serialize, Deserialize, Default)]
pub struct MergeConfig {
    /// Model loader configuration
    pub loader: ModelLoaderConfig,
    /// Validation configuration
    pub validation: ValidationConfig,
    /// Evaluation configuration
    pub evaluation: EvaluationConfig,
    /// Performance optimization settings
    pub performance: PerformanceConfig,
    /// Resource limits
    pub resources: ResourceConfig,
}

/// Configuration for model loading
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ModelLoaderConfig {
    /// Supported model formats
    pub supported_formats: Vec<String>,
    /// Memory mapping for large models
    pub use_memory_mapping: bool,
    /// Streaming threshold in MB
    pub streaming_threshold_mb: usize,
    /// Lazy loading enabled
    pub lazy_loading: bool,
    /// Cache loaded models
    pub cache_models: bool,
    /// Maximum cache size in MB
    pub max_cache_size_mb: usize,
    /// Parallel loading workers
    pub parallel_workers: usize,
}

impl Default for ModelLoaderConfig {
    fn default() -> Self {
        Self {
            supported_formats: vec![
                "safetensors".to_string(),
                "pytorch".to_string(),
                "onnx".to_string(),
                "huggingface".to_string(),
            ],
            use_memory_mapping: true,
            streaming_threshold_mb: 1024,
            lazy_loading: true,
            cache_models: true,
            max_cache_size_mb: 8192,
            parallel_workers: 4,
        }
    }
}

/// Configuration for model validation
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ValidationConfig {
    /// Enable architecture validation
    pub validate_architecture: bool,
    /// Enable parameter compatibility checks
    pub validate_parameters: bool,
    /// Enable numerical stability checks
    pub validate_numerical_stability: bool,
    /// Enable memory usage validation
    pub validate_memory_usage: bool,
    /// Tolerance for numerical comparisons
    pub numerical_tolerance: f32,
    /// Maximum allowed parameter deviation
    pub max_parameter_deviation: f32,
    /// Validation timeout in seconds
    pub timeout_seconds: u64,
}

impl Default for ValidationConfig {
    fn default() -> Self {
        Self {
            validate_architecture: true,
            validate_parameters: true,
            validate_numerical_stability: true,
            validate_memory_usage: true,
            numerical_tolerance: 1e-6,
            max_parameter_deviation: 0.1,
            timeout_seconds: 300,
        }
    }
}

/// Configuration for performance evaluation
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct EvaluationConfig {
    /// Enable performance benchmarking
    pub enable_benchmarking: bool,
    /// Benchmark datasets
    pub benchmark_datasets: Vec<String>,
    /// Evaluation metrics to compute
    pub metrics: Vec<String>,
    /// Evaluation batch size
    pub batch_size: usize,
    /// Number of evaluation samples
    pub num_samples: usize,
    /// Enable A/B testing
    pub enable_ab_testing: bool,
    /// Statistical significance threshold
    pub significance_threshold: f32,
}

impl Default for EvaluationConfig {
    fn default() -> Self {
        Self {
            enable_benchmarking: true,
            benchmark_datasets: vec!["validation".to_string(), "test".to_string()],
            metrics: vec![
                "accuracy".to_string(),
                "loss".to_string(),
                "f1_score".to_string(),
                "perplexity".to_string(),
            ],
            batch_size: 32,
            num_samples: 1000,
            enable_ab_testing: true,
            significance_threshold: 0.05,
        }
    }
}

/// Configuration for performance optimization
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct PerformanceConfig {
    /// Enable GPU acceleration
    pub use_gpu: bool,
    /// GPU device IDs to use
    pub gpu_devices: Vec<usize>,
    /// Enable mixed precision
    pub mixed_precision: bool,
    /// Enable parallel processing
    pub parallel_processing: bool,
    /// Number of worker threads
    pub num_threads: usize,
    /// Batch processing size
    pub batch_size: usize,
    /// Enable vectorized operations
    pub vectorized_ops: bool,
    /// Cache intermediate results
    pub cache_intermediates: bool,
}

impl Default for PerformanceConfig {
    fn default() -> Self {
        Self {
            use_gpu: true,
            gpu_devices: vec![0],
            mixed_precision: true,
            parallel_processing: true,
            num_threads: std::thread::available_parallelism()
                .map(std::num::NonZero::get)
                .unwrap_or(4),
            batch_size: 64,
            vectorized_ops: true,
            cache_intermediates: true,
        }
    }
}

/// Configuration for resource limits
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ResourceConfig {
    /// Maximum memory usage in MB
    pub max_memory_mb: Option<usize>,
    /// Maximum GPU memory usage in MB
    pub max_gpu_memory_mb: Option<usize>,
    /// Maximum disk usage for temporary files in MB
    pub max_disk_usage_mb: Option<usize>,
    /// Operation timeout in seconds
    pub operation_timeout_sec: u64,
    /// Enable resource monitoring
    pub monitor_resources: bool,
    /// Resource check interval in seconds
    pub monitor_interval_sec: u64,
}

impl Default for ResourceConfig {
    fn default() -> Self {
        Self {
            max_memory_mb: None, // No limit by default
            max_gpu_memory_mb: None,
            max_disk_usage_mb: Some(10240), // 10GB default
            operation_timeout_sec: 3600,    // 1 hour
            monitor_resources: true,
            monitor_interval_sec: 30,
        }
    }
}

/// Merge strategy configurations
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(tag = "type", rename_all = "snake_case")]
pub enum MergeStrategy {
    /// TIES (Task-specific Interference Elimination) merging
    Ties(TiesConfig),
    /// DARE (Drop And REscale) merging
    Dare(DareConfig),
    /// SLERP (Spherical Linear Interpolation) merging
    Slerp(SlerpConfig),
    /// Task arithmetic merging
    TaskArithmetic(TaskArithmeticConfig),
    /// Fisher information weighted merging
    Fisher(FisherConfig),
    /// Model soup ensemble merging
    ModelSoup(ModelSoupConfig),
    /// Frankenmerging (layer-wise merging)
    Frankenmerge(FrankenmergeConfig),
    /// Progressive merging with validation
    Progressive(ProgressiveConfig),
}

/// Configuration for TIES merging
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct TiesConfig {
    /// Sign consistency threshold
    pub sign_threshold: f32,
    /// Magnitude threshold for parameter selection
    pub magnitude_threshold: f32,
    /// Rescaling method
    pub rescale_method: RescaleMethod,
    /// Enable sign voting
    pub enable_sign_voting: bool,
    /// Voting threshold
    pub voting_threshold: f32,
    /// Density parameter
    pub density: f32,
}

impl Default for TiesConfig {
    fn default() -> Self {
        Self {
            sign_threshold: 0.5,
            magnitude_threshold: 0.1,
            rescale_method: RescaleMethod::Magnitude,
            enable_sign_voting: true,
            voting_threshold: 0.5,
            density: 0.8,
        }
    }
}

/// Configuration for DARE merging
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct DareConfig {
    /// Drop probability
    pub drop_probability: f32,
    /// Rescaling factor
    pub rescale_factor: f32,
    /// Random seed for reproducibility
    pub seed: Option<u64>,
    /// Enable adaptive dropping
    pub adaptive_dropping: bool,
    /// Importance threshold
    pub importance_threshold: f32,
}

impl Default for DareConfig {
    fn default() -> Self {
        Self {
            drop_probability: 0.1,
            rescale_factor: 1.0,
            seed: None,
            adaptive_dropping: false,
            importance_threshold: 0.01,
        }
    }
}

/// Configuration for SLERP merging
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SlerpConfig {
    /// Interpolation parameter (0.0 to 1.0)
    pub t: f32,
    /// Enable quaternion-based interpolation
    pub use_quaternions: bool,
    /// Normalization method
    pub normalization: NormalizationMethod,
    /// Enable adaptive interpolation
    pub adaptive_t: bool,
    /// Parameter-specific interpolation weights
    pub parameter_weights: Option<HashMap<String, f32>>,
}

impl Default for SlerpConfig {
    fn default() -> Self {
        Self {
            t: 0.5,
            use_quaternions: false,
            normalization: NormalizationMethod::L2,
            adaptive_t: false,
            parameter_weights: None,
        }
    }
}

/// Configuration for task arithmetic
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct TaskArithmeticConfig {
    /// Task vector scaling factors
    pub scaling_factors: Vec<f32>,
    /// Enable sign preservation
    pub preserve_signs: bool,
    /// Magnitude clipping threshold
    pub magnitude_clip: Option<f32>,
    /// Enable selective arithmetic
    pub selective_arithmetic: bool,
    /// Parameter selection criteria
    pub selection_criteria: Vec<String>,
}

impl Default for TaskArithmeticConfig {
    fn default() -> Self {
        Self {
            scaling_factors: vec![1.0],
            preserve_signs: true,
            magnitude_clip: None,
            selective_arithmetic: false,
            selection_criteria: vec![],
        }
    }
}

/// Configuration for Fisher information merging
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct FisherConfig {
    /// Fisher information computation method
    pub fisher_method: FisherMethod,
    /// Regularization strength
    pub regularization: f32,
    /// Enable diagonal approximation
    pub diagonal_fisher: bool,
    /// Number of samples for Fisher estimation
    pub num_samples: usize,
    /// Enable empirical Fisher
    pub empirical_fisher: bool,
}

impl Default for FisherConfig {
    fn default() -> Self {
        Self {
            fisher_method: FisherMethod::Empirical,
            regularization: 1e-6,
            diagonal_fisher: true,
            num_samples: 1000,
            empirical_fisher: true,
        }
    }
}

/// Configuration for model soup merging
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ModelSoupConfig {
    /// Ensemble weights
    pub weights: Vec<f32>,
    /// Soup method
    pub soup_method: SoupMethod,
    /// Enable greedy soup construction
    pub greedy_soup: bool,
    /// Validation metric for greedy selection
    pub validation_metric: String,
    /// Maximum models in soup
    pub max_models: usize,
}

impl Default for ModelSoupConfig {
    fn default() -> Self {
        Self {
            weights: vec![],
            soup_method: SoupMethod::Uniform,
            greedy_soup: false,
            validation_metric: "accuracy".to_string(),
            max_models: 10,
        }
    }
}

/// Configuration for Frankenmerging
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct FrankenmergeConfig {
    /// Layer assignment strategy
    pub layer_assignment: LayerAssignmentStrategy,
    /// Enable cross-architecture merging
    pub cross_architecture: bool,
    /// Layer compatibility checks
    pub compatibility_checks: bool,
    /// Enable layer interpolation
    pub layer_interpolation: bool,
    /// Interpolation weights per layer
    pub layer_weights: HashMap<String, f32>,
}

impl Default for FrankenmergeConfig {
    fn default() -> Self {
        Self {
            layer_assignment: LayerAssignmentStrategy::BestPerformance,
            cross_architecture: false,
            compatibility_checks: true,
            layer_interpolation: false,
            layer_weights: HashMap::new(),
        }
    }
}

/// Configuration for progressive merging
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ProgressiveConfig {
    /// Number of progressive steps
    pub num_steps: usize,
    /// Validation frequency
    pub validation_frequency: usize,
    /// Early stopping patience
    pub early_stopping_patience: Option<usize>,
    /// Validation metric threshold
    pub metric_threshold: f32,
    /// Enable rollback on failure
    pub enable_rollback: bool,
}

impl Default for ProgressiveConfig {
    fn default() -> Self {
        Self {
            num_steps: 10,
            validation_frequency: 1,
            early_stopping_patience: Some(3),
            metric_threshold: 0.01,
            enable_rollback: true,
        }
    }
}

/// Rescaling methods for TIES
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
pub enum RescaleMethod {
    /// Rescale by magnitude
    Magnitude,
    /// Rescale by sign consistency
    SignConsistency,
    /// No rescaling
    None,
}

/// Normalization methods for SLERP
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
pub enum NormalizationMethod {
    /// L2 normalization
    L2,
    /// L1 normalization
    L1,
    /// No normalization
    None,
}

/// Fisher information computation methods
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
pub enum FisherMethod {
    /// Empirical Fisher information
    Empirical,
    /// True Fisher information
    True,
    /// Diagonal Fisher approximation
    Diagonal,
}

/// Model soup methods
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
pub enum SoupMethod {
    /// Uniform averaging
    Uniform,
    /// Weighted averaging
    Weighted,
    /// Greedy soup construction
    Greedy,
}

/// Layer assignment strategies for Frankenmerging
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
pub enum LayerAssignmentStrategy {
    /// Assign based on best performance
    BestPerformance,
    /// Round-robin assignment
    RoundRobin,
    /// Random assignment
    Random,
    /// Performance-weighted assignment
    PerformanceWeighted,
}

/// Load configuration from file
pub fn load_config<P: AsRef<std::path::Path>>(path: P) -> Result<MergeConfig> {
    let content = std::fs::read_to_string(path.as_ref())
        .map_err(|e| MergeError::io(path.as_ref().display().to_string(), e))?;

    let extension = path
        .as_ref()
        .extension()
        .and_then(|s| s.to_str())
        .unwrap_or("");

    match extension {
        "toml" => toml::from_str(&content).map_err(MergeError::from),
        "json" => serde_json::from_str(&content).map_err(MergeError::from),
        _ => Err(MergeError::unsupported_format(format!(
            "Configuration format: {extension}"
        ))),
    }
}

/// Save configuration to file
pub fn save_config<P: AsRef<std::path::Path>>(config: &MergeConfig, path: P) -> Result<()> {
    let extension = path
        .as_ref()
        .extension()
        .and_then(|s| s.to_str())
        .unwrap_or("");

    let content = match extension {
        "toml" => toml::to_string_pretty(config)
            .map_err(|e| MergeError::internal(format!("TOML serialization error: {e}")))?,
        "json" => serde_json::to_string_pretty(config).map_err(MergeError::from)?,
        _ => {
            return Err(MergeError::unsupported_format(format!(
                "Configuration format: {extension}"
            )));
        }
    };

    std::fs::write(path.as_ref(), content)
        .map_err(|e| MergeError::io(path.as_ref().display().to_string(), e))?;

    Ok(())
}

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

    #[test]
    fn test_default_config() {
        let config = MergeConfig::default();
        assert!(config.loader.use_memory_mapping);
        assert!(config.validation.validate_architecture);
        assert!(config.performance.use_gpu);
    }

    #[test]
    fn test_ties_config() {
        let ties = TiesConfig::default();
        assert_eq!(ties.sign_threshold, 0.5);
        assert_eq!(ties.density, 0.8);
    }

    #[test]
    fn test_config_serialization() -> Result<()> {
        let config = MergeConfig::default();
        let json_str = serde_json::to_string(&config)?;
        let deserialized: MergeConfig = serde_json::from_str(&json_str)?;

        assert_eq!(
            config.loader.use_memory_mapping,
            deserialized.loader.use_memory_mapping
        );
        Ok(())
    }

    #[test]
    fn test_config_file_operations() -> Result<()> {
        let config = MergeConfig::default();

        let temp_file = NamedTempFile::new()
            .map_err(|e| MergeError::internal(format!("Failed to create temp file: {}", e)))?;
        let mut temp_path = temp_file.path().to_path_buf();
        temp_path.set_extension("json");

        save_config(&config, &temp_path)?;
        let loaded_config = load_config(&temp_path)?;

        assert_eq!(
            config.loader.parallel_workers,
            loaded_config.loader.parallel_workers
        );
        Ok(())
    }

    #[test]
    fn test_merge_strategy_serialization() -> Result<()> {
        let strategy = MergeStrategy::Ties(TiesConfig::default());
        let json_str = serde_json::to_string(&strategy)?;
        let deserialized: MergeStrategy = serde_json::from_str(&json_str)?;

        match deserialized {
            MergeStrategy::Ties(_) => (),
            _ => panic!("Wrong strategy type"),
        }

        Ok(())
    }
}