//! Fusion pattern detection and analysis
//!
//! This module analyzes operation streams to detect fusion opportunities.
//! It supports multiple fusion patterns including elementwise chains,
//! softmax patterns, layer normalization, and activation chains.

use crate::config::FusionConfig;
use crate::kernel::{DType, KernelSignature, StreamOpKind, StreamOperation, TensorId};
use crate::stream::OperationStream;
use std::collections::{HashMap, HashSet};

/// Pre-computed analysis context for efficient pattern detection
/// This avoids rebuilding data structures for each detection pass
struct AnalysisContext<'a> {
    /// Reference to the operations slice
    ops: &'a [StreamOperation],
    /// Map from output tensor ID to operation index
    output_to_op: HashMap<TensorId, usize>,
    /// Map from operation index to indices of operations that consume its output
    op_consumers: HashMap<usize, Vec<usize>>,
    /// Map from tensor ID to indices of operations that consume it
    tensor_consumers: HashMap<TensorId, Vec<usize>>,
}

impl<'a> AnalysisContext<'a> {
    /// Build analysis context from operations
    fn new(ops: &'a [StreamOperation]) -> Self {
        let capacity = ops.len();

        // Pre-allocate with estimated capacity
        let mut output_to_op = HashMap::with_capacity(capacity);
        let mut op_consumers: HashMap<usize, Vec<usize>> = HashMap::with_capacity(capacity);
        let mut tensor_consumers: HashMap<TensorId, Vec<usize>> =
            HashMap::with_capacity(capacity * 2);

        // Single pass to build all maps
        for (idx, op) in ops.iter().enumerate() {
            output_to_op.insert(op.output, idx);
            for &input in &op.inputs {
                tensor_consumers.entry(input).or_default().push(idx);
            }
        }

        // Build op_consumers from tensor_consumers
        for (idx, op) in ops.iter().enumerate() {
            if let Some(consumers) = tensor_consumers.get(&op.output) {
                op_consumers.insert(idx, consumers.clone());
            }
        }

        Self {
            ops,
            output_to_op,
            op_consumers,
            tensor_consumers,
        }
    }

    /// Get the operation at an index
    #[inline]
    fn get_op(&self, idx: usize) -> &StreamOperation {
        &self.ops[idx]
    }

    /// Get the producer operation index for a tensor
    #[inline]
    fn producer_of(&self, tensor: TensorId) -> Option<usize> {
        self.output_to_op.get(&tensor).copied()
    }

    /// Get consumer operation indices for an operation
    #[inline]
    fn consumers_of_op(&self, op_idx: usize) -> &[usize] {
        self.op_consumers
            .get(&op_idx)
            .map(|v| v.as_slice())
            .unwrap_or(&[])
    }

    /// Get consumer operation indices for a tensor
    #[inline]
    fn consumers_of_tensor(&self, tensor: TensorId) -> &[usize] {
        self.tensor_consumers
            .get(&tensor)
            .map(|v| v.as_slice())
            .unwrap_or(&[])
    }

    /// Check if a tensor is produced by an operation in the stream
    #[inline]
    fn is_internal(&self, tensor: TensorId) -> bool {
        self.output_to_op.contains_key(&tensor)
    }
}

/// Types of fusion patterns that can be detected
#[derive(Debug, Clone, PartialEq)]
pub enum FusionPattern {
    /// Chain of elementwise operations (add, mul, exp, relu, etc.)
    ElementwiseChain,
    /// Softmax pattern: subtract_max -> exp -> sum -> divide
    /// Includes the reduction dimension for proper kernel generation
    SoftmaxPattern { dim: Option<usize> },
    /// Layer normalization pattern: (x - mean) / sqrt(var + eps) * scale + bias
    LayerNormPattern { eps: f32, dim: Option<usize> },
    /// RMS normalization pattern: x / sqrt(mean(x^2) + eps) * scale
    RmsNormPattern { eps: f32, dim: Option<usize> },
    /// Activation chain (e.g., GELU, SiLU with surrounding ops)
    ActivationChain,
    /// Reduction followed by broadcast and element-wise ops
    ReductionBroadcast {
        reduction_op: ReductionType,
        dim: Option<usize>,
    },
    /// Generic fuseable sequence
    Generic,
}

/// Types of reduction operations
#[derive(Debug, Clone, PartialEq)]
pub enum ReductionType {
    Sum,
    Mean,
    Max,
    Min,
}

impl FusionPattern {
    /// Get the expected speedup multiplier for this pattern
    pub fn expected_speedup(&self) -> f64 {
        match self {
            FusionPattern::ElementwiseChain => 2.0,
            FusionPattern::SoftmaxPattern { .. } => 3.0,
            FusionPattern::LayerNormPattern { .. } => 2.5,
            FusionPattern::RmsNormPattern { .. } => 2.5,
            FusionPattern::ActivationChain => 1.8,
            FusionPattern::ReductionBroadcast { .. } => 1.5,
            FusionPattern::Generic => 1.5,
        }
    }

