//! Operation definitions for heterogeneous compute.

use super::Precision;
use serde::{Deserialize, Serialize};

/// Operation types for processor matching.
#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, Serialize, Deserialize)]
pub enum OperationType {
    // Matrix operations
    MatMul,
    Conv2d,
    Conv3d,
    DepthwiseConv,
    BatchNorm,
    LayerNorm,

    // Attention operations
    SelfAttention,
    CrossAttention,
    FlashAttention,

    // Element-wise operations
    Add,
    Mul,
    ReLU,
    GeLU,
    SiLU,
    Softmax,

    // Reduction operations
    Sum,
    Mean,
    Max,
    ArgMax,

    // Data movement
    Transpose,
    Reshape,
    Concat,
    Split,
    Gather,
    Scatter,

    // LLM specific
    Embedding,
    RoPE, // Rotary Position Embedding
    KVCache,
    TopK,
    Sampling,

    // I/O operations
    DataLoad,
    DataPreprocess,
    Tokenization,
    Detokenization,
    Checkpoint,

    // Distributed operations
    AllReduce,
    AllGather,
    ReduceScatter,

    // Training specific
    Backward,
    OptimizerStep,
    GradientClip,
}

/// Concrete operation with parameters.
#[derive(Clone, Debug, Serialize, Deserialize)]
pub enum Operation {
    /// Matrix multiplication.
    MatMul {
        m: usize,
        n: usize,
        k: usize,
        precision: Precision,
    },

    /// 2D Convolution.
    Conv2d {
        batch: usize,
        in_channels: usize,
        out_channels: usize,
        height: usize,
        width: usize,
        kernel_size: usize,
        precision: Precision,
    },

    /// Batch normalization.
    BatchNorm {
        batch: usize,
        channels: usize,
        spatial: usize,
        precision: Precision,
    },

    /// Layer normalization.
    LayerNorm {
        batch: usize,
        seq_len: usize,
        hidden: usize,
        precision: Precision,
    },

    /// Self-attention.
    SelfAttention {
        batch: usize,
        seq_len: usize,
        num_heads: usize,
        head_dim: usize,
        precision: Precision,
    },

    /// Flash attention (fused, memory efficient).
    FlashAttention {
        batch: usize,
        seq_len: usize,
        num_heads: usize,
        head_dim: usize,
        precision: Precision,
    },

    /// Element-wise addition.
    Add {
        elements: usize,
        precision: Precision,
    },

    /// Element-wise multiplication.
    Mul {
        elements: usize,
        precision: Precision,
    },

    /// ReLU activation.
    ReLU { elements: usize },

    /// GeLU activation.
    GeLU { elements: usize },

    /// SiLU (Swish) activation.
    SiLU { elements: usize },

    /// Softmax.
    Softmax {
        batch: usize,
        seq_len: usize,
        precision: Precision,
    },

    /// Embedding lookup.
    Embedding {
        batch: usize,
        seq_len: usize,
        vocab_size: usize,
        embed_dim: usize,
        precision: Precision,
    },

    /// Rotary Position Embedding.
    RoPE {
        batch: usize,
        seq_len: usize,
        head_dim: usize,
        precision: Precision,
    },

    /// KV Cache update.
    KVCache {
        batch: usize,
        seq_len: usize,
        num_heads: usize,
        head_dim: usize,
        precision: Precision,
    },

    /// Top-K sampling.
    TopK {
        batch: usize,
        vocab_size: usize,
        k: usize,
    },

    /// Token sampling.
    Sampling {
        batch: usize,
        vocab_size: usize,
        temperature: f32,
    },

    /// Data loading from storage.
    DataLoad {
        bytes: usize,
        async_: bool,
    },

    /// Data preprocessing.
    DataPreprocess {
        batch: usize,
        transforms: Vec<String>,
    },

    /// Tokenization.
    Tokenization {
        text_bytes: usize,
        vocab_size: usize,
    },

    /// Detokenization.
    Detokenization {
        tokens: usize,
        vocab_size: usize,
    },

    /// Checkpoint save.
    Checkpoint {
        bytes: usize,
        async_: bool,
    },

    /// All-reduce across devices.
    AllReduce {
        elements: usize,
        precision: Precision,
        devices: usize,
    },

    /// Backward pass for a layer.
    Backward {
        forward_op: Box<Operation>,
    },

    /// Optimizer step.
    OptimizerStep {
        parameters: usize,
        optimizer: String,
        precision: Precision,
    },

