//! Backward Pass Implementation for PINN Training
//!
//! This module provides analytical gradient computation for the Helmholtz PINN.
//! Instead of using autograd, we derive and implement closed-form gradients.
//!
//! ## Loss Function
//!
//! L = w_data * L_data + w_pde * L_pde
//!
//! where:
//! - L_data = MSE(u_pred, u_target) = (1/N) Σ ||u_pred - u_target||²
//! - L_pde = MSE(∂²u/∂x² + k²u, 0) (Helmholtz residual)
//!
//! ## Network Architecture
//!
//! u_pred = W_L * tanh(W_{L-1} * ... tanh(W_0 * φ(x, B) + b_0) ... + b_{L-1}) + b_L
//!
//! where φ(x, B) = [sin(2πBx), cos(2πBx)] are Fourier features

use std::f32::consts::PI;
use rtx_tensor::{Tensor, Device};

/// Result type for backward operations
pub type Result<T> = std::result::Result<T, Box<dyn std::error::Error + Send + Sync>>;

/// Workspace for gradient computation
///
/// Stores cached activations from the forward pass needed for backpropagation,
/// plus gradient accumulators for all trainable parameters.
#[derive(Debug)]
pub struct GradientWorkspace {
    // =========================================================================
    // Cached Activations (from forward pass)
    // =========================================================================

    /// Pre-activation values: z_l = W_l @ h_{l-1} + b_l
    /// Shape: [batch, hidden_dim] for hidden layers, [batch, 2] for output
    pub z: Vec<Tensor>,

    /// Post-activation values: h_l = tanh(z_l) for hidden, identity for output
    /// Shape: [batch, hidden_dim] for hidden layers, [batch, 2] for output
    /// Note: h[0] is after first layer, NOT the input features
    pub h: Vec<Tensor>,

    /// Fourier features: φ(x, B) = [sin(2πBx), cos(2πBx)]
    /// Shape: [batch, ff_dim * 2]
    /// This is the input to the first layer
    pub fourier_features: Tensor,

    /// Raw x values for Fourier gradient computation
    /// Shape: [batch, 1]
    pub x_input: Tensor,

    // =========================================================================
    // Gradient Accumulators
    // =========================================================================

    /// Weight gradients: dL/dW_l
    /// Shape: [out_dim, in_dim] for each layer
    pub dW: Vec<Tensor>,

    /// Bias gradients: dL/db_l
    /// Shape: [out_dim] for each layer
    pub db: Vec<Tensor>,

    /// Fourier frequency gradient: dL/dB
    /// Shape: [1, ff_dim]
    pub dB: Tensor,

    // =========================================================================
    // Spatial Derivative Caches (for PDE loss)
    // =========================================================================

    /// First spatial derivative: ∂u/∂x
    /// Shape: [batch, 2] (real and imaginary parts)
    pub du_dx: Tensor,

    /// Second spatial derivative: ∂²u/∂x²
    /// Shape: [batch, 2]
    pub d2u_dx2: Tensor,

    // =========================================================================
    // Backward Pass Intermediate Buffers (pre-allocated to avoid heap allocs)
    // =========================================================================

    /// h² for tanh derivative: 1 - h²
    /// Shape: [batch, hidden_dim] for each layer (except output)
    pub h_sq: Vec<Tensor>,

    /// tanh derivative: 1 - h² = sech²(z)
    /// Shape: [batch, hidden_dim] for each layer (except output)
    pub tanh_deriv: Vec<Tensor>,

    /// dL/dz = dL/dh * tanh'(z) for each layer
    /// Shape: [batch, out_dim] for each layer
    pub dL_dz: Vec<Tensor>,

    /// dL/dz transposed for weight gradient computation
    /// Shape: [out_dim, batch] for each layer
    pub dL_dz_t: Vec<Tensor>,

    /// Gradient w.r.t. previous layer output (passed backward)
    /// Shape: [batch, hidden_dim] for hidden, [batch, ff_dim*2] for first layer
    pub dL_dh_prev: Vec<Tensor>,

    /// MSE gradient output buffer
    /// Shape: [batch, 2]
    pub d_loss: Tensor,

