//! Fusion Pattern Library
//!
//! Provides a comprehensive library of fusion patterns optimized for
//! different model architectures (Transformer, Mamba, MoE, etc.)

use std::collections::HashMap;

/// Standard fusion patterns for common neural network architectures
#[derive(Debug, Clone)]
pub struct FusionPatternLibrary {
    /// Patterns organized by architecture type
    patterns: HashMap<Architecture, Vec<FusionPattern>>,
}

/// Neural network architecture types
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub enum Architecture {
    /// Standard Transformer (BERT, GPT-2 style)
    Transformer,
    /// LLaMA-style Transformer (RMSNorm, SwiGLU, RoPE)
    LLaMA,
    /// Mamba/State Space Model
    Mamba,
    /// Mixture of Experts
    MoE,
    /// Vision Transformer
    ViT,
    /// Convolutional Neural Network
    CNN,
    /// Generic/Unknown
    Generic,
}

/// A fusion pattern with metadata
#[derive(Debug, Clone)]
pub struct FusionPattern {
    /// Name of the pattern
    pub name: String,
    /// Description of what this pattern fuses
    pub description: String,
    /// Sequence of operations to match
    pub operations: Vec<PatternOp>,
    /// Metal kernel name to use
    pub metal_kernel: String,
    /// Estimated performance improvement (1.0 = no improvement)
    pub speedup_factor: f32,
    /// Memory bandwidth reduction factor
    pub memory_reduction: f32,
    /// Whether this pattern is enabled by default
    pub enabled: bool,
}

/// Operations in a pattern (simplified for matching)
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum PatternOp {
    /// Matrix multiplication
    MatMul,
    /// Bias addition
    BiasAdd,
    /// ReLU activation
    ReLU,
    /// GeLU activation
    GeLU,
    /// SiLU/Swish activation
    SiLU,
    /// Residual connection
    Residual,
    /// Layer normalization
    LayerNorm,
    /// RMS normalization
    RMSNorm,
    /// Elementwise add
    Add,
    /// Elementwise multiply
    Mul,
    /// Softmax
    Softmax,
    /// Rotary position embedding
    RoPE,
    /// SwiGLU gated MLP
    SwiGLU,
    /// GeGLU gated MLP
    GeGLU,
    /// Causal mask application
    CausalMask,
    /// Expert routing (MoE)
    ExpertRoute,
    /// Selective scan (Mamba)
    SelectiveScan,
    /// Convolution
    Conv,
    /// Batch normalization
    BatchNorm,
    /// Any activation (wildcard)
    AnyActivation,
}

impl FusionPatternLibrary {
    /// Create a new pattern library with all standard patterns
    pub fn new() -> Self {
        let mut patterns = HashMap::new();

        // Transformer patterns
        patterns.insert(Architecture::Transformer, Self::transformer_patterns());

        // LLaMA patterns
        patterns.insert(Architecture::LLaMA, Self::llama_patterns());

        // Mamba patterns
        patterns.insert(Architecture::Mamba, Self::mamba_patterns());

        // MoE patterns
        patterns.insert(Architecture::MoE, Self::moe_patterns());

        // CNN patterns
        patterns.insert(Architecture::CNN, Self::cnn_patterns());

        // Generic patterns (always applicable)
        patterns.insert(Architecture::Generic, Self::generic_patterns());

        Self { patterns }
    }

    /// Get patterns for a specific architecture
    pub fn patterns_for(&self, arch: Architecture) -> Vec<&FusionPattern> {
        let mut result = Vec::new();

        // Always include generic patterns
        if let Some(generic) = self.patterns.get(&Architecture::Generic) {
            result.extend(generic.iter());
        }

        // Add architecture-specific patterns
        if arch != Architecture::Generic {
            if let Some(specific) = self.patterns.get(&arch) {
                result.extend(specific.iter());
            }
        }

        result
    }

    /// Get all patterns
    pub fn all_patterns(&self) -> impl Iterator<Item = &FusionPattern> {
        self.patterns.values().flatten()
    }

