//! Flow Matching Objective for Flow-Based Policy Optimization (FPO)
//!
//! Computes the clipped surrogate objective for PPO using flow-based log-determinants.
//! All operations are host-side (CPU) since the autograd backend has been removed.

use ml_core::MLError;

/// Configuration for the Flow Matching loss function.
#[derive(Debug, Clone, Copy)]
pub(super) struct FlowMatchingConfig {
    /// The clipping parameter (epsilon) for the PPO surrogate objective.
    pub clip_epsilon: f32,
    /// Minimum value for clipping the log flow ratio.
    pub log_ratio_clip_min: f32,
    /// Maximum value for clipping the log flow ratio.
    pub log_ratio_clip_max: f32,
}

impl Default for FlowMatchingConfig {
    fn default() -> Self {
        Self {
            clip_epsilon: 0.2,
            log_ratio_clip_min: -20.0,
            log_ratio_clip_max: 20.0,
        }
    }
}

/// Computes the PPO clipped surrogate loss using flow matching ratio (host-side).
///
/// Returns a scalar loss value.
pub(super) fn compute_flow_matching_loss_host(
    new_log_dets: &[f32],
    old_log_dets: &[f32],
    advantages: &[f32],
    config: &FlowMatchingConfig,
) -> Result<f32, MLError> {
    if new_log_dets.len() != old_log_dets.len() || new_log_dets.len() != advantages.len() {
        return Err(MLError::InvalidInput(
            "Batch dimensions must match".to_owned(),
        ));
    }

    let batch_size = new_log_dets.len();
    if batch_size == 0 {
        return Err(MLError::InvalidInput("Empty batch".to_owned()));
    }

    let mut sum = 0.0_f32;

    for i in 0..batch_size {
        let new_ld = new_log_dets.get(i).copied().unwrap_or(0.0);
        let old_ld = old_log_dets.get(i).copied().unwrap_or(0.0);
        let adv = advantages.get(i).copied().unwrap_or(0.0);

        let log_ratio = (new_ld - old_ld).clamp(config.log_ratio_clip_min, config.log_ratio_clip_max);
        let ratio = log_ratio.exp();

        let clip_min = 1.0 - config.clip_epsilon;
        let clip_max = 1.0 + config.clip_epsilon;
        let clipped_ratio = ratio.clamp(clip_min, clip_max);

        let surr1 = ratio * adv;
        let surr2 = clipped_ratio * adv;
        sum += surr1.min(surr2);
    }

    let mean = sum / batch_size as f32;
    Ok(-mean) // Negate for minimization
}

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

    #[test]
    fn test_flow_matching_loss_computation() -> Result<(), MLError> {
        let config = FlowMatchingConfig::default();

        let new_log_dets = vec![1.2_f32, 0.8, -0.5, 2.0];
        let old_log_dets = vec![1.0_f32, 1.0, -0.4, 1.5];
        let advantages = vec![1.5_f32, -0.5, 2.0, 0.8];

        let loss = compute_flow_matching_loss_host(&new_log_dets, &old_log_dets, &advantages, &config)?;

        // Result should be around -1.04
        assert!((loss - (-1.040075)).abs() < 0.02, "Expected ~-1.04, got {}", loss);

        Ok(())
    }

    #[test]
    fn test_ratio_clipping() -> Result<(), MLError> {
        let config = FlowMatchingConfig {
            clip_epsilon: 0.2,
            ..Default::default()
        };

        let new_log_dets = vec![2.0_f32];
        let old_log_dets = vec![0.0_f32];
        let advantages = vec![10.0_f32];

        let loss = compute_flow_matching_loss_host(&new_log_dets, &old_log_dets, &advantages, &config)?;
        assert!((loss - (-12.0)).abs() < 0.01, "Expected -12.0, got {}", loss);

        Ok(())
    }
}
