//! Affine Coupling Layer for Flow-Based Policy
//!
//! Implements RealNVP-style affine coupling transformations with context conditioning.
//! All computation is GPU-native via `CudaLinear` + CUDA activation kernels.

use ml_core::MLError;

use crate::cuda_nn::{GpuContext, CudaLinear, cuda_tanh, cuda_from_slice};

/// Affine coupling layer with context conditioning (GPU-native).
pub(super) struct AffineCouplingLayer {
    /// Scale network layers
    scale_fc1: CudaLinear,
    scale_fc2: CudaLinear,
    /// Translation network layers
    trans_fc1: CudaLinear,
    trans_fc2: CudaLinear,
    /// Binary mask (host-side)
    mask: Vec<f32>,
    /// Scale clamping threshold
    scale_clamp: f32,
    /// Action dim
    action_dim: usize,
    /// GPU context
    ctx: GpuContext,
}

impl AffineCouplingLayer {
    /// Create a new affine coupling layer.
    pub(super) fn new(
        action_dim: usize,
        context_dim: usize,
        scale_clamp: f32,
        mask_vec: Vec<f32>,
        ctx: GpuContext,
    ) -> Result<Self, MLError> {
        if mask_vec.len() != action_dim {
            return Err(MLError::ModelError(format!(
                "Mask dimension {} does not match action dimension {}",
                mask_vec.len(), action_dim
            )));
        }

        let hidden_dim = action_dim + context_dim;

        let scale_fc1 = CudaLinear::new(ctx.clone(), action_dim + context_dim, hidden_dim)?;
        let scale_fc2 = CudaLinear::new(ctx.clone(), hidden_dim, action_dim)?;
        let trans_fc1 = CudaLinear::new(ctx.clone(), action_dim + context_dim, hidden_dim)?;
        let trans_fc2 = CudaLinear::new(ctx.clone(), hidden_dim, action_dim)?;

        Ok(Self {
            scale_fc1, scale_fc2, trans_fc1, trans_fc2,
            mask: mask_vec, scale_clamp, action_dim, ctx,
        })
    }

    /// Forward: x -> y (host-side, single sample)
    ///
    /// Returns (y, `log_det`) as host vectors.
    pub(super) fn forward_host(
        &self,
        x: &[f32],
        ctx_vec: &[f32],
        batch_size: usize,
    ) -> Result<(Vec<f32>, Vec<f32>), MLError> {
        // Apply mask
        let masked_x: Vec<f32> = x.iter().zip(self.mask.iter())
            .map(|(&xi, &mi)| xi * mi)
            .collect();

        // Concat masked_x + ctx
        let mut h_input = masked_x;
        h_input.extend_from_slice(ctx_vec);

        // Scale network: tanh(fc2(tanh(fc1(h))))
        let h_gpu = cuda_from_slice(&self.ctx.stream, &h_input)?;
        let s1 = self.scale_fc1.forward(&h_gpu.data, batch_size)?;
        let s1_act = cuda_tanh(&self.ctx.stream, &s1.data, s1.len)?;
        let s2 = self.scale_fc2.forward(&s1_act.data, batch_size)?;
        let s2_act = cuda_tanh(&self.ctx.stream, &s2.data, s2.len)?;
        let scale_host = s2_act.to_vec(&self.ctx.stream)?;

        // Clamp scale
        let scale_clamped: Vec<f32> = scale_host.iter()
            .map(|&s| s.clamp(-self.scale_clamp, self.scale_clamp))
            .collect();

        // Translation network: tanh(fc2(tanh(fc1(h))))
        let t1 = self.trans_fc1.forward(&h_gpu.data, batch_size)?;
        let t1_act = cuda_tanh(&self.ctx.stream, &t1.data, t1.len)?;
        let t2 = self.trans_fc2.forward(&t1_act.data, batch_size)?;
        let t2_act = cuda_tanh(&self.ctx.stream, &t2.data, t2.len)?;
        let translation = t2_act.to_vec(&self.ctx.stream)?;

        // Apply affine transformation: y = x*mask + (1-mask)*(x*exp(s) + t)
        let mut y = vec![0.0_f32; self.action_dim];
        let mut log_det = 0.0_f32;

        for i in 0..self.action_dim {
            let mi = self.mask.get(i).copied().unwrap_or(0.0);
            let si = scale_clamped.get(i).copied().unwrap_or(0.0);
            let ti = translation.get(i).copied().unwrap_or(0.0);
            let xi = x.get(i).copied().unwrap_or(0.0);

            let inv_mask = 1.0 - mi;
            let transformed = xi * si.exp() + ti;
            if let Some(yi) = y.get_mut(i) {
                *yi = xi * mi + inv_mask * transformed;
            }
            log_det += inv_mask * si;
        }

        Ok((y, vec![log_det]))
    }

    /// Inverse: y -> x (host-side, single sample)
    pub(super) fn inverse_host(
        &self,
        y: &[f32],
        ctx_vec: &[f32],
        batch_size: usize,
    ) -> Result<(Vec<f32>, Vec<f32>), MLError> {
        // Apply mask to y
        let masked_y: Vec<f32> = y.iter().zip(self.mask.iter())
            .map(|(&yi, &mi)| yi * mi)
            .collect();

        // Concat
        let mut h_input = masked_y;
        h_input.extend_from_slice(ctx_vec);

        // Scale network
        let h_gpu = cuda_from_slice(&self.ctx.stream, &h_input)?;
        let s1 = self.scale_fc1.forward(&h_gpu.data, batch_size)?;
        let s1_act = cuda_tanh(&self.ctx.stream, &s1.data, s1.len)?;
        let s2 = self.scale_fc2.forward(&s1_act.data, batch_size)?;
        let s2_act = cuda_tanh(&self.ctx.stream, &s2.data, s2.len)?;
        let scale_host = s2_act.to_vec(&self.ctx.stream)?;

        let scale_clamped: Vec<f32> = scale_host.iter()
            .map(|&s| s.clamp(-self.scale_clamp, self.scale_clamp))
            .collect();

        // Translation network
        let t1 = self.trans_fc1.forward(&h_gpu.data, batch_size)?;
        let t1_act = cuda_tanh(&self.ctx.stream, &t1.data, t1.len)?;
        let t2 = self.trans_fc2.forward(&t1_act.data, batch_size)?;
        let t2_act = cuda_tanh(&self.ctx.stream, &t2.data, t2.len)?;
        let translation = t2_act.to_vec(&self.ctx.stream)?;

        // Invert: x = (y - t) / exp(s)
        let mut x = vec![0.0_f32; self.action_dim];
        let mut log_det_inv = 0.0_f32;

        for i in 0..self.action_dim {
            let mi = self.mask.get(i).copied().unwrap_or(0.0);
            let si = scale_clamped.get(i).copied().unwrap_or(0.0);
            let ti = translation.get(i).copied().unwrap_or(0.0);
            let yi = y.get(i).copied().unwrap_or(0.0);

            let inv_mask = 1.0 - mi;
            let inv_transformed = (yi - ti) / si.exp();
            if let Some(xi) = x.get_mut(i) {
                *xi = yi * mi + inv_mask * inv_transformed;
            }
            log_det_inv -= inv_mask * si;
        }

        Ok((x, vec![log_det_inv]))
    }
}