    /// Get the name of this pattern
    pub fn name(&self) -> &'static str {
        match self {
            FusionPattern::ElementwiseChain => "elementwise_chain",
            FusionPattern::SoftmaxPattern { .. } => "softmax",
            FusionPattern::LayerNormPattern { .. } => "layer_norm",
            FusionPattern::RmsNormPattern { .. } => "rms_norm",
            FusionPattern::ActivationChain => "activation_chain",
            FusionPattern::ReductionBroadcast { .. } => "reduction_broadcast",
            FusionPattern::Generic => "generic",
        }
    }

    /// Check if this pattern involves reduction operations
    pub fn has_reduction(&self) -> bool {
        matches!(
            self,
            FusionPattern::SoftmaxPattern { .. }
                | FusionPattern::LayerNormPattern { .. }
                | FusionPattern::RmsNormPattern { .. }
                | FusionPattern::ReductionBroadcast { .. }
        )
    }

    /// Get the reduction dimension if applicable
    pub fn reduction_dim(&self) -> Option<usize> {
        match self {
            FusionPattern::SoftmaxPattern { dim } => *dim,
            FusionPattern::LayerNormPattern { dim, .. } => *dim,
            FusionPattern::RmsNormPattern { dim, .. } => *dim,
            FusionPattern::ReductionBroadcast { dim, .. } => *dim,
            _ => None,
        }
    }
}

/// A detected fusion opportunity
#[derive(Debug, Clone)]
pub struct FusionOpportunity {
    /// The type of fusion pattern detected
    pub pattern: FusionPattern,
    /// Indices of operations in the stream that form this opportunity
    pub operation_indices: Vec<usize>,
    /// The operations themselves
    pub operations: Vec<StreamOpKind>,
    /// External input tensor IDs (not produced within the fusion group)
    pub external_inputs: Vec<TensorId>,
    /// Output tensor ID
    pub output: TensorId,
    /// Estimated speedup from fusing these operations
    pub estimated_speedup: f64,
    /// Data type of the operations
    pub dtype: DType,
}

impl FusionOpportunity {
    /// Get the number of operations in this opportunity
    pub fn operation_count(&self) -> usize {
        self.operation_indices.len()
    }

    /// Check if this opportunity is worthwhile based on minimum ops threshold
    pub fn is_worthwhile(&self, min_ops: usize) -> bool {
        self.operation_indices.len() >= min_ops && self.estimated_speedup > 1.0
    }

    /// Generate a kernel signature for this opportunity
    pub fn to_kernel_signature(&self) -> KernelSignature {
        KernelSignature::new(
            self.operations.clone(),
            self.external_inputs.len(),
            self.dtype,
        )
    }

    /// Calculate the benefit score (ops × speedup)
    pub fn benefit_score(&self) -> f64 {
        self.operation_count() as f64 * self.estimated_speedup
    }
}

/// Analyzer for detecting fusion opportunities in operation streams
#[derive(Debug)]
pub struct FusionAnalyzer {
    config: FusionConfig,
}

impl FusionAnalyzer {
    /// Create a new fusion analyzer with the given configuration
    pub fn new(config: FusionConfig) -> Self {
        Self { config }
    }

    /// Create an analyzer with default configuration
    pub fn with_defaults() -> Self {
        Self::new(FusionConfig::default())
    }

    /// Analyze an operation stream and return fusion opportunities
    pub fn analyze(&self, stream: &OperationStream) -> Vec<FusionOpportunity> {
        if !self.config.enabled || stream.is_empty() {
            return Vec::new();
        }

        // Build analysis context once for all detection passes
        let ops = stream.operations();
        let ctx = AnalysisContext::new(&ops);

        let mut opportunities = Vec::with_capacity(ops.len() / 2);

        // Detect elementwise chains (primary fusion target)
        opportunities.extend(self.detect_elementwise_chains_opt(&ctx));

        // Detect advanced patterns if enabled
        if self.config.detect_advanced_patterns {
            opportunities.extend(self.detect_softmax_patterns_opt(&ctx));
            opportunities.extend(self.detect_layer_norm_patterns_opt(&ctx));
            opportunities.extend(self.detect_activation_chains_opt(&ctx));
        }

        // Sort by benefit score (highest first)
        opportunities.sort_by(|a, b| {
            b.benefit_score()
                .partial_cmp(&a.benefit_score())
                .unwrap_or(std::cmp::Ordering::Equal)
        });

        // Remove overlapping opportunities (greedy selection)
        self.remove_overlaps(opportunities)
    }

    /// Original analyze method that doesn't use the optimized context (for backward compatibility)
    #[allow(dead_code)]
    fn analyze_legacy(&self, stream: &OperationStream) -> Vec<FusionOpportunity> {
        if !self.config.enabled || stream.is_empty() {
            return Vec::new();
        }

        let mut opportunities = Vec::new();

        // Detect elementwise chains (primary fusion target)
        opportunities.extend(self.detect_elementwise_chains(stream));

        // Detect advanced patterns if enabled
        if self.config.detect_advanced_patterns {
            opportunities.extend(self.detect_softmax_patterns(stream));
            opportunities.extend(self.detect_layer_norm_patterns(stream));
            opportunities.extend(self.detect_activation_chains(stream));
        }

        // Sort by benefit score (highest first)
        opportunities.sort_by(|a, b| {
            b.benefit_score()
                .partial_cmp(&a.benefit_score())
                .unwrap_or(std::cmp::Ordering::Equal)
        });

        // Remove overlapping opportunities (greedy selection)
        self.remove_overlaps(opportunities)
    }

