//! Module composition strategies for building complex architectures
//!
//! This module implements various ways to compose modules together,
//! enabling the creation of complex computational graphs from simple building blocks.

use super::*;

/// Strategy for composing multiple modules
#[derive(Debug, Clone, PartialEq)]
pub enum CompositionStrategy {
    /// Sequential composition (pipeline)
    Sequential,

    /// Parallel composition with combination
    Parallel(CombinationStrategy),

    /// Hierarchical composition (multi-level)
    Hierarchical,

    /// Graph-based composition with arbitrary connections
    Graph,

    /// Dynamic composition based on input
    Dynamic,
}

/// Strategy for combining outputs from parallel modules
#[derive(Debug, Clone, PartialEq)]
pub enum CombinationStrategy {
    /// Concatenate outputs along feature dimension
    Concatenate,

    /// Element-wise average
    Average,

    /// Element-wise sum
    Sum,

    /// Weighted sum with fixed weights
    WeightedSum(Vec<f32>),

    /// Attention-weighted combination
    AttentionWeighted,

    /// Maximum pooling across modules
    Max,

    /// Learned combination using a neural network
    Learned,
}

impl CombinationStrategy {
    pub fn combine(&self, tensors: &[Tensor]) -> Result<Tensor> {
        if tensors.is_empty() {
            return Err(ModularError::CompositionError {
                message: "No tensors to combine".to_string(),
            });
        }

        if tensors.len() == 1 {
            return Ok(tensors[0].clone());
        }

        match self {
            CombinationStrategy::Concatenate => {
                // Concatenate along the last dimension (feature dimension)
                Tensor::cat(tensors, -1).map_err(|e| ModularError::CompositionError {
                    message: format!("Concatenation failed: {}", e),
                })
            }

            CombinationStrategy::Average => {
                let sum = tensors
                    .iter()
                    .try_fold(tensors[0].clone(), |acc, t| acc.add(t))?;
                sum.div_scalar(tensors.len() as f32)
                    .map_err(|e| ModularError::CompositionError {
                        message: format!("Average failed: {}", e),
                    })
            }

            CombinationStrategy::Sum => tensors
                .iter()
                .try_fold(tensors[0].clone(), |acc, t| acc.add(t))
                .map_err(|e| ModularError::CompositionError {
                    message: format!("Sum failed: {}", e),
                }),

            CombinationStrategy::WeightedSum(weights) => {
                if weights.len() != tensors.len() {
                    return Err(ModularError::CompositionError {
                        message: format!(
                            "Weight count {} doesn't match tensor count {}",
                            weights.len(),
                            tensors.len()
                        ),
                    });
                }

                let mut result = tensors[0].mul_scalar(weights[0])?;
                for (i, tensor) in tensors.iter().enumerate().skip(1) {
                    let weighted = tensor.mul_scalar(weights[i])?;
                    result = result.add(&weighted)?;
                }
                Ok(result)
            }

            CombinationStrategy::Max => {
                // Stack tensors and take max along the new dimension
                let stacked = Tensor::stack(tensors, 0)?;
                stacked.max().map_err(|e| ModularError::CompositionError {
                    message: format!("Max pooling failed: {}", e),
                })
            }

            CombinationStrategy::AttentionWeighted => {
                // Simplified attention-based combination
                // In practice, this would use a proper attention mechanism
                self.combine(&[tensors[0].clone()]) // Fallback to first tensor
            }

            CombinationStrategy::Learned => {
                // Placeholder - would use a learned combination network
                self.combine(&[tensors[0].clone()]) // Fallback to first tensor
            }
        }
    }
}

/// Result of module composition
#[derive(Debug)]
pub struct CompositionResult {
    /// Final output tensor
    pub output: Tensor,

    /// Intermediate outputs from each stage/module
    pub intermediate_outputs: Vec<Tensor>,

    /// Modules that were actually used
    pub active_modules: Vec<String>,

    /// Total computational cost
    pub total_cost: f32,

    /// Composition metadata
    pub metadata: std::collections::HashMap<String, String>,
}

/// Configuration for module composition
#[derive(Debug, Clone)]
pub struct CompositionConfig {
    /// Whether to store intermediate outputs
    pub store_intermediates: bool,

    /// Maximum depth for hierarchical composition
    pub max_depth: usize,

    /// Whether to enable gradient flow through all paths
    pub enable_full_gradients: bool,

