/***************************************************************************
 * Copyright (c) Johan Mabille, Sylvain Corlay, Wolf Vollprecht and         *
 * Martin Renou                                                             *
 * Copyright (c) QuantStack                                                 *
 * Copyright (c) Serge Guelton                                              *
 *                                                                          *
 * Distributed under the terms of the BSD 3-Clause License.                 *
 *                                                                          *
 * The full license is in the file LICENSE, distributed with this software. *
 ****************************************************************************/

#include "xsimd/xsimd.hpp"

#include <array>
#include <cmath>
#include <complex>
#include <iomanip>
#include <limits>
#include <sstream>
#include <string>
#include <type_traits>
#include <vector>

#include "doctest/doctest.h"

#ifndef XSIMD_TEST_UTILS_HPP
#define XSIMD_TEST_UTILS_HPP

/**************************
 * AppleClang workarounds *
 *************************/

// AppleClang is known for having precision issues
// in the gamma function codegen. It's known to happen
// between AVX and AVX2, but it also happens on SSE4.1
// in GitHub Actions.
// This also seems to happen in M1.
struct precision_t
{
#if defined(__APPLE__) && (XSIMD_WITH_SSE4_1 || XSIMD_WITH_NEON64)
    static constexpr size_t max = 8192;
#else
    static constexpr size_t max = 2048;
#endif
};

/************************
 * Comparison functions *
 ************************/

namespace xsimd
{
    template <class T, class A, size_t N>
    inline bool operator==(const batch<T, A>& lhs, const std::array<T, N>& rhs)
    {
        alignas(A::alignment()) std::array<T, N> tmp;
        lhs.store_aligned(tmp.data());
        return tmp == rhs;
    }

    template <class T, class A, size_t N>
    inline bool operator==(const std::array<T, N>& lhs, const batch<T, A>& rhs)
    {
        return rhs == lhs;
    }
#ifdef XSIMD_ENABLE_XTL_COMPLEX
    template <class T, class A, size_t N, bool i3ec>
    inline bool operator==(const batch<std::complex<T>, A>& lhs, const std::array<xtl::xcomplex<T, T, i3ec>, N>& rhs)
    {
        alignas(A::alignment()) std::array<xtl::xcomplex<T, T, i3ec>, N> tmp;
        lhs.store_aligned(tmp.data());
        return tmp == rhs;
    }

    template <class T, class A, size_t N, bool i3ec>
    inline bool operator==(const std::array<xtl::xcomplex<T, T, i3ec>, N>& lhs, const batch<std::complex<T>, A>& rhs)
    {
        return rhs == lhs;
    }
#endif
}

namespace detail
{
#ifndef __FAST_MATH__
    namespace utils
    {
        // define some overloads here as integer versions do not exist for msvc
        template <class T>
        inline std::enable_if_t<!std::is_integral<T>::value, bool> isinf(const T& c)
        {
            return std::isinf(c);
        }

        template <class T>
        inline std::enable_if_t<std::is_integral<T>::value, bool> isinf(const T&)
        {
            return false;
        }

        template <class T>
        inline std::enable_if_t<!std::is_integral<T>::value, bool> isnan(const T& c)
        {
            return std::isnan(c);
        }

        template <class T>
        inline std::enable_if_t<std::is_integral<T>::value, bool> isnan(const T&)
        {
            return false;
        }
    }
#endif

    inline unsigned char uabs(unsigned char val)
    {
        return val;
    }

    inline unsigned short uabs(unsigned short val)
    {
        return val;
    }

    inline unsigned int uabs(unsigned int val)
    {
        return val;
    }

    inline unsigned long uabs(unsigned long val)
    {
        return val;
    }

    inline unsigned long long uabs(unsigned long long val)
    {
        return val;
    }

    template <class T>
    inline T uabs(T val)
    {
        return std::abs(val);
    }

    template <class T>
    bool check_is_small(const T& value, const T& tolerance)
    {
        return uabs(value) < uabs(tolerance);
    }

