use crate::deep::{Layer, Axon, LayerTy};
use crate::config::FLOAT;
use anyhow::*;

/// Simple macro for creating simple activation functions.
/// ### Example
/// Example implementation of an approximation of the sigmoid activation function.
/// ```
/// activation!{
///    Sigmoid = 
///        |x| x/(x + 1.),
///        |x| 1. / (x.powi(2) + 2. * x + 1.)
/// }
/// ```
#[macro_export]
macro_rules! activation{
    ($name: ident = $f: expr, $d: expr $(,)?) => {
        #[derive(Clone, Copy)]
        pub struct $name<const LEN: usize>;

        impl<const LEN: usize> LayerTy for $name<LEN>{
            type Gradient = [FLOAT; LEN];
            type Output   = [FLOAT; LEN];
        }

        impl<const LEN: usize> Layer<LEN, LEN> for $name<LEN>{
            fn predict(&self, i: impl Axon<LEN>) -> Result<[FLOAT; LEN]>{
                let mut buffer = [0.; LEN];
                for n in 0..LEN{
                    buffer[n] = $f(i[n])
                }
                Ok(buffer)
            }
            fn backpropagate(&mut self, i: impl Axon<LEN>, gradient: impl Axon<LEN>) -> Result<[FLOAT; LEN]>{
                let mut buffer = [0.; LEN];
                for n in 0..LEN{
                    buffer[n] = $d(i[n]) * gradient[n]
                }
                Ok(buffer)
            }
        }
    }
}

pub fn sigmoid(x: FLOAT) -> FLOAT{
    1. / (1. + (-x).exp())
}

activation!{
    Sigmoid =
        |x: FLOAT| sigmoid(x),
        |x: FLOAT| sigmoid(1. - x),
}

activation!{
    Tanh =
        |x: FLOAT| x.tanh(),
        |x: FLOAT| -x.tanh().powi(2) + 1.,
}

// x \cdot \sigma(x)
// \sigma(x) * (1 + x * (1 - \sigma(x)))
activation!{
    Swish =
        |x: FLOAT| x / ((-x).exp() + 1.),
        |x: FLOAT| {
            let sig = sigmoid(x);
            sig * (1. + x * (1. - sig))
        }
}

activation!{
    Relu =
        |x: FLOAT| x.abs(),
        |_| 1.,
}
