#include <iostream>
#include <random>
#include <chrono>

#include "BGN.hpp"

using namespace std;

const size_t NUM_RUNS_PER_TEST = 100;
const size_t MAX_VALUE_IN_TEST = 999;

bool testDecrypt(int x)
{
    bool retval;

    BGN system;

    Scalar testVal(x);
    Scalar one(1);
    Scalar decrypted;

    CurveBipoint curveEnc, curveOne;
    TwistBipoint twistEnc, twistOne;
    Quadripoint quadEncA, quadEncB;

    system.encrypt(curveEnc, testVal);
    system.encrypt(curveOne, one);
    system.encrypt(twistEnc, testVal);
    system.encrypt(twistOne, one);

    quadEncA = system.homomorphic_multiplication(curveEnc, twistOne);
    quadEncB = system.homomorphic_multiplication(curveOne, twistEnc);
    
    decrypted = system.decrypt(curveEnc);
    retval = (decrypted == testVal);

    decrypted = system.decrypt(twistEnc);
    retval = retval && (decrypted == testVal);

    decrypted = system.decrypt(quadEncA);
    retval = retval && (decrypted == testVal);

    decrypted = system.decrypt(quadEncB);
    retval = retval && (decrypted == testVal);

    return retval;
}

double testCurveEncryptSpeed(default_random_engine& generator)
{
    BGN system;

    uniform_int_distribution<int> distribution(0, MAX_VALUE_IN_TEST);

    vector<Scalar> testVals;
    for (size_t i = 0; i < NUM_RUNS_PER_TEST; i++)
        testVals.push_back(Scalar(distribution(generator)));

    vector<CurveBipoint> encryptions(NUM_RUNS_PER_TEST);

    chrono::high_resolution_clock::time_point t0 = chrono::high_resolution_clock::now();
    
    for (size_t i = 0; i < NUM_RUNS_PER_TEST; i++)
        system.encrypt(encryptions[i], testVals[i]);

    chrono::high_resolution_clock::time_point t1 = chrono::high_resolution_clock::now();
    chrono::duration<double> time_span = chrono::duration_cast<chrono::duration<double>>(t1 - t0);

    return time_span.count();
}

double testTwistEncryptSpeed(default_random_engine& generator)
{
    BGN system;

    uniform_int_distribution<int> distribution(0, MAX_VALUE_IN_TEST);

    vector<Scalar> testVals;
    for (size_t i = 0; i < NUM_RUNS_PER_TEST; i++)
        testVals.push_back(Scalar(distribution(generator)));

    vector<TwistBipoint> encryptions(NUM_RUNS_PER_TEST);

    chrono::high_resolution_clock::time_point t0 = chrono::high_resolution_clock::now();
    
    for (size_t i = 0; i < NUM_RUNS_PER_TEST; i++)
        system.encrypt(encryptions[i], testVals[i]);

    chrono::high_resolution_clock::time_point t1 = chrono::high_resolution_clock::now();
    chrono::duration<double> time_span = chrono::duration_cast<chrono::duration<double>>(t1 - t0);

    return time_span.count();
}

double testCurveDecryptSpeed(default_random_engine& generator)
{
    BGN system;

    uniform_int_distribution<int> distribution(0, MAX_VALUE_IN_TEST);

    vector<Scalar> testVals;
    for (size_t i = 0; i < NUM_RUNS_PER_TEST; i++)
        testVals.push_back(Scalar(distribution(generator)));

    vector<CurveBipoint> encryptions(NUM_RUNS_PER_TEST);
    vector<Scalar> decryptions(NUM_RUNS_PER_TEST);
    
    for (size_t i = 0; i < NUM_RUNS_PER_TEST; i++)
        system.encrypt(encryptions[i], testVals[i]);

    chrono::high_resolution_clock::time_point t0 = chrono::high_resolution_clock::now();

    for (size_t i = 0; i < NUM_RUNS_PER_TEST; i++)
        decryptions[i] = system.decrypt(encryptions[i]);

    chrono::high_resolution_clock::time_point t1 = chrono::high_resolution_clock::now();
    chrono::duration<double> time_span = chrono::duration_cast<chrono::duration<double>>(t1 - t0);

    return time_span.count();
}

double testTwistDecryptSpeed(default_random_engine& generator)
{
    BGN system;

    uniform_int_distribution<int> distribution(0, MAX_VALUE_IN_TEST);

    vector<Scalar> testVals;
    for (size_t i = 0; i < NUM_RUNS_PER_TEST; i++)
        testVals.push_back(Scalar(distribution(generator)));

    vector<TwistBipoint> encryptions(NUM_RUNS_PER_TEST);
    vector<Scalar> decryptions(NUM_RUNS_PER_TEST);
    
    for (size_t i = 0; i < NUM_RUNS_PER_TEST; i++)
        system.encrypt(encryptions[i], testVals[i]);

    chrono::high_resolution_clock::time_point t0 = chrono::high_resolution_clock::now();

    for (size_t i = 0; i < NUM_RUNS_PER_TEST; i++)
        decryptions[i] = system.decrypt(encryptions[i]);

    chrono::high_resolution_clock::time_point t1 = chrono::high_resolution_clock::now();
    chrono::duration<double> time_span = chrono::duration_cast<chrono::duration<double>>(t1 - t0);

    return time_span.count();
}