    /// Resource limits for composition
    pub resource_limits: ResourceBudget,
}

impl Default for CompositionConfig {
    fn default() -> Self {
        Self {
            store_intermediates: false,
            max_depth: 5,
            enable_full_gradients: true,
            resource_limits: ResourceBudget::default(),
        }
    }
}

// ============================================================================
// Sequential Composition
// ============================================================================

/// Sequential composition of modules (pipeline)
#[derive(Debug)]
pub struct SequentialComposition {
    modules: Vec<Box<dyn Module>>,
    config: CompositionConfig,
    metadata: ModuleMetadata,
}

impl SequentialComposition {
    pub fn new(modules: Vec<Box<dyn Module>>) -> Self {
        Self {
            modules,
            config: CompositionConfig::default(),
            metadata: ModuleMetadata::default(),
        }
    }

    pub fn with_config(mut self, config: CompositionConfig) -> Self {
        self.config = config;
        self
    }

    pub fn forward(&self, input: &Tensor) -> Result<CompositionResult> {
        if self.modules.is_empty() {
            return Err(ModularError::CompositionError {
                message: "No modules in sequential composition".to_string(),
            });
        }

        let mut current_input = input.clone();
        let mut intermediate_outputs = Vec::new();
        let mut active_modules = Vec::new();
        let mut total_cost = 0.0;

        for module in &self.modules {
            // Check compatibility
            if current_input.dims().last() != Some(&module.input_dim()) {
                return Err(ModularError::IncompatibleDimensions {
                    expected: module.input_dim(),
                    actual: *current_input.dims().last().unwrap(),
                });
            }

            // Forward pass through module
            let output = module.forward(&current_input)?;

            // Update state
            if self.config.store_intermediates {
                intermediate_outputs.push(current_input.clone());
            }
            active_modules.push(module.module_id().to_string());
            total_cost += module.computation_cost(current_input.dims());

            current_input = output;
        }

        let mut metadata = std::collections::HashMap::new();
        metadata.insert("composition_type".to_string(), "sequential".to_string());
        metadata.insert("num_modules".to_string(), self.modules.len().to_string());

        Ok(CompositionResult {
            output: current_input,
            intermediate_outputs,
            active_modules,
            total_cost,
            metadata,
        })
    }
}

impl Module for SequentialComposition {
    fn forward(&self, input: &Tensor) -> Result<Tensor> {
        Ok(self.forward(input)?.output)
    }

    fn input_dim(&self) -> usize {
        self.modules.first().map(|m| m.input_dim()).unwrap_or(0)
    }

    fn output_dim(&self) -> usize {
        self.modules.last().map(|m| m.output_dim()).unwrap_or(0)
    }

    fn module_type(&self) -> &str {
        "sequential_composition"
    }

    fn module_id(&self) -> &str {
        "sequential"
    }

    fn capabilities(&self) -> Vec<String> {
        let mut all_caps = Vec::new();
        for module in &self.modules {
            all_caps.extend(module.capabilities());
        }
        all_caps.sort();
        all_caps.dedup();
        all_caps
    }

    fn metadata(&self) -> &ModuleMetadata {
        self.modules
            .first()
            .map(|m| m.metadata())
            .unwrap_or(&self.metadata)
    }

    fn complexity_score(&self) -> f32 {
        self.modules.iter().map(|m| m.complexity_score()).sum()
    }

    fn memory_footprint(&self) -> usize {
        self.modules.iter().map(|m| m.memory_footprint()).sum()
    }

    fn is_stateful(&self) -> bool {
        self.modules.iter().any(|m| m.is_stateful())
    }

    fn requires_training(&self) -> bool {
        self.modules.iter().any(|m| m.requires_training())
    }

    fn parameters(&self) -> Vec<&Tensor> {
        self.modules.iter().flat_map(|m| m.parameters()).collect()
    }

    fn parameters_mut(&mut self) -> Vec<&mut Tensor> {
        self.modules
            .iter_mut()
            .flat_map(|m| m.parameters_mut())
            .collect()
    }

    fn clone_module(&self) -> Box<dyn Module> {
        let cloned_modules = self.modules.iter().map(|m| m.clone_module()).collect();

        Box::new(SequentialComposition {
            modules: cloned_modules,
            config: self.config.clone(),
            metadata: self.metadata.clone(),
        })
    }

