//! Cross-Entropy Loss Implementation
//!
//! Provides efficient cross-entropy loss computation for language modeling tasks.
//! Supports both standard cross-entropy and sparse cross-entropy with label smoothing.

use crate::losses::{Loss, Reduction};
use crate::{Result, TransformerError};
use rtx_tensor::Tensor;
use std::f32::NEG_INFINITY;
use tracing::{debug, trace};

/// Cross-entropy loss configuration
#[derive(Debug, Clone)]
pub struct CrossEntropyConfig {
    /// Reduction method for loss computation
    pub reduction: Reduction,
    /// Label smoothing factor (0.0 = no smoothing, 0.1 = 10% smoothing)
    pub label_smoothing: f32,
    /// Whether to ignore index -100 (padding tokens)
    pub ignore_index: Option<i64>,
}

impl Default for CrossEntropyConfig {
    fn default() -> Self {
        Self {
            reduction: Reduction::Mean,
            label_smoothing: 0.0,
            ignore_index: Some(-100),
        }
    }
}

/// Cross-entropy loss implementation for language modeling
///
/// Computes the cross-entropy loss between predictions and targets:
/// loss = -sum(target * `log_softmax(predictions)`)
///
/// For language modeling, typically:
/// - predictions: [`batch_size`, `seq_len`, `vocab_size`] logits
/// - targets: [`batch_size`, `seq_len`] token indices
pub struct CrossEntropyLoss {
    config: CrossEntropyConfig,
}

impl CrossEntropyLoss {
    /// Create a new cross-entropy loss with default configuration
    #[must_use]
    pub fn new() -> Self {
        Self {
            config: CrossEntropyConfig::default(),
        }
    }

    /// Create a new cross-entropy loss with custom configuration
    #[must_use]
    pub fn with_config(config: CrossEntropyConfig) -> Self {
        Self { config }
    }

    /// Get the configuration
    #[must_use]
    pub fn config(&self) -> &CrossEntropyConfig {
        &self.config
    }

    /// Compute log-softmax of logits
    ///
    /// For numerical stability, we use the log-sum-exp trick:
    /// `log_softmax(x)` = x - log(sum(exp(x - max(x))))
    fn log_softmax(&self, logits: &Tensor) -> Result<Tensor> {
        let dims = logits.shape().dims();
        let vocab_dim = dims.len() - 1; // Last dimension is vocab_size

        // Find maximum along vocab dimension for numerical stability
        let max_logits = self.max_along_dim(logits, vocab_dim)?;

        // Subtract max for numerical stability: x - max(x)
        let shifted_logits = self.subtract_broadcasted(logits, &max_logits)?;

        // Compute exp(x - max(x))
        let exp_shifted = self.exp_tensor(&shifted_logits)?;

        // Sum along vocab dimension: sum(exp(x - max(x)))
        let sum_exp = self.sum_along_dim(&exp_shifted, vocab_dim)?;

        // Compute log(sum(exp(x - max(x))))
        let log_sum_exp = self.log_tensor(&sum_exp)?;

        // Return log_softmax: (x - max(x)) - log(sum(exp(x - max(x))))
        self.subtract_broadcasted(&shifted_logits, &log_sum_exp)
    }

    /// Compute sparse cross-entropy loss
    ///
    /// For each position (b, t), we compute:
    /// loss[b, t] = -`log_softmax`[b, t, target[b, t]]
    fn sparse_cross_entropy(&self, log_probs: &Tensor, targets: &Tensor) -> Result<Tensor> {
        let batch_size = log_probs.shape()[0];
        let seq_len = log_probs.shape()[1];
        let vocab_size = log_probs.shape()[2];

        // Convert to CPU for indexing operations
        let log_probs_cpu = log_probs.to_cpu()?;
        let targets_cpu = targets.to_cpu()?;

        let mut losses = Vec::with_capacity(batch_size * seq_len);

        for b in 0..batch_size {
            for t in 0..seq_len {
                // Manual indexing into flat vector: [b, t] -> b * seq_len + t
                let target_idx = targets_cpu[b * seq_len + t] as i64;

                let loss_value = if let Some(ignore_idx) = self.config.ignore_index {
                    if target_idx == ignore_idx {
                        0.0 // Ignore this position
                    } else {
                        // Manual indexing: [b, t, c] -> (b * seq_len + t) * vocab_size + c
                        let log_prob =
                            log_probs_cpu[(b * seq_len + t) * vocab_size + target_idx as usize];
                        -log_prob
                    }
                } else {
                    // Manual indexing: [b, t, c] -> (b * seq_len + t) * vocab_size + c
                    let log_prob =
                        log_probs_cpu[(b * seq_len + t) * vocab_size + target_idx as usize];
                    -log_prob
                };

                losses.push(loss_value);
            }
        }

        // Create tensor with per-token losses
        let device = log_probs.device();
        Tensor::from_data(losses, vec![batch_size, seq_len], device)
            .map_err(|e| TransformerError::tensor_op(e.to_string()))
    }

