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

#include "BGN.hpp"
#include "client.hpp"
#include "server.hpp"
#include "serverEntity.hpp"

using namespace std;

// Initialize the classes we use
void initialize_prsona_classes()
{
    Scalar::init();
    PrsonaBase::init();
}

// Quick and dirty mean calculation (used for averaging timings)
double mean(vector<double> xx)
{
    return accumulate(xx.begin(), xx.end(), 0.0) / xx.size(); 
}

void print_user_scores(const vector<PrsonaClient>& users)
{
    std::cout << "<";
    for (size_t i = 0; i < users.size(); i++)
    {
            std::cout << users[i].get_score()
                << (i == users.size() - 1 ? ">" : " ");
    }
    std::cout << std::endl;
}

// Time how long it takes to make a proof of valid votes
vector<double> make_votes(
    default_random_engine& generator,
    vector<vector<TwistBipoint>>& newEncryptedVotes,
    vector<vector<Proof>>& validVoteProofs,
    const vector<PrsonaClient>& users,
    const PrsonaServerEntity& servers,
    size_t numVotes)
{
    vector<double> retval;
    uniform_int_distribution<int> voteDistribution(
        0, PrsonaBase::get_max_allowed_vote());
    size_t numUsers = users.size();
    newEncryptedVotes.clear();

    for (size_t i = 0; i < numUsers; i++)
    {
        // Make the correct number of new votes, but shuffle where they go
        vector<Scalar> votes;
        vector<bool> replaces;
        for (size_t j = 0; j < numUsers; j++)
        {
            votes.push_back(Scalar(voteDistribution(generator)));
            replaces.push_back(j < numVotes);
        }
        shuffle(replaces.begin(), replaces.end(), generator);

        
        Proof baseProof;
        vector<Proof> fullProof;
        Twistpoint shortTermPublicKey =
            users[i].get_short_term_public_key();
        vector<TwistBipoint> currEncryptedVotes =
            servers.get_current_votes_by(baseProof, shortTermPublicKey);
        fullProof.push_back(baseProof);
        servers.get_other_vote_row_commitments(fullProof, shortTermPublicKey);
        
        vector<Proof> currVoteProof;
        chrono::high_resolution_clock::time_point t0 =
            chrono::high_resolution_clock::now();
        currEncryptedVotes = users[i].make_votes(
                                currVoteProof,
                                fullProof,
                                currEncryptedVotes,
                                votes,
                                replaces);
        chrono::high_resolution_clock::time_point t1 =
            chrono::high_resolution_clock::now();

        newEncryptedVotes.push_back(currEncryptedVotes);
        validVoteProofs.push_back(currVoteProof);

        chrono::duration<double> time_span =
            chrono::duration_cast<chrono::duration<double>>(t1 - t0);
        retval.push_back(time_span.count());
    }

    return retval;
}

// Time how long it takes to validate a proof of valid votes
vector<double> transmit_votes_to_servers(
    const vector<vector<TwistBipoint>>& newEncryptedVotes,
    const vector<vector<Proof>>& validVoteProofs,
    const vector<PrsonaClient>& users,
    PrsonaServerEntity& servers)
{
    vector<double> retval;
    size_t numUsers = users.size();
    size_t numServers = servers.get_num_servers();

    for (size_t i = 0; i < numUsers; i++)
    {
        Proof ownerProof;
        Twistpoint shortTermPublicKey =
            users[i].get_short_term_public_key(ownerProof);

        for (size_t j = 0; j < numServers; j++)
        {
            chrono::high_resolution_clock::time_point t0 =
                chrono::high_resolution_clock::now();
            servers.receive_vote(
                validVoteProofs[i],
                newEncryptedVotes[i],
                shortTermPublicKey,
                j);
            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);
            retval.push_back(time_span.count());
        }
            
    }

    return retval;
}

// Time how long it takes to do the operations associated with an epoch
double epoch(PrsonaServerEntity& servers)
{
    // Do the epoch server calculations
    chrono::high_resolution_clock::time_point t0 =
        chrono::high_resolution_clock::now();
    servers.epoch();
    chrono::high_resolution_clock::time_point t1 =
        chrono::high_resolution_clock::now();

    // Return the timing of the epoch server calculations
    chrono::duration<double> time_span =
        chrono::duration_cast<chrono::duration<double>>(t1 - t0);
    return time_span.count();
}