    template <class T>
    T safe_division(const T& lhs, const T& rhs)
    {
        if (rhs == T(0))
        {
            return (std::numeric_limits<T>::max)();
        }
        if (rhs < static_cast<T>(1) && lhs > rhs * (std::numeric_limits<T>::max)())
        {
            return (std::numeric_limits<T>::max)();
        }
        if ((lhs == static_cast<T>(0)) || (rhs > static_cast<T>(1) && lhs < rhs * (std::numeric_limits<T>::min)()))
        {
            return static_cast<T>(0);
        }
        return lhs / rhs;
    }

    template <class T>
    bool check_is_close(const T& lhs, const T& rhs, const T& relative_precision)
    {
        using std::abs;
        T diff = uabs(lhs - rhs);
        T d1 = safe_division(diff, T(uabs(rhs)));
        T d2 = safe_division(diff, T(uabs(lhs)));

        return d1 <= relative_precision && d2 <= relative_precision;
    }

    template <class T>
    struct scalar_comparison_near
    {
        static bool run(const T& lhs, const T& rhs)
        {
            using std::abs;
            using std::max;

            // direct compare integers -- but need tolerance for inexact double conversion
            if (std::is_integral<T>::value && lhs < 10e6 && rhs < 10e6)
            {
                return lhs == rhs;
            }
#ifndef __FAST_MATH__
            if (utils::isnan(lhs))
            {
                return utils::isnan(rhs);
            }

            if (utils::isinf(lhs))
            {
                return utils::isinf(rhs) && (lhs * rhs > 0) /* same sign */;
            }
#endif

            T relative_precision = precision_t::max * std::numeric_limits<T>::epsilon();
            T absolute_zero_prox = precision_t::max * std::numeric_limits<T>::epsilon();

            if (max(uabs(lhs), uabs(rhs)) < T(1e-3))
            {
                using res_type = decltype(lhs - rhs);
                return detail::check_is_small(lhs - rhs, res_type(absolute_zero_prox));
            }
            else
            {
                return detail::check_is_close(lhs, rhs, relative_precision);
            }
        }
    };

    template <class T>
    struct scalar_comparison
    {
        static bool run(const T& lhs, const T& rhs)
        {
            return lhs == rhs;
        }
    };

    template <>
    struct scalar_comparison<float> : scalar_comparison_near<float>
    {
    };

    template <>
    struct scalar_comparison<double> : scalar_comparison_near<double>
    {
    };

    template <class T>
    struct scalar_comparison<std::complex<T>>
    {
        static bool run(const std::complex<T>& lhs, const std::complex<T>& rhs)
        {
            using real_comparison = scalar_comparison<T>;
            return real_comparison::run(lhs.real(), rhs.real()) && real_comparison::run(lhs.imag(), rhs.imag());
        }
    };

#ifdef XSIMD_ENABLE_XTL_COMPLEX
    template <class T, bool i3ec>
    struct scalar_comparison<xtl::xcomplex<T, T, i3ec>>
    {
        static bool run(const xtl::xcomplex<T, T, i3ec>& lhs, const xtl::xcomplex<T, T, i3ec>& rhs)
        {
            using real_comparison = scalar_comparison<T>;
            return real_comparison::run(lhs.real(), rhs.real()) && real_comparison::run(lhs.imag(), rhs.imag());
        }
    };
#endif

    template <class V>
    struct vector_comparison
    {
        static bool run(const V& lhs, const V& rhs)
        {
            using value_type = typename V::value_type;
            for (size_t i = 0; i < lhs.size(); ++i)
            {
                if (!scalar_comparison<value_type>::run(lhs[i], rhs[i]))
                    return false;
            }
            return true;
        }
    };

    template <class T>
    bool expect_scalar_near(const T& lhs, const T& rhs)
    {
        return scalar_comparison<T>::run(lhs, rhs);
    }

    template <class V>
    bool expect_container_near(const V& lhs, const V& rhs)
    {
        return vector_comparison<V>::run(lhs, rhs);
    }

    template <class T, size_t N>
    bool expect_array_near(const std::array<T, N>& lhs, const std::array<T, N>& rhs)
    {
        return expect_container_near(lhs, rhs);
    }

    template <class T, class A>
    bool expect_vector_near(const std::vector<T, A>& lhs, const std::vector<T, A>& rhs)
    {
        return expect_container_near(lhs, rhs);
    }

