//! Advanced model merging strategies
//!
//! This module implements sophisticated merging strategies that build upon
//! the core algorithms:
//! - Model Soups: Ensemble averaging with greedy selection
//! - Frankenmerging: Layer-wise merging across architectures
//! - Progressive: Gradual merging with validation checkpoints

pub mod frankenmerge;
pub mod model_soup;
pub mod progressive;

// Re-export for convenience
pub use frankenmerge::*;
pub use model_soup::*;
pub use progressive::*;

use crate::error::Result;
use crate::types::{MergedModel, Model};
use std::collections::HashMap;

/// Common utilities for advanced merging strategies
pub struct StrategyUtils;

impl StrategyUtils {
    /// Validate that all models in a group have compatible architectures
    pub fn validate_model_group(models: &[Model]) -> Result<()> {
        if models.len() < 2 {
            return Ok(());
        }

        let reference = &models[0];
        for (i, model) in models.iter().enumerate().skip(1) {
            if !reference.is_compatible_with(model) {
                return Err(crate::error::MergeError::compatibility(format!(
                    "Model {i} is incompatible with reference model in group"
                )));
            }
        }

        Ok(())
    }

    /// Compute pairwise similarities between models
    pub fn compute_pairwise_similarities(models: &[Model]) -> Result<Vec<Vec<f32>>> {
        let n = models.len();
        let mut similarities = vec![vec![0.0; n]; n];

        for i in 0..n {
            similarities[i][i] = 1.0; // Self-similarity

            for j in (i + 1)..n {
                let similarity =
                    crate::algorithms::MergeUtils::compute_similarity(&models[i], &models[j])?;
                similarities[i][j] = similarity;
                similarities[j][i] = similarity; // Symmetric
            }
        }

        Ok(similarities)
    }

    /// Group models by similarity threshold
    pub fn group_models_by_similarity(models: &[Model], threshold: f32) -> Result<Vec<Vec<usize>>> {
        let similarities = Self::compute_pairwise_similarities(models)?;
        let n = models.len();

        let mut groups = Vec::new();
        let mut assigned = vec![false; n];

        for i in 0..n {
            if assigned[i] {
                continue;
            }

            let mut group = vec![i];
            assigned[i] = true;

            // Find all models similar to model i
            for j in (i + 1)..n {
                if !assigned[j] && similarities[i][j] >= threshold {
                    group.push(j);
                    assigned[j] = true;
                }
            }

            groups.push(group);
        }

        Ok(groups)
    }

    /// Compute model diversity within a group
    pub fn compute_group_diversity(models: &[Model], group_indices: &[usize]) -> Result<f32> {
        if group_indices.len() < 2 {
            return Ok(0.0);
        }

        let group_models: Vec<&Model> = group_indices.iter().map(|&i| &models[i]).collect();
        let similarities = Self::compute_pairwise_similarities(
            &group_models.into_iter().cloned().collect::<Vec<_>>(),
        )?;

        let mut total_dissimilarity = 0.0;
        let mut pair_count = 0;

        for i in 0..similarities.len() {
            for j in (i + 1)..similarities.len() {
                total_dissimilarity += 1.0 - similarities[i][j];
                pair_count += 1;
            }
        }

        Ok(if pair_count > 0 {
            total_dissimilarity / pair_count as f32
        } else {
            0.0
        })
    }