// Time how long it takes each user to decrypt their new scores
vector<double> transmit_epoch_updates(
    vector<PrsonaClient>& users, const PrsonaServerEntity& servers)
{
    vector<double> retval;
    size_t numUsers = users.size();

    for (size_t i = 0; i < numUsers; i++)
    {
        chrono::high_resolution_clock::time_point t0 =
            chrono::high_resolution_clock::now();
        servers.transmit_updates(users[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);
        retval.push_back(time_span.count());
    }

    return retval;
}

// Test if the proof of reputation level is working as expected
void test_reputation_proof(
    default_random_engine& generator,
    const PrsonaServerEntity& servers,
    const PrsonaClient& a,
    const PrsonaClient& b)
{
    bool flag;
    mpz_class aScore = a.get_score().toInt();
    int i = 0;
    while (i < aScore)
        i++;

    uniform_int_distribution<int> thresholdDistribution(0, i);
    Scalar goodThreshold(thresholdDistribution(generator));
    Scalar badThreshold(aScore + 1);

    Twistpoint shortTermPublicKey = a.get_short_term_public_key();
    vector<Proof> goodRepProof =
        a.generate_reputation_proof(goodThreshold, servers.get_num_clients());

    Proof baseProof;
    vector<Proof> fullProof;
    EGCiphertext currEncryptedScore =
        servers.get_current_user_encrypted_tally(baseProof, shortTermPublicKey);
    fullProof.push_back(baseProof);
    servers.get_other_user_tally_commitments(fullProof, shortTermPublicKey);

    flag = b.verify_reputation_proof(
                goodRepProof,
                shortTermPublicKey,
                goodThreshold,
                fullProof,
                currEncryptedScore);
    cout << "TEST VALID REPUTATION PROOF:     "
        << (flag ? "PASSED (Proof verified)" : "FAILED (Proof not verified)" )
        << endl;
    
    vector<Proof> badRepProof =
        a.generate_reputation_proof(badThreshold, servers.get_num_clients());

    baseProof.clear();
    fullProof.clear();
    currEncryptedScore =
        servers.get_current_user_encrypted_tally(baseProof, shortTermPublicKey);
    fullProof.push_back(baseProof);
    servers.get_other_user_tally_commitments(fullProof, shortTermPublicKey);

    flag = b.verify_reputation_proof(
                badRepProof,
                shortTermPublicKey,
                goodThreshold,
                fullProof,
                currEncryptedScore);
    cout << "TEST INVALID REPUTATION PROOF:   "
        << (flag ? "FAILED (Proof verified)" : "PASSED (Proof not verified)" )
        << endl << endl;
}

// Test if the proof of valid votes is working as expected
void test_vote_proof(
    default_random_engine& generator,
    const PrsonaClient& user,
    PrsonaServerEntity& servers)
{
    size_t numUsers = servers.get_num_clients();
    vector<Scalar> votes;
    vector<bool> replaces;
    bool flag;

    for (size_t i = 0; i < numUsers; i++)
    {
        votes.push_back(Scalar(1));
        replaces.push_back(true); 
    }

    vector<Proof> validVoteProof;

    Proof baseProof;
    vector<Proof> fullProof;
    Twistpoint shortTermPublicKey =
        user.get_short_term_public_key();
    vector<TwistBipoint> encryptedVotes =
        servers.get_current_votes_by(baseProof, shortTermPublicKey);
    fullProof.push_back(baseProof);
    servers.get_other_vote_row_commitments(fullProof, shortTermPublicKey);

    encryptedVotes = user.make_votes(
        validVoteProof,
        fullProof,
        encryptedVotes,
        votes,
        replaces);

    flag = servers.receive_vote(
                validVoteProof, encryptedVotes, shortTermPublicKey);
    cout << "TEST REPLACE VOTE PROOF:         "
        << (flag ? "PASSED (Proof verified)" : "FAILED (Proof not verified)" )
        << endl;

    for (size_t i = 0; i < numUsers; i++)
        replaces[i] = false;

    baseProof.clear();
    fullProof.clear();
    encryptedVotes =
        servers.get_current_votes_by(baseProof, shortTermPublicKey);
    fullProof.push_back(baseProof);
    servers.get_other_vote_row_commitments(fullProof, shortTermPublicKey);
    encryptedVotes = user.make_votes(
        validVoteProof,
        fullProof,
        encryptedVotes,
        votes,
        replaces);

    flag = servers.receive_vote(
                validVoteProof, encryptedVotes, shortTermPublicKey);
    cout << "TEST RERANDOMIZE VOTE PROOF:     "
        << (flag ? "PASSED (Proof verified)" : "FAILED (Proof not verified)" )
        << endl;

    for (size_t i = 0; i < numUsers; i++)
    {
        votes[i] = Scalar(3);
        replaces[i] = true;
    }

    baseProof.clear();
    fullProof.clear();
    encryptedVotes =
        servers.get_current_votes_by(baseProof, shortTermPublicKey);
    fullProof.push_back(baseProof);
    servers.get_other_vote_row_commitments(fullProof, shortTermPublicKey);
    encryptedVotes = user.make_votes(
        validVoteProof,
        fullProof,
        encryptedVotes,
        votes,
        replaces);

    flag = servers.receive_vote(
                validVoteProof, encryptedVotes, shortTermPublicKey);
    cout << "TEST INVALID REPLACE VOTE PROOF: "
        << (flag ? "FAILED (Proof verified)" : "PASSED (Proof not verified)" )
        << endl << endl;
}

void check_vote_matrix_updates()
{
    size_t numServers = 2;
    size_t numUsers = 3;

    cout << "Testing how the vote matrix updates." << endl;

    PrsonaBase::set_client_malicious();

    // Entities we operate with
    PrsonaServerEntity servers(numServers);
    vector<Proof> elGamalBlindGeneratorProof;
    BGNPublicKey bgnPublicKey = servers.get_bgn_public_key();
    Twistpoint elGamalBlindGenerator =
        servers.get_blinding_generator(elGamalBlindGeneratorProof);

    vector<PrsonaClient> users;
    for (size_t i = 0; i < numUsers; i++)
    {
        PrsonaClient currUser(
            elGamalBlindGeneratorProof,
            elGamalBlindGenerator,
            bgnPublicKey,
            numServers);
        users.push_back(currUser);
        servers.add_new_client(users[i]);
    }

    Proof pseudonymsProof;
    vector<Twistpoint> currentPseudonyms =
        servers.get_current_pseudonyms(pseudonymsProof);

    cout << "Making votes." << endl;

    for (size_t i = 0; i < numUsers; i++)
    {   
        Twistpoint shortTermPublicKey =
            users[i].get_short_term_public_key();

        size_t myIndex =
            users[i].binary_search(currentPseudonyms, shortTermPublicKey);

        cout << "User " << i+1 << " has initial index " << myIndex << endl;

        vector<Scalar> votes;
        vector<bool> replaces;
        for (size_t j = 0; j < numUsers; j++)
        {
            if (j == myIndex)
                votes.push_back(Scalar(2));
            else if (j > myIndex)
                votes.push_back(Scalar(1));
            else
                votes.push_back(Scalar(0));
            replaces.push_back(true);
        }

        Proof baseProof;
        vector<Proof> fullProof;
        vector<TwistBipoint> currEncryptedVotes =
            servers.get_current_votes_by(baseProof, shortTermPublicKey);
        fullProof.push_back(baseProof);
        servers.get_other_vote_row_commitments(fullProof, shortTermPublicKey);

        vector<Proof> currVoteProof;
        currEncryptedVotes = users[i].make_votes(
                                currVoteProof,
                                fullProof,
                                currEncryptedVotes,
                                votes,
                                replaces);

        servers.receive_vote(
                currVoteProof,
                currEncryptedVotes,
                shortTermPublicKey);
        cout << "User " << i+1 << " now has the following votes:" << endl;
        servers.print_current_votes_by(shortTermPublicKey);
    }
    
    servers.print_votes();
    
    epoch(servers);

    cout << "First epoch done." << endl;
        
    transmit_epoch_updates(users, servers);

    cout << "Updates given to users." << endl;

    servers.print_votes();

    for (size_t i = 0; i < numUsers; i++)
    {   
        Proof ownerProof;
        Twistpoint shortTermPublicKey =
            users[i].get_short_term_public_key(ownerProof);

        cout << "User " << i+1 << " now has the following votes:" << endl;
        servers.print_current_votes_by(shortTermPublicKey);
    }
}

int main(int argc, char *argv[])
{
    initialize_prsona_classes();

    // Defaults
    size_t numServers = 2;
    size_t numUsers = 5;
    size_t numRounds = 3;
    size_t numVotesPerRound = 3;
    bool maliciousServers = true;
    bool maliciousClients = true;
    string seedStr = "seed";

    // Potentially accept command line inputs
    if (argc > 1)
        numServers = atoi(argv[1]);
    if (argc > 2)
        numUsers = atoi(argv[2]);
    if (argc > 3)
        numRounds = atoi(argv[3]);
    if (argc > 4)
        numVotesPerRound = atoi(argv[4]);
    if (argc > 5)
    {
        bool setting = argv[5][0] == 't' || argv[5][0] == 'T';
        maliciousServers = setting;
    }
    if (argc > 6)
    {
        bool setting = argv[6][0] == 't' || argv[6][0] == 'T';
        maliciousClients = setting;
    }
    if (argc > 7)
        seedStr = argv[7];

    cout << "Running the protocol with the following parameters: " << endl;
    cout << numServers << " PRSONA servers" << endl;
    cout << numUsers << " participants (voters/votees)" << endl;
    cout << numRounds << " epochs" << endl;
    cout << numVotesPerRound << " new (random) votes by each user per epoch" << endl;
    cout << "Servers are set to " << (maliciousServers ? "MALICIOUS" : "HBC") << " security" << endl;
    cout << "Clients are set to " << (maliciousClients ? "MALICIOUS" : "HBC") << " security" << endl;
    cout << "Current randomness seed: \"" << seedStr << "\"" << endl;
    cout << endl;

    // Set malicious flags where necessary
    if (maliciousServers)
        PrsonaBase::set_server_malicious();
    if (maliciousClients)
        PrsonaBase::set_client_malicious();

    // Entities we operate with
    PrsonaServerEntity servers(numServers);
    vector<Proof> elGamalBlindGeneratorProof;
    BGNPublicKey bgnPublicKey = servers.get_bgn_public_key();
    Twistpoint elGamalBlindGenerator =
        servers.get_blinding_generator(elGamalBlindGeneratorProof);

    cout << "Initialization: adding users to system" << endl << endl;
    vector<PrsonaClient> users;
    for (size_t i = 0; i < numUsers; i++)
    {
        PrsonaClient currUser(
            elGamalBlindGeneratorProof,
            elGamalBlindGenerator,
            bgnPublicKey,
            numServers);
        users.push_back(currUser);
        servers.add_new_client(users[i]);
    }

    // Seeded randomness for random votes used in epoch
    seed_seq seed(seedStr.begin(), seedStr.end());
    default_random_engine generator(seed);

    // Do the epoch operations
    for (size_t i = 0; i < numRounds; i++)
    {
        vector<double> timings;

        cout << "Round " << i+1 << " of " << numRounds << ": " << endl;
        
        vector<vector<TwistBipoint>> newEncryptedVotes;
        vector<vector<Proof>> validVoteProofs;
        timings = make_votes(
            generator,
            newEncryptedVotes,
            validVoteProofs,
            users,
            servers,
            numVotesPerRound);
        
        cout << "Vote generation (with proofs): " << mean(timings)
            << " seconds per user" << endl;
        timings.clear();

        timings = transmit_votes_to_servers(
            newEncryptedVotes, validVoteProofs, users, servers);

        cout << "Vote validation: " << mean(timings)
            << " seconds per vote vector/server" << endl;
        timings.clear();

        timings.push_back(epoch(servers));
        
        cout << "Epoch computation: " << mean(timings) << " seconds" << endl;
        timings.clear();

        timings = transmit_epoch_updates(users, servers);

        cout << "Transmit epoch updates: " << mean(timings)
            << " seconds per user" << endl << endl;
    }

    // Pick random users for our tests
    uniform_int_distribution<size_t> userDistribution(0, numUsers - 1);
    size_t user_a = userDistribution(generator);
    size_t user_b = user_a;
    while (user_b == user_a)
        user_b = userDistribution(generator);

    test_reputation_proof(generator, servers, users[user_a], users[user_b]);
    test_vote_proof(generator, users[user_a], servers);

    return 0;
}