    /// Detect chains of elementwise operations
    fn detect_elementwise_chains(&self, stream: &OperationStream) -> Vec<FusionOpportunity> {
        let ops = stream.operations();
        if ops.is_empty() {
            return Vec::new();
        }

        let mut opportunities = Vec::new();
        let mut visited: HashSet<usize> = HashSet::new();

        // Build a map from output tensor to operation index
        let mut output_to_op: HashMap<TensorId, usize> = HashMap::new();
        for (idx, op) in ops.iter().enumerate() {
            output_to_op.insert(op.output, idx);
        }

        // Build consumer map
        let mut consumers: HashMap<usize, Vec<usize>> = HashMap::new();
        for (idx, op) in ops.iter().enumerate() {
            for &input in &op.inputs {
                if let Some(&producer_idx) = output_to_op.get(&input) {
                    consumers.entry(producer_idx).or_default().push(idx);
                }
            }
        }

        // Find chain start points (fuseable ops with no fuseable producer)
        for (idx, op) in ops.iter().enumerate() {
            if visited.contains(&idx) || !op.op.is_fuseable() {
                continue;
            }

            // Check if this is a chain start
            let is_start = op.inputs.iter().all(|input| {
                output_to_op
                    .get(input)
                    .map(|&producer_idx| !ops[producer_idx].op.is_fuseable())
                    .unwrap_or(true)
            });

            if !is_start {
                continue;
            }

            // Extend the chain greedily
            let mut chain = vec![idx];
            let mut chain_outputs: HashSet<TensorId> = HashSet::new();
            chain_outputs.insert(op.output);
            visited.insert(idx);

            let mut current_idx = idx;

            loop {
                // Find the next operation in the chain
                let next_candidates: Vec<_> = consumers
                    .get(&current_idx)
                    .map(|c| c.as_slice())
                    .unwrap_or(&[])
                    .iter()
                    .filter(|&&next_idx| {
                        !visited.contains(&next_idx)
                            && ops[next_idx].op.is_fuseable()
                            && ops[next_idx].inputs.iter().all(|input| {
                                // All inputs must be either external or from our chain
                                chain_outputs.contains(input) || !output_to_op.contains_key(input)
                            })
                    })
                    .copied()
                    .collect();

                if next_candidates.is_empty() {
                    break;
                }

                // Take the first valid candidate (could be smarter about this)
                let next_idx = next_candidates[0];
                let next_op = &ops[next_idx];

                chain.push(next_idx);
                chain_outputs.insert(next_op.output);
                visited.insert(next_idx);
                current_idx = next_idx;

                // Check max chain length
                if chain.len() >= self.config.max_fusion_ops {
                    break;
                }
            }

            // Create opportunity if chain is long enough
            if chain.len() >= self.config.min_fusion_ops {
                let chain_ops: Vec<_> = chain.iter().map(|&i| ops[i].op.clone()).collect();
                let external_inputs = self.find_external_inputs(&chain, ops, &output_to_op);
                let output = ops[*chain.last().unwrap()].output;
                let dtype = ops[chain[0]].dtype;

                // Calculate speedup with bonus for longer chains
                let base_speedup = FusionPattern::ElementwiseChain.expected_speedup();
                let chain_bonus = 1.0 + (chain.len() as f64 - 2.0) * 0.1;
                let estimated_speedup = base_speedup * chain_bonus;

                opportunities.push(FusionOpportunity {
                    pattern: FusionPattern::ElementwiseChain,
                    operation_indices: chain,
                    operations: chain_ops,
                    external_inputs,
                    output,
                    estimated_speedup,
                    dtype,
                });
            }
        }

        opportunities
    }

    /// Detect softmax patterns
    /// Pattern: (sub max) -> exp -> (sum dim) -> div
    /// Or simpler: exp -> sum -> div
    fn detect_softmax_patterns(&self, stream: &OperationStream) -> Vec<FusionOpportunity> {
        let ops = stream.operations();
        let mut opportunities = Vec::new();

        // Build maps for traversal
        let mut output_to_op: HashMap<TensorId, usize> = HashMap::new();
        let mut consumers: HashMap<TensorId, Vec<usize>> = HashMap::new();

        for (idx, op) in ops.iter().enumerate() {
            output_to_op.insert(op.output, idx);
            for &input in &op.inputs {
                consumers.entry(input).or_default().push(idx);
            }
        }

        // Look for Exp operations as anchor points
        for (exp_idx, exp_op) in ops.iter().enumerate() {
            if !matches!(exp_op.op, StreamOpKind::Exp) {
                continue;
            }

            let mut chain = vec![exp_idx];
            let mut dim = None;

            // Look backward for Sub(x, max) pattern (numerical stability)
            if let Some(&input_id) = exp_op.inputs.first() {
                if let Some(&sub_idx) = output_to_op.get(&input_id) {
                    if matches!(ops[sub_idx].op, StreamOpKind::Sub) {
                        // Check if second input could be a max reduction
                        if ops[sub_idx].inputs.len() >= 2 {
                            let second_input = ops[sub_idx].inputs[1];
                            if let Some(&max_idx) = output_to_op.get(&second_input) {
                                if matches!(ops[max_idx].op, StreamOpKind::Max) {
                                    chain.insert(0, max_idx);
                                    chain.insert(1, sub_idx);
                                }
                            }
                        }
                        // Even without max, sub before exp could be part of pattern
                        if chain.len() == 1 {
                            chain.insert(0, sub_idx);
                        }
                    }
                }
            }

            // Look forward for Sum operation
            let exp_output = exp_op.output;
            let sum_consumer = consumers.get(&exp_output).and_then(|cons| {
                cons.iter().find(|&&idx| {
                    matches!(ops[idx].op, StreamOpKind::Sum | StreamOpKind::SumDim { .. })
                })
            });

            if let Some(&sum_idx) = sum_consumer {
                chain.push(sum_idx);

                // Extract dimension from SumDim if present
                if let StreamOpKind::SumDim { dim: d } = &ops[sum_idx].op {
                    dim = Some(*d);
                }

                // Look for Div after Sum
                let sum_output = ops[sum_idx].output;
                let div_consumer = consumers.get(&sum_output).and_then(|cons| {
                    cons.iter()
                        .find(|&&idx| matches!(ops[idx].op, StreamOpKind::Div))
                });

                if let Some(&div_idx) = div_consumer {
                    chain.push(div_idx);
                }
            }

            // Only create opportunity if we have at least exp + sum (core softmax)
            if chain.len() >= 2
                && chain
                    .iter()
                    .any(|&i| matches!(ops[i].op, StreamOpKind::Sum | StreamOpKind::SumDim { .. }))
            {
                let chain_ops: Vec<_> = chain.iter().map(|&i| ops[i].op.clone()).collect();
                let external_inputs = self.find_external_inputs(&chain, ops, &output_to_op);
                let output = ops[*chain.last().unwrap()].output;
                let dtype = exp_op.dtype;

                opportunities.push(FusionOpportunity {
                    pattern: FusionPattern::SoftmaxPattern { dim },
                    operation_indices: chain,
                    operations: chain_ops,
                    external_inputs,
                    output,
                    estimated_speedup: FusionPattern::SoftmaxPattern { dim }.expected_speedup(),
                    dtype,
                });
            }
        }

        opportunities
    }