    /// Select representative models from a group
    pub fn select_representatives(
        models: &[Model],
        group_indices: &[usize],
        max_representatives: usize,
    ) -> Result<Vec<usize>> {
        if group_indices.len() <= max_representatives {
            return Ok(group_indices.to_vec());
        }

        // Use k-means-like selection based on parameter space
        let mut representatives = Vec::new();
        let mut remaining: Vec<usize> = group_indices.to_vec();

        // First representative: model with highest average similarity to others
        let mut best_idx = 0;
        let mut best_avg_sim = 0.0;

        for (i, &model_idx) in remaining.iter().enumerate() {
            let mut total_sim = 0.0;
            for &other_idx in &remaining {
                if model_idx != other_idx {
                    total_sim += crate::algorithms::MergeUtils::compute_similarity(
                        &models[model_idx],
                        &models[other_idx],
                    )?;
                }
            }

            let avg_sim = total_sim / (remaining.len() - 1) as f32;
            if avg_sim > best_avg_sim {
                best_avg_sim = avg_sim;
                best_idx = i;
            }
        }

        representatives.push(remaining.remove(best_idx));

        // Subsequent representatives: maximize diversity
        while representatives.len() < max_representatives && !remaining.is_empty() {
            let mut best_idx = 0;
            let mut best_min_sim = 1.0f32;

            for (i, &candidate_idx) in remaining.iter().enumerate() {
                let mut min_sim = 1.0f32;

                for &rep_idx in &representatives {
                    let sim = crate::algorithms::MergeUtils::compute_similarity(
                        &models[candidate_idx],
                        &models[rep_idx],
                    )?;
                    min_sim = min_sim.min(sim);
                }

                if min_sim < best_min_sim {
                    best_min_sim = min_sim;
                    best_idx = i;
                }
            }

            representatives.push(remaining.remove(best_idx));
        }

        Ok(representatives)
    }

    /// Compute model performance scores (placeholder - would use actual validation)
    pub fn compute_performance_scores(models: &[Model]) -> HashMap<String, f32> {
        let mut scores = HashMap::new();

        for model in models {
            // Placeholder scoring based on model complexity
            let complexity = model.parameter_count() as f32;
            let diversity = Self::compute_parameter_diversity(model);

            // Simple heuristic score (in practice would use validation metrics)
            let score = (1.0 / (1.0 + (complexity / 1000000.0))).max(0.1) * (0.5 + 0.5 * diversity);

            scores.insert(model.name.clone(), score);
        }

        scores
    }

    /// Compute parameter diversity within a single model
    pub fn compute_parameter_diversity(model: &Model) -> f32 {
        let all_params: Vec<f32> = model
            .parameters
            .values()
            .flat_map(|p| p.data.iter())
            .copied()
            .collect();

        if all_params.is_empty() {
            return 0.0;
        }

        let (_, std_dev, _, _) = crate::algorithms::MergeUtils::compute_moments(&all_params);
        std_dev.min(1.0)
    }

    /// Create validation checkpoints for progressive merging
    pub fn create_validation_checkpoints(
        base_model: &Model,
        target_model: &Model,
        num_checkpoints: usize,
    ) -> Result<Vec<MergedModel>> {
        let mut checkpoints = Vec::new();

        for i in 0..=num_checkpoints {
            let t = i as f32 / num_checkpoints as f32;

            // Linear interpolation between base and target
            let mut checkpoint_model = base_model.clone();
            checkpoint_model.id = uuid::Uuid::new_v4();
            checkpoint_model.name = format!("checkpoint_{i}_{t:.2}");

            // Interpolate parameters
            for (param_name, base_param) in &base_model.parameters {
                if let Some(target_param) = target_model.parameters.get(param_name) {
                    let interpolated_data = crate::algorithms::MergeUtils::elementwise_op(
                        &base_param.data,
                        &target_param.data,
                        |base_val, target_val| (1.0 - t) * base_val + t * target_val,
                    )?;

                    if let Some(checkpoint_param) = checkpoint_model.parameters.get_mut(param_name)
                    {
                        checkpoint_param.data = interpolated_data;
                    }
                }
            }

            // Create merged model wrapper
            let merged_checkpoint = MergedModel {
                model: checkpoint_model,
                merge_info: crate::types::MergeInfo {
                    strategy: "ValidationCheckpoint".to_string(),
                    source_models: vec![
                        crate::types::ModelReference {
                            id: base_model.id,
                            name: base_model.name.clone(),
                            path: std::path::PathBuf::from(&base_model.name),
                            weight: Some(1.0 - t),
                        },
                        crate::types::ModelReference {
                            id: target_model.id,
                            name: target_model.name.clone(),
                            path: std::path::PathBuf::from(&target_model.name),
                            weight: Some(t),
                        },
                    ],
                    merged_at: chrono::Utc::now(),
                    config: serde_json::json!({"t": t, "checkpoint": i}),
                    statistics: crate::types::MergeStatistics::default(),
                },
                quality_metrics: crate::types::QualityMetrics::default(),
            };

            checkpoints.push(merged_checkpoint);
        }

        Ok(checkpoints)
    }

