#include <stdint.h>
#include <omp.h>
#include <numeric>
#include <cmath>
#include <vector>
#include <algorithm>
#include <cstdio>

#include "tables.h"
#include "des.h"

#define MAXIMUM_THREADS 8
using namespace std;

float sb0 (uint8_t plainByte, uint8_t key) { return (float)__builtin_popcount(Sb[0][plainByte ^ key]); }
float sb1 (uint8_t plainByte, uint8_t key) { return (float)__builtin_popcount(Sb[1][plainByte ^ key]); }
float sb2 (uint8_t plainByte, uint8_t key) { return (float)__builtin_popcount(Sb[2][plainByte ^ key]); }
float sb3 (uint8_t plainByte, uint8_t key) { return (float)__builtin_popcount(Sb[3][plainByte ^ key]); }
float sb4 (uint8_t plainByte, uint8_t key) { return (float)__builtin_popcount(Sb[4][plainByte ^ key]); }
float sb5 (uint8_t plainByte, uint8_t key) { return (float)__builtin_popcount(Sb[5][plainByte ^ key]); }
float sb6 (uint8_t plainByte, uint8_t key) { return (float)__builtin_popcount(Sb[6][plainByte ^ key]); }
float sb7 (uint8_t plainByte, uint8_t key) { return (float)__builtin_popcount(Sb[7][plainByte ^ key]); }

typedef float (*sboxFunctionPointer) (uint8_t plainByte, uint8_t key);

void readTimeData(char* realtimeFilename, vector<float>& realTime, long numberOfEncryptions){
    FILE *timingFile;
    double currentTiming;

    timingFile = fopen(realtimeFilename,"r");
    printf("%ld timings.\n",numberOfEncryptions);
    for (long i = 0; i < numberOfEncryptions; i++){
        fscanf(timingFile,"%lf\n",&currentTiming);
        realTime[i] = (float)currentTiming;
    }

    fclose(timingFile);
}


void readPlaintextData(char* plaintextFilename,vector< vector<uint8_t> > &rightExpanded, long numberOfEncryptions, int bytesNumber, int keyCandidates){
    FILE *plaintextFile;
    uint32_t Right = 0;
    uint64_t iteratorIP = 0x0ULL;
    uint64_t iteratorExpandedR = 0x0ULL;

    plaintextFile = fopen(plaintextFilename, "r");
    printf("%d plaintexts.\n",numberOfEncryptions);

    for (long i = 0; i < numberOfEncryptions; i++){
        fscanf(plaintextFile,"%llX",&iteratorIP);
        Right = performInitialPermutation(iteratorIP) & 0xffffffff;
        iteratorExpandedR = performExpansion(Right);

        for (int j=0; j<bytesNumber; j++)
            rightExpanded[j][i] = (iteratorExpandedR >> ((bytesNumber - j - 1)*(unsigned int)log2(keyCandidates))) & (keyCandidates-1);

    }
    fclose(plaintextFile);
}

void resizeVectors(vector<vector<uint8_t> > &rightExpanded, vector< vector<double> > &correlationResults, int bytesNumber, int numberOfEncryptions, vector<float> &realTime, vector<float> &model, vector<uint8_t> &key, int keyCandidates){
    model.resize(numberOfEncryptions);
    realTime.resize(numberOfEncryptions);
    key.resize(numberOfEncryptions);


    correlationResults.resize(bytesNumber);
    for (int i = 0; i < bytesNumber; i++)
        correlationResults[i].resize(keyCandidates);

    rightExpanded.resize(bytesNumber);
    for (int i = 0; i < bytesNumber; i++)
        rightExpanded[i].resize(numberOfEncryptions);
}

void populateTimingMode(vector<vector<uint8_t> > &rightExpanded, unsigned int numberOfEncryptions, vector<float> &generatedTimingPredictionModel, vector<uint8_t> &key, int SboxNum){
    int threadCounter = 0;
    sboxFunctionPointer chooseSbox[] = {sb0, sb1, sb2, sb3, sb4, sb5, sb6, sb7};

    omp_set_num_threads(MAXIMUM_THREADS);
    int totalNumberOfThreads = omp_get_num_threads();

    #pragma omp parallel default(none) shared(totalNumberOfThreads, rightExpanded, generatedTimingPredictionModel, key, chooseSbox)
    {
            int tid     = omp_get_thread_num();
            int nitems  = numberOfEncryptions/totalNumberOfThreads;
            int start   = tid*nitems;
            int end     = start + nitems;

            if (tid == totalNumberOfThreads-1) end = numberOfEncryptions;

            #pragma omp for
            for (threadCounter=0; threadCounter<totalNumberOfThreads; threadCounter++)
                std::transform (rightExpanded[SboxNum].begin()+start, rightExpanded[SboxNum].begin()+end, key.begin()+start, generatedTimingPredictionModel.begin()+start, chooseSbox[SboxNum]);
    }
}