    /// Transformer-specific fusion patterns
    fn transformer_patterns() -> Vec<FusionPattern> {
        vec![
            FusionPattern {
                name: "attention_qkv_projection".to_string(),
                description: "Fused Q, K, V projections into single GEMM".to_string(),
                operations: vec![PatternOp::MatMul, PatternOp::MatMul, PatternOp::MatMul],
                metal_kernel: "qkv_projection_fused_f32".to_string(),
                speedup_factor: 1.3,
                memory_reduction: 0.33,
                enabled: true,
            },
            FusionPattern {
                name: "attention_scores_softmax".to_string(),
                description: "Fused QK^T / sqrt(d) + causal mask + softmax".to_string(),
                operations: vec![PatternOp::MatMul, PatternOp::Softmax],
                metal_kernel: "attention_scores_causal_f32".to_string(),
                speedup_factor: 1.5,
                memory_reduction: 0.50,
                enabled: true,
            },
            FusionPattern {
                name: "ffn_gelu".to_string(),
                description: "Fused FFN with GeLU activation".to_string(),
                operations: vec![PatternOp::MatMul, PatternOp::BiasAdd, PatternOp::GeLU],
                metal_kernel: "gemm_bias_gelu_f32".to_string(),
                speedup_factor: 1.4,
                memory_reduction: 0.50,
                enabled: true,
            },
            FusionPattern {
                name: "residual_layernorm".to_string(),
                description: "Fused residual add + layer normalization".to_string(),
                operations: vec![PatternOp::Residual, PatternOp::LayerNorm],
                metal_kernel: "residual_layernorm_f32".to_string(),
                speedup_factor: 1.3,
                memory_reduction: 0.50,
                enabled: true,
            },
        ]
    }

    /// LLaMA-specific fusion patterns
    fn llama_patterns() -> Vec<FusionPattern> {
        vec![
            FusionPattern {
                name: "llama_rmsnorm".to_string(),
                description: "Optimized RMS normalization for LLaMA".to_string(),
                operations: vec![PatternOp::RMSNorm],
                metal_kernel: "rmsnorm_f32".to_string(),
                speedup_factor: 1.2,
                memory_reduction: 0.0,
                enabled: true,
            },
            FusionPattern {
                name: "llama_residual_rmsnorm".to_string(),
                description: "Fused residual + RMS normalization".to_string(),
                operations: vec![PatternOp::Residual, PatternOp::RMSNorm],
                metal_kernel: "residual_rmsnorm_f32".to_string(),
                speedup_factor: 1.4,
                memory_reduction: 0.50,
                enabled: true,
            },
            FusionPattern {
                name: "llama_swiglu".to_string(),
                description: "Fused SwiGLU gated MLP (gate * up projections)".to_string(),
                operations: vec![PatternOp::SwiGLU],
                metal_kernel: "swiglu_fused_f32".to_string(),
                speedup_factor: 1.6,
                memory_reduction: 0.50,
                enabled: true,
            },
            FusionPattern {
                name: "llama_rope".to_string(),
                description: "Fused rotary position embedding".to_string(),
                operations: vec![PatternOp::RoPE],
                metal_kernel: "rope_fused_f32".to_string(),
                speedup_factor: 1.2,
                memory_reduction: 0.0,
                enabled: true,
            },
            FusionPattern {
                name: "llama_attention_rope".to_string(),
                description: "Fused QK projection + RoPE".to_string(),
                operations: vec![PatternOp::MatMul, PatternOp::RoPE],
                metal_kernel: "qk_rope_fused_f32".to_string(),
                speedup_factor: 1.3,
                memory_reduction: 0.33,
                enabled: true,
            },
        ]
    }

    /// Mamba-specific fusion patterns
    fn mamba_patterns() -> Vec<FusionPattern> {
        vec![
            FusionPattern {
                name: "mamba_selective_scan".to_string(),
                description: "Optimized selective scan kernel".to_string(),
                operations: vec![PatternOp::SelectiveScan],
                metal_kernel: "mamba_selective_scan_f32".to_string(),
                speedup_factor: 2.0,
                memory_reduction: 0.0,
                enabled: true,
            },
            FusionPattern {
                name: "mamba_conv_silu".to_string(),
                description: "Fused causal conv1d + SiLU".to_string(),
                operations: vec![PatternOp::Conv, PatternOp::SiLU],
                metal_kernel: "mamba_causal_conv_silu_f32".to_string(),
                speedup_factor: 1.4,
                memory_reduction: 0.50,
                enabled: true,
            },
            FusionPattern {
                name: "mamba_ssm_update".to_string(),
                description: "Fused SSM state update".to_string(),
                operations: vec![PatternOp::MatMul, PatternOp::Add, PatternOp::Mul],
                metal_kernel: "mamba_ssm_update_f32".to_string(),
                speedup_factor: 1.5,
                memory_reduction: 0.50,
                enabled: true,
            },
        ]
    }

