//! Operation traits for backends.
//!
//! These traits define operations that can be specialized per backend.
//! The main Backend trait provides basic operations, while these traits
//! provide additional specialized operations.

use crate::Backend;

/// Tensor operations that can be specialized per backend.
///
/// Most operations are defined on the main Backend trait. This trait
/// provides additional operations that may need more customization.
pub trait TensorOps<B: Backend, const D: usize>: Clone + Send + Sync {
    /// Check if the tensor is contiguous in memory.
    fn is_contiguous(&self) -> bool;

    /// Make the tensor contiguous (copy if necessary).
    fn contiguous(self) -> Self;

    /// Get the number of elements.
    fn numel(&self) -> usize;

    /// Get the strides.
    fn strides(&self) -> [usize; D];
}

/// Module operations (convolution, pooling, etc.).
pub trait ModuleOps<B: Backend>: Sized {
    /// 2D convolution.
    fn conv2d(
        input: &B::TensorPrimitive<4>,  // [batch, in_channels, height, width]
        weight: &B::TensorPrimitive<4>, // [out_channels, in_channels, kH, kW]
        bias: Option<&B::TensorPrimitive<1>>,
        stride: [usize; 2],
        padding: [usize; 2],
        dilation: [usize; 2],
        groups: usize,
    ) -> B::TensorPrimitive<4>;

    /// 2D max pooling.
    fn max_pool2d(
        input: &B::TensorPrimitive<4>,
        kernel_size: [usize; 2],
        stride: [usize; 2],
        padding: [usize; 2],
    ) -> B::TensorPrimitive<4>;

    /// 2D average pooling.
    fn avg_pool2d(
        input: &B::TensorPrimitive<4>,
        kernel_size: [usize; 2],
        stride: [usize; 2],
        padding: [usize; 2],
    ) -> B::TensorPrimitive<4>;

    /// Batch normalization.
    #[allow(clippy::too_many_arguments)] // Inherent to batch_norm API
    fn batch_norm(
        input: &B::TensorPrimitive<4>,
        running_mean: &B::TensorPrimitive<1>,
        running_var: &B::TensorPrimitive<1>,
        weight: Option<&B::TensorPrimitive<1>>,
        bias: Option<&B::TensorPrimitive<1>>,
        training: bool,
        momentum: f64,
        eps: f64,
    ) -> B::TensorPrimitive<4>;

    /// Dropout (returns unchanged tensor if not training).
    fn dropout(input: &B::TensorPrimitive<4>, prob: f64, training: bool) -> B::TensorPrimitive<4>;

    /// Linear projection (matrix multiply + optional bias).
    fn linear(
        input: &B::TensorPrimitive<2>,
        weight: &B::TensorPrimitive<2>,
        bias: Option<&B::TensorPrimitive<1>>,
    ) -> B::TensorPrimitive<2>;

    /// Embedding lookup.
    fn embedding(
        indices: &B::TensorPrimitive<2>, // [batch, seq_len] as int
        weight: &B::TensorPrimitive<2>,  // [vocab_size, embed_dim]
    ) -> B::TensorPrimitive<3>; // [batch, seq_len, embed_dim]
}

/// Activation function operations.
pub trait ActivationOps<B: Backend>: Sized {
    /// ReLU activation.
    fn relu<const D: usize>(tensor: &B::TensorPrimitive<D>) -> B::TensorPrimitive<D>;

    /// Leaky ReLU activation.
    fn leaky_relu<const D: usize>(
        tensor: &B::TensorPrimitive<D>,
        negative_slope: f64,
    ) -> B::TensorPrimitive<D>;

    /// Sigmoid activation.
    fn sigmoid<const D: usize>(tensor: &B::TensorPrimitive<D>) -> B::TensorPrimitive<D>;

    /// Tanh activation.
    fn tanh<const D: usize>(tensor: &B::TensorPrimitive<D>) -> B::TensorPrimitive<D>;

    /// GELU activation (Gaussian Error Linear Unit).
    fn gelu<const D: usize>(tensor: &B::TensorPrimitive<D>) -> B::TensorPrimitive<D>;

    /// SiLU (Swish) activation: x * sigmoid(x).
    fn silu<const D: usize>(tensor: &B::TensorPrimitive<D>) -> B::TensorPrimitive<D>;

    /// Mish activation: x * tanh(softplus(x)).
    fn mish<const D: usize>(tensor: &B::TensorPrimitive<D>) -> B::TensorPrimitive<D>;

    /// Softplus activation: log(1 + exp(x)).
    fn softplus<const D: usize>(tensor: &B::TensorPrimitive<D>, beta: f64)
    -> B::TensorPrimitive<D>;

    /// ELU activation (Exponential Linear Unit).
    fn elu<const D: usize>(tensor: &B::TensorPrimitive<D>, alpha: f64) -> B::TensorPrimitive<D>;
}