double calculateCorrelation(vector<float> &x, vector<float> &y){
    double sumX = 0.0,  sumY = 0.0;
    double sumX2 = 0.0, sumY2 = 0.0;
    double sumXY = 0.0;
    const unsigned int num = x.size();
    const double init = 0;

    double correlationValue = 0.0;
    int nthreads = 0;

    omp_set_dynamic(0);
    omp_set_num_threads(MAXIMUM_THREADS);

    #pragma omp parallel
    #pragma omp single
    nthreads = omp_get_num_threads();

    #pragma omp parallel default(none) reduction(+:sumX, sumX2, sumY, sumY2, sumXY) shared(nthreads,x,y)
    {
            int tid     = omp_get_thread_num();
            int nitems  = num/nthreads;
            int start   = tid*nitems;
            int end     = start + nitems;

            if (tid == nthreads-1) end = num;

            #pragma omp for
            for (int i=0; i<nthreads; i++){
                sumX2   += std::inner_product( x.begin() + start, x.begin()+end, x.begin()+start, init);
                sumX    += std::accumulate(x.begin()+start, x.begin()+end, init);
                sumY2   += std::inner_product( y.begin() + start,y.begin()+end, y.begin()+start, init);
                sumY    += std::accumulate(y.begin()+start, y.begin()+end, init);
                sumXY   += std::inner_product( x.begin() + start, x.begin()+end, y.begin()+start, init);
            }
    }

    correlationValue = (sumXY - sumX*sumY/num) / sqrt((sumX2 - sumX * sumX / num) * (sumY2 - sumY * sumY / num));

    if (!isnan(correlationValue))
        return correlationValue;
    else{
        return 0;
    }

    return 0;
}

double saveResultsToFile(const char * resultFilename, int keyCandidates, vector< vector<double> > &pcc, uint8_t SboxNum){
    FILE *resultsFile;
    resultsFile = fopen(resultFilename,"w");
    fprintf(resultsFile,"#\t\tPrediction\n");
    for (int i = 0; i < keyCandidates; i++)
        fprintf(resultsFile,"%d\t%12f\n",i, pcc[SboxNum][i]*pcc[SboxNum][i]*100000);

    fclose(resultsFile);
}

void singleSboxTimingAttack(unsigned int numberOfEncryptions, char* realtimeFilename, char* plaintextFilename, const char* resultFilename, uint8_t SboxNum){
    const int keyCandidates = 64;
    const int bytesNumber = 8;

    vector< vector<double> > correlationResults;                            /*The resulted likelihood information*/
    vector<float> timingData;                               /*Explicit real time*/
    vector<float> generatedTimingPredictionModel;           /*Timing model that is used as a prediction*/
    vector<uint8_t> key;                                /*Key value*/
    vector<vector<uint8_t> > rightExpanded;                             /*Expanded R, i.e. E(R) permutation of DES*/

    resizeVectors(rightExpanded, correlationResults, bytesNumber, numberOfEncryptions, timingData, generatedTimingPredictionModel, key, keyCandidates);
    readTimeData(realtimeFilename, timingData, numberOfEncryptions);
    readPlaintextData(plaintextFilename, rightExpanded, numberOfEncryptions, bytesNumber, keyCandidates);

    //create model
    for (int i = 0; i < keyCandidates; i++){
        std::fill(key.begin(),key.end(),i); //Entire key is composed of the same key candidate

        populateTimingMode(rightExpanded, numberOfEncryptions, generatedTimingPredictionModel, key, SboxNum);
        correlationResults[SboxNum][i] = calculateCorrelation(timingData, generatedTimingPredictionModel);

        printf("Working...\t %.2lf%%\r\n",(double)(i+1)*(double)100/(double)(keyCandidates));
        fflush(stdout);
    }

    saveResultsToFile(resultFilename, keyCandidates, correlationResults, SboxNum); //write to file
}