double testQuadDecryptSpeed(default_random_engine& generator)
{
    BGN system;

    uniform_int_distribution<int> distribution(0, MAX_VALUE_IN_TEST);

    vector<Scalar> testVals;
    for (size_t i = 0; i < NUM_RUNS_PER_TEST; i++)
        testVals.push_back(Scalar(distribution(generator)));
    
    Scalar one(1);

    TwistBipoint oneEncryption;
    vector<CurveBipoint> firstEncryptions(NUM_RUNS_PER_TEST);
    vector<Quadripoint> realEncryptions(NUM_RUNS_PER_TEST);
    vector<Scalar> decryptions(NUM_RUNS_PER_TEST);
    
    system.encrypt(oneEncryption, one);

    for (size_t i = 0; i < NUM_RUNS_PER_TEST; i++)
    {
        system.encrypt(firstEncryptions[i], testVals[i]);
        realEncryptions[i] = system.homomorphic_multiplication(firstEncryptions[i], oneEncryption);
    }   

    chrono::high_resolution_clock::time_point t0 = chrono::high_resolution_clock::now();

    for (size_t i = 0; i < NUM_RUNS_PER_TEST; i++)
        decryptions[i] = system.decrypt(realEncryptions[i]);

    chrono::high_resolution_clock::time_point t1 = chrono::high_resolution_clock::now();
    chrono::duration<double> time_span = chrono::duration_cast<chrono::duration<double>>(t1 - t0);

    return time_span.count();
}

bool testAddition(int x, int y)
{
    bool retval;

    BGN system;

    Scalar testX(x);
    Scalar testY(y);
    Scalar testSum(x + y);
    Scalar one(1);
    Scalar decrypted;

    CurveBipoint curveX, curveY, curveSum, curveOne;
    TwistBipoint twistX, twistY, twistSum, twistOne;
    Quadripoint quadXA, quadXB, quadYA, quadYB,
        quadSumAA, quadSumAB, quadSumBA, quadSumBB;

    system.encrypt(curveX, testX);
    system.encrypt(curveY, testY);
    system.encrypt(curveOne, one);
    system.encrypt(twistX, testX);
    system.encrypt(twistY, testY);
    system.encrypt(twistOne, one);
    
    curveSum = system.homomorphic_addition(curveX, curveY);
    twistSum = system.homomorphic_addition(twistX, twistY);

    quadXA = system.homomorphic_multiplication(curveX, twistOne);
    quadXB = system.homomorphic_multiplication(curveOne, twistX);
    quadYA = system.homomorphic_multiplication(curveY, twistOne);
    quadYB = system.homomorphic_multiplication(curveOne, twistY);

    quadSumAA = system.homomorphic_addition(quadXA, quadYA);
    quadSumAB = system.homomorphic_addition(quadXA, quadYB);
    quadSumBA = system.homomorphic_addition(quadXB, quadYA);
    quadSumBB = system.homomorphic_addition(quadXB, quadYB);

    decrypted = system.decrypt(curveSum);
    retval = (decrypted == testSum);

    decrypted = system.decrypt(twistSum);
    retval = retval && (decrypted == testSum);

    decrypted = system.decrypt(quadSumAA);
    retval = retval && (decrypted == testSum);

    decrypted = system.decrypt(quadSumAB);
    retval = retval && (decrypted == testSum);

    decrypted = system.decrypt(quadSumBA);
    retval = retval && (decrypted == testSum);

    decrypted = system.decrypt(quadSumBB);
    retval = retval && (decrypted == testSum);

    return retval;
}