/// Attention operations - LLM performance critical.
///
/// These operations delegate to hand-optimized kernels:
/// - CUDA: FlashAttention-3 kernels
/// - Metal: Custom MSL shaders
/// - WebGPU: WGSL compute shaders
/// - CPU: Optimized BLAS-based fallback
pub trait AttentionOps<B: Backend>: Sized {
    /// Flash Attention v3 (optimized scaled dot-product attention).
    ///
    /// # Arguments
    /// - `query`: `[batch, heads, seq_len, head_dim]`
    /// - `key`: `[batch, kv_heads, kv_len, head_dim]`
    /// - `value`: `[batch, kv_heads, kv_len, head_dim]`
    /// - `mask`: Optional attention mask
    /// - `scale`: Attention scale (typically 1/sqrt(head_dim))
    /// - `causal`: Whether to apply causal masking
    ///
    /// # Performance
    /// - CUDA: Uses FlashAttention-3 for 5-8x speedup over naive
    /// - Metal: Uses fused MSL kernel for Apple Silicon
    /// - WebGPU: Uses tiled WGSL kernel
    fn flash_attention(
        query: &B::TensorPrimitive<4>,
        key: &B::TensorPrimitive<4>,
        value: &B::TensorPrimitive<4>,
        mask: Option<&B::TensorPrimitive<4>>,
        scale: B::FloatElem,
        causal: bool,
    ) -> B::TensorPrimitive<4>;

    /// Flash Attention with FP8 for reduced memory.
    fn flash_attention_fp8(
        query: &B::TensorPrimitive<4>,
        key: &B::TensorPrimitive<4>,
        value: &B::TensorPrimitive<4>,
        mask: Option<&B::TensorPrimitive<4>>,
        scale: B::FloatElem,
        causal: bool,
    ) -> B::TensorPrimitive<4> {
        // Default: fallback to regular flash attention
        Self::flash_attention(query, key, value, mask, scale, causal)
    }

    /// Grouped Query Attention (GQA) - used in LLaMA 2+, Mistral.
    ///
    /// # Arguments
    /// - `query`: `[batch, q_heads, seq_len, head_dim]`
    /// - `key`: `[batch, kv_heads, kv_len, head_dim]` where kv_heads < q_heads
    /// - `value`: `[batch, kv_heads, kv_len, head_dim]`
    ///
    /// Note: kv_heads divides q_heads evenly.
    fn grouped_query_attention(
        query: &B::TensorPrimitive<4>,
        key: &B::TensorPrimitive<4>,
        value: &B::TensorPrimitive<4>,
        mask: Option<&B::TensorPrimitive<4>>,
        scale: B::FloatElem,
        causal: bool,
        num_kv_groups: usize,
    ) -> B::TensorPrimitive<4>;

    /// Multi-Query Attention (MQA) - single KV head.
    fn multi_query_attention(
        query: &B::TensorPrimitive<4>,
        key: &B::TensorPrimitive<4>,
        value: &B::TensorPrimitive<4>,
        mask: Option<&B::TensorPrimitive<4>>,
        scale: B::FloatElem,
        causal: bool,
    ) -> B::TensorPrimitive<4> {
        Self::grouped_query_attention(query, key, value, mask, scale, causal, 1)
    }

    /// Ring Attention for extremely long context (16M+ tokens).
    ///
    /// Distributes attention computation across devices using ring topology.
    fn ring_attention(
        query: &B::TensorPrimitive<4>,
        key: &B::TensorPrimitive<4>,
        value: &B::TensorPrimitive<4>,
        mask: Option<&B::TensorPrimitive<4>>,
        scale: B::FloatElem,
        causal: bool,
        chunk_size: usize,
    ) -> B::TensorPrimitive<4>;

    /// Sliding Window Attention (used in Mistral, LongT5).
    fn sliding_window_attention(
        query: &B::TensorPrimitive<4>,
        key: &B::TensorPrimitive<4>,
        value: &B::TensorPrimitive<4>,
        window_size: usize,
        scale: B::FloatElem,
    ) -> B::TensorPrimitive<4>;
}

/// KV-Cache operations for LLM inference.
///
/// RustyTorch++ unique feature: entropy-guided cache eviction.
pub trait KVCacheOps<B: Backend>: Sized {
    /// KV-Cache handle type.
    type CacheHandle: Clone + Send + Sync;

    /// Create a new KV-cache.
    fn create_cache(
        batch_size: usize,
        num_heads: usize,
        max_seq_len: usize,
        head_dim: usize,
        device: &B::Device,
    ) -> Self::CacheHandle;

    /// Update cache with new key-value pairs.
    fn update_cache(
        cache: &mut Self::CacheHandle,
        key: &B::TensorPrimitive<4>,
        value: &B::TensorPrimitive<4>,
        position: usize,
    );

    /// Get cached key-value pairs.
    fn get_cache(cache: &Self::CacheHandle) -> (B::TensorPrimitive<4>, B::TensorPrimitive<4>);

    /// Get current sequence length in cache.
    fn cache_length(cache: &Self::CacheHandle) -> usize;

    /// Apply entropy-guided eviction (RustyTorch++ unique feature).
    ///
    /// Evicts low-entropy (uninformative) tokens while preserving
    /// high-entropy (important) tokens like proper nouns, numbers.
    ///
    /// # Arguments
    /// - `cache`: The cache to evict from
    /// - `entropy_threshold`: Tokens below this entropy are candidates
    /// - `min_retention`: Keep at least this fraction of tokens
    ///
    /// # Returns
    /// Number of tokens evicted
    fn evict_by_entropy(
        cache: &mut Self::CacheHandle,
        entropy_threshold: f32,
        min_retention: f32,
    ) -> usize;

    /// Clear the cache.
    fn clear_cache(cache: &mut Self::CacheHandle);

    /// Clone the cache (for speculative decoding).
    fn clone_cache(cache: &Self::CacheHandle) -> Self::CacheHandle;

    /// Quantize cache to FP8 for memory savings.
    fn quantize_cache_fp8(cache: &mut Self::CacheHandle);
}

#[cfg(test)]
mod tests {
    // Tests would go here but require a concrete backend implementation
}