/// @file dpf/fp61.hpp
/// @brief Prime field of order `2^61 - 1`, as an additive output type.
/// @copyright Copyright (c) 2019-2026 Ryan Henry and [others](@ref authors)
/// @license Released under a GNU General Public v2.0 (GPLv2) license;
///          see [LICENSE.md](@ref license) for details.

#ifndef LIBDPF_INCLUDE_DPF_FP61_HPP__
#define LIBDPF_INCLUDE_DPF_FP61_HPP__

#include <cstddef>
#include <cstdint>
#include <cstring>
#include <ostream>
#include <stdexcept>
#include <type_traits>

#include "hedley/hedley.h"

#include "dpf/utils.hpp"
#include "dpf/leaf_arithmetic.hpp"
#include "dpf/random.hpp"

namespace dpf
{

/// @brief Modulus \f$p = 2^{61}-1\f$. `p` itself reduces to 0.
inline constexpr std::uint64_t fp61_mod = (std::uint64_t{1} << 61) - 1;

/// @brief Additive element of the field of order `2^61 - 1`.
class fp61
{
  public:
    /// @brief Underlying unsigned word. Values are stored already reduced.
    using integral_type = std::uint64_t;
    static constexpr std::size_t num_bits = 61;
    static constexpr bool dpf_modint = true;
    static constexpr bool dpf_fp61 = true;

    /// @brief Reduce `v` into the field.
    /// @param v the integer to reduce. Defaults to 0
    HEDLEY_NO_THROW
    HEDLEY_ALWAYS_INLINE
    constexpr fp61(integral_type v = 0) noexcept
        : val{reduce(v)}
    { }

    /// @brief Copy constructor.
    HEDLEY_NO_THROW
    HEDLEY_ALWAYS_INLINE
    constexpr fp61(const fp61 &) noexcept = default;
    /// @brief Move constructor.
    HEDLEY_NO_THROW
    HEDLEY_ALWAYS_INLINE
    constexpr fp61(fp61 &&) noexcept = default;
    /// @brief Copy assignment.
    /// @return `*this`
    HEDLEY_NO_THROW
    HEDLEY_ALWAYS_INLINE
    constexpr fp61 & operator=(const fp61 &) noexcept = default;
    /// @brief Move assignment.
    /// @return `*this`
    HEDLEY_NO_THROW
    HEDLEY_ALWAYS_INLINE
    constexpr fp61 & operator=(fp61 &&) noexcept = default;

    /// @brief The reduced representative in `[0, p)`.
    /// @return the stored field element
    HEDLEY_NO_THROW
    HEDLEY_ALWAYS_INLINE
    HEDLEY_PURE
    constexpr integral_type raw() const noexcept { return reduce(val); }

    /// @brief Build a field element from PRG bytes (Mersenne reduction).
    /// @param bytes the PRG output
    /// @param n the number of bytes available
    /// @return the field element
    HEDLEY_ALWAYS_INLINE
    static fp61 from_seed(const void * bytes, std::size_t n) noexcept
    {
        unsigned char buf[16]{};
        if (n > sizeof(buf))
            n = sizeof(buf);
        std::memcpy(buf, bytes, n);
        std::uint64_t w[2]{};
        std::memcpy(w, buf, sizeof(w));
        using u128 = unsigned __int128;
        const u128 wide = static_cast<u128>(w[0])
            | (static_cast<u128>(w[1]) << 64);
        const auto lo = static_cast<integral_type>(wide) & fp61_mod;
        const auto mid = static_cast<integral_type>(wide >> 61) & fp61_mod;
        const auto hi = static_cast<integral_type>(wide >> 122);
        return fp61{lo + mid + hi};
    }

    /// @brief Same value as `raw()`.
    /// @return the stored field element
    HEDLEY_NO_THROW
    HEDLEY_ALWAYS_INLINE
    HEDLEY_PURE
    explicit constexpr operator integral_type() const noexcept { return val; }