    /// MoE-specific fusion patterns
    fn moe_patterns() -> Vec<FusionPattern> {
        vec![
            FusionPattern {
                name: "moe_topk_gating".to_string(),
                description: "Fused top-k expert selection with softmax".to_string(),
                operations: vec![
                    PatternOp::MatMul,
                    PatternOp::Softmax,
                    PatternOp::ExpertRoute,
                ],
                metal_kernel: "moe_topk_gating_f32".to_string(),
                speedup_factor: 1.5,
                memory_reduction: 0.33,
                enabled: true,
            },
            FusionPattern {
                name: "moe_expert_gemm".to_string(),
                description: "Batched GEMM for expert computation".to_string(),
                operations: vec![PatternOp::ExpertRoute, PatternOp::MatMul],
                metal_kernel: "moe_expert_gemm_f32".to_string(),
                speedup_factor: 1.4,
                memory_reduction: 0.25,
                enabled: true,
            },
            FusionPattern {
                name: "moe_combine_experts".to_string(),
                description: "Fused weighted expert combination".to_string(),
                operations: vec![PatternOp::Mul, PatternOp::Add],
                metal_kernel: "moe_combine_experts_f32".to_string(),
                speedup_factor: 1.3,
                memory_reduction: 0.50,
                enabled: true,
            },
        ]
    }

    /// CNN-specific fusion patterns
    fn cnn_patterns() -> Vec<FusionPattern> {
        vec![
            FusionPattern {
                name: "conv_bn_relu".to_string(),
                description: "Fused convolution + batch norm + ReLU".to_string(),
                operations: vec![PatternOp::Conv, PatternOp::BatchNorm, PatternOp::ReLU],
                metal_kernel: "conv_bn_relu_f32".to_string(),
                speedup_factor: 1.6,
                memory_reduction: 0.66,
                enabled: true,
            },
            FusionPattern {
                name: "conv_relu".to_string(),
                description: "Fused convolution + ReLU".to_string(),
                operations: vec![PatternOp::Conv, PatternOp::ReLU],
                metal_kernel: "conv_relu_f32".to_string(),
                speedup_factor: 1.3,
                memory_reduction: 0.50,
                enabled: true,
            },
        ]
    }

    /// Generic fusion patterns (applicable to any architecture)
    fn generic_patterns() -> Vec<FusionPattern> {
        vec![
            FusionPattern {
                name: "gemm_relu".to_string(),
                description: "Fused GEMM + ReLU".to_string(),
                operations: vec![PatternOp::MatMul, PatternOp::ReLU],
                metal_kernel: "gemm_relu_f32".to_string(),
                speedup_factor: 1.3,
                memory_reduction: 0.50,
                enabled: true,
            },
            FusionPattern {
                name: "gemm_gelu".to_string(),
                description: "Fused GEMM + GeLU".to_string(),
                operations: vec![PatternOp::MatMul, PatternOp::GeLU],
                metal_kernel: "gemm_gelu_f32".to_string(),
                speedup_factor: 1.4,
                memory_reduction: 0.50,
                enabled: true,
            },
            FusionPattern {
                name: "gemm_silu".to_string(),
                description: "Fused GEMM + SiLU".to_string(),
                operations: vec![PatternOp::MatMul, PatternOp::SiLU],
                metal_kernel: "gemm_silu_f32".to_string(),
                speedup_factor: 1.4,
                memory_reduction: 0.50,
                enabled: true,
            },
            FusionPattern {
                name: "gemm_bias_relu".to_string(),
                description: "Fused GEMM + bias + ReLU".to_string(),
                operations: vec![PatternOp::MatMul, PatternOp::BiasAdd, PatternOp::ReLU],
                metal_kernel: "gemm_bias_relu_f32".to_string(),
                speedup_factor: 1.5,
                memory_reduction: 0.66,
                enabled: true,
            },
            FusionPattern {
                name: "gemm_bias_gelu".to_string(),
                description: "Fused GEMM + bias + GeLU".to_string(),
                operations: vec![PatternOp::MatMul, PatternOp::BiasAdd, PatternOp::GeLU],
                metal_kernel: "gemm_bias_gelu_f32".to_string(),
                speedup_factor: 1.5,
                memory_reduction: 0.66,
                enabled: true,
            },
            FusionPattern {
                name: "add_relu".to_string(),
                description: "Fused elementwise add + ReLU".to_string(),
                operations: vec![PatternOp::Add, PatternOp::ReLU],
                metal_kernel: "add_relu_f32".to_string(),
                speedup_factor: 1.3,
                memory_reduction: 0.50,
                enabled: true,
            },
            FusionPattern {
                name: "add_gelu".to_string(),
                description: "Fused elementwise add + GeLU".to_string(),
                operations: vec![PatternOp::Add, PatternOp::GeLU],
                metal_kernel: "add_gelu_f32".to_string(),
                speedup_factor: 1.3,
                memory_reduction: 0.50,
                enabled: true,
            },
            FusionPattern {
                name: "fma_relu".to_string(),
                description: "Fused FMA + ReLU (a*b+c then ReLU)".to_string(),
                operations: vec![PatternOp::Mul, PatternOp::Add, PatternOp::ReLU],
                metal_kernel: "fma_relu_f32".to_string(),
                speedup_factor: 1.4,
                memory_reduction: 0.66,
                enabled: true,
            },
            FusionPattern {
                name: "geglu".to_string(),
                description: "Fused GeGLU gated activation".to_string(),
                operations: vec![PatternOp::GeGLU],
                metal_kernel: "geglu_fused_f32".to_string(),
                speedup_factor: 1.5,
                memory_reduction: 0.50,
                enabled: true,
            },
        ]
    }

