mod coupling_layer;
#[cfg(test)]
mod flow_matching;

use coupling_layer::AffineCouplingLayer;

use rand::thread_rng;
use rand_distr::{Distribution, Normal};
use serde::{Deserialize, Serialize};

use ml_core::MLError;

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

/// Configuration for the normalizing flow policy network.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct FlowPolicyConfig {
    /// Dimensionality of the input state.
    pub state_dim: usize,
    /// Dimensionality of the action space.
    pub action_dim: usize,
    /// Dimensionality of the context encoding.
    pub context_dim: usize,
    /// Number of affine coupling layers.
    pub num_layers: usize,
    /// Clamping threshold for scale parameters.
    pub scale_clamp: f32,
}

impl Default for FlowPolicyConfig {
    fn default() -> Self {
        Self {
            state_dim: 225,
            action_dim: 1,
            context_dim: 128,
            num_layers: 4,
            scale_clamp: 5.0,
        }
    }
}

/// Normalizing flow-based policy network for continuous action spaces.
///
/// All computation is GPU-native via `cuda_nn` primitives.
pub struct FlowPolicy {
    /// Context encoder
    context_enc: CudaLinear,
    /// Stack of affine coupling layers
    layers: Vec<AffineCouplingLayer>,
    /// Configuration
    config: FlowPolicyConfig,
    /// GPU context
    ctx: GpuContext,
}

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

impl FlowPolicy {
    /// Creates a new `FlowPolicy`.
    pub fn new(config: FlowPolicyConfig) -> Result<Self, MLError> {
        let ctx = GpuContext::new()?;

        let context_enc = CudaLinear::new(ctx.clone(), config.state_dim, config.context_dim)?;

        let mut layers = Vec::with_capacity(config.num_layers);
        for i in 0..config.num_layers {
            let layer = AffineCouplingLayer::new(
                config.action_dim,
                config.context_dim,
                config.scale_clamp,
                Self::make_mask(config.action_dim, i),
                ctx.clone(),
            )?;
            layers.push(layer);
        }

        Ok(Self {
            context_enc,
            layers,
            config,
            ctx,
        })
    }

    /// Sample action from the flow policy (host-side, single sample).
    ///
    /// Returns (action, `log_prob`) as host scalars.
    pub fn sample_action_host(&self, state: &[f32]) -> Result<(Vec<f32>, Vec<f32>), MLError> {
        let batch_size = 1;

        // Encode context
        let ctx_vec = self.encode_context_host(state, batch_size)?;

        // Sample z ~ N(0, 1)
        let z = self.sample_base_noise();

        // Flow forward
        let (y, log_det_flow) = self.flow_forward_host(&z, &ctx_vec, batch_size)?;

        // Squash: a = tanh(y)
        let a: Vec<f32> = y.iter().map(|&yi| yi.tanh()).collect();

        // Log-det correction for tanh squashing
        let log_det_squash: f32 = a.iter().map(|&ai| {
            let ai_clamped = ai.clamp(-0.999, 0.999);
            ((1.0 - ai_clamped) * (1.0 + ai_clamped)).ln()
        }).sum();

        // log p(z) = -0.5 * (z^2 + log(2pi))
        let log_2pi = (2.0_f64 * std::f64::consts::PI).ln() as f32;
        let log_pz: f32 = z.iter().map(|&zi| -0.5 * (zi * zi + log_2pi)).sum();

        let log_det_flow_scalar = log_det_flow.first().copied().unwrap_or(0.0);
        let log_prob = log_pz - log_det_flow_scalar - log_det_squash;

        Ok((a, vec![log_prob]))
    }

    /// Evaluate log probabilities for given state-action pairs (host-side).
    pub fn evaluate_actions_host(
        &self,
        state: &[f32],
        action: &[f32],
    ) -> Result<f32, MLError> {
        let batch_size = 1;
        let ctx_vec = self.encode_context_host(state, batch_size)?;

        // Unsquash: y = atanh(a) = 0.5 * ln((1+a)/(1-a))
        let y: Vec<f32> = action.iter().map(|&ai| {
            let ai_clamped = ai.clamp(-0.999, 0.999);
            0.5 * ((1.0 + ai_clamped) / (1.0 - ai_clamped)).ln()
        }).collect();

        // Inverse flow
        let (z, log_det_flow) = self.flow_inverse_host(&y, &ctx_vec, batch_size)?;

        // Log-det correction for tanh
        let a_clamped: Vec<f32> = action.iter().map(|&ai| ai.clamp(-0.999, 0.999)).collect();
        let log_det_squash: f32 = a_clamped.iter().map(|&ai| {
            ((1.0 - ai) * (1.0 + ai)).ln()
        }).sum();

        // log p(z)
        let log_2pi = (2.0_f64 * std::f64::consts::PI).ln() as f32;
        let log_pz: f32 = z.iter().map(|&zi| -0.5 * (zi * zi + log_2pi)).sum();

        let log_det_flow_scalar = log_det_flow.first().copied().unwrap_or(0.0);
        Ok(log_pz - log_det_flow_scalar - log_det_squash)
    }