    /// @brief Mersenne reduction of a 64-bit word.
    /// @param x the integer to reduce
    /// @return `x` modulo `2^61-1`, with `p` itself represented as 0
    HEDLEY_NO_THROW
    HEDLEY_ALWAYS_INLINE
    HEDLEY_CONST
    static constexpr integral_type reduce(integral_type x) noexcept
    {
        x = (x & fp61_mod) + (x >> 61);
        if (x >= fp61_mod)
            x -= fp61_mod;
        return x;
    }

    /// @brief Field addition.
    /// @param a left addend
    /// @param b right addend
    /// @return `a + b` in the field
    HEDLEY_NO_THROW
    HEDLEY_ALWAYS_INLINE
    HEDLEY_CONST
    friend constexpr fp61 operator+(fp61 a, fp61 b) noexcept
    {
        return fp61{reduce(a.val) + reduce(b.val)};
    }

    /// @brief Field subtraction.
    /// @param a minuend
    /// @param b subtrahend
    /// @return `a - b` in the field
    HEDLEY_NO_THROW
    HEDLEY_ALWAYS_INLINE
    HEDLEY_CONST
    friend constexpr fp61 operator-(fp61 a, fp61 b) noexcept
    {
        return fp61{reduce(a.val) + fp61_mod - reduce(b.val)};
    }

    /// @brief Field negation.
    /// @param a the element to negate
    /// @return `-a`, with `-0 = 0`
    HEDLEY_NO_THROW
    HEDLEY_ALWAYS_INLINE
    HEDLEY_CONST
    friend constexpr fp61 operator-(fp61 a) noexcept
    {
        const auto v = reduce(a.val);
        return fp61{v == 0 ? 0 : fp61_mod - v};
    }

    /// @brief Field multiplication.
    /// @param a left factor
    /// @param b right factor
    /// @return `a * b` in the field
    HEDLEY_NO_THROW
    HEDLEY_ALWAYS_INLINE
    HEDLEY_CONST
    friend constexpr fp61 operator*(fp61 a, fp61 b) noexcept
    {
        using u128 = unsigned __int128;
        const u128 p = static_cast<u128>(reduce(a.val)) * static_cast<u128>(reduce(b.val));
        const auto lo = static_cast<integral_type>(p) & fp61_mod;
        const auto mid = static_cast<integral_type>(p >> 61) & fp61_mod;
        const auto hi = static_cast<integral_type>(p >> 122);
        return fp61{lo + mid + hi};
    }

    /// @brief Field equality.
    /// @param a left element
    /// @param b right element
    /// @return `true` when the reduced values match
    HEDLEY_NO_THROW
    HEDLEY_ALWAYS_INLINE
    HEDLEY_CONST
    friend constexpr bool operator==(fp61 a, fp61 b) noexcept
    {
        return reduce(a.val) == reduce(b.val);
    }

    /// @brief Field inequality.
    /// @param a left element
    /// @param b right element
    /// @return `true` when the reduced values differ
    HEDLEY_NO_THROW
    HEDLEY_ALWAYS_INLINE
    HEDLEY_CONST
    friend constexpr bool operator!=(fp61 a, fp61 b) noexcept
    {
        return reduce(a.val) != reduce(b.val);
    }

    /// @brief Write the reduced representative in decimal.
    /// @param os the output stream
    /// @param a the element to write
    /// @return `os`
    friend std::ostream & operator<<(std::ostream & os, fp61 a)
    {
        return os << a.val;
    }