    template <class T, size_t N, class A>
    bool expect_batch_near(const ::xsimd::batch<T, A>& lhs, const std::array<T, N>& rhs)
    {
        alignas(A::alignment()) std::array<T, N> tmp;
        lhs.store_aligned(tmp.data());
        return expect_array_near(tmp, rhs);
    }

    template <class T, size_t N, class A>
    bool expect_batch_near(const std::array<T, N>& lhs, const ::xsimd::batch<T, A>& rhs)
    {
        alignas(A::alignment()) std::array<T, N> tmp;
        rhs.store_aligned(tmp.data());
        return expect_array_near(lhs, tmp);
    }

    template <class T, class A>
    bool expect_batch_near(const ::xsimd::batch<T, A>& lhs, const ::xsimd::batch<T, A>& rhs)
    {
        constexpr auto N = xsimd::batch<T, A>::size;
        alignas(A::alignment()) std::array<T, N> tmp;
        lhs.store_aligned(tmp.data());
        return expect_batch_near(tmp, rhs);
    }

    template <class T, size_t N, class A>
    bool expect_batch_near(const ::xsimd::batch_bool<T, A>& lhs, const std::array<bool, N>& rhs)
    {
        alignas(A::alignment()) std::array<bool, N> tmp;
        lhs.store_aligned(tmp.data());
        return expect_array_near(tmp, rhs);
    }

    template <class T, size_t N, class A>
    bool expect_batch_near(const std::array<bool, N>& lhs, const ::xsimd::batch_bool<T, A>& rhs)
    {
        alignas(A::alignment()) std::array<bool, N> tmp;
        rhs.store_aligned(tmp.data());
        return expect_array_near(lhs, tmp);
    }

    template <class T, class A>
    bool expect_batch_near(const ::xsimd::batch_bool<T, A>& lhs, const ::xsimd::batch_bool<T, A>& rhs)
    {
        constexpr auto N = xsimd::batch<T, A>::size;
        alignas(A::alignment()) std::array<bool, N> tmp;
        lhs.store_aligned(tmp.data());
        return expect_batch_near(tmp, rhs);
    }

    template <class T, class A>
    size_t get_nb_diff_near(const std::vector<T, A>& lhs, const std::vector<T, A>& rhs, float precision)
    {
        size_t i = 0;
        for (size_t i = 0; i < lhs.size(); i++)
        {
            if (std::abs(lhs[i] - rhs[i]) > precision)
            {
                i++;
            }
        }
        return i;
    }

    template <class T, size_t N>
    size_t get_nb_diff_near(const std::array<T, N>& lhs, const std::array<T, N>& rhs, float precision)
    {
        size_t i = 0;
        for (size_t i = 0; i < lhs.size(); i++)
        {
            if (std::abs(lhs[i] - rhs[i]) > precision)
            {
                i++;
            }
        }
        return i;
    }

    template <class B, class S>
    void load_batch(B& b, const S& src, size_t i = 0)
    {
        b = B::load_unaligned(src.data() + i);
    }

    template <class B, class D>
    void store_batch(const B& b, D& dst, size_t i = 0)
    {
        b.store_unaligned(dst.data() + i);
    }

    // Non-template context scope to avoid per-instantiation vtable issues with MinGW GCC.
    // INFO() creates a ContextScope<Lambda> with a unique vtable per template instantiation.
    // This concrete class has a single vtable definition shared across all instantiations.
    struct StringContextScope : doctest::detail::ContextScopeBase
    {
        std::string msg_;
        explicit StringContextScope(std::string msg)
            : msg_(std::move(msg))
        {
        }
        void stringify(std::ostream* os) const override { *os << msg_; }
    };

    template <typename T, typename std::enable_if<std::is_floating_point<T>::value, int>::type = 0>
    std::string to_string_full_precision(T value)
    {
        // TODO(C++17): use std::to_chars
        std::ostringstream ss;
        ss << std::setprecision(std::numeric_limits<T>::max_digits10) << value;
        return ss.str();
    }

    template <typename T, typename std::enable_if<!std::is_floating_point<T>::value, int>::type = 0>
    std::string to_string_full_precision(T value)
    {
        return doctest::toString(value).c_str();
    }

