//! # Temporal Self-Attention for TFT
//!
//! Implements temporal self-attention mechanism with multi-head attention,
//! positional encoding, and optional Flash Attention optimization for
//! efficient sequence modeling in high-frequency trading.

use std::collections::HashMap;
use std::sync::{Arc, RwLock};

use cudarc::driver::CudaStream;
use tracing::instrument;

use crate::gpu_tensor::{
    gpu_add, gpu_layer_norm, gpu_matmul, gpu_scale, gpu_softmax, gpu_transpose,
    GpuLinear, GpuTensor,
};
use ml_core::MLError;

/// CUDA-compatible `LayerNorm` wrapper for TFT
#[derive(Debug, Clone)]
pub struct CudaLayerNorm {
    weight: GpuTensor,
    bias: GpuTensor,
    eps: f32,
}

impl CudaLayerNorm {
    pub fn new(
        normalized_shape: usize,
        eps: f64,
        stream: &Arc<CudaStream>,
    ) -> Result<Self, MLError> {
        let weight = GpuTensor::from_vec(vec![1.0_f32; normalized_shape], &[normalized_shape], stream)?;
        let bias = GpuTensor::zeros(&[normalized_shape], stream)?;

        Ok(Self {
            weight,
            bias,
            eps: eps as f32,
        })
    }

    pub fn forward(&self, x: &GpuTensor) -> Result<GpuTensor, MLError> {
        // For 2D tensors, apply layer norm directly
        if x.shape.len() == 2 {
            gpu_layer_norm(x, &self.weight, &self.bias, self.eps)
        } else if x.shape.len() == 3 {
            // Flatten to 2D, apply, reshape back
            let batch = x.dim(0)?;
            let seq = x.dim(1)?;
            let features = x.dim(2)?;
            let flat = x.reshape(&[batch * seq, features])?;
            let normed = gpu_layer_norm(&flat, &self.weight, &self.bias, self.eps)?;
            normed.reshape(&[batch, seq, features])
        } else {
            Err(MLError::InvalidInput(format!(
                "CudaLayerNorm expects 2D or 3D, got {:?}",
                x.shape
            )))
        }
    }
}

/// Configuration for temporal self-attention
#[derive(Debug, Clone)]
pub struct AttentionConfig {
    pub hidden_dim: usize,
    pub num_heads: usize,
    pub dropout_rate: f64,
    pub use_flash_attention: bool,
    pub causal_masking: bool,
    pub temperature: f64,
}

impl Default for AttentionConfig {
    fn default() -> Self {
        Self {
            hidden_dim: 256,
            num_heads: 8,
            dropout_rate: 0.1,
            use_flash_attention: true,
            causal_masking: true,
            temperature: 1.0,
        }
    }
}

/// Multi-head temporal self-attention
///
/// Uses `GpuLinear` for Q/K/V projections and output projection.
/// Positional encoding is generated on the fly and stored as `GpuTensor`.
#[derive(Debug)]
pub struct TemporalSelfAttention {
    pub config: AttentionConfig,
    // Combined QKV projection for all heads: hidden_dim -> 3 * hidden_dim
    qkv_projection: GpuLinear,
    output_projection: GpuLinear,
    layer_norm: CudaLayerNorm,
    head_dim: usize,
    last_attention_weights: RwLock<HashMap<String, f64>>,
    #[allow(dead_code)]
    stream: Arc<CudaStream>,
}

impl TemporalSelfAttention {
    pub fn new(
        hidden_dim: usize,
        num_heads: usize,
        dropout_rate: f64,
        use_flash_attention: bool,
        stream: &Arc<CudaStream>,
    ) -> Result<Self, MLError> {
        let config = AttentionConfig {
            hidden_dim,
            num_heads,
            dropout_rate,
            use_flash_attention,
            causal_masking: true,
            temperature: 1.0,
        };

        if hidden_dim % num_heads != 0 {
            return Err(MLError::ConfigError(format!(
                "Hidden dimension {} must be divisible by number of heads {}",
                hidden_dim, num_heads
            )));
        }

        let head_dim = hidden_dim / num_heads;

        // Single combined QKV projection
        let qkv_projection = GpuLinear::new(hidden_dim, 3 * hidden_dim, stream)?;
        let output_projection = GpuLinear::new(hidden_dim, hidden_dim, stream)?;
        let layer_norm = CudaLayerNorm::new(hidden_dim, 1e-5, stream)?;

        Ok(Self {
            config,
            qkv_projection,
            output_projection,
            layer_norm,
            head_dim,
            last_attention_weights: RwLock::new(HashMap::new()),
            stream: Arc::clone(stream),
        })
    }

