//! Flash Attention configuration

/// Configuration for Flash Attention
#[derive(Clone, Debug)]
pub struct FlashAttentionConfig {
    /// Block size for Q dimension tiling (threads per threadgroup)
    pub block_q: usize,
    /// Block size for K/V dimension tiling
    pub block_kv: usize,
    /// Maximum head dimension supported
    pub max_head_dim: usize,
    /// Softmax scaling factor (typically 1/sqrt(head_dim))
    /// If None, will be computed automatically from head_dim
    pub softmax_scale: Option<f32>,
    /// Enable causal masking (for autoregressive models)
    pub causal: bool,
    /// Dropout probability (0.0 = no dropout)
    pub dropout_p: f32,
    /// Use float16 (half precision) for Q/K/V tensors
    /// Accumulation is always done in float32 for numerical stability
    pub use_f16: bool,
}

impl Default for FlashAttentionConfig {
    fn default() -> Self {
        // Apple Silicon GPUs have 32KB threadgroup memory limit
        // Threadgroup memory includes:
        //   - Q_shared: BLOCK_Q * HEAD_DIM * bytes_per_element
        //   - K_shared: BLOCK_KV * HEAD_DIM * bytes_per_element
        //   - V_shared: BLOCK_KV * HEAD_DIM * bytes_per_element
        //   - Per-thread spill (o_acc array): BLOCK_Q * HEAD_DIM * 4 (always f32)
        // With block_q=16, block_kv=16, head_dim=64: ~16KB total (f32), ~12KB (f16)
        Self {
            block_q: 16,
            block_kv: 16,
            max_head_dim: 64,
            softmax_scale: None,
            causal: false,
            dropout_p: 0.0,
            use_f16: false,
        }
    }
}

impl FlashAttentionConfig {
    /// Create a new configuration with default settings
    pub fn new() -> Self {
        Self::default()
    }

    /// Create a configuration for causal (autoregressive) attention
    pub fn causal() -> Self {
        Self {
            causal: true,
            ..Default::default()
        }
    }

    /// Create a configuration optimized for large head dimensions (128)
    /// Uses smaller block sizes to fit within 32KB threadgroup memory
    pub fn large_head_dim() -> Self {
        // For head_dim=128 with BLOCK_Q=8, BLOCK_KV=8:
        // Memory = (8 + 16) * 128 * 4 = 12,288 bytes (~12KB)
        Self {
            block_q: 8,
            block_kv: 8,
            max_head_dim: 128,
            softmax_scale: None,
            causal: false,
            dropout_p: 0.0,
            use_f16: false,
        }
    }

    /// Create a configuration using float16 (half precision)
    ///
    /// Uses f16 for Q/K/V storage but f32 accumulation for numerical stability.
    /// This reduces memory bandwidth and allows larger block sizes.
    pub fn f16() -> Self {
        // With f16, we can use larger blocks:
        // Memory = (16 + 32) * 64 * 2 = 6,144 bytes (~6KB for shared)
        // Plus f32 accumulators: 16 * 64 * 4 = 4KB
        Self {
            block_q: 16,
            block_kv: 16,
            max_head_dim: 64,
            softmax_scale: None,
            causal: false,
            dropout_p: 0.0,
            use_f16: true,
        }
    }

    /// Create a causal configuration using float16
    pub fn f16_causal() -> Self {
        Self {
            causal: true,
            ..Self::f16()
        }
    }

    /// Create a large head dimension config with float16
    pub fn f16_large_head_dim() -> Self {
        // For f16 with head_dim=128:
        // Memory = (8 + 16) * 128 * 2 = 6,144 bytes (~6KB for shared)
        Self {
            block_q: 8,
            block_kv: 8,
            max_head_dim: 128,
            softmax_scale: None,
            causal: false,
            dropout_p: 0.0,
            use_f16: true,
        }
    }

    /// Create a causal configuration for large head dimensions (128)
    pub fn large_head_dim_causal() -> Self {
        Self {
            causal: true,
            ..Self::large_head_dim()
        }
    }

    /// Set maximum head dimension
    pub fn with_max_head_dim(mut self, max_head_dim: usize) -> Self {
        self.max_head_dim = max_head_dim;
        self
    }