double testCurveAdditionSpeed(default_random_engine& generator)
{
    BGN system;

    uniform_int_distribution<int> distribution(0, MAX_VALUE_IN_TEST);

    vector<Scalar> testXs;
    for (size_t i = 0; i < NUM_RUNS_PER_TEST; i++)
        testXs.push_back(Scalar(distribution(generator)));

    vector<Scalar> testYs;
    for (size_t i = 0; i < NUM_RUNS_PER_TEST; i++)
        testYs.push_back(Scalar(distribution(generator)));

    vector<CurveBipoint> encXs(NUM_RUNS_PER_TEST);
    vector<CurveBipoint> encYs(NUM_RUNS_PER_TEST);
    vector<CurveBipoint> encSums(NUM_RUNS_PER_TEST);
    
    for (size_t i = 0; i < NUM_RUNS_PER_TEST; i++)
    {
        system.encrypt(encXs[i], testXs[i]);
        system.encrypt(encYs[i], testYs[i]);
    }

    chrono::high_resolution_clock::time_point t0 = chrono::high_resolution_clock::now();

    for (size_t i = 0; i < NUM_RUNS_PER_TEST; i++)
        encSums[i] = system.homomorphic_addition(encXs[i], encYs[i]);

    chrono::high_resolution_clock::time_point t1 = chrono::high_resolution_clock::now();
    chrono::duration<double> time_span = chrono::duration_cast<chrono::duration<double>>(t1 - t0);

    return time_span.count();
}

double testTwistAdditionSpeed(default_random_engine& generator)
{
    BGN system;

    uniform_int_distribution<int> distribution(0, MAX_VALUE_IN_TEST);

    vector<Scalar> testXs;
    for (size_t i = 0; i < NUM_RUNS_PER_TEST; i++)
        testXs.push_back(Scalar(distribution(generator)));

    vector<Scalar> testYs;
    for (size_t i = 0; i < NUM_RUNS_PER_TEST; i++)
        testYs.push_back(Scalar(distribution(generator)));

    vector<TwistBipoint> encXs(NUM_RUNS_PER_TEST);
    vector<TwistBipoint> encYs(NUM_RUNS_PER_TEST);
    vector<TwistBipoint> encSums(NUM_RUNS_PER_TEST);
    
    for (size_t i = 0; i < NUM_RUNS_PER_TEST; i++)
    {
        system.encrypt(encXs[i], testXs[i]);
        system.encrypt(encYs[i], testYs[i]);
    }

    chrono::high_resolution_clock::time_point t0 = chrono::high_resolution_clock::now();

    for (size_t i = 0; i < NUM_RUNS_PER_TEST; i++)
        encSums[i] = system.homomorphic_addition(encXs[i], encYs[i]);

    chrono::high_resolution_clock::time_point t1 = chrono::high_resolution_clock::now();
    chrono::duration<double> time_span = chrono::duration_cast<chrono::duration<double>>(t1 - t0);

    return time_span.count();
}

double testQuadAdditionSpeed(default_random_engine& generator)
{
    BGN system;

    uniform_int_distribution<int> distribution(0, MAX_VALUE_IN_TEST);

    vector<Scalar> testXs;
    for (size_t i = 0; i < NUM_RUNS_PER_TEST; i++)
        testXs.push_back(Scalar(distribution(generator)));

    vector<Scalar> testYs;
    for (size_t i = 0; i < NUM_RUNS_PER_TEST; i++)
        testYs.push_back(Scalar(distribution(generator)));

    Scalar one(1);

    TwistBipoint oneEncryption;
    vector<CurveBipoint> firstEncXs(NUM_RUNS_PER_TEST);
    vector<Quadripoint> realEncXs(NUM_RUNS_PER_TEST);
    vector<CurveBipoint> firstEncYs(NUM_RUNS_PER_TEST);
    vector<Quadripoint> realEncYs(NUM_RUNS_PER_TEST);
    vector<Quadripoint> encSums(NUM_RUNS_PER_TEST);
    
    system.encrypt(oneEncryption, one);

    for (size_t i = 0; i < NUM_RUNS_PER_TEST; i++)
    {
        system.encrypt(firstEncXs[i], testXs[i]);
        system.encrypt(firstEncYs[i], testYs[i]);
        realEncXs[i] = system.homomorphic_multiplication(firstEncXs[i], oneEncryption);
        realEncYs[i] = system.homomorphic_multiplication(firstEncYs[i], oneEncryption);
    }

    chrono::high_resolution_clock::time_point t0 = chrono::high_resolution_clock::now();

    for (size_t i = 0; i < NUM_RUNS_PER_TEST; i++)
        encSums[i] = system.homomorphic_addition(realEncXs[i], realEncYs[i]);

    chrono::high_resolution_clock::time_point t1 = chrono::high_resolution_clock::now();
    chrono::duration<double> time_span = chrono::duration_cast<chrono::duration<double>>(t1 - t0);

    return time_span.count();
}

bool testMultiplication(int x, int y)
{
    bool retval;

    BGN system;

    Scalar testX(x);
    Scalar testY(y);
    Scalar testProduct(x * y);
    Scalar decrypted;

    CurveBipoint curveX, curveY;
    TwistBipoint twistX, twistY;
    Quadripoint productA, productB;

    system.encrypt(curveX, testX);
    system.encrypt(curveY, testY);
    system.encrypt(twistX, testX);
    system.encrypt(twistY, testY);
    
    productA = system.homomorphic_multiplication(curveX, twistY);
    productB = system.homomorphic_multiplication(curveY, twistX);

    decrypted = system.decrypt(productA);
    retval = (decrypted == testProduct);

    decrypted = system.decrypt(productB);
    retval = retval && (decrypted == testProduct);

    return retval;
}