    /// Transpose.
    Transpose {
        shape: Vec<usize>,
        axes: Vec<usize>,
    },

    /// Reshape.
    Reshape {
        from: Vec<usize>,
        to: Vec<usize>,
    },

    /// Concatenate tensors.
    Concat {
        shapes: Vec<Vec<usize>>,
        axis: usize,
    },

    /// Generic operation.
    Generic {
        op_type: OperationType,
        flops: f64,
        memory: u64,
    },
}

impl Operation {
    /// Returns the operation type.
    pub fn op_type(&self) -> OperationType {
        match self {
            Operation::MatMul { .. } => OperationType::MatMul,
            Operation::Conv2d { .. } => OperationType::Conv2d,
            Operation::BatchNorm { .. } => OperationType::BatchNorm,
            Operation::LayerNorm { .. } => OperationType::LayerNorm,
            Operation::SelfAttention { .. } => OperationType::SelfAttention,
            Operation::FlashAttention { .. } => OperationType::FlashAttention,
            Operation::Add { .. } => OperationType::Add,
            Operation::Mul { .. } => OperationType::Mul,
            Operation::ReLU { .. } => OperationType::ReLU,
            Operation::GeLU { .. } => OperationType::GeLU,
            Operation::SiLU { .. } => OperationType::SiLU,
            Operation::Softmax { .. } => OperationType::Softmax,
            Operation::Embedding { .. } => OperationType::Embedding,
            Operation::RoPE { .. } => OperationType::RoPE,
            Operation::KVCache { .. } => OperationType::KVCache,
            Operation::TopK { .. } => OperationType::TopK,
            Operation::Sampling { .. } => OperationType::Sampling,
            Operation::DataLoad { .. } => OperationType::DataLoad,
            Operation::DataPreprocess { .. } => OperationType::DataPreprocess,
            Operation::Tokenization { .. } => OperationType::Tokenization,
            Operation::Detokenization { .. } => OperationType::Detokenization,
            Operation::Checkpoint { .. } => OperationType::Checkpoint,
            Operation::AllReduce { .. } => OperationType::AllReduce,
            Operation::Backward { .. } => OperationType::Backward,
            Operation::OptimizerStep { .. } => OperationType::OptimizerStep,
            Operation::Transpose { .. } => OperationType::Transpose,
            Operation::Reshape { .. } => OperationType::Reshape,
            Operation::Concat { .. } => OperationType::Concat,
            Operation::Generic { op_type, .. } => *op_type,
        }
    }

    /// Returns the precision used.
    pub fn precision(&self) -> Precision {
        match self {
            Operation::MatMul { precision, .. }
            | Operation::Conv2d { precision, .. }
            | Operation::BatchNorm { precision, .. }
            | Operation::LayerNorm { precision, .. }
            | Operation::SelfAttention { precision, .. }
            | Operation::FlashAttention { precision, .. }
            | Operation::Add { precision, .. }
            | Operation::Mul { precision, .. }
            | Operation::Softmax { precision, .. }
            | Operation::Embedding { precision, .. }
            | Operation::RoPE { precision, .. }
            | Operation::KVCache { precision, .. }
            | Operation::AllReduce { precision, .. }
            | Operation::OptimizerStep { precision, .. } => *precision,
            Operation::Backward { forward_op } => forward_op.precision(),
            _ => Precision::Fp32, // Default
        }
    }

    /// Estimates FLOPS for the operation.
    pub fn estimated_flops(&self) -> f64 {
        match self {
            // MatMul: 2 * M * N * K (multiply-add)
            Operation::MatMul { m, n, k, .. } => 2.0 * (*m as f64) * (*n as f64) * (*k as f64),

            // Conv2d: 2 * batch * out * H * W * in * K * K
            Operation::Conv2d {
                batch,
                in_channels,
                out_channels,
                height,
                width,
                kernel_size,
                ..
            } => {
                2.0 * (*batch as f64)
                    * (*out_channels as f64)
                    * (*height as f64)
                    * (*width as f64)
                    * (*in_channels as f64)
                    * (*kernel_size as f64)
                    * (*kernel_size as f64)
            }

            // Self-attention: 4 * batch * seq * seq * head_dim * heads
            Operation::SelfAttention {
                batch,
                seq_len,
                num_heads,
                head_dim,
                ..
            }
            | Operation::FlashAttention {
                batch,
                seq_len,
                num_heads,
                head_dim,
                ..
            } => {
                4.0 * (*batch as f64)
                    * (*seq_len as f64)
                    * (*seq_len as f64)
                    * (*head_dim as f64)
                    * (*num_heads as f64)
            }

            // Element-wise: 1 FLOP per element
            Operation::Add { elements, .. }
            | Operation::Mul { elements, .. }
            | Operation::ReLU { elements }
            | Operation::GeLU { elements }
            | Operation::SiLU { elements } => *elements as f64,

            // Softmax: ~5 ops per element (exp, sum, div)
            Operation::Softmax {
                batch, seq_len, ..
            } => 5.0 * (*batch as f64) * (*seq_len as f64),

            // Embedding: just lookup, minimal FLOPS
            Operation::Embedding {
                batch,
                seq_len,
                embed_dim,
                ..
            } => (*batch as f64) * (*seq_len as f64) * (*embed_dim as f64) * 0.1,

            // Backward: ~2x forward
            Operation::Backward { forward_op } => forward_op.estimated_flops() * 2.0,

            // Generic
            Operation::Generic { flops, .. } => *flops,

            // I/O operations: minimal compute
            _ => 1000.0,
        }
    }