    /// Apply label smoothing
    ///
    /// Smoothed loss = (1 - α) * `ce_loss` + α * `uniform_loss`
    /// where α is the smoothing factor
    fn apply_label_smoothing(&self, ce_loss: &Tensor, log_probs: &Tensor) -> Result<Tensor> {
        if self.config.label_smoothing == 0.0 {
            return Ok(ce_loss.clone());
        }

        let vocab_size = log_probs.shape()[2] as f32;
        let smoothing = self.config.label_smoothing;

        // Uniform distribution loss: -mean(log_probs)
        let uniform_loss = self.mean_along_dim(log_probs, 2)?;
        let uniform_loss = self.negate_tensor(&uniform_loss)?;

        // Combine: (1 - smoothing) * ce_loss + smoothing * uniform_loss
        let ce_weight = 1.0 - smoothing;
        let uniform_weight = smoothing / vocab_size;

        let weighted_ce = self.scale_tensor(ce_loss, ce_weight)?;
        let weighted_uniform = self.scale_tensor(&uniform_loss, uniform_weight)?;

        self.add_tensors(&weighted_ce, &weighted_uniform)
    }

    /// Apply reduction to loss tensor
    fn apply_reduction(&self, loss: &Tensor) -> Result<Tensor> {
        match self.config.reduction {
            Reduction::None => Ok(loss.clone()),
            Reduction::Mean => {
                let total_elements = loss.numel() as f32;
                let sum_loss = self.sum_all(loss)?;
                self.scale_tensor(&sum_loss, 1.0 / total_elements)
            }
            Reduction::Sum => self.sum_all(loss),
        }
    }

    // Helper methods for tensor operations
    // These would ideally be implemented as proper tensor operations

    fn max_along_dim(&self, tensor: &Tensor, dim: usize) -> Result<Tensor> {
        // Simplified implementation - in practice, use proper tensor max operation
        let data = tensor.to_cpu()?;
        let shape = tensor.shape().dims();

        if dim >= shape.len() {
            return Err(TransformerError::tensor_op(
                "Invalid dimension for max operation".to_string(),
            ));
        }

        let reduced_shape: Vec<usize> = shape
            .iter()
            .enumerate()
            .filter_map(|(i, &size)| if i == dim { None } else { Some(size) })
            .collect();

        let max_val = data.iter().fold(NEG_INFINITY, |a, &b| a.max(b));
        let result_size = reduced_shape.iter().product();
        let result_data = vec![max_val; result_size];

        Tensor::from_data(result_data, reduced_shape, tensor.device())
            .map_err(|e| TransformerError::tensor_op(e.to_string()))
    }

    fn subtract_broadcasted(&self, a: &Tensor, b: &Tensor) -> Result<Tensor> {
        // Simplified - in practice, use proper broadcasting
        (a.clone() - b.clone()).map_err(|e| TransformerError::tensor_op(e.to_string()))
    }

    fn exp_tensor(&self, tensor: &Tensor) -> Result<Tensor> {
        tensor
            .exp()
            .map_err(|e| TransformerError::tensor_op(e.to_string()))
    }

    fn log_tensor(&self, tensor: &Tensor) -> Result<Tensor> {
        tensor
            .log()
            .map_err(|e| TransformerError::tensor_op(e.to_string()))
    }

    fn sum_along_dim(&self, tensor: &Tensor, dim: usize) -> Result<Tensor> {
        tensor
            .sum(Some(dim))
            .map_err(|e| TransformerError::tensor_op(e.to_string()))
    }

    fn mean_along_dim(&self, tensor: &Tensor, dim: usize) -> Result<Tensor> {
        let sum = self.sum_along_dim(tensor, dim)?;
        let size = tensor.shape()[dim] as f32;
        self.scale_tensor(&sum, 1.0 / size)
    }

    fn negate_tensor(&self, tensor: &Tensor) -> Result<Tensor> {
        tensor
            .mul_scalar(-1.0)
            .map_err(|e| TransformerError::tensor_op(e.to_string()))
    }

    fn scale_tensor(&self, tensor: &Tensor, scale: f32) -> Result<Tensor> {
        tensor
            .mul_scalar(scale)
            .map_err(|e| TransformerError::tensor_op(e.to_string()))
    }

    fn add_tensors(&self, a: &Tensor, b: &Tensor) -> Result<Tensor> {
        a.add(b)
            .map_err(|e| TransformerError::tensor_op(e.to_string()))
    }

    fn sum_all(&self, tensor: &Tensor) -> Result<Tensor> {
        tensor
            .sum(None)
            .map_err(|e| TransformerError::tensor_op(e.to_string()))
    }
}

