//! LSTM Encoder for TFT
//!
//! Standard LSTM implementation for temporal feature encoding in TFT.
//! 2-layer LSTM with `hidden_dim=128` (configurable).
//!
//! ## Architecture
//! - Input gate: `i_t` = `σ(W_ii` * `x_t` + `W_hi` * h_(t-1) + `b_i`)
//! - Forget gate: `f_t` = `σ(W_if` * `x_t` + `W_hf` * h_(t-1) + `b_f`)
//! - Cell gate: `g_t` = `tanh(W_ig` * `x_t` + `W_hg` * h_(t-1) + `b_g`)
//! - Output gate: `o_t` = `σ(W_io` * `x_t` + `W_ho` * h_(t-1) + `b_o`)
//! - Cell state: `c_t` = `f_t` ⊙ c_(t-1) + `i_t` ⊙ `g_t`
//! - Hidden state: `h_t` = `o_t` ⊙ `tanh(c_t)`
//!
//! ## Weight Matrices (per layer)
//! - `W_ii`, `W_if`, `W_ig`, `W_io`: Input-to-hidden weights [`hidden_size`, `input_size`]
//! - `W_hi`, `W_hf`, `W_hg`, `W_ho`: Hidden-to-hidden weights [`hidden_size`, `hidden_size`]

use std::sync::Arc;

use cudarc::driver::CudaStream;
use ml_core::MLError;

use crate::gpu_tensor::{
    gpu_add, gpu_mul, gpu_select_dim1, gpu_sigmoid, gpu_stack_2d, gpu_tanh, GpuLinear, GpuTensor,
};

/// Single LSTM layer with 4 gates
#[derive(Debug, Clone)]
pub struct LSTMLayer {
    /// Input gate: input-to-hidden
    w_ii: GpuLinear,
    /// Forget gate: input-to-hidden
    w_if: GpuLinear,
    /// Cell gate: input-to-hidden
    w_ig: GpuLinear,
    /// Output gate: input-to-hidden
    w_io: GpuLinear,

    /// Input gate: hidden-to-hidden
    w_hi: GpuLinear,
    /// Forget gate: hidden-to-hidden
    w_hf: GpuLinear,
    /// Cell gate: hidden-to-hidden
    w_hg: GpuLinear,
    /// Output gate: hidden-to-hidden
    w_ho: GpuLinear,

    hidden_size: usize,
    stream: Arc<CudaStream>,
}

impl LSTMLayer {
    /// Create new LSTM layer
    pub fn new(
        input_size: usize,
        hidden_size: usize,
        stream: &Arc<CudaStream>,
    ) -> Result<Self, MLError> {
        // Input-to-hidden weights
        let w_ii = GpuLinear::new(input_size, hidden_size, stream)?;
        let w_if = GpuLinear::new(input_size, hidden_size, stream)?;
        let w_ig = GpuLinear::new(input_size, hidden_size, stream)?;
        let w_io = GpuLinear::new(input_size, hidden_size, stream)?;

        // Hidden-to-hidden weights
        let w_hi = GpuLinear::new(hidden_size, hidden_size, stream)?;
        let w_hf = GpuLinear::new(hidden_size, hidden_size, stream)?;
        let w_hg = GpuLinear::new(hidden_size, hidden_size, stream)?;
        let w_ho = GpuLinear::new(hidden_size, hidden_size, stream)?;

        Ok(Self {
            w_ii,
            w_if,
            w_ig,
            w_io,
            w_hi,
            w_hf,
            w_hg,
            w_ho,
            hidden_size,
            stream: Arc::clone(stream),
        })
    }