    /// Entropy estimate via Monte Carlo (host-side).
    pub fn entropy_host(&self, state: &[f32]) -> Result<f32, MLError> {
        let (_action, log_probs) = self.sample_action_host(state)?;
        let log_prob = log_probs.first().copied().unwrap_or(0.0);
        Ok(-log_prob)
    }

    fn encode_context_host(&self, state: &[f32], batch_size: usize) -> Result<Vec<f32>, MLError> {
        let state_gpu = cuda_from_slice(&self.ctx.stream, state)?;
        let encoded = self.context_enc.forward(&state_gpu.data, batch_size)?;
        let tanh_out = cuda_tanh(&self.ctx.stream, &encoded.data, encoded.len)?;
        tanh_out.to_vec(&self.ctx.stream)
    }

    fn flow_forward_host(
        &self,
        z: &[f32],
        ctx: &[f32],
        batch_size: usize,
    ) -> Result<(Vec<f32>, Vec<f32>), MLError> {
        let mut x = z.to_vec(); // cpu-side: FixedPoint slice clone
        let mut total_log_det = 0.0_f32;

        for layer in &self.layers {
            let (x_new, log_det) = layer.forward_host(&x, ctx, batch_size)?;
            total_log_det += log_det.first().copied().unwrap_or(0.0);
            x = x_new;
        }

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

    fn flow_inverse_host(
        &self,
        y: &[f32],
        ctx: &[f32],
        batch_size: usize,
    ) -> Result<(Vec<f32>, Vec<f32>), MLError> {
        let mut x = y.to_vec(); // cpu-side: FixedPoint slice clone
        let mut total_log_det = 0.0_f32;

        for layer in self.layers.iter().rev() {
            let (x_new, log_det) = layer.inverse_host(&x, ctx, batch_size)?;
            total_log_det += log_det.first().copied().unwrap_or(0.0);
            x = x_new;
        }

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

    fn sample_base_noise(&self) -> Vec<f32> {
        // SAFETY: Normal::new(0.0, 1.0) is infallible (std > 0), so this never panics.
        #[allow(clippy::unwrap_used)]
        let normal = Normal::new(0.0_f32, 1.0_f32).unwrap_or_else(|_| Normal::new(0.0, 1.0).unwrap());
        let mut rng = thread_rng();
        (0..self.config.action_dim)
            .map(|_| normal.sample(&mut rng))
            .collect()
    }

    fn make_mask(action_dim: usize, layer_idx: usize) -> Vec<f32> {
        let mask_value = if layer_idx % 2 == 0 { 0.0_f32 } else { 1.0_f32 };
        vec![mask_value; action_dim]
    }

    /// Returns reference to configuration.
    pub const fn config(&self) -> &FlowPolicyConfig {
        &self.config
    }
}

#[cfg(test)]
#[allow(clippy::manual_range_contains, clippy::unnecessary_wraps)]
mod tests {
    use super::*;

    #[test]
    fn test_flow_policy_creation() -> Result<(), MLError> {
        let config = FlowPolicyConfig::default();
        let policy = FlowPolicy::new(config)?;

        assert_eq!(policy.config.state_dim, 225);
        assert_eq!(policy.config.action_dim, 1);
        assert_eq!(policy.config.num_layers, 4);
        assert_eq!(policy.layers.len(), 4);

        Ok(())
    }

    #[test]
    fn test_sample_action() -> Result<(), MLError> {
        let config = FlowPolicyConfig {
            state_dim: 8,
            action_dim: 1,
            context_dim: 16,
            num_layers: 2,
            scale_clamp: 5.0,
        };
        let policy = FlowPolicy::new(config)?;

        let state = vec![0.1_f32; 8];
        let (action, log_prob) = policy.sample_action_host(&state)?;

        assert_eq!(action.len(), 1);
        assert_eq!(log_prob.len(), 1);

        let a = action.first().copied().unwrap_or(0.0);
        assert!(a >= -1.0 && a <= 1.0, "Action {} out of bounds", a);

        let lp = log_prob.first().copied().unwrap_or(0.0);
        assert!(lp.is_finite(), "Log prob must be finite, got {}", lp);

        Ok(())
    }

    #[test]
    fn test_mask_alternation() -> Result<(), MLError> {
        let mask0 = FlowPolicy::make_mask(1, 0);
        let mask1 = FlowPolicy::make_mask(1, 1);

        assert_eq!(mask0[0], 0.0);
        assert_eq!(mask1[0], 1.0);

        Ok(())
    }
}
