#include "openfhe.h"

using namespace lbcrypto;
using namespace std;

double binomial_coeff(float k, int a) {
    double b=1;
    for(int p=1; p<=a; p++) {
        b=b*(k+1-p)/p;
    }
    return b;
}

int main() {
    CCParams<CryptoContextCKKSRNS> parameters;

    parameters.SetSecretKeyDist(UNIFORM_TERNARY);
    parameters.SetSecurityLevel(lbcrypto::HEStd_NotSet);
    parameters.SetRingDim(1 << 15);

    ScalingTechnique rescaleTech = FLEXIBLEAUTO;
    usint dcrtBits               = 59;
    usint firstMod               = 60;

    parameters.SetScalingModSize(dcrtBits);
    parameters.SetScalingTechnique(rescaleTech);
    parameters.SetFirstModSize(firstMod);

    std::vector<uint32_t> levelBudget = {3, 3};
    uint32_t approxBootstrapDepth     = 9;

    uint32_t levelsUsedBeforeBootstrap = 10;
    usint depth =
            levelsUsedBeforeBootstrap + FHECKKSRNS::GetBootstrapDepth(approxBootstrapDepth, levelBudget, UNIFORM_TERNARY);

    parameters.SetMultiplicativeDepth(depth);

    CryptoContext<DCRTPoly> crypto_context = GenCryptoContext(parameters);

    crypto_context->Enable(PKE);
    crypto_context->Enable(KEYSWITCH);
    crypto_context->Enable(LEVELEDSHE);
    crypto_context->Enable(ADVANCEDSHE);
    crypto_context->Enable(FHE);

    usint ringDim = crypto_context->GetRingDimension();
    usint numSlots = ringDim / 2;

    crypto_context->EvalBootstrapSetup(levelBudget);
    KeyPair<DCRTPoly> key_pair = crypto_context->KeyGen();

    crypto_context->EvalMultKeyGen(key_pair.secretKey);
    crypto_context->EvalBootstrapKeyGen(key_pair.secretKey, numSlots);

    vector<double> val = {0.1, 0.2, 0.3, 0.4, 0.5};
    Ciphertext<DCRTPoly> ciphertext = crypto_context->Encrypt(
            key_pair.publicKey, crypto_context->MakeCKKSPackedPlaintext(val));

    vector<double> minus_half_vec = {-0.5, -0.5, -0.5, -0.5, -0.5};
    Plaintext minus_half = crypto_context->MakeCKKSPackedPlaintext(minus_half_vec);

    ciphertext = crypto_context->EvalAdd(ciphertext, minus_half);

    vector<double> coefficients;

    int degree = 8;

    for (int i = 0; i < degree; i++) {
        coefficients.push_back(pow(2, -0.5 + i) * binomial_coeff(0.5, i));
    }

    cout << "Coefficients: " << coefficients << endl;

    Ciphertext<DCRTPoly> output = crypto_context->EvalPoly(ciphertext, coefficients);

    Plaintext plain_result;
    crypto_context->Decrypt(key_pair.secretKey, output, &plain_result);


    cout << "Result: " << plain_result << endl;


}