    /// Enable or disable a pattern by name
    pub fn set_pattern_enabled(&mut self, name: &str, enabled: bool) -> bool {
        for patterns in self.patterns.values_mut() {
            for pattern in patterns.iter_mut() {
                if pattern.name == name {
                    pattern.enabled = enabled;
                    return true;
                }
            }
        }
        false
    }

    /// Get pattern by name
    pub fn get_pattern(&self, name: &str) -> Option<&FusionPattern> {
        for patterns in self.patterns.values() {
            for pattern in patterns {
                if pattern.name == name {
                    return Some(pattern);
                }
            }
        }
        None
    }

    /// Estimate total speedup for an architecture
    pub fn estimate_speedup(&self, arch: Architecture) -> f32 {
        let patterns = self.patterns_for(arch);
        if patterns.is_empty() {
            return 1.0;
        }

        // Geometric mean of speedup factors (more realistic than arithmetic mean)
        let product: f32 = patterns
            .iter()
            .filter(|p| p.enabled)
            .map(|p| p.speedup_factor)
            .product();

        let count = patterns.iter().filter(|p| p.enabled).count();
        if count == 0 {
            1.0
        } else {
            product.powf(1.0 / count as f32)
        }
    }
}

impl Default for FusionPatternLibrary {
    fn default() -> Self {
        Self::new()
    }
}

/// Pattern matching result
#[derive(Debug, Clone)]
pub struct PatternMatch {
    /// The matched pattern
    pub pattern: FusionPattern,
    /// Start index in the operation sequence
    pub start_idx: usize,
    /// End index in the operation sequence
    pub end_idx: usize,
    /// Confidence score (0.0 to 1.0)
    pub confidence: f32,
}

/// Pattern matcher for finding fusion opportunities
#[derive(Debug)]
pub struct PatternMatcher {
    library: FusionPatternLibrary,
    architecture: Architecture,
}

impl PatternMatcher {
    /// Create a new pattern matcher
    pub fn new(architecture: Architecture) -> Self {
        Self {
            library: FusionPatternLibrary::new(),
            architecture,
        }
    }

    /// Create with a custom pattern library
    pub fn with_library(library: FusionPatternLibrary, architecture: Architecture) -> Self {
        Self {
            library,
            architecture,
        }
    }

    /// Find all pattern matches in an operation sequence
    pub fn find_matches(&self, operations: &[PatternOp]) -> Vec<PatternMatch> {
        let mut matches = Vec::new();
        let patterns = self.library.patterns_for(self.architecture);

        for pattern in patterns {
            if !pattern.enabled {
                continue;
            }

            // Try to match pattern at each position
            for start_idx in 0..operations.len() {
                if let Some((end_idx, confidence)) =
                    self.try_match_at(operations, &pattern.operations, start_idx)
                {
                    matches.push(PatternMatch {
                        pattern: pattern.clone(),
                        start_idx,
                        end_idx,
                        confidence,
                    });
                }
            }
        }

        // Sort by confidence and position
        matches.sort_by(|a, b| {
            b.confidence
                .partial_cmp(&a.confidence)
                .unwrap_or(std::cmp::Ordering::Equal)
                .then(a.start_idx.cmp(&b.start_idx))
        });

        matches
    }

