#pragma once

#include <bit>
#include <cmath>
#include <limits>
#include <span>
#include <type_traits>

#include "cloudkey.hpp"
#include "detwfa.hpp"
#include "keyswitch.hpp"
#include "params.hpp"
#include "trlwe.hpp"
#include "utils.hpp"

#ifdef USE_KEY_BUNDLE
#include "keybundle.hpp"
#endif

namespace TFHEpp {

// Most TFHEpp parameter families use all 2N automorphisms of the target
// ring.  Shallow bootstrap parameters may instead use a smaller power-of-two
// LWE modulus q, provided q divides 2N; this is the setting of ePrint
// 2026/1730 (q = 512, N = 1024).
template <class P>
inline constexpr uint32_t BlindRotationModulus = [] {
    if constexpr (requires { P::blind_rotation_modulus; })
        return P::blind_rotation_modulus;
    else
        return 2 * P::targetP::n;
}();

template <class P>
void EvalIdentityKeySwitch(TLWE<typename P::targetP>& res,
                           const TLWE<typename P::domainP>& tlwe,
                           const EvalKey& ek)
{
    if constexpr (std::is_same_v<P, lvl10param> ||
                  std::is_same_v<P, lvl21param>) {
        static_assert(P::domainP::k * P::domainP::n >=
                      P::targetP::k * P::targetP::n);
        SubsetIdentityKeySwitch<P>(res, tlwe, ek.getsubiksk<P>());
    }
    else {
        IdentityKeySwitch<P>(res, tlwe, ek.getiksk<P>());
    }
}

// https://eprint.iacr.org/2025/809
template <class P, uint32_t num_out>
void BRModSwitch(ModswitchTLWE<typename P::domainP>& moded,
                 const TLWE<typename P::domainP>& tlwe)
{
    using domainT = typename P::domainP::T;
    static_assert(std::numeric_limits<domainT>::digits <= 64,
                  "BRModSwitch correction requires at most a 64-bit torus");
    using correctionT =
        std::conditional_t<(std::numeric_limits<domainT>::digits <= 32),
                           int64_t, __int128_t>;
    constexpr uint32_t modulus = BlindRotationModulus<P>;
    static_assert(std::has_single_bit(modulus),
                  "blind rotation modulus must be a power of two");
    static_assert(modulus <= 2 * P::targetP::n,
                  "blind rotation modulus cannot exceed the ring automorphisms");
    static_assert((2 * P::targetP::n) % modulus == 0,
                  "blind rotation modulus must divide the ring automorphisms");
    constexpr uint32_t modulusbit = std::countr_zero(modulus);
    constexpr uint32_t exponent_scale = 2 * P::targetP::n / modulus;
    constexpr uint32_t bitwidth = bits_needed<num_out - 1>();
    correctionT c = 0;
    constexpr domainT roundoffset =
        1ULL << (std::numeric_limits<domainT>::digits - 1 - modulusbit +
                 bitwidth);
    for (int i = 0; i < P::domainP::k * P::domainP::n; i++) {
        const domainT reduced =
            (tlwe[i] + roundoffset) >>
            (std::numeric_limits<domainT>::digits - modulusbit + bitwidth)
                << bitwidth;
        moded[i] = reduced * exponent_scale;
        const domainT rounded =
            reduced
            << (std::numeric_limits<domainT>::digits - modulusbit);
        c += static_cast<std::make_signed_t<domainT>>(tlwe[i] - rounded);
    }
    const domainT reduced_b =
        modulus -
        (static_cast<domainT>(
             static_cast<correctionT>(
                 tlwe[P::domainP::k * P::domainP::n]) -
             c / 2 + roundoffset) >>
         (std::numeric_limits<domainT>::digits - modulusbit + bitwidth)
             << bitwidth);
    moded[P::domainP::k * P::domainP::n] = reduced_b * exponent_scale;
}

template <class P, uint32_t num_out = 1>
void BlindRotate(TRLWE<typename P::targetP>& res,
                 const TLWE<typename P::domainP>& tlwe,
                 const BootstrappingKeyFFT<P>& bkfft,
                 const Polynomial<typename P::targetP>& testvector)
{
    ModswitchTLWE<typename P::domainP> moded;
    BRModSwitch<P, num_out>(moded, tlwe);
    res = {};
    PolynomialMulByXai<typename P::targetP>(
        res[P::targetP::k], testvector, moded[P::domainP::k * P::domainP::n]);
#ifdef USE_KEY_BUNDLE
    for (int i = 0; i < P::domainP::k * P::domainP::n / P::Addends; i++) {
        alignas(64) TRGSWFFT<typename P::targetP> BKadded;
        KeyBundleFFT<P>(BKadded, bkfft[i],
                        std::span(moded)
                            .subspan(P::Addends * i, P::Addends)
                            .template first<P::Addends>());
        ExternalProduct<typename P::targetP>(res, res, BKadded);
    }
#else
    if constexpr (requires { P::domainP::ell; }) {
        static_assert(P::domainP::k == 1,
                      "block-binary blind rotation expects domain k = 1");
        static_assert(P::domainP::n % P::domainP::ell == 0,
                      "block-binary dimension must be divisible by ell");
        static_assert(P::Addends == 1,
                      "block-binary blind rotation does not use key bundling");

        constexpr uint32_t ell = P::domainP::ell;
        constexpr uint32_t blocks = P::domainP::n / ell;
        for (uint32_t block = 0; block < blocks; block++) {
            const uint32_t base = block * ell;
            if (block + 1 < blocks) {
                const char* next_bk =
                    reinterpret_cast<const char*>(&bkfft[base + ell]);
                for (int p = 0; p < 8; p++)
                    __builtin_prefetch(next_bk + p * 4096, 0, 1);
            }
            std::span<const BootstrappingKeyElementFFT<P>, ell> bkfft_block(
                bkfft.begin() + base, ell);
            std::span<const typename P::domainP::T, ell> bara(
                moded.begin() + base, ell);
            CMUXFFTwithBlockBinaryPolynomialMulByXaiMinusOne<P>(
                res, bkfft_block, bara);
        }
    }
    else {
        for (int i = 0; i < P::domainP::k * P::domainP::n; i++) {
            // Prefetch the next BK element (128KB ahead) to overlap memory
            // latency with computation.  Spread prefetch hints across multiple
            // cache lines at the start of the next TRGSW element.
            if (i + 1 < P::domainP::k * P::domainP::n) {
                const char* next_bk =
                    reinterpret_cast<const char*>(&bkfft[i + 1]);
                for (int p = 0; p < 8; p++)
                    __builtin_prefetch(next_bk + p * 4096, 0, 1);
            }
            if (moded[i] == 0) continue;
            CMUXwithPolynomialMulByXaiMinusOne<P>(res, bkfft[i], moded[i]);
        }
    }
#endif
}

template <class P, uint32_t num_out = 1>
void BlindRotate(TRLWE<typename P::targetP>& res,
                 const TLWE<typename P::domainP>& tlwe,
                 const BootstrappingKeyFFT<P>& bkfft,
                 const TRLWE<typename P::targetP>& testvector)
{
    ModswitchTLWE<typename P::domainP> moded;
    BRModSwitch<P, num_out>(moded, tlwe);
    for (int k = 0; k < P::targetP::k + 1; k++)
        PolynomialMulByXai<typename P::targetP>(
            res[k], testvector[k],
            moded[P::domainP::k * P::domainP::n]);
#ifdef USE_KEY_BUNDLE
    alignas(64) std::array<TRGSWFFT<typename P::targetP>,
                           P::domainP::k * P::domainP::n / P::Addends>
        BKadded;
#pragma omp parallel for num_threads(4)
    for (int i = 0; i < P::domainP::k * P::domainP::n / P::Addends; i++) {
        std::array<typename P::domainP::T, P::Addends> bara;
        bara[0] = moded[2 * i];
        bara[1] = moded[2 * i + 1];
        KeyBundleFFT<P>(BKadded[i], bkfft[i], bara);
    }
    for (int i = 0; i < P::domainP::k * P::domainP::n / P::Addends; i++) {
        ExternalProduct<typename P::targetP>(res, res, BKadded[i]);
    }
#else
    for (int i = 0; i < P::domainP::k * P::domainP::n; i++) {
        const uint32_t ā = moded[i];
        if (ā == 0) continue;
        if (i + 1 < P::domainP::k * P::domainP::n) {
            const char* next_bk = reinterpret_cast<const char*>(&bkfft[i + 1]);
            for (int p = 0; p < 8; p++)
                __builtin_prefetch(next_bk + p * 4096, 0, 1);
        }
        CMUXwithPolynomialMulByXaiMinusOne<P>(res, bkfft[i], ā);
    }
#endif
}

template <class P, uint32_t num_out = 1>
void BlindRotate(TRLWE<typename P::targetP>& res,
                 const TLWE<typename P::domainP>& tlwe,
                 const BootstrappingKeyNTT<P>& bkntt,
                 const Polynomial<typename P::targetP>& testvector)
{
    ModswitchTLWE<typename P::domainP> moded;
    BRModSwitch<P, num_out>(moded, tlwe);
    res = {};
    PolynomialMulByXai<typename P::targetP>(
        res[P::targetP::k], testvector,
        moded[P::domainP::k * P::domainP::n]);
    for (int i = 0; i < P::domainP::k * P::domainP::n; i++) {
        const uint32_t ā = moded[i];
        if (ā == 0) continue;
        // Do not use CMUXNTT to avoid unnecessary copy.
        CMUXwithPolynomialMulByXaiMinusOne<typename P::targetP>(res, bkntt[i],
                                                                ā);
    }
}

template <class P, uint32_t num_out = 1>
void BlindRotate(TRLWE<typename P::targetP>& res,
                 const TLWE<typename P::domainP>& tlwe,
                 const BootstrappingKeyRAINTT<P>& bkraintt,
                 const Polynomial<typename P::targetP>& testvector)
{
    ModswitchTLWE<typename P::domainP> moded;
    BRModSwitch<P, num_out>(moded, tlwe);
    res = {};
    PolynomialMulByXai<typename P::targetP>(
        res[P::targetP::k], testvector,
        moded[P::domainP::k * P::domainP::n]);
    for (int i = 0; i < P::domainP::k * P::domainP::n; i++) {
        const uint32_t ā = moded[i];
        if (ā == 0) continue;
        // Do not use CMUXNTT to avoid unnecessary copy.
        CMUXwithPolynomialMulByXaiMinusOne<typename P::targetP>(res,
                                                                bkraintt[i], ā);
    }
}

template <class P, uint32_t num_out = 1>
void BlindRotate(TRLWE<typename P::targetP>& res,
                 const TLWE<typename P::domainP>& tlwe,
                 const BootstrappingKeyFNT<P>& bkfnt,
                 const Polynomial<typename P::targetP>& testvector)
{
    ModswitchTLWE<typename P::domainP> moded;
    BRModSwitch<P, num_out>(moded, tlwe);
    res = {};
    PolynomialMulByXai<typename P::targetP>(
        res[P::targetP::k], testvector, moded[P::domainP::k * P::domainP::n]);
    for (int i = 0; i < P::domainP::k * P::domainP::n; i++) {
        if (moded[i] == 0) continue;
        CMUXwithPolynomialMulByXaiMinusOne<typename P::targetP>(res, bkfnt[i],
                                                                moded[i]);
    }
}

template <class P, uint32_t num_out = 1>
void BlindRotate(TRLWE<typename P::targetP>& res,
                 const TLWE<typename P::domainP>& tlwe,
                 const BootstrappingKeyFNT<P>& bkfnt,
                 const TRLWE<typename P::targetP>& testvector)
{
    ModswitchTLWE<typename P::domainP> moded;
    BRModSwitch<P, num_out>(moded, tlwe);
    for (int k = 0; k < P::targetP::k + 1; k++)
        PolynomialMulByXai<typename P::targetP>(
            res[k], testvector[k],
            moded[P::domainP::k * P::domainP::n]);
    for (int i = 0; i < P::domainP::k * P::domainP::n; i++) {
        const uint32_t ā = moded[i];
        if (ā == 0) continue;
        CMUXwithPolynomialMulByXaiMinusOne<typename P::targetP>(res, bkfnt[i],
                                                                ā);
    }
}

template <class P>
void GateBootstrappingTLWE2TLWE(
    TLWE<typename P::targetP>& res, const TLWE<typename P::domainP>& tlwe,
    const BootstrappingKeyFFT<P>& bkfft,
    const Polynomial<typename P::targetP>& testvector)
{
    alignas(64) TRLWE<typename P::targetP> acc;
    BlindRotate<P>(acc, tlwe, bkfft, testvector);
    SampleExtractIndex<typename P::targetP>(res, acc, 0);
}

template <class P>
void GateBootstrappingTLWE2TLWE(TLWE<typename P::targetP>& res,
                                const TLWE<typename P::domainP>& tlwe,
                                const BootstrappingKeyFFT<P>& bkfft,
                                const TRLWE<typename P::targetP>& testvector)
{
    alignas(64) TRLWE<typename P::targetP> acc;
    BlindRotate<P>(acc, tlwe, bkfft, testvector);
    SampleExtractIndex<typename P::targetP>(res, acc, 0);
}

template <class P>
void GateBootstrappingTLWE2TLWE(
    TLWE<typename P::targetP>& res, const TLWE<typename P::domainP>& tlwe,
    const BootstrappingKeyNTT<P>& bkntt,
    const Polynomial<typename P::targetP>& testvector)
{
    alignas(64) TRLWE<typename P::targetP> acc;
    BlindRotate<P>(acc, tlwe, bkntt, testvector);
    SampleExtractIndex<typename P::targetP>(res, acc, 0);
}

template <class P>
void GateBootstrappingTLWE2TLWE(
    TLWE<typename P::targetP>& res, const TLWE<typename P::domainP>& tlwe,
    const BootstrappingKeyRAINTT<P>& bkraintt,
    const Polynomial<typename P::targetP>& testvector)
{
    alignas(64) TRLWE<typename P::targetP> acc;
    BlindRotate<P>(acc, tlwe, bkraintt, testvector);
    SampleExtractIndex<typename P::targetP>(res, acc, 0);
}

template <class P>
void GateBootstrappingTLWE2TLWE(
    TLWE<typename P::targetP>& res, const TLWE<typename P::domainP>& tlwe,
    const BootstrappingKeyFNT<P>& bkfnt,
    const Polynomial<typename P::targetP>& testvector)
{
    alignas(64) TRLWE<typename P::targetP> acc;
    BlindRotate<P>(acc, tlwe, bkfnt, testvector);
    SampleExtractIndex<typename P::targetP>(res, acc, 0);
}

template <class P>
void GateBootstrappingTLWE2TLWE(TLWE<typename P::targetP>& res,
                                const TLWE<typename P::domainP>& tlwe,
                                const BootstrappingKeyFNT<P>& bkfnt,
                                const TRLWE<typename P::targetP>& testvector)
{
    alignas(64) TRLWE<typename P::targetP> acc;
    BlindRotate<P>(acc, tlwe, bkfnt, testvector);
    SampleExtractIndex<typename P::targetP>(res, acc, 0);
}

template <class P, uint32_t num_out>
void GateBootstrappingManyLUT(
    std::array<TLWE<typename P::targetP>, num_out>& res,
    const TLWE<typename P::domainP>& tlwe, const BootstrappingKeyFFT<P>& bkfft,
    const Polynomial<typename P::targetP>& testvector)
{
    alignas(64) TRLWE<typename P::targetP> acc;
    BlindRotate<P, num_out>(acc, tlwe, bkfft, testvector);
    for (int i = 0; i < num_out; i++)
        SampleExtractIndex<typename P::targetP>(res[i], acc, i);
}

template <class P, uint32_t num_out>
void GateBootstrappingManyLUT(
    std::array<TLWE<typename P::targetP>, num_out>& res,
    const TLWE<typename P::domainP>& tlwe, const BootstrappingKeyFFT<P>& bkfft,
    const TRLWE<typename P::targetP>& testvector)
{
    alignas(64) TRLWE<typename P::targetP> acc;
    BlindRotate<P, num_out>(acc, tlwe, bkfft, testvector);
    for (int i = 0; i < num_out; i++)
        SampleExtractIndex<typename P::targetP>(res[i], acc, i);
}

template <class P, typename P::T μ>
constexpr Polynomial<P> μpolygen()
{
    Polynomial<P> poly;
    for (typename P::T& p : poly) p = μ;
    return poly;
}

template <class bkP, typename bkP::targetP::T μ, class iksP>
void GateBootstrapping(TLWE<typename iksP::targetP>& res,
                       const TLWE<typename bkP::domainP>& tlwe,
                       const EvalKey& ek)
{
    alignas(64) TLWE<typename bkP::targetP> tlwelvl1;
    GateBootstrappingTLWE2TLWE<bkP>(tlwelvl1, tlwe, ek.getbkfft<bkP>(),
                                    μpolygen<typename bkP::targetP, μ>());
    EvalIdentityKeySwitch<iksP>(res, tlwelvl1, ek);
}

template <class iksP, class bkP, typename bkP::targetP::T μ>
void GateBootstrapping(TLWE<typename bkP::targetP>& res,
                       const TLWE<typename iksP::domainP>& tlwe,
                       const EvalKey& ek)
{
    alignas(64) TLWE<typename iksP::targetP> tlwelvl0;
    EvalIdentityKeySwitch<iksP>(tlwelvl0, tlwe, ek);
    GateBootstrappingTLWE2TLWE<bkP>(res, tlwelvl0, ek.getbkfft<bkP>(),
                                    μpolygen<typename bkP::targetP, μ>());
}

template <class bkP, typename bkP::targetP::T μ, class iksP>
void GateBootstrappingNTT(TLWE<typename iksP::tagetP>& res,
                          const TLWE<typename bkP::domainP>& tlwe,
                          const EvalKey& ek)
{
    alignas(64) TLWE<typename bkP::targetP> tlwelvl1;
    GateBootstrappingTLWE2TLWE<bkP>(tlwelvl1, tlwe, ek.getbkntt<bkP>(),
                                    μpolygen<typename bkP::targetP, μ>());
    EvalIdentityKeySwitch<iksP>(res, tlwelvl1, ek);
}

template <class iksP, class bkP, typename bkP::targetP::T μ>
void GateBootstrappingNTT(TLWE<typename bkP::targetP>& res,
                          const TLWE<typename iksP::domainP>& tlwe,
                          const EvalKey& ek)
{
    alignas(64) TLWE<typename iksP::targetP> tlwelvl0;
    EvalIdentityKeySwitch<iksP>(tlwelvl0, tlwe, ek);
    GateBootstrappingTLWE2TLWE<bkP>(res, tlwelvl0, ek.getbkntt<bkP>(),
                                    μpolygen<typename bkP::targetP, μ>());
}

}  // namespace TFHEpp