%%cython -3
# distutils: extra_compile_args = -O3
# cython: cdivision=True
# cython: boundscheck=False
# cython: wraparound=False


import cython
from cython.parallel import prange
import numpy as np
from libc.math cimport exp, log, log1p
cimport numpy as np

np.import_array()


# Numerically stable log(1 + exp(x))
cdef inline double log1pexp(double x) nogil:
    if x <= -37:
        return exp(x)
    elif x <= 18:
        return log1p(exp(x))
    elif x <= 33.3:
        return x + exp(-x)
    else:
        return x


# Faster version of numerically stable log(1 + exp(x))
cdef inline double log1pexp_fast(double x) nogil:
    if x <= -37:
        return exp(x)
    elif x <= -2:
        return log1p(exp(x))
    elif x <= 18:
        return log(1 + exp(x))
    elif x <= 33.3:
        return x + exp(-x)
    else:
        return x



cdef inline double c_logloss(double y_true, double raw) nogil:
    return log1pexp(raw) - y_true * raw
    

cdef inline double c_logloss_fast(double y_true, double raw) nogil:
    return log1pexp_fast(raw) - y_true * raw


def cy_logloss_stable(double[::1] y_true, double[::1] raw):
    cdef:
        int n_samples
        int i
        double[::1] out = np.empty_like(y_true)
    
    n_samples = y_true.shape[0]
    for i in range(n_samples):
        out[i] = c_logloss(y_true[i], raw[i])
        
    return np.asarray(out)


def cy_logloss_stable_fast(double[::1] y_true, double[::1] raw):
    cdef:
        int n_samples
        int i
        double[::1] out = np.empty_like(y_true)
    
    n_samples = y_true.shape[0]
    for i in range(n_samples):
        out[i] = c_logloss_fast(y_true[i], raw[i])
        
    return np.asarray(out)