    // =========================================================================
    // Fourier Backward Intermediate Buffers
    // =========================================================================

    /// x * scale for Fourier backward
    /// Shape: [batch, 1]
    pub x_scaled: Tensor,

    /// B @ x_scaled = 2πBx
    /// Shape: [batch, ff_dim]
    pub Bx: Tensor,

    /// cos(2πBx) for Fourier backward
    /// Shape: [batch, ff_dim]
    pub cos_Bx: Tensor,

    /// sin(2πBx) for Fourier backward
    /// Shape: [batch, ff_dim]
    pub sin_Bx: Tensor,

    /// dphi_sin/dB = 2πx * cos(2πBx)
    /// Shape: [batch, ff_dim]
    pub dphi_sin_dB: Tensor,

    /// dphi_cos/dB = -2πx * sin(2πBx)
    /// Shape: [batch, ff_dim]
    pub dphi_cos_dB: Tensor,

    /// Slice of dL_dphi for sin part
    /// Shape: [batch, ff_dim]
    pub dL_dphi_sin: Tensor,

    /// Slice of dL_dphi for cos part
    /// Shape: [batch, ff_dim]
    pub dL_dphi_cos: Tensor,

    /// dL_dB from sin contribution
    /// Shape: [batch, ff_dim]
    pub dL_dB_from_sin: Tensor,

    /// dL_dB from cos contribution
    /// Shape: [batch, ff_dim]
    pub dL_dB_from_cos: Tensor,

    /// Device for tensor allocation
    device: Device,

    /// Number of layers
    num_layers: usize,

    /// Batch size
    batch_size: usize,

    /// Hidden dimension
    hidden_dim: usize,

    /// Fourier feature dimension
    ff_dim: usize,
}