    /// Estimates memory usage (bytes).
    pub fn estimated_memory(&self) -> u64 {
        let precision_bytes = match self.precision() {
            Precision::Fp64 => 8,
            Precision::Fp32 => 4,
            Precision::Fp16 | Precision::Bf16 => 2,
            Precision::Int8 => 1,
            Precision::Int4 => 1, // Rounded up
        };

        match self {
            Operation::MatMul { m, n, k, .. } => {
                // Input A (m×k) + Input B (k×n) + Output (m×n)
                ((*m * *k) + (*k * *n) + (*m * *n)) as u64 * precision_bytes
            }

            Operation::SelfAttention {
                batch,
                seq_len,
                num_heads,
                head_dim,
                ..
            } => {
                // Q, K, V, Output, intermediate attention
                5 * (*batch as u64)
                    * (*seq_len as u64)
                    * (*num_heads as u64)
                    * (*head_dim as u64)
                    * precision_bytes
            }

            Operation::FlashAttention {
                batch,
                seq_len,
                num_heads,
                head_dim,
                ..
            } => {
                // FlashAttention uses much less memory
                2 * (*batch as u64)
                    * (*seq_len as u64)
                    * (*num_heads as u64)
                    * (*head_dim as u64)
                    * precision_bytes
            }

            Operation::KVCache {
                batch,
                seq_len,
                num_heads,
                head_dim,
                ..
            } => {
                // K and V caches
                2 * (*batch as u64)
                    * (*seq_len as u64)
                    * (*num_heads as u64)
                    * (*head_dim as u64)
                    * precision_bytes
            }

            Operation::Generic { memory, .. } => *memory,

            _ => 1024 * 1024, // 1 MB default
        }
    }

    /// Creates the backward operation for this operation.
    pub fn backward(&self) -> Option<Operation> {
        match self {
            Operation::MatMul { .. }
            | Operation::Conv2d { .. }
            | Operation::SelfAttention { .. }
            | Operation::FlashAttention { .. }
            | Operation::LayerNorm { .. }
            | Operation::BatchNorm { .. } => Some(Operation::Backward {
                forward_op: Box::new(self.clone()),
            }),
            _ => None,
        }
    }
}

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

    #[test]
    fn test_matmul_flops() {
        let op = Operation::MatMul {
            m: 1024,
            n: 1024,
            k: 1024,
            precision: Precision::Fp32,
        };

        let flops = op.estimated_flops();
        // 2 * 1024^3 = ~2.1 billion FLOPS
        assert!(flops > 2e9 && flops < 2.2e9);
    }

    #[test]
    fn test_attention_memory() {
        let regular = Operation::SelfAttention {
            batch: 1,
            seq_len: 4096,
            num_heads: 32,
            head_dim: 128,
            precision: Precision::Fp16,
        };

        let flash = Operation::FlashAttention {
            batch: 1,
            seq_len: 4096,
            num_heads: 32,
            head_dim: 128,
            precision: Precision::Fp16,
        };

        // FlashAttention should use less memory
        assert!(flash.estimated_memory() < regular.estimated_memory());
    }

    #[test]
    fn test_backward_creation() {
        let forward = Operation::MatMul {
            m: 1024,
            n: 1024,
            k: 1024,
            precision: Precision::Fp32,
        };

        let backward = forward.backward();
        assert!(backward.is_some());

        if let Some(Operation::Backward { forward_op }) = backward {
            assert!(matches!(*forward_op, Operation::MatMul { .. }));
        }
    }
}