    fn specialize_for_task(
        &self,
        task_type: TaskType,
        training_data: &[TestData],
    ) -> Result<Box<dyn Module>> {
        let specialized_modules = self
            .modules
            .iter()
            .map(|m| m.specialize_for_task(task_type.clone(), training_data))
            .collect::<Result<Vec<_>>>()?;

        Ok(Box::new(SequentialComposition::new(specialized_modules)))
    }

    fn continually_specialize(&mut self, training_data: &[TestData], task_id: usize) -> Result<()> {
        for module in &mut self.modules {
            module.continually_specialize(training_data, task_id)?;
        }
        Ok(())
    }
}

// ============================================================================
// Parallel Composition
// ============================================================================

/// Parallel composition of modules with combination strategy
#[derive(Debug)]
pub struct ParallelComposition {
    modules: Vec<Box<dyn Module>>,
    combination_strategy: CombinationStrategy,
    config: CompositionConfig,
    metadata: ModuleMetadata,
}

impl ParallelComposition {
    pub fn new(modules: Vec<Box<dyn Module>>, combination_strategy: CombinationStrategy) -> Self {
        Self {
            modules,
            combination_strategy,
            config: CompositionConfig::default(),
            metadata: ModuleMetadata::default(),
        }
    }

    pub fn with_config(mut self, config: CompositionConfig) -> Self {
        self.config = config;
        self
    }

    pub fn forward(&self, input: &Tensor) -> Result<CompositionResult> {
        if self.modules.is_empty() {
            return Err(ModularError::CompositionError {
                message: "No modules in parallel composition".to_string(),
            });
        }

        let mut outputs = Vec::new();
        let mut active_modules = Vec::new();
        let mut total_cost = 0.0;

        // Execute all modules in parallel
        for module in &self.modules {
            // Check input compatibility
            if input.dims().last() != Some(&module.input_dim()) {
                return Err(ModularError::IncompatibleDimensions {
                    expected: module.input_dim(),
                    actual: *input.dims().last().unwrap(),
                });
            }

            let output = module.forward(input)?;
            outputs.push(output);
            active_modules.push(module.module_id().to_string());
            total_cost += module.computation_cost(input.dims());
        }

        // Combine outputs
        let combined_output = self.combination_strategy.combine(&outputs)?;

        let mut metadata = std::collections::HashMap::new();
        metadata.insert("composition_type".to_string(), "parallel".to_string());
        metadata.insert("num_modules".to_string(), self.modules.len().to_string());
        metadata.insert(
            "combination_strategy".to_string(),
            format!("{:?}", self.combination_strategy),
        );

        Ok(CompositionResult {
            output: combined_output,
            intermediate_outputs: if self.config.store_intermediates {
                outputs
            } else {
                Vec::new()
            },
            active_modules,
            total_cost,
            metadata,
        })
    }
}

impl Module for ParallelComposition {
    fn forward(&self, input: &Tensor) -> Result<Tensor> {
        Ok(self.forward(input)?.output)
    }

    fn input_dim(&self) -> usize {
        // All modules should have the same input dimension
        self.modules.first().map(|m| m.input_dim()).unwrap_or(0)
    }

    fn output_dim(&self) -> usize {
        match &self.combination_strategy {
            CombinationStrategy::Concatenate => self.modules.iter().map(|m| m.output_dim()).sum(),
            _ => {
                // For other strategies, output dimension equals input dimension of modules
                self.modules.first().map(|m| m.output_dim()).unwrap_or(0)
            }
        }
    }

    fn module_type(&self) -> &str {
        "parallel_composition"
    }

    fn module_id(&self) -> &str {
        "parallel"
    }

    fn capabilities(&self) -> Vec<String> {
        let mut all_caps = Vec::new();
        for module in &self.modules {
            all_caps.extend(module.capabilities());
        }
        all_caps.sort();
        all_caps.dedup();
        all_caps
    }

    fn metadata(&self) -> &ModuleMetadata {
        self.modules
            .first()
            .map(|m| m.metadata())
            .unwrap_or(&self.metadata)
    }

    fn complexity_score(&self) -> f32 {
        self.modules
            .iter()
            .map(|m| m.complexity_score())
            .fold(0.0, f32::max) // Max complexity since they run in parallel
    }

    fn memory_footprint(&self) -> usize {
        self.modules.iter().map(|m| m.memory_footprint()).sum() // All modules loaded simultaneously
    }

