#include <gtest/gtest.h>
#include <tuple>

#include "grotto/closed_form.hpp"

#include <cmath>
#include <cstdint>

namespace
{

std::int64_t raw_of(long double x, unsigned k)
{
    return std::llround(std::ldexp(x, static_cast<int>(k)));
}

long double truth(grotto::closed which, long double x)
{
    switch (which)
    {
    case grotto::closed::atanh: return std::atanh(x);
    case grotto::closed::asinh: return std::asinh(x);
    case grotto::closed::acosh: return std::acosh(x);
    case grotto::closed::atan: return std::atan(x);
    case grotto::closed::acot: return 1.57079632679489661923L - std::atan(x);
    case grotto::closed::asec: return std::acos(1.0L / x);
    case grotto::closed::acsc: return std::asin(1.0L / x);
    case grotto::closed::asech: return std::acosh(1.0L / x);
    case grotto::closed::acsch: return std::asinh(1.0L / x);
    case grotto::closed::acoth: return std::atanh(1.0L / x);
    case grotto::closed::selu:
        return 1.0507009873554804934L * (x > 0 ? x : 1.6732632423543772848L * std::expm1(x));
    case grotto::closed::elu:
    case grotto::closed::celu:
        return x > 0 ? x : std::expm1(x);
    case grotto::closed::softsign: return x / (1.0L + std::fabsl(x));
    case grotto::closed::tanhshrink: return x - std::tanh(x);
    case grotto::closed::logistic: return std::log(x / (1.0L - x));
    case grotto::closed::exponential: return -std::log(1.0L - x);
    case grotto::closed::laplace:
        return x <= 0.5L ? std::log(2.0L * x) : -std::log(2.0L * (1.0L - x));
    case grotto::closed::cauchy: return std::tan(3.14159265358979323846L * (x - 0.5L));
    case grotto::closed::sinc: return x == 0 ? 1.0L : std::sin(x) / x;
    case grotto::closed::cbrt: return std::cbrt(x);
    case grotto::closed::qtrt: return std::sqrt(std::sqrt(x));
    case grotto::closed::icbrt: return 1.0L / std::cbrt(x);
    case grotto::closed::iqtrt: return 1.0L / std::sqrt(std::sqrt(x));
    case grotto::closed::pow_m01: return std::exp(-0.1L * std::log(x));
    case grotto::closed::pow_p15: return x * std::sqrt(x);
    case grotto::closed::pow_m3: return 1.0L / (x * x * x);
    }
    return 0;
}

void expect_ulps(grotto::closed which, unsigned k, long double x, long double ulps)
{
    const std::int64_t raw = raw_of(x, k);
    std::int64_t got = 0;
    ASSERT_NO_THROW(got = grotto::eval_closed(which, k, raw)) << static_cast<int>(which) << " x=" << static_cast<double>(x);
    const long double xr = std::ldexp(static_cast<long double>(raw), -static_cast<int>(k));
    const long double want = truth(which, xr) * std::ldexp(1.0L, static_cast<int>(k));
    EXPECT_LE(std::fabsl(static_cast<long double>(got) - want), ulps)
        << static_cast<int>(which) << " k=" << k << " x=" << static_cast<double>(x)
        << " got=" << got << " want=" << static_cast<double>(want);
}

}  // namespace