impl Loss for CrossEntropyLoss {
    fn forward(&self, predictions: &Tensor, targets: &Tensor) -> Result<Tensor> {
        trace!(
            "Computing cross-entropy loss: predictions={:?}, targets={:?}",
            predictions.shape(),
            targets.shape()
        );

        // Validate input shapes
        let pred_shape = predictions.shape().dims();
        let target_shape = targets.shape().dims();

        if pred_shape.len() != 3 {
            return Err(TransformerError::tensor_op(format!(
                "Predictions must be 3D [batch, seq, vocab], got shape {pred_shape:?}"
            )));
        }

        if target_shape.len() != 2 {
            return Err(TransformerError::tensor_op(format!(
                "Targets must be 2D [batch, seq], got shape {target_shape:?}"
            )));
        }

        if pred_shape[0] != target_shape[0] || pred_shape[1] != target_shape[1] {
            return Err(TransformerError::tensor_op(format!(
                "Batch/sequence dimensions mismatch: pred {:?} vs target {:?}",
                &pred_shape[..2],
                target_shape
            )));
        }

        // Compute log-softmax
        let log_probs = self.log_softmax(predictions)?;

        // Compute sparse cross-entropy loss
        let mut ce_loss = self.sparse_cross_entropy(&log_probs, targets)?;

        // Apply label smoothing if configured
        if self.config.label_smoothing > 0.0 {
            ce_loss = self.apply_label_smoothing(&ce_loss, &log_probs)?;
        }

        // Apply reduction
        let final_loss = self.apply_reduction(&ce_loss)?;

        debug!(
            "Cross-entropy loss computed: reduction={:?}, value={:.4}",
            self.config.reduction,
            final_loss
                .to_cpu()
                .unwrap_or_else(|_| vec![0.0])
                .first()
                .copied()
                .unwrap_or(0.0)
        );

        Ok(final_loss)
    }

    fn reduction(&self) -> Reduction {
        self.config.reduction
    }

    fn name(&self) -> &'static str {
        "CrossEntropyLoss"
    }
}

#[cfg(all(test, feature = "disabled_tests"))]
mod tests {
    use super::*;
    use rtx_tensor::Device;

    #[test]
    fn test_cross_entropy_basic() {
        let device = Device::cuda(0).unwrap_or(Device::default());
        let loss_fn = CrossEntropyLoss::new();

        // Create simple test case
        let predictions = Tensor::from_data(
            vec![
                // Batch 1, Token 1: high prob for class 0
                1.0, 0.0, 0.0, // Batch 1, Token 2: high prob for class 1
                0.0, 1.0, 0.0,
            ],
            vec![1, 2, 3], // [batch=1, seq=2, vocab=3]
            &device,
        )
        .unwrap();

        let targets = Tensor::from_data(
            vec![0i64, 1i64], // Correct predictions
            vec![1, 2],       // [batch=1, seq=2]
            &device,
        )
        .unwrap();

        let loss = loss_fn.forward(&predictions, &targets).unwrap();
        let loss_value: f32 = loss.to_cpu().unwrap()[0];

        // Loss should be low for correct predictions
        assert!(loss_value >= 0.0, "Loss should be non-negative");
        assert!(
            loss_value < 2.0,
            "Loss should be reasonable for correct predictions"
        );
    }

    #[test]
    fn test_cross_entropy_shapes() {
        let device = Device::cuda(0).unwrap_or(Device::default());
        let loss_fn = CrossEntropyLoss::new();

        // Wrong prediction shape (2D instead of 3D)
        let predictions = Tensor::zeros(&[2, 3], &device).unwrap();
        let targets = Tensor::zeros_typed(&[2], DType::I64, &device).unwrap();

        let result = loss_fn.forward(&predictions, &targets);
        assert!(result.is_err(), "Should fail with wrong prediction shape");

        // Wrong target shape (3D instead of 2D)
        let predictions = Tensor::zeros(&[2, 3, 4], &device).unwrap();
        let targets = Tensor::zeros_typed(&[2, 3, 1], DType::I64, &device).unwrap();

        let result = loss_fn.forward(&predictions, &targets);
        assert!(result.is_err(), "Should fail with wrong target shape");
    }

    #[test]
    fn test_label_smoothing() {
        let device = Device::cuda(0).unwrap_or(Device::default());
        let config = CrossEntropyConfig {
            label_smoothing: 0.1,
            ..Default::default()
        };
        let loss_fn = CrossEntropyLoss::with_config(config);

        let predictions =
            Tensor::from_data(vec![1.0, 0.0, 0.0, 0.0, 1.0, 0.0], vec![1, 2, 3], &device).unwrap();

        let targets = Tensor::from_data(vec![0i64, 1i64], vec![1, 2], &device).unwrap();

        let loss = loss_fn.forward(&predictions, &targets);
        assert!(loss.is_ok(), "Label smoothing should work");
    }
}