//! Activation Functions for Liquid Neural Networks
//!
//! Fixed-point implementations of activation functions optimized for ultra-low latency.

use super::{FixedPoint, LiquidError, Result, PRECISION};
use serde::{Deserialize, Serialize};

/// Activation function types
#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq)]
pub enum ActivationType {
    Sigmoid,
    Tanh,
    ReLU,
    LeakyReLU,
    Swish,
    GELU,
    Linear,
}

/// Fast sigmoid approximation using fixed-point arithmetic
/// Uses polynomial approximation for ultra-low latency
pub fn sigmoid(x: FixedPoint) -> Result<FixedPoint> {
    // Clamp input to prevent overflow
    let clamped_x = if x.0 > 8 * PRECISION {
        FixedPoint(8 * PRECISION)
    } else if x.0 < -8 * PRECISION {
        FixedPoint(-8 * PRECISION)
    } else {
        x
    };

    // Fast sigmoid approximation: 0.5 * (x / (1 + abs(x))) + 0.5
    let abs_x = FixedPoint(clamped_x.0.abs());
    let one = FixedPoint::one();
    let half = FixedPoint(PRECISION / 2);

    let denominator = (one + abs_x)?;
    let fraction = (clamped_x / denominator)?;
    

    (half * fraction)? + half
}

/// Fast tanh approximation using fixed-point arithmetic
pub fn tanh(x: FixedPoint) -> Result<FixedPoint> {
    // Clamp input to prevent overflow
    let clamped_x = if x.0 > 4 * PRECISION {
        return Ok(FixedPoint::one());
    } else if x.0 < -4 * PRECISION {
        return Ok(FixedPoint(-PRECISION));
    } else {
        x
    };

    // Fast tanh approximation: x / (1 + abs(x))
    let abs_x = FixedPoint(clamped_x.0.abs());
    let one = FixedPoint::one();
    let denominator = (one + abs_x)?;
    let result = (clamped_x / denominator)?;

    Ok(result)
}

/// `ReLU` activation function
pub const fn relu(x: FixedPoint) -> FixedPoint {
    if x.0 > 0 {
        x
    } else {
        FixedPoint::zero()
    }
}

/// Leaky `ReLU` activation function
pub fn leaky_relu(x: FixedPoint, alpha: FixedPoint) -> Result<FixedPoint> {
    if x.0 > 0 {
        Ok(x)
    } else {
        alpha * x
    }
}

/// Swish activation function: x * sigmoid(x)
pub fn swish(x: FixedPoint) -> Result<FixedPoint> {
    let sig_x = sigmoid(x)?;
    x * sig_x
}

/// GELU approximation using tanh
pub fn gelu(x: FixedPoint) -> Result<FixedPoint> {
    let half = FixedPoint(PRECISION / 2);
    let sqrt_2_over_pi = FixedPoint((0.7978845608 * PRECISION as f64) as i64); // sqrt(2/π)
    let coeff = FixedPoint((0.044715 * PRECISION as f64) as i64); // 0.044715

    // x³ calculation
    let x_squared = (x * x)?;
    let x_cubed = (x_squared * x)?;

    // 0.044715 * x³
    let cubic_term = (coeff * x_cubed)?;

    // x + 0.044715 * x³
    let inner_sum = (x + cubic_term)?;

    // sqrt(2/π) * (x + 0.044715 * x³)
    let tanh_input = (sqrt_2_over_pi * inner_sum)?;

    // tanh(sqrt(2/π) * (x + 0.044715 * x³))
    let tanh_result = tanh(tanh_input)?;

    // 1 + tanh(...)
    let one_plus_tanh = (FixedPoint::one() + tanh_result)?;

    // 0.5 * x * (1 + tanh(...))
    
    (half * x)? * one_plus_tanh
}

/// Linear activation (identity function)
pub const fn linear(x: FixedPoint) -> FixedPoint {
    x
}

/// Apply activation function based on type
pub fn apply_activation(x: FixedPoint, activation_type: ActivationType) -> Result<FixedPoint> {
    match activation_type {
        ActivationType::Sigmoid => sigmoid(x),
        ActivationType::Tanh => tanh(x),
        ActivationType::ReLU => Ok(relu(x)),
        ActivationType::LeakyReLU => leaky_relu(x, FixedPoint(PRECISION / 100)), // α = 0.01
        ActivationType::Swish => swish(x),
        ActivationType::GELU => gelu(x),
        ActivationType::Linear => Ok(linear(x)),
    }
}

/// Activation function derivatives for backpropagation
pub mod derivatives {
    use super::{FixedPoint, Result, sigmoid, tanh, ActivationType, PRECISION, LiquidError};