double testMultiplicationSpeed(default_random_engine& generator)
{
    BGN system;

    uniform_int_distribution<int> distribution(0, MAX_VALUE_IN_TEST);

    vector<Scalar> testXs;
    for (size_t i = 0; i < NUM_RUNS_PER_TEST; i++)
        testXs.push_back(Scalar(distribution(generator)));

    vector<Scalar> testYs;
    for (size_t i = 0; i < NUM_RUNS_PER_TEST; i++)
        testYs.push_back(Scalar(distribution(generator)));

    vector<CurveBipoint> encXs(NUM_RUNS_PER_TEST);
    vector<TwistBipoint> encYs(NUM_RUNS_PER_TEST);
    vector<Quadripoint> encProducts(NUM_RUNS_PER_TEST);
    
    for (size_t i = 0; i < NUM_RUNS_PER_TEST; i++)
    {
        system.encrypt(encXs[i], testXs[i]);
        system.encrypt(encYs[i], testYs[i]);
    }

    chrono::high_resolution_clock::time_point t0 = chrono::high_resolution_clock::now();

    for (size_t i = 0; i < NUM_RUNS_PER_TEST; i++)
        encProducts[i] = system.homomorphic_multiplication(encXs[i], encYs[i]);

    chrono::high_resolution_clock::time_point t1 = chrono::high_resolution_clock::now();
    chrono::duration<double> time_span = chrono::duration_cast<chrono::duration<double>>(t1 - t0);

    return time_span.count();
}

int main(int argc, char *argv[])
{
    string seedStr("default");
    if (argc > 1)
        seedStr = argv[1];

    seed_seq seed(seedStr.begin(), seedStr.end());
    default_random_engine generator(seed);
    uniform_int_distribution<int> distribution(0, MAX_VALUE_IN_TEST);

    cout << "test_PointAtInfinity: ";
    if (testDecrypt(0))
        cout << "PASS" << endl;
    else
        cout << "FAIL" << endl;
    
    cout << "test_GeneratorPoint: ";
    if (testDecrypt(1))
        cout << "PASS" << endl;
    else
        cout << "FAIL" << endl;

    int randomPoint = distribution(generator);
    cout << "test_RandomPoint (" << randomPoint << "): ";
    if (testDecrypt(randomPoint))
        cout << "PASS" << endl;
    else
        cout << "FAIL" << endl;

    cout << "test_CurveEncryptSpeed (" << NUM_RUNS_PER_TEST << " runs): ";
    cout << testCurveEncryptSpeed(generator) << " seconds" << endl;

    cout << "test_TwistEncryptSpeed (" << NUM_RUNS_PER_TEST << " runs): ";
    cout << testTwistEncryptSpeed(generator) << " seconds" << endl;

    cout << "test_CurveDecryptSpeed (" << NUM_RUNS_PER_TEST << " runs): ";
    cout << testCurveDecryptSpeed(generator) << " seconds" << endl;

    cout << "test_TwistDecryptSpeed (" << NUM_RUNS_PER_TEST << " runs): ";
    cout << testTwistDecryptSpeed(generator) << " seconds" << endl;

    cout << "test_QuadDecryptSpeed (" << NUM_RUNS_PER_TEST << " runs): ";
    cout << testQuadDecryptSpeed(generator) << " seconds" << endl;

    int addX = distribution(generator);
    int addY = distribution(generator);
    cout << "test_Addition (" << addX << ", " << addY << "): ";
    if (testAddition(addX, addY))
        cout << "PASS" << endl;
    else
        cout << "FAIL" << endl;

    cout << "test_CurveAdditionSpeed (" << NUM_RUNS_PER_TEST << " runs): ";
    cout << testCurveAdditionSpeed(generator) << " seconds" << endl;

    cout << "test_TwistAdditionSpeed (" << NUM_RUNS_PER_TEST << " runs): ";
    cout << testTwistAdditionSpeed(generator) << " seconds" << endl;

    cout << "test_QuadAdditionSpeed (" << NUM_RUNS_PER_TEST << " runs): ";
    cout << testQuadAdditionSpeed(generator) << " seconds" << endl;

    int multX = distribution(generator);
    int multY = distribution(generator);
    cout << "test_Multiplication (" << multX << ", " << multY << "): ";
    if (testMultiplication(multX, multY))
        cout << "PASS" << endl;
    else
        cout << "FAIL" << endl;

    cout << "test_MultiplicationSpeed (" << NUM_RUNS_PER_TEST << " runs): ";
    cout << testMultiplicationSpeed(generator) << " seconds" << endl;

    return 0;
}