//! Gated Residual Network (GRN) for TFT
//!
//! Implements gated linear units with residual connections for improved
//! gradient flow and feature learning in temporal fusion transformers.

use std::sync::Arc;

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

use crate::gpu_tensor::{
    gpu_add, gpu_elu, gpu_layer_norm, gpu_mul, gpu_ones, gpu_sigmoid, GpuLinear, GpuTensor,
};

/// Gated Linear Unit for feature gating
#[derive(Debug)]
pub struct GatedLinearUnit {
    pub output_dim: usize,
    linear: GpuLinear,
    gate: GpuLinear,
}

impl GatedLinearUnit {
    pub fn new(input_dim: usize, output_dim: usize, stream: &Arc<CudaStream>) -> Result<Self, MLError> {
        let linear = GpuLinear::new(input_dim, output_dim, stream)?;
        let gate = GpuLinear::new(input_dim, output_dim, stream)?;

        Ok(Self {
            output_dim,
            linear,
            gate,
        })
    }

    pub fn forward(&self, x: &GpuTensor) -> Result<GpuTensor, MLError> {
        let linear_out = self.linear.forward(x)?;
        let gate_out = gpu_sigmoid(&self.gate.forward(x)?)?;
        gpu_mul(&linear_out, &gate_out)
    }
}

/// Gated Residual Network with optional context
#[derive(Debug)]
pub struct GatedResidualNetwork {
    pub input_dim: usize,
    pub output_dim: usize,
    // Primary processing layers
    linear1: GpuLinear,
    linear2: GpuLinear,
    // Gating mechanism
    glu: GatedLinearUnit,
    // Layer normalization parameters
    ln_gamma: GpuTensor,
    ln_beta: GpuTensor,
    // Optional skip connection projection
    skip_projection: Option<GpuLinear>,
    // Context integration
    context_projection: Option<GpuLinear>,
}

impl GatedResidualNetwork {
    pub fn new(input_dim: usize, output_dim: usize, stream: &Arc<CudaStream>) -> Result<Self, MLError> {
        let linear1 = GpuLinear::new(input_dim, output_dim, stream)?;
        let linear2 = GpuLinear::new(output_dim, output_dim, stream)?;
        let glu = GatedLinearUnit::new(output_dim, output_dim, stream)?;
        let ln_gamma = gpu_ones(&[output_dim], stream)?;
        let ln_beta = GpuTensor::zeros(&[output_dim], stream)?;

        let skip_projection = (input_dim != output_dim)
            .then(|| GpuLinear::new(input_dim, output_dim, stream))
            .transpose()?;

        let context_projection = Some(GpuLinear::new(output_dim, output_dim, stream)?);

        Ok(Self {
            input_dim,
            output_dim,
            linear1,
            linear2,
            glu,
            ln_gamma,
            ln_beta,
            skip_projection,
            context_projection,
        })
    }

    pub fn forward(&self, x: &GpuTensor, context: Option<&GpuTensor>) -> Result<GpuTensor, MLError> {
        // For 3D input: flatten batch*seq_len, apply, reshape back
        if x.shape.len() == 3 {
            let batch = x.dim(0)?;
            let seq_len = x.dim(1)?;
            let feat = x.dim(2)?;
            let flat = x.reshape(&[batch * seq_len, feat])?;
            let ctx_flat = context.map(|c| {
                if c.shape.len() == 3 {
                    let cb = c.dim(0).unwrap_or(batch);
                    let cs = c.dim(1).unwrap_or(seq_len);
                    let cf = c.dim(2).unwrap_or(feat);
                    c.reshape(&[cb * cs, cf]).unwrap_or_else(|_| c.clone())
                } else {
                    c.clone()
                }
            });
            let result_flat = self.forward_2d(&flat, ctx_flat.as_ref())?;
            let out_dim = result_flat.dim(1)?;
            return result_flat.reshape(&[batch, seq_len, out_dim]);
        }

        self.forward_2d(x, context)
    }

    fn forward_2d(&self, x: &GpuTensor, context: Option<&GpuTensor>) -> Result<GpuTensor, MLError> {
        // First linear transformation + ELU
        let mut hidden = gpu_elu(&self.linear1.forward(x)?, 1.0)?;

        // Apply context if provided
        if let (Some(ctx), Some(ctx_proj)) = (context, &self.context_projection) {
            let ctx_out = ctx_proj.forward(ctx)?;
            hidden = gpu_add(&hidden, &ctx_out)?;
        }

        // Second linear transformation
        hidden = self.linear2.forward(&hidden)?;

        // Apply gating
        let gated = self.glu.forward(&hidden)?;

        // Skip connection
        let output = if let Some(proj) = &self.skip_projection {
            let skip = proj.forward(x)?;
            gpu_add(&gated, &skip)?
        } else {
            gpu_add(&gated, x)?
        };

        // Layer normalization
        gpu_layer_norm(&output, &self.ln_gamma, &self.ln_beta, 1e-5)
    }
}