impl GradientWorkspace {
    /// Create a new gradient workspace
    ///
    /// # Arguments
    /// * `batch_size` - Number of samples in a batch
    /// * `ff_dim` - Fourier feature dimension
    /// * `hidden_dim` - Hidden layer dimension
    /// * `num_hidden_layers` - Number of hidden layers (not including output)
    /// * `device` - Device for tensor allocation
    pub fn new(
        batch_size: usize,
        ff_dim: usize,
        hidden_dim: usize,
        num_hidden_layers: usize,
        device: &Device,
    ) -> Result<Self> {
        let num_layers = num_hidden_layers + 1; // +1 for output layer

        // Pre-allocate activation caches
        let mut z = Vec::with_capacity(num_layers);
        let mut h = Vec::with_capacity(num_layers);

        // Hidden layers
        for _ in 0..num_hidden_layers {
            z.push(Tensor::zeros(&[batch_size, hidden_dim], device)?);
            h.push(Tensor::zeros(&[batch_size, hidden_dim], device)?);
        }
        // Output layer
        z.push(Tensor::zeros(&[batch_size, 2], device)?);
        h.push(Tensor::zeros(&[batch_size, 2], device)?);

        // Pre-allocate gradient accumulators
        let mut dW = Vec::with_capacity(num_layers);
        let mut db = Vec::with_capacity(num_layers);

        // First hidden layer: [hidden_dim, ff_dim * 2]
        dW.push(Tensor::zeros(&[hidden_dim, ff_dim * 2], device)?);
        db.push(Tensor::zeros(&[hidden_dim], device)?);

        // Middle hidden layers: [hidden_dim, hidden_dim]
        for _ in 1..num_hidden_layers {
            dW.push(Tensor::zeros(&[hidden_dim, hidden_dim], device)?);
            db.push(Tensor::zeros(&[hidden_dim], device)?);
        }

        // Output layer: [2, hidden_dim]
        dW.push(Tensor::zeros(&[2, hidden_dim], device)?);
        db.push(Tensor::zeros(&[2], device)?);

        // Fourier frequency gradient
        let dB = Tensor::zeros(&[1, ff_dim], device)?;

        // Fourier features cache
        let fourier_features = Tensor::zeros(&[batch_size, ff_dim * 2], device)?;

        // Input cache
        let x_input = Tensor::zeros(&[batch_size, 1], device)?;

        // Spatial derivatives
        let du_dx = Tensor::zeros(&[batch_size, 2], device)?;
        let d2u_dx2 = Tensor::zeros(&[batch_size, 2], device)?;

        // =========================================================================
        // Backward Pass Intermediate Buffers
        // =========================================================================

        // Buffers for tanh derivative computation (only for hidden layers)
        let mut h_sq = Vec::with_capacity(num_hidden_layers);
        let mut tanh_deriv = Vec::with_capacity(num_hidden_layers);
        for _ in 0..num_hidden_layers {
            h_sq.push(Tensor::zeros(&[batch_size, hidden_dim], device)?);
            tanh_deriv.push(Tensor::zeros(&[batch_size, hidden_dim], device)?);
        }

        // dL/dz for each layer
        let mut dL_dz = Vec::with_capacity(num_layers);
        let mut dL_dz_t = Vec::with_capacity(num_layers);
        for i in 0..num_layers {
            let out_dim = if i < num_hidden_layers { hidden_dim } else { 2 };
            dL_dz.push(Tensor::zeros(&[batch_size, out_dim], device)?);
            dL_dz_t.push(Tensor::zeros(&[out_dim, batch_size], device)?);
        }

        // dL/dh_prev for each layer (gradient passed backward)
        let mut dL_dh_prev = Vec::with_capacity(num_layers);
        // First layer: gradient goes to Fourier features [batch, ff_dim*2]
        dL_dh_prev.push(Tensor::zeros(&[batch_size, ff_dim * 2], device)?);
        // Hidden layers: [batch, hidden_dim]
        for _ in 1..num_layers {
            dL_dh_prev.push(Tensor::zeros(&[batch_size, hidden_dim], device)?);
        }

        // MSE gradient buffer
        let d_loss = Tensor::zeros(&[batch_size, 2], device)?;

        // =========================================================================
        // Fourier Backward Intermediate Buffers
        // =========================================================================

        let x_scaled = Tensor::zeros(&[batch_size, 1], device)?;
        let Bx = Tensor::zeros(&[batch_size, ff_dim], device)?;
        let cos_Bx = Tensor::zeros(&[batch_size, ff_dim], device)?;
        let sin_Bx = Tensor::zeros(&[batch_size, ff_dim], device)?;
        let dphi_sin_dB = Tensor::zeros(&[batch_size, ff_dim], device)?;
        let dphi_cos_dB = Tensor::zeros(&[batch_size, ff_dim], device)?;
        let dL_dphi_sin = Tensor::zeros(&[batch_size, ff_dim], device)?;
        let dL_dphi_cos = Tensor::zeros(&[batch_size, ff_dim], device)?;
        let dL_dB_from_sin = Tensor::zeros(&[batch_size, ff_dim], device)?;
        let dL_dB_from_cos = Tensor::zeros(&[batch_size, ff_dim], device)?;

        Ok(Self {
            z,
            h,
            fourier_features,
            x_input,
            dW,
            db,
            dB,
            du_dx,
            d2u_dx2,
            // Backward intermediates
            h_sq,
            tanh_deriv,
            dL_dz,
            dL_dz_t,
            dL_dh_prev,
            d_loss,
            // Fourier backward intermediates
            x_scaled,
            Bx,
            cos_Bx,
            sin_Bx,
            dphi_sin_dB,
            dphi_cos_dB,
            dL_dphi_sin,
            dL_dphi_cos,
            dL_dB_from_sin,
            dL_dB_from_cos,
            // Metadata
            device: device.clone(),
            num_layers,
            batch_size,
            hidden_dim,
            ff_dim,
        })
    }

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

    /// Zero all gradients (call before backward pass)
    pub fn zero_grad(&mut self) -> Result<()> {
        for dw in &mut self.dW {
            *dw = Tensor::zeros(dw.shape(), &self.device)?;
        }
        for db in &mut self.db {
            *db = Tensor::zeros(db.shape(), &self.device)?;
        }
        self.dB = Tensor::zeros(self.dB.shape(), &self.device)?;
        Ok(())
    }
}