    fn is_stateful(&self) -> bool {
        self.modules.iter().any(|m| m.is_stateful())
    }

    fn requires_training(&self) -> bool {
        self.modules.iter().any(|m| m.requires_training())
    }

    fn parameters(&self) -> Vec<&Tensor> {
        self.modules.iter().flat_map(|m| m.parameters()).collect()
    }

    fn parameters_mut(&mut self) -> Vec<&mut Tensor> {
        self.modules
            .iter_mut()
            .flat_map(|m| m.parameters_mut())
            .collect()
    }

    fn clone_module(&self) -> Box<dyn Module> {
        let cloned_modules = self.modules.iter().map(|m| m.clone_module()).collect();

        Box::new(ParallelComposition {
            modules: cloned_modules,
            combination_strategy: self.combination_strategy.clone(),
            config: self.config.clone(),
            metadata: ModuleMetadata::default(),
        })
    }

    fn specialize_for_task(
        &self,
        task_type: TaskType,
        training_data: &[TestData],
    ) -> Result<Box<dyn Module>> {
        let specialized_modules = self
            .modules
            .iter()
            .map(|m| m.specialize_for_task(task_type.clone(), training_data))
            .collect::<Result<Vec<_>>>()?;

        Ok(Box::new(ParallelComposition::new(
            specialized_modules,
            self.combination_strategy.clone(),
        )))
    }

    fn continually_specialize(&mut self, training_data: &[TestData], task_id: usize) -> Result<()> {
        for module in &mut self.modules {
            module.continually_specialize(training_data, task_id)?;
        }
        Ok(())
    }
}

// ============================================================================
// Hierarchical Composition
// ============================================================================

/// Hierarchical composition with multiple levels
#[derive(Debug)]
pub struct HierarchicalComposition {
    levels: Vec<Vec<Box<dyn Module>>>,
    level_combination_strategies: Vec<CombinationStrategy>,
    config: CompositionConfig,
    metadata: ModuleMetadata,
}

impl HierarchicalComposition {
    pub fn new() -> Self {
        Self {
            levels: Vec::new(),
            level_combination_strategies: Vec::new(),
            config: CompositionConfig::default(),
            metadata: ModuleMetadata::default(),
        }
    }

    pub fn add_level(mut self, modules: Vec<Box<dyn Module>>) -> Self {
        self.levels.push(modules);
        // Default to averaging for each level
        self.level_combination_strategies
            .push(CombinationStrategy::Average);
        self
    }

    pub fn with_level_strategy(mut self, level: usize, strategy: CombinationStrategy) -> Self {
        if level < self.level_combination_strategies.len() {
            self.level_combination_strategies[level] = strategy;
        }
        self
    }

    pub fn with_config(mut self, config: CompositionConfig) -> Self {
        self.config = config;
        self
    }

    pub fn forward(&self, input: &Tensor) -> Result<CompositionResult> {
        if self.levels.is_empty() {
            return Err(ModularError::CompositionError {
                message: "No levels in hierarchical composition".to_string(),
            });
        }

        let mut current_input = input.clone();
        let mut all_intermediate_outputs = Vec::new();
        let mut all_active_modules = Vec::new();
        let mut total_cost = 0.0;

        for (level_idx, level_modules) in self.levels.iter().enumerate() {
            if level_modules.is_empty() {
                continue;
            }

            // Process all modules in this level (parallel)
            let mut level_outputs = Vec::new();
            let mut level_modules_used = Vec::new();

            for module in level_modules {
                let output = module.forward(&current_input)?;
                level_outputs.push(output);
                level_modules_used.push(module.module_id().to_string());
                total_cost += module.computation_cost(current_input.dims());
            }

            // Combine outputs from this level
            let combination_strategy = self
                .level_combination_strategies
                .get(level_idx)
                .unwrap_or(&CombinationStrategy::Average);

            let level_output = combination_strategy.combine(&level_outputs)?;

            if self.config.store_intermediates {
                all_intermediate_outputs.extend(level_outputs);
            }
            all_active_modules.extend(level_modules_used);

            // Output of this level becomes input to next level
            current_input = level_output;
        }

        let mut metadata = std::collections::HashMap::new();
        metadata.insert("composition_type".to_string(), "hierarchical".to_string());
        metadata.insert("num_levels".to_string(), self.levels.len().to_string());

        Ok(CompositionResult {
            output: current_input,
            intermediate_outputs: all_intermediate_outputs,
            active_modules: all_active_modules,
            total_cost,
            metadata,
        })
    }
}