  private:
    integral_type val;
};

namespace utils
{

template <>
struct bitlength_of<fp61>
    : std::integral_constant<std::size_t, 61>
{ };

template <>
struct has_characteristic_two<fp61> : std::false_type
{ };

}  // namespace utils

namespace leaf_arithmetic
{

namespace detail
{

template <std::size_t Lanes>
HEDLEY_ALWAYS_INLINE
void fp61_lanes(const void * a, const void * b, void * out,
    fp61 (*op)(fp61, fp61)) noexcept
{
    std::uint64_t aa[Lanes], bb[Lanes], cc[Lanes];
    std::memcpy(aa, a, sizeof(aa));
    std::memcpy(bb, b, sizeof(bb));
    for (std::size_t i = 0; i < Lanes; ++i)
        cc[i] = op(fp61{aa[i]}, fp61{bb[i]}).raw();
    std::memcpy(out, cc, sizeof(cc));
}

template <std::size_t Lanes>
HEDLEY_ALWAYS_INLINE
void fp61_scale(const void * a, fp61 b, void * out) noexcept
{
    std::uint64_t aa[Lanes], cc[Lanes];
    std::memcpy(aa, a, sizeof(aa));
    for (std::size_t i = 0; i < Lanes; ++i)
        cc[i] = (fp61{aa[i]} * b).raw();
    std::memcpy(out, cc, sizeof(cc));
}

}  // namespace detail

HEDLEY_PRAGMA(GCC diagnostic push)
HEDLEY_PRAGMA(GCC diagnostic ignored "-Wignored-attributes")
template <>
struct add_t<fp61, simde__m128i>
{
    auto operator()(const simde__m128i & a, const simde__m128i & b) const
    {
        simde__m128i out;
        detail::fp61_lanes<2>(&a, &b, &out, [](fp61 x, fp61 y) {
            return x + y;
        });
        return out;
    }
};

template <>
struct subtract_t<fp61, simde__m128i>
{
    auto operator()(const simde__m128i & a, const simde__m128i & b) const
    {
        simde__m128i out;
        detail::fp61_lanes<2>(&a, &b, &out, [](fp61 x, fp61 y) {
            return x - y;
        });
        return out;
    }
};

template <>
struct multiply_t<fp61, simde__m128i>
{
    auto operator()(const simde__m128i & a, fp61 b) const
    {
        simde__m128i out;
        detail::fp61_scale<2>(&a, b, &out);
        return out;
    }
};

template <>
struct add_t<fp61, simde__m256i>
{
    auto operator()(const simde__m256i & a, const simde__m256i & b) const
    {
        simde__m256i out;
        detail::fp61_lanes<4>(&a, &b, &out, [](fp61 x, fp61 y) {
            return x + y;
        });
        return out;
    }
};

template <>
struct subtract_t<fp61, simde__m256i>
{
    auto operator()(const simde__m256i & a, const simde__m256i & b) const
    {
        simde__m256i out;
        detail::fp61_lanes<4>(&a, &b, &out, [](fp61 x, fp61 y) {
            return x - y;
        });
        return out;
    }
};

template <>
struct multiply_t<fp61, simde__m256i>
{
    auto operator()(const simde__m256i & a, fp61 b) const
    {
        simde__m256i out;
        detail::fp61_scale<4>(&a, b, &out);
        return out;
    }
};
HEDLEY_PRAGMA(GCC diagnostic pop)

}  // namespace leaf_arithmetic

/// @brief Sample a uniformly reduced field element by rejection.
/// @return an element of the field
template <>
HEDLEY_NO_THROW
inline auto uniform_sample<fp61>() noexcept
{
    for (;;)
    {
        const auto v = uniform_sample<std::uint64_t>() & fp61_mod;
        if (v < fp61_mod)
            return fp61{v};
    }
}

namespace detail
{

template <>
struct shamir_field<fp61> : std::true_type
{
    /// @brief `a^{-1}` by Fermat, `a^{p-2}`.
    /// @param a a non-zero field element
    /// @return `a^{-1}`
    /// @throws std::invalid_argument if `a` is zero
    static fp61 inv(fp61 a)
    {
        if (a.raw() == 0)
            throw std::invalid_argument("shamir: inverse of zero");
        fp61 base = a;
        fp61 out{1};
        auto e = fp61_mod - 2;
        while (e != 0)
        {
            if (e & 1u)
                out = out * base;
            base = base * base;
            e >>= 1;
        }
        return out;
    }
};

}  // namespace detail

}  // namespace dpf

#endif  // LIBDPF_INCLUDE_DPF_FP61_HPP__