TEST(ClosedForm, InverseHyperbolicsAndTrig)
{
    for (unsigned k : {16u, 32u})
    {
        for (long double x : {-0.6L, -0.2L, 0.2L, 0.6L})
            expect_ulps(grotto::closed::atanh, k, x, 8.0L);
        for (long double x : {-2.0L, -0.5L, 0.3L, 1.5L, 4.0L})
            expect_ulps(grotto::closed::asinh, k, x, 16.0L);
        for (long double x : {1.0L, 1.5L, 3.0L})
            expect_ulps(grotto::closed::acosh, k, x, 16.0L);
        for (long double x : {-2.0L, -0.4L, 0.0L, 0.5L, 1.0L, 3.0L})
            expect_ulps(grotto::closed::atan, k, x, 8.0L);
        for (long double x : {-1.5L, 0.4L, 2.0L})
            expect_ulps(grotto::closed::acot, k, x, 8.0L);
        for (long double x : {-2.0L, -1.2L, 1.2L, 3.0L})
        {
            expect_ulps(grotto::closed::asec, k, x, 8.0L);
            expect_ulps(grotto::closed::acsc, k, x, 8.0L);
        }
        for (long double x : {0.2L, 0.5L, 1.0L})
            expect_ulps(grotto::closed::asech, k, x, 16.0L);
        for (long double x : {-1.5L, -0.4L, 0.4L, 2.0L})
            expect_ulps(grotto::closed::acsch, k, x, 16.0L);
        for (long double x : {-2.0L, 1.4L, 3.0L})
            expect_ulps(grotto::closed::acoth, k, x, 12.0L);
    }
}

TEST(ClosedForm, ActivationsQuantilesAndSinc)
{
    for (unsigned k : {16u, 32u})
    {
        for (long double x : {-1.5L, -0.2L, 0.0L, 0.4L, 2.0L})
        {
            expect_ulps(grotto::closed::elu, k, x, 8.0L);
            expect_ulps(grotto::closed::celu, k, x, 8.0L);
            expect_ulps(grotto::closed::selu, k, x, 16.0L);
            expect_ulps(grotto::closed::softsign, k, x, 4.0L);
            expect_ulps(grotto::closed::tanhshrink, k, x, 8.0L);
            expect_ulps(grotto::closed::sinc, k, x, 1.0L);
        }
        for (long double p : {0.1L, 0.3L, 0.5L, 0.8L})
        {
            expect_ulps(grotto::closed::logistic, k, p, 12.0L);
            expect_ulps(grotto::closed::exponential, k, p, 12.0L);
            expect_ulps(grotto::closed::laplace, k, p, 12.0L);
        }
        for (long double p : {0.2L, 0.4L, 0.6L, 0.8L})
            expect_ulps(grotto::closed::cauchy, k, p, 16.0L);
    }
}

TEST(ClosedForm, ExtraPowers)
{
    for (unsigned k : {16u, 32u})
    {
        for (long double x : {-8.0L, -1.0L, 0.5L, 1.0L, 2.0L, 8.0L})
        {
            expect_ulps(grotto::closed::cbrt, k, x, 24.0L);
            expect_ulps(grotto::closed::icbrt, k, x, 1.0L);
            expect_ulps(grotto::closed::pow_m3, k, x, 16.0L);
        }
        for (long double x : {0.25L, 0.5L, 1.0L, 2.0L, 9.0L, 81.0L})
        {
            expect_ulps(grotto::closed::qtrt, k, x, 24.0L);
            expect_ulps(grotto::closed::iqtrt, k, x, 32.0L);
            expect_ulps(grotto::closed::pow_m01, k, x, 24.0L);
            expect_ulps(grotto::closed::pow_p15, k, x, 1.0L);
        }
        expect_ulps(grotto::closed::icbrt, k, std::ldexp(1.0L, -8), 1.0L);
        expect_ulps(grotto::closed::icbrt, k, 100.0L, 1.0L);
    }
}

TEST(ClosedForm, RejectsPolesAndBadPrecision)
{
    EXPECT_THROW(grotto::eval_closed(grotto::closed::atanh, 16, 1 << 16), std::domain_error);
    EXPECT_THROW(grotto::eval_closed(grotto::closed::acosh, 16, 0), std::domain_error);
    EXPECT_THROW(grotto::eval_closed(grotto::closed::qtrt, 16, -4), std::domain_error);
    EXPECT_THROW(grotto::eval_closed(grotto::closed::pow_m3, 16, 0), std::domain_error);
    EXPECT_THROW(grotto::eval_closed(grotto::closed::logistic, 16, 0), std::domain_error);
    EXPECT_THROW(grotto::eval_closed(grotto::closed::sinc, 7, 1), std::invalid_argument);
}