    /// Try to match a pattern at a specific position
    fn try_match_at(
        &self,
        operations: &[PatternOp],
        pattern: &[PatternOp],
        start_idx: usize,
    ) -> Option<(usize, f32)> {
        if start_idx + pattern.len() > operations.len() {
            return None;
        }

        let mut confidence = 1.0;

        for (i, pattern_op) in pattern.iter().enumerate() {
            let op = &operations[start_idx + i];

            match (pattern_op, op) {
                // Exact match
                (a, b) if a == b => {}
                // Wildcard activation
                (PatternOp::AnyActivation, PatternOp::ReLU)
                | (PatternOp::AnyActivation, PatternOp::GeLU)
                | (PatternOp::AnyActivation, PatternOp::SiLU) => {
                    confidence *= 0.9; // Slightly lower confidence for wildcards
                }
                // No match
                _ => return None,
            }
        }

        Some((start_idx + pattern.len(), confidence))
    }

    /// Find the best non-overlapping matches
    pub fn find_best_matches(&self, operations: &[PatternOp]) -> Vec<PatternMatch> {
        let all_matches = self.find_matches(operations);
        let mut best = Vec::new();
        let mut covered = vec![false; operations.len()];

        for m in all_matches {
            // Check if this match overlaps with existing ones
            let overlaps = (m.start_idx..m.end_idx).any(|i| covered[i]);

            if !overlaps {
                // Mark as covered
                for i in m.start_idx..m.end_idx {
                    covered[i] = true;
                }
                best.push(m);
            }
        }

        best
    }
}

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

    #[test]
    fn test_pattern_library_creation() {
        let lib = FusionPatternLibrary::new();
        assert!(lib.patterns_for(Architecture::Transformer).len() > 0);
        assert!(lib.patterns_for(Architecture::LLaMA).len() > 0);
        assert!(lib.patterns_for(Architecture::Mamba).len() > 0);
    }

    #[test]
    fn test_pattern_matching() {
        let matcher = PatternMatcher::new(Architecture::Generic);

        let ops = vec![PatternOp::MatMul, PatternOp::BiasAdd, PatternOp::ReLU];

        let matches = matcher.find_matches(&ops);
        assert!(!matches.is_empty());
    }

    #[test]
    fn test_llama_patterns() {
        let lib = FusionPatternLibrary::new();
        let llama_patterns = lib.patterns_for(Architecture::LLaMA);

        // Should have SwiGLU pattern
        assert!(llama_patterns.iter().any(|p| p.name == "llama_swiglu"));

        // Should have RMSNorm pattern
        assert!(llama_patterns.iter().any(|p| p.name == "llama_rmsnorm"));
    }

    #[test]
    fn test_speedup_estimation() {
        let lib = FusionPatternLibrary::new();

        let speedup = lib.estimate_speedup(Architecture::LLaMA);
        assert!(speedup > 1.0, "LLaMA should have speedup > 1.0");
    }

    #[test]
    fn test_pattern_enable_disable() {
        let mut lib = FusionPatternLibrary::new();

        assert!(lib.set_pattern_enabled("gemm_relu", false));

        let pattern = lib.get_pattern("gemm_relu").unwrap();
        assert!(!pattern.enabled);
    }

    #[test]
    fn test_best_matches_no_overlap() {
        let matcher = PatternMatcher::new(Architecture::Generic);

        // Sequence that could have overlapping patterns
        let ops = vec![
            PatternOp::MatMul,
            PatternOp::ReLU,
            PatternOp::MatMul,
            PatternOp::GeLU,
        ];

        let best = matcher.find_best_matches(&ops);

        // Check no overlapping matches
        for i in 0..best.len() {
            for j in (i + 1)..best.len() {
                let range_i = best[i].start_idx..best[i].end_idx;
                let range_j = best[j].start_idx..best[j].end_idx;

                // Ranges should not overlap
                assert!(
                    range_i.end <= range_j.start || range_j.end <= range_i.start,
                    "Matches should not overlap"
                );
            }
        }
    }
}