    /// Forward pass through single LSTM layer
    ///
    /// # Arguments
    /// * `input` - Input tensor [batch, `seq_len`, `input_size`] (flattened to 2D per timestep)
    /// * `h0` - Initial hidden state [batch, `hidden_size`]
    /// * `c0` - Initial cell state [batch, `hidden_size`]
    ///
    /// # Returns
    /// - output: [batch, `seq_len`, `hidden_size`]
    /// - `h_final`: [batch, `hidden_size`]
    /// - `c_final`: [batch, `hidden_size`]
    #[allow(clippy::too_many_lines)]
    pub fn forward(
        &self,
        input: &GpuTensor,
        h0: Option<&GpuTensor>,
        c0: Option<&GpuTensor>,
    ) -> Result<(GpuTensor, GpuTensor, GpuTensor), MLError> {
        // Input shape: [batch, seq_len, input_size]
        if input.shape.len() != 3 {
            return Err(MLError::DimensionMismatch {
                expected: 3,
                actual: input.shape.len(),
            });
        }
        let batch_size = input.dim(0)?;
        let seq_len = input.dim(1)?;

        // Initialize hidden and cell states if not provided
        let mut h_t = match h0 {
            Some(h) => h.clone(),
            None => GpuTensor::zeros(&[batch_size, self.hidden_size], &self.stream)?,
        };

        let mut c_t = match c0 {
            Some(c) => c.clone(),
            None => GpuTensor::zeros(&[batch_size, self.hidden_size], &self.stream)?,
        };

        // Collect outputs for each timestep — all on GPU
        let mut output_steps: Vec<GpuTensor> = Vec::with_capacity(seq_len);

        for t in 0..seq_len {
            // Extract timestep t on GPU: input[batch, t, input_size] -> [batch, input_size]
            let x_t = gpu_select_dim1(input, t)?;

            // Input gate: i_t = sigma(W_ii * x_t + W_hi * h_(t-1))
            let i_input = self.w_ii.forward(&x_t)?;
            let i_hidden = self.w_hi.forward(&h_t)?;
            let i_sum = gpu_add(&i_input, &i_hidden)?;
            let i_t = gpu_sigmoid(&i_sum)?;

            // Forget gate: f_t = sigma(W_if * x_t + W_hf * h_(t-1))
            let f_input = self.w_if.forward(&x_t)?;
            let f_hidden = self.w_hf.forward(&h_t)?;
            let f_sum = gpu_add(&f_input, &f_hidden)?;
            let f_t = gpu_sigmoid(&f_sum)?;

            // Cell gate: g_t = tanh(W_ig * x_t + W_hg * h_(t-1))
            let g_input = self.w_ig.forward(&x_t)?;
            let g_hidden = self.w_hg.forward(&h_t)?;
            let g_sum = gpu_add(&g_input, &g_hidden)?;
            let g_t = gpu_tanh(&g_sum)?;

            // Output gate: o_t = sigma(W_io * x_t + W_ho * h_(t-1))
            let o_input = self.w_io.forward(&x_t)?;
            let o_hidden = self.w_ho.forward(&h_t)?;
            let o_sum = gpu_add(&o_input, &o_hidden)?;
            let o_t = gpu_sigmoid(&o_sum)?;

            // Cell state: c_t = f_t * c_(t-1) + i_t * g_t
            let fc = gpu_mul(&f_t, &c_t)?;
            let ig = gpu_mul(&i_t, &g_t)?;
            c_t = gpu_add(&fc, &ig)?;

            // Hidden state: h_t = o_t * tanh(c_t)
            let c_tanh = gpu_tanh(&c_t)?;
            h_t = gpu_mul(&o_t, &c_tanh)?;

            output_steps.push(h_t.clone());
        }

        // Assemble output on GPU: [batch, seq_len, hidden_size]
        // output_steps[t] is [batch, hidden_size]. For each batch, extract row b
        // from each timestep and concatenate to get [seq_len, hidden_size], then
        // stack across batches.
        let mut batch_sequences: Vec<GpuTensor> = Vec::with_capacity(batch_size);
        for b in 0..batch_size {
            let mut steps_for_batch: Vec<GpuTensor> = Vec::with_capacity(seq_len);
            for step in &output_steps {
                // step: [batch, hidden_size] -> extract row b: [1, hidden_size]
                let row = crate::gpu_tensor::gpu_narrow_2d(step, 0, b, 1)?;
                steps_for_batch.push(row);
            }
            // Stack [seq_len] items of [1, hidden] -> cat dim0 gives [seq_len, hidden]
            let batch_seq = crate::gpu_tensor::gpu_cat_dim0(&steps_for_batch, &self.stream)?;
            batch_sequences.push(batch_seq);
        }
        // Stack [batch] items of [seq_len, hidden] -> [batch, seq_len, hidden]
        let output = gpu_stack_2d(&batch_sequences)?;

        Ok((output, h_t, c_t))
    }
}

/// Multi-layer LSTM encoder
pub struct LSTMEncoder {
    layers: Vec<LSTMLayer>,
    num_layers: usize,
    hidden_size: usize,
}

impl std::fmt::Debug for LSTMEncoder {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("LSTMEncoder")
            .field("num_layers", &self.num_layers)
            .field("hidden_size", &self.hidden_size)
            .finish()
    }
}

impl LSTMEncoder {
    /// Create new LSTM encoder
    ///
    /// # Arguments
    /// * `num_layers` - Number of LSTM layers (typically 2)
    /// * `input_size` - Input feature dimension
    /// * `hidden_size` - Hidden state dimension (typically 128)
    /// * `stream` - CUDA stream for GPU operations
    pub fn new(
        num_layers: usize,
        input_size: usize,
        hidden_size: usize,
        stream: &Arc<CudaStream>,
    ) -> Result<Self, MLError> {
        let mut layers = Vec::new();

        for i in 0..num_layers {
            let layer_input_size = if i == 0 { input_size } else { hidden_size };
            let layer = LSTMLayer::new(layer_input_size, hidden_size, stream)?;
            layers.push(layer);
        }

        Ok(Self {
            layers,
            num_layers,
            hidden_size,
        })
    }

    /// Forward pass through all LSTM layers
    pub fn forward(
        &self,
        input: &GpuTensor,
        _states: Option<(&GpuTensor, &GpuTensor)>,
    ) -> Result<(GpuTensor, GpuTensor, GpuTensor), MLError> {
        let mut layer_input = input.clone();
        let mut h_final_last = GpuTensor::zeros(&[1], &layer_input.stream)?;
        let mut c_final_last = GpuTensor::zeros(&[1], &layer_input.stream)?;

        for layer in &self.layers {
            let (output, h_final, c_final) = layer.forward(&layer_input, None, None)?;
            layer_input = output;
            h_final_last = h_final;
            c_final_last = c_final;
        }

        Ok((layer_input, h_final_last, c_final_last))
    }

    /// Get number of layers
    pub const fn num_layers(&self) -> usize {
        self.num_layers
    }

    /// Get hidden size
    pub const fn hidden_size(&self) -> usize {
        self.hidden_size
    }

    /// Estimate memory usage in MB (FP32)
    pub fn estimate_memory_mb(&self) -> f64 {
        let params_per_layer = 8 * self.hidden_size * self.hidden_size;
        let total_params = params_per_layer * self.num_layers;
        let bytes = total_params * 4;
        bytes as f64 / (1024.0 * 1024.0)
    }
}