    template <class T>
    StringContextScope make_context_info(const char* name, const T& val)
    {
        return StringContextScope(std::string(name) + ":" + to_string_full_precision(val));
    }
}

// We use make_context_info instead of INFO() to avoid MinGW GCC vtable issues
// (see StringContextScope above). Unlike INFO(), make_context_info is eager:
// it stringifies its operands at construction. To keep the happy path cheap
// (these macros are called in tight loops and QEMU-emulated CI targets like
// ppc64le blow past wall-clock otherwise), we first evaluate the predicate
// and only build the context when it fails. On failure the predicate is
// re-evaluated inside CHECK_UNARY so doctest records the expression text;
// this requires the operands to be side-effect-free, which holds for all
// call sites here.
#define CHECK_BATCH_EQ(b1, b2)                                             \
    do                                                                     \
    {                                                                      \
        const bool batches_are_near = ::detail::expect_batch_near(b1, b2); \
        if (!batches_are_near)                                             \
        {                                                                  \
            auto _ctx1 = ::detail::make_context_info(#b1, b1);             \
            auto _ctx2 = ::detail::make_context_info(#b2, b2);             \
            CHECK_UNARY(batches_are_near);                                 \
        }                                                                  \
    } while (0)

#define CHECK_SCALAR_EQ(s1, s2)                                             \
    do                                                                      \
    {                                                                       \
        const bool scalars_are_near = ::detail::expect_scalar_near(s1, s2); \
        if (!scalars_are_near)                                              \
        {                                                                   \
            auto _ctx1 = ::detail::make_context_info(#s1, s1);              \
            auto _ctx2 = ::detail::make_context_info(#s2, s2);              \
            CHECK_UNARY(scalars_are_near);                                  \
        }                                                                   \
    } while (0)

#define CHECK_VECTOR_EQ(v1, v2)                                             \
    do                                                                      \
    {                                                                       \
        const bool vectors_are_near = ::detail::expect_vector_near(v1, v2); \
        if (!vectors_are_near)                                              \
        {                                                                   \
            auto _ctx1 = ::detail::make_context_info(#v1, v1);              \
            auto _ctx2 = ::detail::make_context_info(#v2, v2);              \
            CHECK_UNARY(vectors_are_near);                                  \
        }                                                                   \
    } while (0)

/***********************
 * Testing types lists *
 ***********************/

#define BATCH_INT_TYPES xsimd::batch<uint8_t>,  \
                        xsimd::batch<int8_t>,   \
                        xsimd::batch<uint16_t>, \
                        xsimd::batch<int16_t>,  \
                        xsimd::batch<uint32_t>, \
                        xsimd::batch<int32_t>,  \
                        xsimd::batch<uint64_t>, \
                        xsimd::batch<int64_t>

#if XSIMD_WITH_NEON64 || !XSIMD_WITH_NEON
#define BATCH_FLOAT_TYPES xsimd::batch<float>, xsimd::batch<double>
#else
#define BATCH_FLOAT_TYPES xsimd::batch<float>
#endif
#if XSIMD_WITH_NEON64 || !XSIMD_WITH_NEON
#define BATCH_COMPLEX_TYPES xsimd::batch<std::complex<float>>, xsimd::batch<std::complex<double>>
#else
#define BATCH_COMPLEX_TYPES xsimd::batch<std::complex<float>>
#endif

#define BATCH_TYPES BATCH_INT_TYPES, BATCH_FLOAT_TYPES
#define BATCH_MATH_TYPES xsimd::batch<int32_t>, BATCH_FLOAT_TYPES

#define BATCH_SWIZZLE_TYPES BATCH_FLOAT_TYPES, BATCH_COMPLEX_TYPES, BATCH_INT_TYPES

/********************
 * conversion utils *
 ********************/
template <size_t N, size_t A>
struct conversion_param
{
    static constexpr size_t size = N;
    static constexpr size_t alignment = A;
};

#define CONVERSION_TYPES conversion_param<sizeof(xsimd::types::simd_register<int, xsimd::default_arch>) / sizeof(double), xsimd::default_arch::alignment()>

#endif // XXSIMD_TEST_UTILS_HPP