    /// Sigmoid derivative: σ(x) * (1 - σ(x))
    pub fn sigmoid_derivative(x: FixedPoint) -> Result<FixedPoint> {
        let sig_x = sigmoid(x)?;
        let one_minus_sig = (FixedPoint::one() - sig_x)?;
        sig_x * one_minus_sig
    }

    /// Tanh derivative: 1 - tanh²(x)
    pub fn tanh_derivative(x: FixedPoint) -> Result<FixedPoint> {
        let tanh_x = tanh(x)?;
        let tanh_squared = (tanh_x * tanh_x)?;
        FixedPoint::one() - tanh_squared
    }

    /// `ReLU` derivative
    pub const fn relu_derivative(x: FixedPoint) -> FixedPoint {
        if x.0 > 0 {
            FixedPoint::one()
        } else {
            FixedPoint::zero()
        }
    }

    /// Leaky `ReLU` derivative
    pub const fn leaky_relu_derivative(x: FixedPoint, alpha: FixedPoint) -> FixedPoint {
        if x.0 > 0 {
            FixedPoint::one()
        } else {
            alpha
        }
    }

    /// Apply activation derivative based on type
    pub fn apply_activation_derivative(
        x: FixedPoint,
        activation_type: ActivationType,
    ) -> Result<FixedPoint> {
        match activation_type {
            ActivationType::Sigmoid => sigmoid_derivative(x),
            ActivationType::Tanh => tanh_derivative(x),
            ActivationType::ReLU => Ok(relu_derivative(x)),
            ActivationType::LeakyReLU => Ok(leaky_relu_derivative(x, FixedPoint(PRECISION / 100))),
            ActivationType::Linear => Ok(FixedPoint::one()),
            ActivationType::Swish | ActivationType::GELU => {
                Err(LiquidError::InvalidConfiguration(
                    "Derivative not implemented for this activation".to_owned(),
                ))
            },
        }
    }
}

#[cfg(test)]
#[allow(clippy::unnecessary_wraps)]
mod tests {
    use super::*;
    use crate::liquid::PRECISION;


    #[test]
    fn test_sigmoid() -> Result<()> {
        let zero = FixedPoint::zero();
        let result = sigmoid(zero)?;
        // sigmoid(0) should be approximately 0.5
        assert!((result.to_f64() - 0.5).abs() < 0.1);

        let positive = FixedPoint(2 * PRECISION);
        let result = sigmoid(positive)?;
        assert!(result.to_f64() > 0.5);

        let negative = FixedPoint(-2 * PRECISION);
        let result = sigmoid(negative)?;
        assert!(result.to_f64() < 0.5);
        Ok(())
    }

    #[test]
    fn test_tanh() -> Result<()> {
        let zero = FixedPoint::zero();
        let result = tanh(zero)?;
        // tanh(0) should be approximately 0
        assert!(result.to_f64().abs() < 0.1);

        let positive = FixedPoint(PRECISION);
        let result = tanh(positive)?;
        assert!(result.to_f64() > 0.0);

        let negative = FixedPoint(-PRECISION);
        let result = tanh(negative)?;
        assert!(result.to_f64() < 0.0);
        Ok(())
    }

    #[test]
    fn test_relu() -> Result<()> {
        let positive = FixedPoint(PRECISION);
        let result = relu(positive);
        assert_eq!(result.0, PRECISION);

        let negative = FixedPoint(-PRECISION);
        let result = relu(negative);
        assert_eq!(result.0, 0);

        let zero = FixedPoint::zero();
        let result = relu(zero);
        assert_eq!(result.0, 0);
        Ok(())
    }

    #[test]
    fn test_leaky_relu() -> Result<()> {
        let alpha = FixedPoint(PRECISION / 100); // 0.01

        let positive = FixedPoint(PRECISION);
        let result = leaky_relu(positive, alpha)?;
        assert_eq!(result.0, PRECISION);

        let negative = FixedPoint(-PRECISION);
        let result = leaky_relu(negative, alpha)?;
        assert_eq!(result.0, -PRECISION / 100);
        Ok(())
    }

    #[test]
    fn test_activation_derivatives() -> Result<()> {
        let x = FixedPoint(PRECISION / 2); // 0.5

        let sig_deriv = derivatives::sigmoid_derivative(x)?;
        assert!(sig_deriv.to_f64() > 0.0);

        let tanh_deriv = derivatives::tanh_derivative(x)?;
        assert!(tanh_deriv.to_f64() > 0.0);

        let relu_deriv = derivatives::relu_derivative(x);
        assert_eq!(relu_deriv.0, PRECISION);
        Ok(())
    }
}