    /// Detect layer normalization patterns
    /// Pattern: mean -> sub -> square -> mean -> (add eps) -> sqrt -> div -> (mul scale) -> (add bias)
    fn detect_layer_norm_patterns(&self, stream: &OperationStream) -> Vec<FusionOpportunity> {
        let ops = stream.operations();
        let mut opportunities = Vec::new();

        // Build maps
        let mut output_to_op: HashMap<TensorId, usize> = HashMap::new();
        let mut consumers: HashMap<TensorId, Vec<usize>> = HashMap::new();

        for (idx, op) in ops.iter().enumerate() {
            output_to_op.insert(op.output, idx);
            for &input in &op.inputs {
                consumers.entry(input).or_default().push(idx);
            }
        }

        // Look for Mean operations as potential pattern starts
        for (mean_idx, mean_op) in ops.iter().enumerate() {
            if !matches!(
                mean_op.op,
                StreamOpKind::Mean | StreamOpKind::MeanDim { .. }
            ) {
                continue;
            }

            let mut chain = vec![mean_idx];
            let mut dim = None;
            let eps = 1e-5f32; // Default eps

            // Extract dimension
            if let StreamOpKind::MeanDim { dim: d } = &mean_op.op {
                dim = Some(*d);
            }

            // Look for Sub(x, mean) consumer
            let mean_output = mean_op.output;
            let sub_consumer = consumers.get(&mean_output).and_then(|cons| {
                cons.iter().find(|&&idx| {
                    matches!(ops[idx].op, StreamOpKind::Sub)
                        && ops[idx].inputs.get(1) == Some(&mean_output)
                })
            });

            if let Some(&sub_idx) = sub_consumer {
                chain.push(sub_idx);

                // Look for Square(sub_result)
                let sub_output = ops[sub_idx].output;
                let square_consumer = consumers.get(&sub_output).and_then(|cons| {
                    cons.iter()
                        .find(|&&idx| matches!(ops[idx].op, StreamOpKind::Square))
                });

                if let Some(&square_idx) = square_consumer {
                    chain.push(square_idx);

                    // Look for second Mean (variance computation)
                    let square_output = ops[square_idx].output;
                    let mean2_consumer = consumers.get(&square_output).and_then(|cons| {
                        cons.iter().find(|&&idx| {
                            matches!(
                                ops[idx].op,
                                StreamOpKind::Mean | StreamOpKind::MeanDim { .. }
                            )
                        })
                    });

                    if let Some(&mean2_idx) = mean2_consumer {
                        chain.push(mean2_idx);

                        // Continue looking for add_eps -> sqrt -> div pattern
                        let var_output = ops[mean2_idx].output;

                        // Look for AddScalar (eps) or Sqrt directly
                        if let Some(add_eps_idx) = consumers.get(&var_output).and_then(|cons| {
                            cons.iter()
                                .find(|&&idx| matches!(ops[idx].op, StreamOpKind::AddScalar(_)))
                        }) {
                            chain.push(*add_eps_idx);
                        }

                        // Continue extending if we find sqrt, div, etc.
                        // This is a simplified version - full detection would trace the complete path
                    }
                }
            }

            // Create opportunity if we found enough of the pattern
            if chain.len() >= 3 {
                let chain_ops: Vec<_> = chain.iter().map(|&i| ops[i].op.clone()).collect();
                let external_inputs = self.find_external_inputs(&chain, ops, &output_to_op);
                let output = ops[*chain.last().unwrap()].output;
                let dtype = mean_op.dtype;

                opportunities.push(FusionOpportunity {
                    pattern: FusionPattern::LayerNormPattern { eps, dim },
                    operation_indices: chain,
                    operations: chain_ops,
                    external_inputs,
                    output,
                    estimated_speedup: FusionPattern::LayerNormPattern { eps, dim }
                        .expected_speedup(),
                    dtype,
                });
            }
        }

        opportunities
    }