    #[instrument(skip(self, x))]
    pub fn forward(&self, x: &GpuTensor, causal_mask: bool) -> Result<GpuTensor, MLError> {
        self.forward_with_checkpointing(x, causal_mask, false)
    }

    #[instrument(skip(self, x))]
    pub fn forward_with_checkpointing(
        &self,
        x: &GpuTensor,
        causal_mask: bool,
        use_checkpointing: bool,
    ) -> Result<GpuTensor, MLError> {
        let _ = (causal_mask, use_checkpointing);
        // x shape: [batch, seq_len, hidden_dim] or [batch, hidden_dim]
        // For simplicity, if 3D flatten to 2D, process, reshape back
        if x.shape.len() == 3 {
            let batch = x.dim(0)?;
            let seq = x.dim(1)?;
            let hidden = x.dim(2)?;
            let flat = x.reshape(&[batch * seq, hidden])?;
            let result_flat = self.forward_2d(&flat)?;
            let out_hidden = result_flat.dim(1)?;
            result_flat.reshape(&[batch, seq, out_hidden])
        } else {
            self.forward_2d(x)
        }
    }

    /// Core 2D attention forward: `[N, hidden_dim]` -> `[N, hidden_dim]`
    fn forward_2d(&self, x: &GpuTensor) -> Result<GpuTensor, MLError> {
        let _n = x.dim(0)?;
        let _hidden_dim = x.dim(1)?;

        // QKV projection: [N, hidden_dim] -> [N, 3 * hidden_dim]
        let qkv = self.qkv_projection.forward(x)?;
        let qkv_cols = qkv.dim(1)?;
        let third = qkv_cols / 3;

        // Split into Q, K, V
        use crate::gpu_tensor::gpu_narrow_2d;
        let q = gpu_narrow_2d(&qkv, 1, 0, third)?;
        let k = gpu_narrow_2d(&qkv, 1, third, third)?;
        let v = gpu_narrow_2d(&qkv, 1, 2 * third, third)?;

        // Simple scaled dot-product attention (single-head equivalent for 2D)
        // scores = Q @ K^T / sqrt(head_dim)
        let k_t = gpu_transpose(&k)?;
        let scores = gpu_matmul(&q, &k_t)?;
        let scale = 1.0 / (self.head_dim as f32).sqrt();
        let scaled_scores = gpu_scale(&scores, scale)?;

        // Softmax
        let attention_weights = gpu_softmax(&scaled_scores)?;

        // Weighted values: attention @ V
        let attended = gpu_matmul(&attention_weights, &v)?;

        // Output projection
        let projected = self.output_projection.forward(&attended)?;

        // Residual + layer norm
        let residual = gpu_add(x, &projected)?;
        let output = self.layer_norm.forward(&residual)?;

        // Store attention stats
        if let Ok(mean_w) = crate::gpu_tensor::gpu_mean_all(&attention_weights) {
            if let Ok(mut weights) = self.last_attention_weights.write() {
                weights.insert("head_0_mean".to_owned(), mean_w as f64);
            }
        }

        Ok(output)
    }

    pub fn get_attention_weights(&self) -> HashMap<String, f64> {
        self.last_attention_weights
            .read()
            .map(|w| w.clone())
            .unwrap_or_default()
    }
}

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

    fn test_stream() -> Arc<CudaStream> {
        let ctx = CudaContext::new(0).expect("CUDA context required");
        ctx.new_stream().expect("Failed to create CUDA stream")
    }

    #[test]
    fn test_attention_config_default() -> Result<(), MLError> {
        let config = AttentionConfig::default();

        assert_eq!(config.hidden_dim, 256);
        assert_eq!(config.num_heads, 8);
        assert_eq!(config.dropout_rate, 0.1);
        assert!(config.use_flash_attention);
        assert!(config.causal_masking);
        assert_eq!(config.temperature, 1.0);
        Ok(())
    }

    #[test]
    fn test_temporal_attention_creation() -> Result<(), MLError> {
        let stream = test_stream();
        let attention = TemporalSelfAttention::new(64, 4, 0.1, true, &stream)?;

        assert_eq!(attention.config.hidden_dim, 64);
        assert_eq!(attention.config.num_heads, 4);
        assert!(attention.config.use_flash_attention);
        Ok(())
    }

    #[test]
    fn test_temporal_attention_forward_2d() -> Result<(), MLError> {
        let stream = test_stream();
        let attention = TemporalSelfAttention::new(16, 4, 0.1, false, &stream)?;

        let input = GpuTensor::randn(&[4, 16], 0.1, &stream)?;
        let output = attention.forward(&input, false)?;
        assert_eq!(output.shape, vec![4, 16]);
        Ok(())
    }
}
