#include <pybind11/complex.h>
#include <pybind11/iostream.h>
#include <pybind11/operators.h>
#include <pybind11/pybind11.h>
#include <pybind11/stl.h>

#include <fstream>

#include "seal/seal.h"
#include "seal/util/clipnormal.h"
#include "seal/util/croots.h"
#include "seal/util/hash.h"
#include "seal/util/hestdparms.h"
#include "seal/util/iterator.h"
#include "seal/util/numth.h"
#include "seal/util/pointer.h"
#include "seal/util/polyarithsmallmod.h"
#include "seal/util/rlwe.h"
#include "seal/util/scalingvariant.h"
#include "tenseal/sealapi/bindings.h"

using namespace seal;
using namespace seal::util;
using namespace std;
namespace py = pybind11;

/***
 *Notes:
 * Some methods, like add_poly_poly, have dedicated lambda implementation
 *because uint64* is not handled as the start of an array by pybind, but as a
 *value. We need to use std::vector explicitly.
 * **/

template <typename T>
void bind_seal_pointer(py::module &m, const std::string &name) {
    /*******************
     * "util/pointer.h" {
     ***/
    using type = Pointer<T>;
    std::string class_name = "Pointer" + name;

    py::class_<type, std::shared_ptr<type>>(m, class_name.c_str())
        .def(py::init<>())
        .def("__getitem__",
             py::overload_cast<std::size_t>(&type::operator[], py::const_))
        .def("is_set", &type::is_set)
        .def("get", py::overload_cast<>(&type::get, py::const_),
             py::return_value_policy::reference);

    using const_type = ConstPointer<T>;
    std::string const_class_name = "ConstPointer" + name;

    py::class_<const_type>(m, const_class_name.c_str())
        .def(py::init<>())
        .def("__getitem__", py::overload_cast<std::size_t>(
                                &const_type::operator[], py::const_))
        .def("is_set", &const_type::is_set)
        .def("get", py::overload_cast<>(&const_type::get, py::const_),
             py::return_value_policy::reference);
    /***
     * } "util/pointer.h"
     *******************/
}

