//! Loss functions for Mamba2 SSM training
//!
//! Includes directional MSE that penalizes wrong-sign predictions
//! more heavily -- critical for trading where direction matters more
//! than magnitude.
//!
//! Uses GPU ops (`gpu_sub`, `gpu_sqr`, `gpu_mean_all`) instead of downloading
//! tensors to host. Only a single scalar result is returned.

use ml_core::{MLError, MLResult};

use crate::gpu_tensor::{gpu_abs, gpu_add, gpu_mean_all, gpu_mul, gpu_scale, gpu_sqr, gpu_sub, GpuTensor};

/// Standard MSE loss (baseline for comparison).
///
/// Computed entirely on GPU. Returns a scalar f32.
pub fn mse_loss(predictions: &GpuTensor, targets: &GpuTensor) -> MLResult<f32> {
    if predictions.shape != targets.shape {
        return Err(MLError::DimensionMismatch {
            expected: predictions.numel(),
            actual: targets.numel(),
        });
    }

    if predictions.numel() == 0 {
        return Ok(0.0);
    }

    // diff = pred - target
    let diff = gpu_sub(predictions, targets)?;
    // sq = diff^2
    let sq = gpu_sqr(&diff)?;
    // mean(sq) -- single scalar download
    gpu_mean_all(&sq)
}

/// Directional MSE loss.
///
/// When prediction and target have opposite signs (wrong direction),
/// the squared error is multiplied by `direction_weight` (default: 2.0).
/// When signs agree, normal MSE is used.
///
/// `direction_weight = 1.0` produces identical results to standard MSE.
/// `direction_weight = 2.0` penalizes wrong-direction errors 2x.
///
/// Implementation: `weight = 1 + (direction_weight - 1) * wrong_mask`
/// where `wrong_mask = 1` when signs differ, 0 when same.
///
/// Uses the identity: wrong direction means `pred * target < 0`.
/// Computes `sign_agree = sign(pred * target)` via `product / (|product| + eps)`,
/// then derives `wrong_mask = (1 - sign_agree) / 2`.
///
/// For production correctness, we compute both terms on GPU and combine.
pub fn directional_mse_loss(
    predictions: &GpuTensor,
    targets: &GpuTensor,
    direction_weight: f64,
) -> MLResult<f32> {
    if predictions.shape != targets.shape {
        return Err(MLError::DimensionMismatch {
            expected: predictions.numel(),
            actual: targets.numel(),
        });
    }

    if predictions.numel() == 0 {
        return Ok(0.0);
    }

    // When direction_weight == 1.0, directional loss == MSE
    if (direction_weight - 1.0).abs() < 1e-10 {
        return mse_loss(predictions, targets);
    }

    // diff = pred - target, sq = diff^2
    let diff = gpu_sub(predictions, targets)?;
    let sq = gpu_sqr(&diff)?;

    // product = pred * target (negative when signs differ)
    let product = gpu_mul(predictions, targets)?;
    // abs_product = |pred * target|
    let abs_product = gpu_abs(&product)?;
    // sign_indicator = product / |product| gives +1 when same sign, -1 when different
    // But we need to handle zeros. Use: wrong_mask = (1 - product / (|product| + eps)) / 2
    // wrong_mask ~ 1 when signs differ, ~ 0 when same
    let eps_tensor = GpuTensor::from_vec(
        vec![1e-7_f32; predictions.numel()],
        &predictions.shape,
        &predictions.stream,
    )?;
    let safe_abs = gpu_add(&abs_product, &eps_tensor)?;
    // ratio = product / safe_abs, in [-1, 1]
    let ratio = crate::gpu_tensor::gpu_div(&product, &safe_abs)?;
    // wrong_mask = (1 - ratio) / 2, in [0, 1]
    let one_minus_ratio = crate::gpu_tensor::gpu_scalar_sub(1.0, &ratio)?;
    let wrong_mask = gpu_scale(&one_minus_ratio, 0.5)?;

    // weight = 1 + (direction_weight - 1) * wrong_mask
    let extra_weight = direction_weight as f32 - 1.0;
    let penalty = gpu_scale(&wrong_mask, extra_weight)?;
    let ones = crate::gpu_tensor::gpu_ones(&predictions.shape, &predictions.stream)?;
    let weights = gpu_add(&ones, &penalty)?;

    // weighted_sq = sq * weights
    let weighted_sq = gpu_mul(&sq, &weights)?;
    gpu_mean_all(&weighted_sq)
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::sync::Arc;
    use cudarc::driver::{CudaContext, CudaStream};

    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_directional_loss_same_as_mse_when_correct_direction() {
        let stream = test_stream();
        let preds = GpuTensor::from_vec(vec![0.3_f32, 0.5, 0.8], &[3], &stream).unwrap();
        let targets = GpuTensor::from_vec(vec![0.2_f32, 0.6, 0.9], &[3], &stream).unwrap();

        let dir_loss = directional_mse_loss(&preds, &targets, 2.0).unwrap();
        let mse = mse_loss(&preds, &targets).unwrap();

        assert!(
            (dir_loss - mse).abs() < 1e-6,
            "same-direction: dir={} vs mse={}",
            dir_loss,
            mse
        );
    }

    #[test]
    fn test_directional_loss_higher_for_wrong_direction() {
        let stream = test_stream();
        let preds = GpuTensor::from_vec(vec![0.3_f32], &[1], &stream).unwrap();
        let targets = GpuTensor::from_vec(vec![-0.2_f32], &[1], &stream).unwrap();

        let dir_loss = directional_mse_loss(&preds, &targets, 2.0).unwrap();
        let mse = mse_loss(&preds, &targets).unwrap();

        assert!(
            dir_loss > mse,
            "wrong-direction loss {} should be > mse {}",
            dir_loss,
            mse
        );
    }

    #[test]
    fn test_directional_loss_weight_1_equals_mse() {
        let stream = test_stream();
        let preds = GpuTensor::from_vec(vec![0.3_f32, -0.1], &[2], &stream).unwrap();
        let targets = GpuTensor::from_vec(vec![-0.2_f32, 0.5], &[2], &stream).unwrap();

        let dir_loss = directional_mse_loss(&preds, &targets, 1.0).unwrap();
        let mse = mse_loss(&preds, &targets).unwrap();

        assert!(
            (dir_loss - mse).abs() < 1e-6,
            "weight=1.0: dir={} should == mse={}",
            dir_loss,
            mse
        );
    }
}