impl Module for HierarchicalComposition {
    fn forward(&self, input: &Tensor) -> Result<Tensor> {
        Ok(self.forward(input)?.output)
    }

    fn input_dim(&self) -> usize {
        self.levels
            .first()
            .and_then(|level| level.first())
            .map(|m| m.input_dim())
            .unwrap_or(0)
    }

    fn output_dim(&self) -> usize {
        // For hierarchical composition, the output dimension depends on the last level
        // and its combination strategy
        if let Some(last_level) = self.levels.last() {
            if let Some(last_strategy) = self.level_combination_strategies.last() {
                match last_strategy {
                    CombinationStrategy::Concatenate => {
                        last_level.iter().map(|m| m.output_dim()).sum()
                    }
                    _ => last_level.first().map(|m| m.output_dim()).unwrap_or(0),
                }
            } else {
                last_level.first().map(|m| m.output_dim()).unwrap_or(0)
            }
        } else {
            0
        }
    }

    fn module_type(&self) -> &str {
        "hierarchical_composition"
    }

    fn module_id(&self) -> &str {
        "hierarchical"
    }

    fn capabilities(&self) -> Vec<String> {
        let mut all_caps = Vec::new();
        for level in &self.levels {
            for module in level {
                all_caps.extend(module.capabilities());
            }
        }
        all_caps.sort();
        all_caps.dedup();
        all_caps
    }

    fn metadata(&self) -> &ModuleMetadata {
        &self.metadata
    }

    fn complexity_score(&self) -> f32 {
        let mut total_complexity = 0.0;
        for level in &self.levels {
            let level_complexity = level
                .iter()
                .map(|m| m.complexity_score())
                .fold(0.0, f32::max); // Max within level (parallel)
            total_complexity += level_complexity; // Sum across levels (sequential)
        }
        total_complexity
    }

    fn memory_footprint(&self) -> usize {
        self.levels
            .iter()
            .flat_map(|level| level.iter())
            .map(|m| m.memory_footprint())
            .sum()
    }

    fn is_stateful(&self) -> bool {
        self.levels
            .iter()
            .any(|level| level.iter().any(|m| m.is_stateful()))
    }

    fn requires_training(&self) -> bool {
        self.levels
            .iter()
            .any(|level| level.iter().any(|m| m.requires_training()))
    }

    fn parameters(&self) -> Vec<&Tensor> {
        self.levels
            .iter()
            .flat_map(|level| level.iter())
            .flat_map(|m| m.parameters())
            .collect()
    }

    fn parameters_mut(&mut self) -> Vec<&mut Tensor> {
        self.levels
            .iter_mut()
            .flat_map(|level| level.iter_mut())
            .flat_map(|m| m.parameters_mut())
            .collect()
    }

    fn clone_module(&self) -> Box<dyn Module> {
        let cloned_levels = self
            .levels
            .iter()
            .map(|level| level.iter().map(|m| m.clone_module()).collect())
            .collect();

        Box::new(HierarchicalComposition {
            levels: cloned_levels,
            level_combination_strategies: self.level_combination_strategies.clone(),
            config: self.config.clone(),
            metadata: ModuleMetadata::default(),
        })
    }

    fn specialize_for_task(
        &self,
        task_type: TaskType,
        training_data: &[TestData],
    ) -> Result<Box<dyn Module>> {
        let specialized_levels = self
            .levels
            .iter()
            .map(|level| {
                level
                    .iter()
                    .map(|m| m.specialize_for_task(task_type.clone(), training_data))
                    .collect::<Result<Vec<_>>>()
            })
            .collect::<Result<Vec<_>>>()?;

        Ok(Box::new(HierarchicalComposition {
            levels: specialized_levels,
            level_combination_strategies: self.level_combination_strategies.clone(),
            config: self.config.clone(),
            metadata: ModuleMetadata::default(),
        }))
    }

    fn continually_specialize(&mut self, training_data: &[TestData], task_id: usize) -> Result<()> {
        for level in &mut self.levels {
            for module in level {
                module.continually_specialize(training_data, task_id)?;
            }
        }
        Ok(())
    }
}