void bind_seal_util_namespace(pybind11::module &m) {
    m.doc() = "SEAL::util bindings for Python";

    py::class_<MemoryPoolHandle>(m, "MemoryPoolHandle", py::module_local());

    /*******************
     * "util/rns.h" {
     ***/
    py::class_<RNSBase, std::shared_ptr<RNSBase>>(m, "RNSBase",
                                                  py::module_local())
        .def(py::init<const RNSBase &>())
        .def(py::init<const std::vector<Modulus> &, MemoryPoolHandle>(),
             py::arg(), py::arg() = MemoryManager::GetPool())
        .def("__getitem__",
             py::overload_cast<std::size_t>(&RNSBase::operator[], py::const_))
        .def("size", &RNSBase::size)
        .def("contains", &RNSBase::contains)
        .def("is_subbase_of", &RNSBase::is_subbase_of)
        .def("is_superbase_of", &RNSBase::is_superbase_of)
        .def("is_superbase_of", &RNSBase::is_superbase_of)
        .def("is_proper_subbase_of", &RNSBase::is_proper_subbase_of)
        .def("is_proper_superbase_of", &RNSBase::is_proper_superbase_of)
        .def("extend",
             py::overload_cast<const Modulus &>(&RNSBase::extend, py::const_))
        .def("extend",
             py::overload_cast<const RNSBase &>(&RNSBase::extend, py::const_))
        .def("drop", py::overload_cast<>(&RNSBase::drop, py::const_))
        .def("drop",
             py::overload_cast<const Modulus &>(&RNSBase::drop, py::const_))
        .def("decompose",
             [](const RNSBase &obj, std::vector<std::uint64_t> &value) {
                 std::vector<std::uint64_t> out = value;
                 obj.decompose(out.data(), MemoryManager::GetPool());
                 return out;
             })
        .def("decompose_array",
             [](const RNSBase &obj, std::vector<std::uint64_t> &value,
                std::size_t count) {
                 std::vector<std::uint64_t> out = value;
                 obj.decompose_array(out.data(), count,
                                     MemoryManager::GetPool());
                 return out;
             })
        .def("compose",
             [](const RNSBase &obj, std::vector<std::uint64_t> &value) {
                 std::vector<std::uint64_t> out = value;
                 obj.compose(out.data(), MemoryManager::GetPool());
                 return out;
             })
        .def("compose_array",
             [](const RNSBase &obj, std::vector<std::uint64_t> &value,
                std::size_t count) {
                 std::vector<std::uint64_t> out = value;
                 obj.compose_array(out.data(), count, MemoryManager::GetPool());
                 return out;
             })
        .def("base_prod", &RNSBase::base_prod,
             py::return_value_policy::reference);

    py::class_<BaseConverter>(m, "BaseConverter", py::module_local())
        .def(py::init<const RNSBase &, const RNSBase &, MemoryPoolHandle>(),
             py::arg(), py::arg(), py::arg() = MemoryManager::GetPool())
        .def("ibase_size", &BaseConverter::ibase_size)
        .def("obase_size", &BaseConverter::obase_size)
        .def("ibase", &BaseConverter::ibase)
        .def("obase", &BaseConverter::obase)
        .def(
            "fast_convert",
            [](const BaseConverter &obj, const std::vector<std::uint64_t> &in) {
                std::vector<std::uint64_t> out(obj.obase_size());
                obj.fast_convert(in.data(), out.data(),
                                 MemoryManager::GetPool());
                return out;
            })
        .def("fast_convert_array", [](const BaseConverter &obj,
                                      const std::vector<std::uint64_t> &in,
                                      std::size_t poly_modulus_degree) {
            std::vector<std::uint64_t> out(obj.obase_size() *
                                           poly_modulus_degree);
            obj.fast_convert_array(ConstRNSIter(in.data(), poly_modulus_degree),
                                   RNSIter(out.data(), poly_modulus_degree),
                                   MemoryManager::GetPool());

            return out;
        });

    py::class_<RNSTool, std::shared_ptr<RNSTool>>(m, "RNSTool",
                                                  py::module_local())
        .def(py::init<std::size_t, const RNSBase &, const Modulus &,
                      MemoryPoolHandle>(),
             py::arg(), py::arg(), py::arg(),
             py::arg() = MemoryManager::GetPool())
        .def("divide_and_round_q_last_inplace",
             [](const RNSTool &obj, std::vector<uint64_t> &input,
                std::size_t poly_modulus_degree) {
                 obj.divide_and_round_q_last_inplace(
                     RNSIter(input.data(), poly_modulus_degree),
                     MemoryManager::GetPool());
                 return input;
             })
        .def("divide_and_round_q_last_ntt_inplace",
             [](const RNSTool &obj, std::vector<uint64_t> &input,
                std::size_t poly_modulus_degree,
                const NTTTables *rns_ntt_tables) {
                 obj.divide_and_round_q_last_ntt_inplace(
                     RNSIter(input.data(), poly_modulus_degree), rns_ntt_tables,
                     MemoryManager::GetPool());
                 return input;
             })
        .def("fastbconv_sk",
             [](const RNSTool &obj, const std::vector<uint64_t> &input,
                std::size_t poly_modulus_degree) {
                 std::vector<uint64_t> out(poly_modulus_degree *
                                           obj.base_q()->size());
                 obj.fastbconv_sk(
                     ConstRNSIter(input.data(), poly_modulus_degree),
                     RNSIter(out.data(), poly_modulus_degree),
                     MemoryManager::GetPool());
                 return out;
             })
        .def("sm_mrq",
             [](const RNSTool &obj, const std::vector<uint64_t> &input,
                std::size_t poly_modulus_degree) {
                 std::vector<uint64_t> out(poly_modulus_degree *
                                           obj.base_Bsk()->size());
                 obj.sm_mrq(ConstRNSIter(input.data(), poly_modulus_degree),
                            RNSIter(out.data(), poly_modulus_degree),
                            MemoryManager::GetPool());
                 return out;
             })
        .def("fast_floor",
             [](const RNSTool &obj, const std::vector<uint64_t> &input,
                std::size_t poly_modulus_degree) {
                 std::vector<uint64_t> out(poly_modulus_degree *
                                           obj.base_Bsk()->size());
                 obj.fast_floor(ConstRNSIter(input.data(), poly_modulus_degree),
                                RNSIter(out.data(), poly_modulus_degree),
                                MemoryManager::GetPool());
                 return out;
             })
        .def("fastbconv_m_tilde",
             [](const RNSTool &obj, const std::vector<uint64_t> &input,
                std::size_t poly_modulus_degree) {
                 std::vector<uint64_t> out(poly_modulus_degree *
                                           obj.base_Bsk_m_tilde()->size());
                 obj.fastbconv_m_tilde(
                     ConstRNSIter(input.data(), poly_modulus_degree),
                     RNSIter(out.data(), poly_modulus_degree),
                     MemoryManager::GetPool());
                 return out;
             })
        .def("decrypt_scale_and_round",
             [](const RNSTool &obj, const std::vector<uint64_t> &input,
                std::size_t poly_modulus_degree) {
                 std::vector<uint64_t> out(poly_modulus_degree *
                                           obj.base_q()->size());
                 obj.decrypt_scale_and_round(
                     ConstRNSIter(input.data(), poly_modulus_degree),
                     CoeffIter(out.data()), MemoryManager::GetPool());
                 return out;
             })
        .def("inv_q_last_mod_q", &RNSTool::inv_q_last_mod_q)
        .def("base_Bsk_ntt_tables", &RNSTool::base_Bsk_ntt_tables)
        .def("base_q", &RNSTool::base_q)
        .def("base_B", &RNSTool::base_B)
        .def("base_Bsk", &RNSTool::base_Bsk)
        .def("base_Bsk_m_tilde", &RNSTool::base_Bsk_m_tilde)
        .def("base_t_gamma", &RNSTool::base_t_gamma)
        .def("m_tilde", &RNSTool::m_tilde)
        .def("m_sk", &RNSTool::m_sk)
        .def("t", &RNSTool::t)
        .def("gamma", &RNSTool::gamma);
    /***
     * } "util/rns.h"
     *******************/

    /*******************
     * "util/ntt.h" {
     ***/

    py::class_<NTTTables>(m, "NTTTables", py::module_local())
        .def(py::init<NTTTables &>())
        .def(py::init<int, const Modulus &, MemoryPoolHandle>(), py::arg(),
             py::arg(), py::arg() = MemoryManager::GetPool())
        .def("get_root", &NTTTables::get_root)
        .def("get_from_root_powers",
             py::overload_cast<>(&NTTTables::get_from_root_powers, py::const_))
        .def("get_from_root_powers",
             py::overload_cast<std::size_t>(&NTTTables::get_from_root_powers,
                                            py::const_))
        .def("get_from_inv_root_powers",
             py::overload_cast<>(&NTTTables::get_from_inv_root_powers,
                                 py::const_))
        .def("get_from_inv_root_powers",
             py::overload_cast<std::size_t>(
                 &NTTTables::get_from_inv_root_powers, py::const_))
        .def("modulus", &NTTTables::modulus)
        .def("coeff_count_power", &NTTTables::coeff_count_power)
        .def("coeff_count", &NTTTables::coeff_count);
    m.def("CreateNTTTables",
          [](int coeff_count_power, const std::vector<Modulus> &modulus) {
              Pointer<NTTTables> tables;
              CreateNTTTables(coeff_count_power, modulus, tables,
                              MemoryManager::GetPool());
              return tables;
          })
        .def("ntt_negacyclic_harvey_lazy",
             [](std::vector<uint64_t> &op, const NTTTables &tables) {
                 ntt_negacyclic_harvey_lazy(op.data(), tables);
                 return op;
             })
        .def("ntt_negacyclic_harvey",
             [](std::vector<uint64_t> &op, const NTTTables &tables) {
                 ntt_negacyclic_harvey(op.data(), tables);
                 return op;
             })
        .def("inverse_ntt_negacyclic_harvey_lazy",
             [](std::vector<uint64_t> &op, const NTTTables &tables) {
                 inverse_ntt_negacyclic_harvey_lazy(op.data(), tables);
                 return op;
             })
        .def("inverse_ntt_negacyclic_harvey",
             [](std::vector<uint64_t> &op, const NTTTables &tables) {
                 inverse_ntt_negacyclic_harvey_lazy(op.data(), tables);
                 return op;
             });
    /***
     * } "util/ntt.h"
     *******************/

    /*******************
     * "util/galois.h" {
     ***/
    py::class_<GaloisTool>(m, "GaloisTool")
        .def(py::init<int, MemoryPoolHandle>(), py::arg(),
             py::arg() = MemoryManager::GetPool())
        .def(
            "apply_galois",
            [](const GaloisTool &obj, const std::vector<std::uint64_t> &operand,
               std::uint32_t galois_elt, const Modulus &modulus,
               std::size_t coeff_count) {
                std::vector<uint64_t> out(galois_elt * coeff_count);
                obj.apply_galois(operand.data(), galois_elt, modulus,
                                 out.data());
                return out;
            })
        .def("apply_galois_ntt",
             [](GaloisTool &obj, const std::vector<std::uint64_t> &operand,
                std::uint32_t galois_elt, std::size_t coeff_count) {
                 std::vector<uint64_t> out(coeff_count * galois_elt);
                 obj.apply_galois_ntt(operand.data(), galois_elt, out.data());
                 return out;
             })
        .def("get_elt_from_step", &GaloisTool::get_elt_from_step)
        .def("get_elts_from_steps", &GaloisTool::get_elts_from_steps)
        .def("get_elts_all", &GaloisTool::get_elts_all)
        .def_static("GetIndexFromElt", &GaloisTool::GetIndexFromElt);
    /***
     * } "util/galois.h"
     *******************/

    /*******************
     * "util/rlwe.h" {
     ***/
    m.def("sample_poly_ternary",
          [](std::shared_ptr<UniformRandomGenerator> rng,
             const EncryptionParameters &parms) {
              auto coeff_modulus = parms.coeff_modulus();
              std::vector<uint64_t> out(coeff_modulus.size() *
                                        parms.poly_modulus_degree());
              sample_poly_ternary(rng, parms, out.data());
              return out;
          })
        .def("sample_poly_normal",
             [](std::shared_ptr<UniformRandomGenerator> rng,
                const EncryptionParameters &parms) {
                 auto coeff_modulus = parms.coeff_modulus();
                 std::vector<uint64_t> out(coeff_modulus.size() *
                                           parms.poly_modulus_degree());
                 sample_poly_normal(rng, parms, out.data());
                 return out;
             })
        .def("sample_poly_uniform",
             [](std::shared_ptr<UniformRandomGenerator> rng,
                const EncryptionParameters &parms) {
                 auto coeff_modulus = parms.coeff_modulus();
                 std::vector<uint64_t> out(coeff_modulus.size() *
                                           parms.poly_modulus_degree());
                 sample_poly_uniform(rng, parms, out.data());
                 return out;
             });
    /***
     * } "util/rlwe.h"
     *******************/

    /*******************
     * "util/polyarithsmallmod.h" {
     ***/
    m.def("modulo_poly_coeffs",
          [](const std::vector<uint64_t> &poly, std::size_t coeff_count,
             const Modulus &modulus) {
              std::vector<uint64_t> out(coeff_count);
              modulo_poly_coeffs(poly.data(), coeff_count, modulus, out.data());
              return out;
          })
        .def("negate_poly_coeffmod",
             [](const std::vector<uint64_t> &poly, std::size_t coeff_count,
                const Modulus &modulus) {
                 std::vector<uint64_t> out(coeff_count);
                 negate_poly_coeffmod(poly.data(), coeff_count, modulus,
                                      out.data());
                 return out;
             })
        .def("add_poly_coeffmod",
             [](const std::vector<uint64_t> &operand1,
                const std::vector<uint64_t> &operand2, std::size_t coeff_count,
                const Modulus &modulus) {
                 std::vector<uint64_t> out(coeff_count);
                 add_poly_coeffmod(operand1.data(), operand2.data(),
                                   coeff_count, modulus, out.data());
                 return out;
             })
        .def("sub_poly_coeffmod",
             [](const std::vector<uint64_t> &operand1,
                const std::vector<uint64_t> &operand2, std::size_t coeff_count,
                const Modulus &modulus) {
                 std::vector<uint64_t> out(coeff_count);
                 sub_poly_coeffmod(operand1.data(), operand2.data(),
                                   coeff_count, modulus, out.data());
                 return out;
             })
        .def("multiply_poly_scalar_coeffmod",
             [](const std::vector<uint64_t> &poly, std::size_t coeff_count,
                std::uint64_t scalar, const Modulus &modulus) {
                 std::vector<uint64_t> out(coeff_count);
                 multiply_poly_scalar_coeffmod(poly.data(), coeff_count, scalar,
                                               modulus, out.data());
                 return out;
             })
        .def("dyadic_product_coeffmod",
             [](const std::vector<std::uint64_t> &operand1,
                const std::vector<std::uint64_t> &operand2,
                std::size_t coeff_count, const Modulus &modulus,
                std::size_t result_coeff_count) {
                 std::vector<uint64_t> out(coeff_count * coeff_count);
                 dyadic_product_coeffmod(operand1.data(), operand2.data(),
                                         coeff_count, modulus, out.data());
                 return out;
             })
        .def("poly_infty_norm_coeffmod",
             [](const std::vector<uint64_t> &poly, std::size_t coeff_count,
                const Modulus &modulus) {
                 return poly_infty_norm_coeffmod(poly.data(), coeff_count,
                                                 modulus);
             })
        .def("negacyclic_shift_poly_coeffmod",
             [](const std::vector<uint64_t> &poly, std::size_t coeff_count,
                std::size_t shift, const Modulus &modulus) {
                 std::vector<uint64_t> out(coeff_count);
                 negacyclic_shift_poly_coeffmod(poly.data(), coeff_count, shift,
                                                modulus, out.data());
                 return out;
             })
        .def("negacyclic_multiply_poly_mono_coeffmod",
             [](const std::vector<uint64_t> &poly, std::size_t coeff_count,
                std::uint64_t mono_coeff, std::size_t mono_exponent,
                const Modulus &modulus) {
                 std::vector<uint64_t> out(coeff_count);
                 negacyclic_multiply_poly_mono_coeffmod(
                     poly.data(), coeff_count, mono_coeff, mono_exponent,
                     modulus, out.data(), MemoryManager::GetPool());
                 return out;
             });
    /***
     * } "util/polyarithsmallmod.h"
     *******************/

    /*******************
     * "util/polycore.h" {
     ***/
    m.def("poly_to_hex_string",
          [](const std::vector<uint64_t> &poly, std::size_t coeff_count,
             std::size_t coeff_uint64_count) {
              return poly_to_hex_string(poly.data(), coeff_count,
                                        coeff_uint64_count);
          })
        .def("poly_to_dec_string",
             [](const std::vector<uint64_t> &poly, std::size_t coeff_count,
                std::size_t coeff_uint64_count) {
                 return poly_to_dec_string(poly.data(), coeff_count,
                                           coeff_uint64_count,
                                           MemoryManager::GetPool());
             });
    /***
     * } "util/polycore.h"
     *******************/

    /*******************
     * "util/croots.h" {
     ***/
    py::class_<ComplexRoots>(m, "ComplexRoots", py::module_local())
        .def(py::init<std::size_t, MemoryPoolHandle>(), py::arg(),
             py::arg() = MemoryManager::GetPool())
        .def("get_root", &ComplexRoots::get_root);

    /***
     * } "util/croots.h"
     *******************/

    /*******************
     * "util/numth.h" {
     ***/
    m.def("naf", &naf)
        .def("gcd",
             [](std::uint64_t a, std::uint64_t b) -> uint64_t {
                 return gcd(a, b);
             })
        .def("xgcd", &xgcd)
        .def("are_coprime", &are_coprime)
        .def("is_prime", &is_prime)
        .def("get_primes", &get_primes)
        .def("get_prime", &get_prime)
        .def("is_primitive_root", &is_primitive_root)
        .def("try_minimal_primitive_root", &try_minimal_primitive_root);

    /***
     * } "util/numth.h"
     *******************/

    /*******************
     * "util/hestdparms.h" {
     ***/
    m.def("seal_he_std_parms_128_tc", &seal_he_std_parms_128_tc)
        .def("seal_he_std_parms_192_tc", &seal_he_std_parms_192_tc)
        .def("seal_he_std_parms_256_tc", &seal_he_std_parms_256_tc)
        .def("seal_he_std_parms_128_tq", &seal_he_std_parms_128_tq)
        .def("seal_he_std_parms_192_tq", &seal_he_std_parms_192_tq)
        .def("seal_he_std_parms_256_tq", &seal_he_std_parms_256_tq);
    /***
     * } "util/hestdparms.h"
     *******************/

    /*******************
     * "util/clipnormal.h" {
     ***/
    py::class_<ClippedNormalDistribution>(m, "ClippedNormalDistribution",
                                          py::module_local())
        .def(py::init<double /*mean*/, double /*standard_deviation*/,
                      double /*max_deviation*/>())
        .def("mean", &ClippedNormalDistribution::mean)
        .def("standard_deviation",
             &ClippedNormalDistribution::standard_deviation)
        .def("max_deviation", &ClippedNormalDistribution::max_deviation)
        .def("min", &ClippedNormalDistribution::min)
        .def("max", &ClippedNormalDistribution::max)
        .def("reset", &ClippedNormalDistribution::reset)
        .def("__call__", py::overload_cast<RandomToStandardAdapter &>(
                             &ClippedNormalDistribution::operator()<
                                 RandomToStandardAdapter>));
    /***
     * } "util/clipnormal.h"
     *******************/

    /*******************
     * "util/hash.h" {
     ***/
    py::class_<HashFunction>(m, "HashFunction", py::module_local())
        .def_static("hash", [](const std::vector<uint64_t> &in) {
            HashFunction::hash_block_type out;
            HashFunction::hash(in.data(), in.size(), out);
            return out;
        });

    /***
     * } "util/hash.h"
     *******************/

    py::class_<MultiplyUIntModOperand>(m, "MultiplyUIntModOperand",
                                       py::module_local())
        .def(py::init<>())
        .def_readwrite("operand", &MultiplyUIntModOperand::operand)
        .def_readwrite("quotient", &MultiplyUIntModOperand::quotient)
        .def("set_quotient", &MultiplyUIntModOperand::set_quotient)
        .def("set", &MultiplyUIntModOperand::set);

    /*******************
     * "util/pointer.h" {
     ***/
    bind_seal_pointer<Modulus>(m, "Modulus");
    bind_seal_pointer<std::complex<double>>(m, "ComplexDouble");
    bind_seal_pointer<std::uint64_t>(m, "UInt64");
    bind_seal_pointer<GaloisTool>(m, "GaloisTool");
    bind_seal_pointer<NTTTables>(m, "NTTTables");
    bind_seal_pointer<RNSTool>(m, "RNSTool");
    bind_seal_pointer<RNSBase>(m, "RNSBase");
    bind_seal_pointer<MultiplyUIntModOperand>(m, "MultiplyUIntModOperand");

    /***
     * } "util/pointer.h"
     *******************/
}