    /// Detect activation chains (e.g., mul -> sigmoid -> mul for SiLU)
    fn detect_activation_chains(&self, stream: &OperationStream) -> Vec<FusionOpportunity> {
        let ops = stream.operations();
        let mut opportunities = Vec::new();
        let mut visited: HashSet<usize> = HashSet::new();

        // Build output to op map
        let mut output_to_op: HashMap<TensorId, usize> = HashMap::new();
        for (idx, op) in ops.iter().enumerate() {
            output_to_op.insert(op.output, idx);
        }

        // Look for activation functions
        for (idx, op) in ops.iter().enumerate() {
            if visited.contains(&idx) {
                continue;
            }

            let is_activation = matches!(
                op.op,
                StreamOpKind::GELU
                    | StreamOpKind::SiLU
                    | StreamOpKind::Sigmoid
                    | StreamOpKind::Tanh
                    | StreamOpKind::ReLU
            );

            if !is_activation || !op.op.is_fuseable() {
                continue;
            }

            let mut chain = vec![idx];
            visited.insert(idx);

            // Look backward for preceding elementwise op
            if let Some(&input_id) = op.inputs.first() {
                if let Some(&prev_idx) = output_to_op.get(&input_id) {
                    if ops[prev_idx].op.is_fuseable() && !visited.contains(&prev_idx) {
                        chain.insert(0, prev_idx);
                        visited.insert(prev_idx);
                    }
                }
            }

            // Look forward for following elementwise op
            for (next_idx, next_op) in ops.iter().enumerate() {
                if next_op.inputs.contains(&op.output)
                    && next_op.op.is_fuseable()
                    && !visited.contains(&next_idx)
                {
                    chain.push(next_idx);
                    visited.insert(next_idx);
                    break;
                }
            }

            if chain.len() >= self.config.min_fusion_ops {
                let chain_ops: Vec<_> = chain.iter().map(|&i| ops[i].op.clone()).collect();
                let external_inputs = self.find_external_inputs(&chain, ops, &output_to_op);
                let output = ops[*chain.last().unwrap()].output;
                let dtype = op.dtype;

                opportunities.push(FusionOpportunity {
                    pattern: FusionPattern::ActivationChain,
                    operation_indices: chain,
                    operations: chain_ops,
                    external_inputs,
                    output,
                    estimated_speedup: FusionPattern::ActivationChain.expected_speedup(),
                    dtype,
                });
            }
        }

        opportunities
    }

    // ========================================================================
    // Optimized detection methods using AnalysisContext
    // ========================================================================

    /// Optimized elementwise chain detection using pre-computed context
    fn detect_elementwise_chains_opt(&self, ctx: &AnalysisContext<'_>) -> Vec<FusionOpportunity> {
        if ctx.ops.is_empty() {
            return Vec::new();
        }

        let mut opportunities = Vec::with_capacity(ctx.ops.len() / 4);
        let mut visited: HashSet<usize> = HashSet::with_capacity(ctx.ops.len());

        // Find chain start points (fuseable ops with no fuseable producer)
        for (idx, op) in ctx.ops.iter().enumerate() {
            if visited.contains(&idx) || !op.op.is_fuseable() {
                continue;
            }

            // Check if this is a chain start
            let is_start = op.inputs.iter().all(|input| {
                ctx.producer_of(*input)
                    .map(|producer_idx| !ctx.get_op(producer_idx).op.is_fuseable())
                    .unwrap_or(true)
            });

            if !is_start {
                continue;
            }

            // Extend the chain greedily
            let mut chain = Vec::with_capacity(self.config.max_fusion_ops);
            chain.push(idx);
            let mut chain_outputs: HashSet<TensorId> =
                HashSet::with_capacity(self.config.max_fusion_ops);
            chain_outputs.insert(op.output);
            visited.insert(idx);

            let mut current_idx = idx;

            loop {
                // Find the next operation in the chain
                let consumers = ctx.consumers_of_op(current_idx);
                let next_candidate = consumers
                    .iter()
                    .filter(|&&next_idx| {
                        !visited.contains(&next_idx)
                            && ctx.get_op(next_idx).op.is_fuseable()
                            && ctx.get_op(next_idx).inputs.iter().all(|input| {
                                chain_outputs.contains(input) || !ctx.is_internal(*input)
                            })
                    })
                    .next();

                match next_candidate {
                    Some(&next_idx) => {
                        let next_op = ctx.get_op(next_idx);
                        chain.push(next_idx);
                        chain_outputs.insert(next_op.output);
                        visited.insert(next_idx);
                        current_idx = next_idx;

                        if chain.len() >= self.config.max_fusion_ops {
                            break;
                        }
                    }
                    None => break,
                }
            }

            // Create opportunity if chain is long enough
            if chain.len() >= self.config.min_fusion_ops {
                let chain_ops: Vec<_> = chain.iter().map(|&i| ctx.get_op(i).op.clone()).collect();
                let external_inputs = self.find_external_inputs_opt(&chain, ctx);
                let output = ctx.get_op(*chain.last().unwrap()).output;
                let dtype = ctx.get_op(chain[0]).dtype;

                let base_speedup = FusionPattern::ElementwiseChain.expected_speedup();
                let chain_bonus = 1.0 + (chain.len() as f64 - 2.0) * 0.1;

                opportunities.push(FusionOpportunity {
                    pattern: FusionPattern::ElementwiseChain,
                    operation_indices: chain,
                    operations: chain_ops,
                    external_inputs,
                    output,
                    estimated_speedup: base_speedup * chain_bonus,
                    dtype,
                });
            }
        }

        opportunities
    }