/// Stack of Gated Residual Networks
#[derive(Debug)]
pub struct GRNStack {
    pub num_layers: usize,
    layers: Vec<GatedResidualNetwork>,
}

impl GRNStack {
    pub fn new(
        input_dim: usize,
        hidden_dim: usize,
        output_dim: usize,
        num_layers: usize,
        stream: &Arc<CudaStream>,
    ) -> Result<Self, MLError> {
        let mut layers = Vec::with_capacity(num_layers);

        for i in 0..num_layers {
            let layer_input_dim = if i == 0 { input_dim } else { hidden_dim };
            let layer_output_dim = if i == num_layers - 1 {
                output_dim
            } else {
                hidden_dim
            };

            let grn = GatedResidualNetwork::new(layer_input_dim, layer_output_dim, stream)?;
            layers.push(grn);
        }

        Ok(Self { num_layers, layers })
    }

    pub fn forward(&self, x: &GpuTensor, context: Option<&GpuTensor>) -> Result<GpuTensor, MLError> {
        let mut output = if let Some(first_layer) = self.layers.first() {
            first_layer.forward(x, context)?
        } else {
            return Ok(x.clone());
        };

        for layer in self.layers.iter().skip(1) {
            output = layer.forward(&output, context)?;
        }

        Ok(output)
    }
}

#[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_grn_creation() -> Result<(), MLError> {
        let stream = test_stream();
        let grn = GatedResidualNetwork::new(64, 32, &stream)?;
        assert_eq!(grn.input_dim, 64);
        assert_eq!(grn.output_dim, 32);
        Ok(())
    }

    #[test]
    fn test_grn_forward_same_dims() -> Result<(), MLError> {
        let stream = test_stream();
        let grn = GatedResidualNetwork::new(32, 32, &stream)?;
        let input = GpuTensor::from_vec(vec![1.0_f32; 64], &[2, 32], &stream)?;
        let output = grn.forward(&input, None)?;
        assert_eq!(output.shape, vec![2, 32]);
        Ok(())
    }

    #[test]
    fn test_grn_forward_different_dims() -> Result<(), MLError> {
        let stream = test_stream();
        let grn = GatedResidualNetwork::new(64, 32, &stream)?;
        let input = GpuTensor::from_vec(vec![1.0_f32; 128], &[2, 64], &stream)?;
        let output = grn.forward(&input, None)?;
        assert_eq!(output.shape, vec![2, 32]);
        Ok(())
    }

    #[test]
    fn test_grn_forward_with_context() -> Result<(), MLError> {
        let stream = test_stream();
        let grn = GatedResidualNetwork::new(32, 32, &stream)?;
        let input = GpuTensor::from_vec(vec![1.0_f32; 64], &[2, 32], &stream)?;
        let context = GpuTensor::from_vec(vec![0.5_f32; 64], &[2, 32], &stream)?;
        let output = grn.forward(&input, Some(&context))?;
        assert_eq!(output.shape, vec![2, 32]);
        Ok(())
    }

    #[test]
    fn test_grn_forward_3d() -> Result<(), MLError> {
        let stream = test_stream();
        let grn = GatedResidualNetwork::new(16, 16, &stream)?;
        let input = GpuTensor::from_vec(vec![1.0_f32; 160], &[2, 5, 16], &stream)?;
        let output = grn.forward(&input, None)?;
        assert_eq!(output.shape, vec![2, 5, 16]);
        Ok(())
    }

    #[test]
    fn test_glu_creation() -> Result<(), MLError> {
        let stream = test_stream();
        let glu = GatedLinearUnit::new(64, 32, &stream)?;
        assert_eq!(glu.output_dim, 32);
        Ok(())
    }

    #[test]
    fn test_glu_forward() -> Result<(), MLError> {
        let stream = test_stream();
        let glu = GatedLinearUnit::new(32, 16, &stream)?;
        let input = GpuTensor::from_vec(vec![1.0_f32; 64], &[2, 32], &stream)?;
        let output = glu.forward(&input)?;
        assert_eq!(output.shape, vec![2, 16]);
        Ok(())
    }

    #[test]
    fn test_grn_stack() -> Result<(), MLError> {
        let stream = test_stream();
        let stack = GRNStack::new(64, 32, 16, 3, &stream)?;
        assert_eq!(stack.num_layers, 3);
        let input = GpuTensor::from_vec(vec![1.0_f32; 128], &[2, 64], &stream)?;
        let output = stack.forward(&input, None)?;
        assert_eq!(output.shape, vec![2, 16]);
        Ok(())
    }
}