/// Compute the gradient of a single layer's backward pass
///
/// Given dL/dh (gradient w.r.t. layer output), computes:
/// - dL/dW (weight gradient)
/// - dL/db (bias gradient)
/// - dL/dh_prev (gradient to propagate backward)
///
/// For a layer with h = tanh(W @ h_prev + b):
/// - dL/dz = dL/dh * (1 - h²)  (tanh derivative)
/// - dL/dW = dL/dz.T @ h_prev
/// - dL/db = sum(dL/dz, axis=0)
/// - dL/dh_prev = dL/dz @ W
///
/// # Arguments
/// * `dL_dh` - Gradient w.r.t. layer output [batch, out_dim]
/// * `h` - Cached activation (tanh output) [batch, out_dim]
/// * `h_prev` - Previous layer activation (input to this layer) [batch, in_dim]
/// * `W` - Layer weights [out_dim, in_dim]
/// * `apply_tanh_grad` - Whether to apply tanh derivative (false for output layer)
///
/// # Returns
/// Tuple of (dL/dW, dL/db, dL/dh_prev)
pub fn layer_backward(
    dL_dh: &Tensor,
    h: &Tensor,
    h_prev: &Tensor,
    W: &Tensor,
    apply_tanh_grad: bool,
) -> Result<(Tensor, Tensor, Tensor)> {
    // Compute dL/dz
    let dL_dz = if apply_tanh_grad {
        // tanh'(z) = 1 - tanh²(z) = 1 - h²
        let h_sq = h.mul(h)?;
        let one_minus_h_sq = h_sq.mul_scalar(-1.0)?.add_scalar(1.0)?;
        dL_dh.mul(&one_minus_h_sq)?
    } else {
        // No activation (output layer)
        dL_dh.clone()
    };

    // dL/dW = dL/dz.T @ h_prev = [out_dim, batch] @ [batch, in_dim] = [out_dim, in_dim]
    let dL_dz_t = dL_dz.transpose(0, 1)?;
    let dL_dW = dL_dz_t.matmul(h_prev)?;

    // dL/db = sum(dL/dz, axis=0) = [out_dim]
    let dL_db = dL_dz.sum(Some(0))?;

    // dL/dh_prev = dL/dz @ W = [batch, out_dim] @ [out_dim, in_dim] = [batch, in_dim]
    let dL_dh_prev = dL_dz.matmul(W)?;

    Ok((dL_dW, dL_db, dL_dh_prev))
}

/// Compute gradient of Fourier features w.r.t. B
///
/// φ(x, B) = [sin(2πBx), cos(2πBx)]
/// ∂φ/∂B = [2πx * cos(2πBx), -2πx * sin(2πBx)]
///
/// # Arguments
/// * `dL_dphi` - Gradient w.r.t. Fourier features [batch, ff_dim * 2]
/// * `x` - Input positions [batch, 1]
/// * `B` - Fourier frequencies [1, ff_dim]
///
/// # Returns
/// dL/dB with shape [1, ff_dim]
pub fn fourier_backward(
    dL_dphi: &Tensor,
    x: &Tensor,
    B: &Tensor,
) -> Result<Tensor> {
    let scale = 2.0 * PI;
    let ff_dim = B.shape()[1] as usize;

    // Compute 2πBx
    let x_scaled = x.mul_scalar(scale)?; // [batch, 1]
    let Bx = x_scaled.matmul(B)?;        // [batch, ff_dim]

    // ∂sin(2πBx)/∂B = 2πx * cos(2πBx)
    // ∂cos(2πBx)/∂B = -2πx * sin(2πBx)
    let cos_Bx = Bx.cos()?;
    let sin_Bx = Bx.sin()?;

    // dphi/dB for sin part: 2πx * cos(2πBx)
    let dphi_sin_dB = x_scaled.mul(&cos_Bx)?; // Broadcasting: [batch, 1] * [batch, ff_dim]

    // dphi/dB for cos part: -2πx * sin(2πBx)
    let dphi_cos_dB = x_scaled.mul(&sin_Bx)?.mul_scalar(-1.0)?;

    // Split dL_dphi into sin and cos parts
    // dL_dphi has shape [batch, ff_dim * 2] = [batch, sin_part | cos_part]
    // Use slice(dim, start, end) API
    let dL_dphi_sin = dL_dphi.slice(1, 0, ff_dim)?;
    let dL_dphi_cos = dL_dphi.slice(1, ff_dim, ff_dim * 2)?;

    // Chain rule: dL/dB = dL/dphi_sin * dphi_sin/dB + dL/dphi_cos * dphi_cos/dB
    let dL_dB_from_sin = dL_dphi_sin.mul(&dphi_sin_dB)?;
    let dL_dB_from_cos = dL_dphi_cos.mul(&dphi_cos_dB)?;

    // Sum over batch dimension to get final gradient
    // sum(dim) where dim=0 for batch dimension
    let dL_dB = dL_dB_from_sin.add(&dL_dB_from_cos)?.sum(Some(0))?;

    Ok(dL_dB)
}