    /// Optimized softmax pattern detection
    fn detect_softmax_patterns_opt(&self, ctx: &AnalysisContext<'_>) -> Vec<FusionOpportunity> {
        let mut opportunities = Vec::new();

        for (exp_idx, exp_op) in ctx.ops.iter().enumerate() {
            if !matches!(exp_op.op, StreamOpKind::Exp) {
                continue;
            }

            let mut chain = vec![exp_idx];
            let mut dim = None;

            // Look backward for Sub(x, max) pattern
            if let Some(&input_id) = exp_op.inputs.first() {
                if let Some(sub_idx) = ctx.producer_of(input_id) {
                    if matches!(ctx.get_op(sub_idx).op, StreamOpKind::Sub) {
                        if ctx.get_op(sub_idx).inputs.len() >= 2 {
                            let second_input = ctx.get_op(sub_idx).inputs[1];
                            if let Some(max_idx) = ctx.producer_of(second_input) {
                                if matches!(ctx.get_op(max_idx).op, StreamOpKind::Max) {
                                    chain.insert(0, max_idx);
                                    chain.insert(1, sub_idx);
                                }
                            }
                        }
                        if chain.len() == 1 {
                            chain.insert(0, sub_idx);
                        }
                    }
                }
            }

            // Look forward for Sum operation
            let exp_output = exp_op.output;
            let sum_consumer = ctx.consumers_of_tensor(exp_output).iter().find(|&&idx| {
                matches!(
                    ctx.get_op(idx).op,
                    StreamOpKind::Sum | StreamOpKind::SumDim { .. }
                )
            });

            if let Some(&sum_idx) = sum_consumer {
                chain.push(sum_idx);

                if let StreamOpKind::SumDim { dim: d } = &ctx.get_op(sum_idx).op {
                    dim = Some(*d);
                }

                let sum_output = ctx.get_op(sum_idx).output;
                let div_consumer = ctx
                    .consumers_of_tensor(sum_output)
                    .iter()
                    .find(|&&idx| matches!(ctx.get_op(idx).op, StreamOpKind::Div));

                if let Some(&div_idx) = div_consumer {
                    chain.push(div_idx);
                }
            }

            if chain.len() >= 2
                && chain.iter().any(|&i| {
                    matches!(
                        ctx.get_op(i).op,
                        StreamOpKind::Sum | StreamOpKind::SumDim { .. }
                    )
                })
            {
                let chain_ops: Vec<_> = chain.iter().map(|&i| ctx.get_op(i).op.clone()).collect();
                let external_inputs = self.find_external_inputs_opt(&chain, ctx);
                let output = ctx.get_op(*chain.last().unwrap()).output;
                let dtype = exp_op.dtype;

                opportunities.push(FusionOpportunity {
                    pattern: FusionPattern::SoftmaxPattern { dim },
                    operation_indices: chain,
                    operations: chain_ops,
                    external_inputs,
                    output,
                    estimated_speedup: FusionPattern::SoftmaxPattern { dim }.expected_speedup(),
                    dtype,
                });
            }
        }

        opportunities
    }

    /// Optimized layer norm pattern detection
    fn detect_layer_norm_patterns_opt(&self, ctx: &AnalysisContext<'_>) -> Vec<FusionOpportunity> {
        let mut opportunities = Vec::new();
        let eps = 1e-5f32;

        for (mean_idx, mean_op) in ctx.ops.iter().enumerate() {
            if !matches!(
                mean_op.op,
                StreamOpKind::Mean | StreamOpKind::MeanDim { .. }
            ) {
                continue;
            }

            let mut chain = vec![mean_idx];
            let mut dim = None;

            if let StreamOpKind::MeanDim { dim: d } = &mean_op.op {
                dim = Some(*d);
            }

            let mean_output = mean_op.output;
            let sub_consumer = ctx.consumers_of_tensor(mean_output).iter().find(|&&idx| {
                matches!(ctx.get_op(idx).op, StreamOpKind::Sub)
                    && ctx.get_op(idx).inputs.get(1) == Some(&mean_output)
            });

            if let Some(&sub_idx) = sub_consumer {
                chain.push(sub_idx);

                let sub_output = ctx.get_op(sub_idx).output;
                let square_consumer = ctx
                    .consumers_of_tensor(sub_output)
                    .iter()
                    .find(|&&idx| matches!(ctx.get_op(idx).op, StreamOpKind::Square));

                if let Some(&square_idx) = square_consumer {
                    chain.push(square_idx);

                    let square_output = ctx.get_op(square_idx).output;
                    let mean2_consumer =
                        ctx.consumers_of_tensor(square_output).iter().find(|&&idx| {
                            matches!(
                                ctx.get_op(idx).op,
                                StreamOpKind::Mean | StreamOpKind::MeanDim { .. }
                            )
                        });

                    if let Some(&mean2_idx) = mean2_consumer {
                        chain.push(mean2_idx);

                        let var_output = ctx.get_op(mean2_idx).output;
                        if let Some(&add_eps_idx) = ctx
                            .consumers_of_tensor(var_output)
                            .iter()
                            .find(|&&idx| matches!(ctx.get_op(idx).op, StreamOpKind::AddScalar(_)))
                        {
                            chain.push(add_eps_idx);
                        }
                    }
                }
            }

            if chain.len() >= 3 {
                let chain_ops: Vec<_> = chain.iter().map(|&i| ctx.get_op(i).op.clone()).collect();
                let external_inputs = self.find_external_inputs_opt(&chain, ctx);
                let output = ctx.get_op(*chain.last().unwrap()).output;
                let dtype = mean_op.dtype;

                opportunities.push(FusionOpportunity {
                    pattern: FusionPattern::LayerNormPattern { eps, dim },
                    operation_indices: chain,
                    operations: chain_ops,
                    external_inputs,
                    output,
                    estimated_speedup: FusionPattern::LayerNormPattern { eps, dim }
                        .expected_speedup(),
                    dtype,
                });
            }
        }

        opportunities
    }

