# Losses
from scipy.special import gammaln
import numpy as np
from numba import njit

@njit
def MAE(t, cal):
    total = 0.0
    for i in range(len(cal)):
        total += abs(cal[i] - t[i])
    return total / len(cal)

@njit
def MSE(t, cal):
    total = 0.0
    for i in range(len(cal)):
        diff = cal[i] - t[i]
        total += diff * diff
    return total / len(cal)

@njit
def RMSE(t, cal):
    total = 0.0
    for i in range(len(cal)):
        diff = cal[i] - t[i]
        total += diff * diff
    return np.sqrt(total / len(cal))

@njit
def R2(t, cal):
    mean_t = 0.0
    for i in range(len(t)):
        mean_t += t[i]
    mean_t /= len(t)

    ss_res = 0.0
    ss_tot = 0.0
    for i in range(len(t)):
        diff_res = t[i] - cal[i]
        diff_tot = t[i] - mean_t
        ss_res += diff_res * diff_res
        ss_tot += diff_tot * diff_tot

    return 1 - ss_res / ss_tot if ss_tot != 0 else 0.0

@njit
def log_cosh(t, cal):
    total = 0.0
    n = len(cal)
    for i in range(n):
        diff = cal[i] - t[i]
        total += diff * np.tanh(diff / 2.0)
    return total / n

@njit
def binary_cross_entropy(t, cal):
    total = 0.0
    n = len(cal)
    for i in range(n):
        x = cal[i]
        x = max(min(x, 50.0), -50.0)         # clip for stability
        prob = 1.0 / (1.0 + np.exp(-x))    # sigmoid
        prob = min(max(prob, 1e-16), 1 - 1e-16)  # clip to prevent log(0)
        total += -t[i] * np.log(prob) - (1 - t[i]) * np.log(1 - prob)
    return total / n

@njit
def exp_loss(t, cal):
    total = 0.0
    n = len(cal)
    for i in range(n):
        total += np.exp(-cal[i] * t[i])
    return total / n

@njit
def poisson_deviance(t, cal):
    eps = 1e-8
    total = 0.0
    n = len(cal)
    for i in range(n):
        c = cal[i]
        ti = t[i]
        if c < eps:
            c = eps
        total += ti * np.log((ti + eps) / c) - (ti - c)
    return total / n

losses_dict = {'MAE': MAE, 'MSE': MSE, 'RMSE': RMSE, 'R2': R2, 'LogCosh': log_cosh,
                'BinaryCrossEntropy': binary_cross_entropy, 'ExpLoss': exp_loss, 
                'PoissonDeviance': poisson_deviance}

def get_loss(name: str):
    if name not in losses_dict: 
        raise KeyError(f"'{name}' is not a valid loss name. Valid names are: {list(losses_dict.keys())}")
    loss_fn = losses_dict[name]
    cal = np.random.rand(1000)
    t = np.random.rand(1000)
    loss_fn(cal, t)
    return loss_fn