    /// Estimate memory requirements for merging strategy
    pub fn estimate_memory_requirements(models: &[Model], strategy: &str) -> Result<usize> {
        let total_model_memory: usize = models.iter().map(super::types::Model::memory_size).sum();

        let overhead_factor = match strategy {
            "ModelSoup" => 1.2,    // Small overhead for ensemble averaging
            "Frankenmerge" => 1.5, // Moderate overhead for layer-wise operations
            "Progressive" => 2.0,  // Higher overhead for checkpoints
            _ => 1.1,
        };

        let estimated_memory = (total_model_memory as f32 * overhead_factor) as usize;
        Ok(estimated_memory)
    }

    /// Check resource constraints
    pub fn check_resource_constraints(
        models: &[Model],
        strategy: &str,
        max_memory_mb: Option<usize>,
    ) -> Result<()> {
        let estimated_memory = Self::estimate_memory_requirements(models, strategy)?;
        let estimated_memory_mb = estimated_memory / (1024 * 1024);

        if let Some(max_mb) = max_memory_mb
            && estimated_memory_mb > max_mb
        {
            return Err(crate::error::MergeError::memory(estimated_memory_mb));
        }

        // Check for extremely large memory requirements
        if estimated_memory_mb > 32768 {
            // 32GB
            return Err(crate::error::MergeError::resource_exhausted(format!(
                "Memory requirement: {estimated_memory_mb}MB"
            )));
        }

        Ok(())
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::types::{DataType, ModelArchitecture, ParameterTensor};

    fn create_test_model(name: &str, params: Vec<(String, Vec<f32>)>) -> Model {
        let arch = ModelArchitecture {
            arch_type: "test".to_string(),
            num_layers: 1,
            hidden_dim: params.len(),
            params: std::collections::HashMap::new(),
        };

        let mut model = Model::new(name.to_string(), arch);

        for (param_name, data) in params {
            let param = ParameterTensor::new(param_name, vec![data.len()], DataType::Float32, data);
            model.add_parameter(param);
        }

        model
    }

    #[test]
    fn test_model_group_validation() -> Result<()> {
        let model1 = create_test_model("model1", vec![("weight".to_string(), vec![1.0, 2.0])]);

        let model2 = create_test_model("model2", vec![("weight".to_string(), vec![3.0, 4.0])]);

        let result = StrategyUtils::validate_model_group(&[model1, model2]);
        assert!(result.is_ok());

        Ok(())
    }

    #[test]
    fn test_pairwise_similarities() -> Result<()> {
        let model1 = create_test_model("model1", vec![("weight".to_string(), vec![1.0, 2.0, 3.0])]);

        let model2 = create_test_model(
            "model2",
            vec![
                ("weight".to_string(), vec![1.0, 2.0, 3.0]), // Identical
            ],
        );

        let model3 = create_test_model(
            "model3",
            vec![
                ("weight".to_string(), vec![4.0, 5.0, 6.0]), // Different
            ],
        );

        let similarities = StrategyUtils::compute_pairwise_similarities(&[model1, model2, model3])?;

        assert_eq!(similarities.len(), 3);
        assert_eq!(similarities[0].len(), 3);

        // Self-similarities should be 1.0
        assert_eq!(similarities[0][0], 1.0);
        assert_eq!(similarities[1][1], 1.0);
        assert_eq!(similarities[2][2], 1.0);

        // Model1 and Model2 should be very similar (identical)
        assert!(similarities[0][1] > 0.9);

        // Model1/Model2 and Model3 should be less similar
        assert!(similarities[0][2] < similarities[0][1]);

        Ok(())
    }

    #[test]
    fn test_similarity_grouping() -> Result<()> {
        let model1 = create_test_model("model1", vec![("weight".to_string(), vec![1.0, 2.0])]);

        let model2 = create_test_model(
            "model2",
            vec![
                ("weight".to_string(), vec![1.1, 2.1]), // Very similar to model1
            ],
        );

        let model3 = create_test_model(
            "model3",
            vec![
                ("weight".to_string(), vec![10.0, 20.0]), // Very different
            ],
        );

        let groups = StrategyUtils::group_models_by_similarity(&[model1, model2, model3], 0.8)?;

        // Should have at least 2 groups (similar models together, different model separate)
        assert!(groups.len() >= 2);

        // Total models should be preserved
        let total_models: usize = groups.iter().map(|g| g.len()).sum();
        assert_eq!(total_models, 3);

        Ok(())
    }

    #[test]
    fn test_representative_selection() -> Result<()> {
        let models = (0..5)
            .map(|i| {
                create_test_model(
                    &format!("model{}", i),
                    vec![("weight".to_string(), vec![i as f32, (i + 1) as f32])],
                )
            })
            .collect::<Vec<_>>();

        let group_indices = vec![0, 1, 2, 3, 4];
        let representatives = StrategyUtils::select_representatives(&models, &group_indices, 3)?;

        assert_eq!(representatives.len(), 3);
        assert!(representatives.iter().all(|&i| i < 5));

        // Should select diverse representatives
        assert!(representatives[0] != representatives[1]);
        assert!(representatives[1] != representatives[2]);
        assert!(representatives[0] != representatives[2]);

        Ok(())
    }

    #[test]
    fn test_parameter_diversity() {
        let model1 = create_test_model(
            "uniform",
            vec![
                ("weight".to_string(), vec![1.0, 1.0, 1.0, 1.0]), // No diversity
            ],
        );

        let model2 = create_test_model(
            "diverse",
            vec![
                ("weight".to_string(), vec![1.0, 5.0, 10.0, 2.0]), // High diversity
            ],
        );

        let diversity1 = StrategyUtils::compute_parameter_diversity(&model1);
        let diversity2 = StrategyUtils::compute_parameter_diversity(&model2);

        assert!(diversity1 < diversity2);
        assert_eq!(diversity1, 0.0); // Uniform values should have zero diversity
        assert!(diversity2 > 0.0); // Diverse values should have positive diversity
    }

    #[test]
    fn test_memory_estimation() -> Result<()> {
        let model1 = create_test_model("model1", vec![("weight".to_string(), vec![1.0; 1000])]);

        let model2 = create_test_model("model2", vec![("weight".to_string(), vec![1.0; 1000])]);

        let models = vec![model1, model2];

        let soup_memory = StrategyUtils::estimate_memory_requirements(&models, "ModelSoup")?;
        let progressive_memory =
            StrategyUtils::estimate_memory_requirements(&models, "Progressive")?;

        // Progressive should require more memory than ModelSoup
        assert!(progressive_memory > soup_memory);

        Ok(())
    }

    #[test]
    fn test_resource_constraints() -> Result<()> {
        let small_model = create_test_model("small", vec![("weight".to_string(), vec![1.0; 10])]);

        let models = vec![small_model];

        // Should pass with reasonable memory limit
        let result = StrategyUtils::check_resource_constraints(&models, "ModelSoup", Some(1024));
        assert!(result.is_ok());

        // Should fail with very small memory limit
        let result = StrategyUtils::check_resource_constraints(&models, "ModelSoup", Some(1));
        assert!(result.is_err());

        Ok(())
    }

    #[test]
    fn test_validation_checkpoints() -> Result<()> {
        let base_model = create_test_model("base", vec![("weight".to_string(), vec![0.0, 0.0])]);

        let target_model =
            create_test_model("target", vec![("weight".to_string(), vec![2.0, 4.0])]);

        let checkpoints =
            StrategyUtils::create_validation_checkpoints(&base_model, &target_model, 4)?;

        assert_eq!(checkpoints.len(), 5); // 0, 1, 2, 3, 4 (inclusive)

        // First checkpoint should be base model
        let first_weight = checkpoints[0].model.get_parameter("weight").unwrap();
        assert_eq!(first_weight.data, vec![0.0, 0.0]);

        // Last checkpoint should be target model
        let last_weight = checkpoints[4].model.get_parameter("weight").unwrap();
        assert_eq!(last_weight.data, vec![2.0, 4.0]);

        // Middle checkpoint should be interpolated
        let middle_weight = checkpoints[2].model.get_parameter("weight").unwrap();
        assert_eq!(middle_weight.data, vec![1.0, 2.0]); // 50% interpolation

        Ok(())
    }
}