/// Compute MSE loss gradient
///
/// L = (1/N) Σ ||pred - target||²
/// dL/dpred = (2/N) * (pred - target)
///
/// # Arguments
/// * `pred` - Predictions [batch, 2]
/// * `target` - Targets [batch, 2]
///
/// # Returns
/// dL/dpred with shape [batch, 2]
pub fn mse_backward(pred: &Tensor, target: &Tensor) -> Result<Tensor> {
    let batch_size = pred.shape()[0] as f32;
    let diff = pred.sub(target)?;
    let grad = diff.mul_scalar(2.0 / batch_size)?;
    Ok(grad)
}

// =============================================================================
// GPU-NATIVE ZERO-ALLOCATION BACKWARD PASS
// =============================================================================
//
// These functions use PinnStreamContext GPU ops to avoid tensor allocations.
// All intermediate results are written to pre-allocated GradientWorkspace buffers.

#[cfg(feature = "cuda")]
use crate::cuda_stream_context::PinnStreamContext;

/// Compute MSE loss gradient into pre-allocated buffer (zero-alloc version)
///
/// # Arguments
/// * `pred` - Predictions [batch, 2]
/// * `target` - Targets [batch, 2]
/// * `ctx` - CUDA stream context for GPU ops
/// * `d_loss` - Pre-allocated output buffer [batch, 2]
#[cfg(feature = "cuda")]
pub fn mse_backward_inplace(
    pred: &Tensor,
    target: &Tensor,
    ctx: &PinnStreamContext,
    d_loss: &mut Tensor,
) -> Result<()> {
    let batch_size = pred.shape()[0] as f32;
    let scale = 2.0 / batch_size;

    // d_loss = pred - target
    ctx.sub_out(pred, target, d_loss)?;

    // d_loss = d_loss * scale (in-place via copy to same buffer)
    // Note: mul_scalar_out writes to output, we need in-place
    // For now, use existing mul_scalar_ if available
    d_loss.mul_scalar_(scale)?;

    Ok(())
}