    /// Optimized activation chain detection
    fn detect_activation_chains_opt(&self, ctx: &AnalysisContext<'_>) -> Vec<FusionOpportunity> {
        let mut opportunities = Vec::new();
        let mut visited: HashSet<usize> = HashSet::with_capacity(ctx.ops.len());

        for (idx, op) in ctx.ops.iter().enumerate() {
            if visited.contains(&idx) {
                continue;
            }

            let is_activation = matches!(
                op.op,
                StreamOpKind::GELU
                    | StreamOpKind::SiLU
                    | StreamOpKind::Sigmoid
                    | StreamOpKind::Tanh
                    | StreamOpKind::ReLU
            );

            if !is_activation || !op.op.is_fuseable() {
                continue;
            }

            let mut chain = vec![idx];
            visited.insert(idx);

            // Look backward
            if let Some(&input_id) = op.inputs.first() {
                if let Some(prev_idx) = ctx.producer_of(input_id) {
                    if ctx.get_op(prev_idx).op.is_fuseable() && !visited.contains(&prev_idx) {
                        chain.insert(0, prev_idx);
                        visited.insert(prev_idx);
                    }
                }
            }

            // Look forward
            for (next_idx, next_op) in ctx.ops.iter().enumerate() {
                if next_op.inputs.contains(&op.output)
                    && next_op.op.is_fuseable()
                    && !visited.contains(&next_idx)
                {
                    chain.push(next_idx);
                    visited.insert(next_idx);
                    break;
                }
            }

            if chain.len() >= self.config.min_fusion_ops {
                let chain_ops: Vec<_> = chain.iter().map(|&i| ctx.get_op(i).op.clone()).collect();
                let external_inputs = self.find_external_inputs_opt(&chain, ctx);
                let output = ctx.get_op(*chain.last().unwrap()).output;
                let dtype = op.dtype;

                opportunities.push(FusionOpportunity {
                    pattern: FusionPattern::ActivationChain,
                    operation_indices: chain,
                    operations: chain_ops,
                    external_inputs,
                    output,
                    estimated_speedup: FusionPattern::ActivationChain.expected_speedup(),
                    dtype,
                });
            }
        }

        opportunities
    }

    /// Optimized external input detection using context
    fn find_external_inputs_opt(
        &self,
        chain: &[usize],
        ctx: &AnalysisContext<'_>,
    ) -> Vec<TensorId> {
        let chain_set: HashSet<_> = chain.iter().copied().collect();
        let mut external = Vec::with_capacity(chain.len());

        for &idx in chain {
            for &input in &ctx.get_op(idx).inputs {
                let is_external = ctx
                    .producer_of(input)
                    .map(|producer_idx| !chain_set.contains(&producer_idx))
                    .unwrap_or(true);

                if is_external && !external.contains(&input) {
                    external.push(input);
                }
            }
        }

        external
    }

    /// Find external inputs for a chain of operations
    fn find_external_inputs(
        &self,
        chain: &[usize],
        ops: &[StreamOperation],
        output_to_op: &HashMap<TensorId, usize>,
    ) -> Vec<TensorId> {
        let chain_set: HashSet<_> = chain.iter().copied().collect();
        let mut external = Vec::new();

        for &idx in chain {
            for &input in &ops[idx].inputs {
                // External if not produced by an op in the chain
                let is_external = output_to_op
                    .get(&input)
                    .map(|&producer_idx| !chain_set.contains(&producer_idx))
                    .unwrap_or(true);

                if is_external && !external.contains(&input) {
                    external.push(input);
                }
            }
        }

        external
    }

    /// Remove overlapping opportunities (greedy: keep highest benefit)
    fn remove_overlaps(&self, mut opportunities: Vec<FusionOpportunity>) -> Vec<FusionOpportunity> {
        let mut result = Vec::new();
        let mut used_ops: HashSet<usize> = HashSet::new();

        for opp in opportunities {
            // Check if any operation in this opportunity is already used
            let has_overlap = opp
                .operation_indices
                .iter()
                .any(|idx| used_ops.contains(idx));

            if !has_overlap {
                // Mark operations as used
                for &idx in &opp.operation_indices {
                    used_ops.insert(idx);
                }
                result.push(opp);
            }
        }

        result
    }

    /// Analyze a single operation for fusion potential
    pub fn analyze_operation(&self, op: &StreamOpKind) -> OperationAnalysis {
        OperationAnalysis {
            is_fuseable: op.is_fuseable(),
            is_sync_point: op.is_sync_point(),
            fusion_benefit: op.fusion_benefit(),
            preferred_patterns: if op.is_fuseable() {
                vec![
                    FusionPattern::ElementwiseChain,
                    FusionPattern::ActivationChain,
                ]
            } else {
                vec![]
            },
        }
    }
}

