//! Differentiable loss functions (MSE, CrossEntropy, etc.) used to drive
//! training by comparing predictions against ground-truth targets.

use crate::error::{TuphiumError, TuphiumResult};

/// Small epsilon used to clamp probabilities away from `0` and `1` so the
/// logarithms and divisions in the cross-entropy losses stay finite.
const EPS: f32 = 1e-7;

/// A differentiable loss function comparing predictions against targets.
pub trait Loss {
    /// Returns the scalar loss between targets `y_true` and predictions `y_pred`.
    ///
    /// Both slices must have the same length.
    fn forward(&self, y_true: &[f32], y_pred: &[f32]) -> TuphiumResult<f32>;

    /// Returns the gradient of the loss with respect to each prediction.
    ///
    /// The returned vector has the same length as `y_pred` and is fed into the
    /// model's backward pass.
    fn backward(&self, y_true: &[f32], y_pred: &[f32]) -> TuphiumResult<Vec<f32>>;
}

fn validate_loss_inputs(
    operation: &'static str,
    y_true: &[f32],
    y_pred: &[f32],
) -> TuphiumResult<()> {
    if y_true.len() != y_pred.len() {
        return Err(TuphiumError::ShapeMismatch {
            operation,
            expected: vec![y_true.len()],
            actual: vec![y_pred.len()],
        });
    }
    if y_true.is_empty() {
        return Err(TuphiumError::EmptyData {
            operation,
            target: "loss inputs".to_string(),
        });
    }
    if let Some(&value) = y_true
        .iter()
        .chain(y_pred.iter())
        .find(|value| !value.is_finite())
    {
        return Err(TuphiumError::NonFiniteValue { operation, value });
    }
    Ok(())
}

/// Clamps a probability into `[EPS, 1 - EPS]` so its logarithm and the
/// division by `p(1 - p)` remain finite.
fn clamp_probability(probability: f32) -> f32 {
    probability.clamp(EPS, 1.0 - EPS)
}

/// Binary Cross-Entropy. Expects sigmoid probabilities, averaged over outputs.
/// backward gives (p - t)/(p(1-p))/n; multiplied by sigmoid'(p)=p(1-p) inside
/// the output layer this recovers the classic (p - t), clamp-guarded.
pub struct BinaryCrossEntropy;
impl Loss for BinaryCrossEntropy {
    /// Mean binary cross-entropy. Predictions are clamped away from `0` and `1`
    /// to keep the logarithms finite.
    ///
    /// # Panics
    ///
    /// Panics if `y_true` and `y_pred` differ in length.
    fn forward(&self, y_true: &[f32], y_pred: &[f32]) -> TuphiumResult<f32> {
        validate_loss_inputs("BinaryCrossEntropy::forward", y_true, y_pred)?;
        let total: f32 = y_true
            .iter()
            .zip(y_pred)
            .map(|(&target, &prediction)| {
                let prediction = clamp_probability(prediction);
                -(target * prediction.ln() + (1.0 - target) * (1.0 - prediction).ln())
            })
            .sum();
        Ok(total / y_true.len() as f32)
    }
    /// Per-prediction gradient `(p - t) / (p(1 - p)) / n`, clamp-guarded.
    fn backward(&self, y_true: &[f32], y_pred: &[f32]) -> TuphiumResult<Vec<f32>> {
        validate_loss_inputs("BinaryCrossEntropy::backward", y_true, y_pred)?;
        let count = y_true.len() as f32;
        Ok(y_true
            .iter()
            .zip(y_pred)
            .map(|(&target, &prediction)| {
                let prediction = clamp_probability(prediction);
                ((prediction - target) / (prediction * (1.0 - prediction))) / count
            })
            .collect())
    }
}

/// Mean Squared Error.
pub struct MeanSquaredError;
impl Loss for MeanSquaredError {
    /// Mean of the squared differences between predictions and targets.
    ///
    /// # Panics
    ///
    /// Panics if `y_true` and `y_pred` differ in length.
    fn forward(&self, y_true: &[f32], y_pred: &[f32]) -> TuphiumResult<f32> {
        validate_loss_inputs("MeanSquaredError::forward", y_true, y_pred)?;
        let count = y_true.len() as f32;
        let total: f32 = y_true
            .iter()
            .zip(y_pred)
            .map(|(&target, &prediction)| {
                let difference = prediction - target;
                difference * difference
            })
            .sum();
        Ok(total / count)
    }
    /// Per-prediction gradient `2 (p - t) / n`.
    fn backward(&self, y_true: &[f32], y_pred: &[f32]) -> TuphiumResult<Vec<f32>> {
        validate_loss_inputs("MeanSquaredError::backward", y_true, y_pred)?;
        let count = y_true.len() as f32;
        Ok(y_true
            .iter()
            .zip(y_pred)
            .map(|(&target, &prediction)| 2.0 * (prediction - target) / count)
            .collect())
    }
}

/// Categorical Cross-Entropy. Expects probabilities (e.g. post-softmax) and a
/// one-hot target; summed over classes. Note: a fused Softmax+CE output layer
/// (giving the clean p - t gradient) is left for a later phase.
pub struct CrossEntropy;
impl Loss for CrossEntropy {
    /// Categorical cross-entropy `-Σ t_i · ln(p_i)`, with predictions clamped to
    /// keep the logarithm finite. Expects a one-hot `y_true`.
    ///
    /// # Panics
    ///
    /// Panics if `y_true` and `y_pred` differ in length.
    fn forward(&self, y_true: &[f32], y_pred: &[f32]) -> TuphiumResult<f32> {
        validate_loss_inputs("CrossEntropy::forward", y_true, y_pred)?;
        Ok(-y_true
            .iter()
            .zip(y_pred)
            .map(|(&target, &prediction)| target * clamp_probability(prediction).ln())
            .sum::<f32>())
    }
    /// Per-class gradient `-t_i / p_i`, clamp-guarded.
    fn backward(&self, y_true: &[f32], y_pred: &[f32]) -> TuphiumResult<Vec<f32>> {
        validate_loss_inputs("CrossEntropy::backward", y_true, y_pred)?;
        Ok(y_true
            .iter()
            .zip(y_pred)
            .map(|(&target, &prediction)| -target / clamp_probability(prediction))
            .collect())
    }
}