/// Compute single layer backward pass using pre-allocated buffers (zero-alloc version)
///
/// # Arguments
/// * `dL_dh` - Gradient w.r.t. layer output [batch, out_dim]
/// * `h` - Cached activation (tanh output) [batch, out_dim]
/// * `h_prev` - Previous layer activation [batch, in_dim]
/// * `W` - Layer weights [out_dim, in_dim]
/// * `layer_idx` - Index of this layer (for accessing workspace buffers)
/// * `apply_tanh_grad` - Whether to apply tanh derivative
/// * `ctx` - CUDA stream context for GPU ops
/// * `ws` - GradientWorkspace with pre-allocated buffers
#[cfg(feature = "cuda")]
pub fn layer_backward_inplace(
    dL_dh: &Tensor,
    h: &Tensor,
    h_prev: &Tensor,
    W: &Tensor,
    layer_idx: usize,
    apply_tanh_grad: bool,
    ctx: &PinnStreamContext,
    ws: &mut GradientWorkspace,
) -> Result<()> {
    // Compute dL/dz into ws.dL_dz[layer_idx]
    if apply_tanh_grad {
        // dL/dz = dL/dh * (1 - h²)
        // Use fused tanh_deriv_mul kernel: out = (1 - h*h) * dL_dh
        ctx.tanh_deriv_mul_out(h, dL_dh, &mut ws.dL_dz[layer_idx])?;
    } else {
        // No activation - copy dL_dh to dL_dz
        // For output layer, dL_dz = dL_dh directly
        // We can use add with zero, or just assign (need copy)
        // Use a simple copy via mul_scalar(1.0)
        ctx.mul_scalar_out(dL_dh, 1.0, &mut ws.dL_dz[layer_idx])?;
    }

    // dL/dW = dL/dz.T @ h_prev
    // First transpose dL_dz: [batch, out] -> [out, batch]
    // Then matmul: [out, batch] @ [batch, in] -> [out, in]
    //
    // For matmul, we use the Tensor API since we have matmul_out
    let dL_dz_t = ws.dL_dz[layer_idx].transpose(0, 1)?;
    dL_dz_t.matmul_out(h_prev, &mut ws.dW[layer_idx])?;

    // dL/db = sum(dL/dz, axis=0)
    ctx.sum_rows_out(&ws.dL_dz[layer_idx], &mut ws.db[layer_idx])?;

    // dL/dh_prev = dL/dz @ W
    ws.dL_dz[layer_idx].matmul_out(W, &mut ws.dL_dh_prev[layer_idx])?;

    Ok(())
}

/// Compute Fourier backward gradient using fused GPU kernel (zero-alloc version)
///
/// This replaces the allocating `fourier_backward()` with a single fused CUDA kernel
/// that eliminates ~10 D2H copies per training step.
///
/// # Arguments
/// * `dL_dphi` - Gradient from layer 0 [batch, 2*ff_dim]
/// * `x` - Input coordinates [batch, 1]
/// * `B` - Fourier frequencies [1, ff_dim]
/// * `ctx` - CUDA stream context for GPU ops
/// * `dL_dB` - Pre-allocated output buffer [1, ff_dim]
#[cfg(feature = "cuda")]
pub fn fourier_backward_inplace(
    dL_dphi: &Tensor,
    x: &Tensor,
    B: &Tensor,
    ctx: &PinnStreamContext,
    dL_dB: &mut Tensor,
) -> Result<()> {
    use std::f32::consts::PI;

    // Zero the output buffer before atomic accumulation
    ctx.memset_zero_out(dL_dB)?;

    // Launch fused kernel: computes entire Fourier gradient in one pass
    ctx.fourier_grad(x, dL_dphi, B, dL_dB, 2.0 * PI)?;

    Ok(())
}

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

    #[test]
    fn test_gradient_workspace_creation() {
        let device = Device::Cpu;
        let ws = GradientWorkspace::new(32, 16, 64, 3, &device).unwrap();

        // Check shapes
        assert_eq!(ws.z.len(), 4); // 3 hidden + 1 output
        assert_eq!(ws.h.len(), 4);
        assert_eq!(ws.dW.len(), 4);
        assert_eq!(ws.db.len(), 4);

        // Check dimensions
        assert_eq!(ws.dW[0].shape(), &[64, 32]); // hidden_dim x ff_dim*2
        assert_eq!(ws.dW[1].shape(), &[64, 64]); // hidden x hidden
        assert_eq!(ws.dW[3].shape(), &[2, 64]);  // output x hidden
    }

    #[test]
    fn test_mse_backward() {
        let device = Device::Cpu;
        let pred = Tensor::from_vec(vec![1.0f32, 2.0, 3.0, 4.0], &[2, 2], &device).unwrap();
        let target = Tensor::from_vec(vec![0.0f32, 0.0, 0.0, 0.0], &[2, 2], &device).unwrap();

        let grad = mse_backward(&pred, &target).unwrap();

        // dL/dpred = (2/2) * (pred - 0) = pred
        let grad_cpu = grad.to_cpu().unwrap();
        assert!((grad_cpu[0] - 1.0).abs() < 1e-5);
        assert!((grad_cpu[1] - 2.0).abs() < 1e-5);
    }
}