    /// Set block sizes for Q and K/V tiling
    pub fn with_block_sizes(mut self, block_q: usize, block_kv: usize) -> Self {
        self.block_q = block_q;
        self.block_kv = block_kv;
        self
    }

    /// Set explicit softmax scale
    pub fn with_softmax_scale(mut self, scale: f32) -> Self {
        self.softmax_scale = Some(scale);
        self
    }

    /// Enable or disable causal masking
    pub fn with_causal(mut self, causal: bool) -> Self {
        self.causal = causal;
        self
    }

    /// Set dropout probability
    pub fn with_dropout(mut self, p: f32) -> Self {
        self.dropout_p = p.clamp(0.0, 1.0);
        self
    }

    /// Enable or disable float16 (half precision) mode
    pub fn with_f16(mut self, use_f16: bool) -> Self {
        self.use_f16 = use_f16;
        self
    }

    /// Get the softmax scale, computing from head_dim if not explicitly set
    pub fn get_softmax_scale(&self, head_dim: usize) -> f32 {
        self.softmax_scale
            .unwrap_or_else(|| 1.0 / (head_dim as f32).sqrt())
    }

    /// Validate configuration against tensor dimensions
    pub fn validate(&self, head_dim: usize) -> Result<(), String> {
        if head_dim > self.max_head_dim {
            return Err(format!(
                "head_dim {} exceeds max_head_dim {}",
                head_dim, self.max_head_dim
            ));
        }
        if self.block_q == 0 || self.block_kv == 0 {
            return Err("Block sizes must be non-zero".to_string());
        }
        if self.dropout_p < 0.0 || self.dropout_p > 1.0 {
            return Err("Dropout probability must be in [0, 1]".to_string());
        }
        Ok(())
    }
}

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

    #[test]
    fn test_default_config() {
        let config = FlashAttentionConfig::default();
        assert_eq!(config.block_q, 16);
        assert_eq!(config.block_kv, 16);
        assert_eq!(config.max_head_dim, 64);
        assert!(!config.causal);
        assert_eq!(config.dropout_p, 0.0);
        assert!(!config.use_f16);
    }

    #[test]
    fn test_causal_config() {
        let config = FlashAttentionConfig::causal();
        assert!(config.causal);
    }

    #[test]
    fn test_large_head_dim_config() {
        let config = FlashAttentionConfig::large_head_dim();
        assert_eq!(config.block_q, 8);
        assert_eq!(config.block_kv, 8);
        assert_eq!(config.max_head_dim, 128);
        assert!(!config.causal);
    }

    #[test]
    fn test_large_head_dim_causal_config() {
        let config = FlashAttentionConfig::large_head_dim_causal();
        assert_eq!(config.max_head_dim, 128);
        assert!(config.causal);
    }

    #[test]
    fn test_f16_config() {
        let config = FlashAttentionConfig::f16();
        assert!(config.use_f16);
        assert!(!config.causal);
        assert_eq!(config.max_head_dim, 64);

        let config_causal = FlashAttentionConfig::f16_causal();
        assert!(config_causal.use_f16);
        assert!(config_causal.causal);

        let config_large = FlashAttentionConfig::f16_large_head_dim();
        assert!(config_large.use_f16);
        assert_eq!(config_large.max_head_dim, 128);
    }

    #[test]
    fn test_softmax_scale() {
        let config = FlashAttentionConfig::default();
        let scale = config.get_softmax_scale(64);
        assert!((scale - 0.125).abs() < 1e-6); // 1/sqrt(64) = 0.125

        let scale_128 = config.get_softmax_scale(128);
        assert!((scale_128 - 0.08838835).abs() < 1e-6); // 1/sqrt(128)
    }

    #[test]
    fn test_validation() {
        let config = FlashAttentionConfig::default();
        assert!(config.validate(64).is_ok());
        assert!(config.validate(128).is_err()); // exceeds max_head_dim for default

        let large_config = FlashAttentionConfig::large_head_dim();
        assert!(large_config.validate(128).is_ok());
        assert!(large_config.validate(256).is_err());
    }
}