/// Analysis result for a single operation
#[derive(Debug)]
pub struct OperationAnalysis {
    /// Whether the operation can be fused
    pub is_fuseable: bool,
    /// Whether this is a sync point (triggers flush)
    pub is_sync_point: bool,
    /// Benefit score for fusion (0.0-1.0)
    pub fusion_benefit: f64,
    /// Preferred fusion patterns for this operation
    pub preferred_patterns: Vec<FusionPattern>,
}

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

    fn make_op(op: StreamOpKind, inputs: Vec<TensorId>, output: TensorId) -> StreamOperation {
        StreamOperation::new(op, inputs, output, vec![1024], DType::F32)
    }

    #[test]
    fn test_pattern_speedup() {
        assert!(FusionPattern::ElementwiseChain.expected_speedup() > 1.0);
        assert!(
            FusionPattern::SoftmaxPattern { dim: None }.expected_speedup()
                > FusionPattern::ElementwiseChain.expected_speedup()
        );
        assert!(
            FusionPattern::LayerNormPattern {
                eps: 1e-5,
                dim: None
            }
            .expected_speedup()
                > FusionPattern::ElementwiseChain.expected_speedup()
        );
    }

    #[test]
    fn test_analyzer_creation() {
        let analyzer = FusionAnalyzer::with_defaults();
        assert!(analyzer.config.enabled);
    }

    #[test]
    fn test_empty_stream() {
        let analyzer = FusionAnalyzer::with_defaults();
        let stream = OperationStream::with_defaults();

        let opportunities = analyzer.analyze(&stream);
        assert!(opportunities.is_empty());
    }

    #[test]
    fn test_elementwise_chain_detection() {
        let analyzer = FusionAnalyzer::new(FusionConfig::default().with_min_fusion_ops(2));
        let mut stream = OperationStream::with_defaults();

        let a = TensorId::new();
        let b = TensorId::new();
        let c = TensorId::new();
        let d = TensorId::new();
        let e = TensorId::new();

        stream.register_input(a);
        stream.register_input(b);

        // Chain: add(a, b) -> mul(c, a) -> relu(d)
        stream.record(make_op(StreamOpKind::Add, vec![a, b], c));
        stream.record(make_op(StreamOpKind::Mul, vec![c, a], d));
        stream.record(make_op(StreamOpKind::ReLU, vec![d], e));

        let opportunities = analyzer.analyze(&stream);
        assert!(!opportunities.is_empty());

        let first = &opportunities[0];
        assert_eq!(first.pattern, FusionPattern::ElementwiseChain);
        assert!(first.operation_count() >= 2);
    }

    #[test]
    fn test_sync_point_breaks_chain() {
        let analyzer = FusionAnalyzer::new(FusionConfig::default().with_min_fusion_ops(2));
        let mut stream = OperationStream::with_defaults();

        let a = TensorId::new();
        let b = TensorId::new();
        let c = TensorId::new();
        let d = TensorId::new();

        stream.register_input(a);

        // add -> matmul (sync point) -> mul
        stream.record(make_op(StreamOpKind::Add, vec![a, a], b));
        stream.record(make_op(StreamOpKind::MatMul, vec![b], c));
        stream.record(make_op(StreamOpKind::Mul, vec![c, c], d));

        let opportunities = analyzer.analyze(&stream);

        // MatMul should break the chain
        for opp in &opportunities {
            assert!(
                !opp.operations
                    .iter()
                    .any(|op| matches!(op, StreamOpKind::MatMul))
            );
        }
    }

    #[test]
    fn test_opportunity_worthwhile() {
        let opp = FusionOpportunity {
            pattern: FusionPattern::ElementwiseChain,
            operation_indices: vec![0, 1],
            operations: vec![StreamOpKind::Add, StreamOpKind::Mul],
            external_inputs: vec![TensorId::new()],
            output: TensorId::new(),
            estimated_speedup: 2.0,
            dtype: DType::F32,
        };

        assert!(opp.is_worthwhile(2));
        assert!(!opp.is_worthwhile(3));
    }

    #[test]
    fn test_kernel_signature_generation() {
        let opp = FusionOpportunity {
            pattern: FusionPattern::ElementwiseChain,
            operation_indices: vec![0, 1, 2],
            operations: vec![StreamOpKind::Add, StreamOpKind::Mul, StreamOpKind::ReLU],
            external_inputs: vec![TensorId::new(), TensorId::new()],
            output: TensorId::new(),
            estimated_speedup: 2.5,
            dtype: DType::F32,
        };

        let sig = opp.to_kernel_signature();
        assert_eq!(sig.ops.len(), 3);
        assert_eq!(sig.num_inputs, 2);
        assert_eq!(sig.dtype, DType::F32);
    }

    #[test]
    fn test_operation_analysis() {
        let analyzer = FusionAnalyzer::with_defaults();

        let add_analysis = analyzer.analyze_operation(&StreamOpKind::Add);
        assert!(add_analysis.is_fuseable);
        assert!(!add_analysis.is_sync_point);

        let matmul_analysis = analyzer.analyze_operation(&StreamOpKind::MatMul);
        assert!(!matmul_analysis.is_fuseable);
        assert!(matmul_analysis.is_sync_point);
    }
}