#include<iostream>
#include<iomanip>
#include<fstream>
#include<sstream>
#include<string>
#include<vector>
#include<random>

int M = 0;
int N = 0;


//-----------------------------------------------------------------------------
//source for time: https://stackoverflow.com/questions/17432502/how-can-i-measure-cpu-time-and-wall-clock-time-on-both-linux-windows
//  Windows
#ifdef _WIN32
#include <Windows.h>
double get_wall_time(){
    LARGE_INTEGER time,freq;
    if (!QueryPerformanceFrequency(&freq)){
        //  Handle error
        return 0;
    }
    if (!QueryPerformanceCounter(&time)){
        //  Handle error
        return 0;
    }
    return (double)time.QuadPart / freq.QuadPart;
}
double get_cpu_time(){
    FILETIME a,b,c,d;
    if (GetProcessTimes(GetCurrentProcess(),&a,&b,&c,&d) != 0){
        //  Returns total user time.
        //  Can be tweaked to include kernel times as well.
        return
            (double)(d.dwLowDateTime |
                    ((unsigned long long)d.dwHighDateTime << 32)) * 0.0000001;
    }else{
        //  Handle error
        return 0;
    }
}

//  Posix/Linux
#else
#include <time.h>
#include <sys/time.h>
double get_wall_time(){
    struct timeval time;
    if (gettimeofday(&time,NULL)){
        //  Handle error
        return 0;
    }
    return (double)time.tv_sec + (double)time.tv_usec * .000001;
}
double get_cpu_time(){
    return (double)clock() / CLOCKS_PER_SEC;
}
#endif
//-----------------------------------------------------------------------------

struct dpbwtnode{
    dpbwtnode *below, *above, *u, *v;
    int divergence, id, originalid;
};

struct dpbwt{
    dpbwtnode bottomfirstcol;
    std::vector<dpbwtnode*> firstcol;
    std::vector<std::vector<bool>> panel;
};

void ReadVCF(std::string inFile, bool ** & panel){
    using namespace std;
    {//count M and N
        string line = "1#";
        ifstream in(inFile);
        while (line[1] == '#')
            getline(in, line);
        stringstream linestr;

        linestr.str(line);
        linestr.clear();
        for (int i = 0; i<9; i++)
            linestr >> line;
        N = M = 0;
        while (!linestr.eof()){
            ++M;
            ++M;
            linestr >> line;
        }
        while (getline(in,line)){
            ++N;
        }
        in.close();
    }
    bool *temp = new bool[(long long)M * N];
    panel = new bool*[M];
    for (long long i = 0; i<M; i++){
        panel[i] = &(temp[i*N]);
    }
    string line = "##";
    ifstream in(inFile);
    stringstream linestr;
    int x = 0;
    char y = 0;

    while (line[1] == '#')
        getline(in, line);
    for (int j = 0; j<N; ++j){
        getline(in, line);
        linestr.str(line);
        linestr.clear();
        for (int i = 0; i<9; ++i){
            linestr >> line; 
        }
        for (int i = 0; i<M/2; ++i){
            linestr >> x >> y;
            panel[i*2][j] = (bool)x;
            linestr >> x;
            panel[i*2+1][j] = (bool)x;

        }
    }
    in.close();
}

void OutputPanel(const char* outFile, bool** panel, int num){
    std::ofstream out(outFile);
    for (int i = 0; i<num; ++i){
        for (int j = 0; j<N; ++j)
            out << std::setw(3) << (int)panel[i][j];
        out << std::endl;
    }
    out.close();
}

void OutputAll(const char * outFilep, const char * outFiled, const char * outFileu,
        const char * outFilev, int numHap, dpbwt x){
    std::ofstream outp(outFilep);
    std::ofstream outd(outFiled);
    std::ofstream outu(outFileu);
    std::ofstream outv(outFilev);

    dpbwtnode *top[N+1];
    top[0] = &x.bottomfirstcol;
    while (top[0]->above!=nullptr)
        top[0] = top[0]->above;
    for (int j = 0; j<N; ++j)
        top[j+1] = top[j]->u;
    for (int i=0; i<numHap; ++i){
        for (int j=1; j<N+1; ++j){
            outp << std::setw(7) << (top[j]->originalid);
            outd << std::setw(7) << top[j]->divergence;
            if (j!=N){
                if (top[j]->u->id != -1)
                    outu << std::setw(7) << top[j]->u->originalid;
                else
                    outu << std::setw(7) << -1;
                if (top[j]->v->id != -1)
                    outv << std::setw(7) << top[j]->v->originalid;
                else
                    outv << std::setw(7) << -1;
            }

            top[j] = top[j]->below;
        }
        outp << std::endl;
        outd << std::endl;
        outu << std::endl;
        outv << std::endl;
    }
    outp.close();
    outd.close();
    outu.close();
    outv.close();
}

void initializedpbwt(dpbwt &x){
    x.bottomfirstcol.divergence = 0;
    x.bottomfirstcol.id = -1;
    x.bottomfirstcol.below = x.bottomfirstcol.above = nullptr;
    dpbwtnode *prev = &x.bottomfirstcol, *current;
    for (int j = 0; j<N; ++j){
        current = new dpbwtnode;
        current->divergence = j+1;
        current->id = -1;
        current->below = current->above = nullptr;
        prev->u = prev->v = current;
        prev = current;
    }
}

void insertdpbwt(double *time[], int *uv, const int & timeIndex, const int  & insertIndex, bool **panel, dpbwt & x, 
        bool debugging){
    double oldtime = 0, oldcputime = 0;
    int uvUpdated = 0;

    oldtime = get_wall_time();
    oldcputime = get_cpu_time();

    x.panel.push_back(std::vector<bool>());
    x.panel[x.panel.size()-1].resize(N);
    dpbwtnode *t = &(x.bottomfirstcol),
              *botk = &(x.bottomfirstcol),
              *z = new dpbwtnode[N+1];
    x.firstcol.push_back(&z[0]);

    int newindex = x.panel.size()-1;

    z[0].id = newindex;
    z[0].originalid = insertIndex; //!!!! different from indel benchmark
    z[0].divergence = 0;
    z[0].below = t;
    z[0].above = t->above;
    t->above = &z[0];
    if (z[0].above != nullptr)
        z[0].above->below = &z[0];
    for (int k = 0; k<N; ++k){
        dpbwtnode* temp = z[k].above;
        while (temp != nullptr && x.panel[temp->id][k] != panel[insertIndex][k]){
            if (!panel[insertIndex][k])
                temp->u = &z[k+1];
            else 
                temp->v = &z[k+1];
            temp = temp->above;
            uvUpdated++;
        }
        if (temp == nullptr && panel[insertIndex][k]){
            temp = botk->above;
            botk->u = &z[k+1];
            uvUpdated++;
            while (temp != &z[k] && x.panel[temp->id][k]){
                temp->u = &z[k+1];
                temp = temp->above;
                uvUpdated++;
            }
        }
        if (!panel[insertIndex][k]){
            z[k].u = &z[k+1];
            z[k].v = t->v;
        }
        else{
            z[k].u = t->u;
            z[k].v = &z[k+1];
        }
        t = (panel[insertIndex][k])? t->v : t->u;
        z[k+1].id = newindex;
        z[k+1].originalid = insertIndex;
        z[k+1].below = t;
        z[k+1].above = t->above;
        t->above = &z[k+1];
        if (z[k+1].above != nullptr)
            z[k+1].above->below = &z[k+1];
        botk = botk->v;
        x.panel[newindex][k] = panel[insertIndex][k];
    }

    int zdtemp = N,
        bdtemp = N;

    for (int k = N; k>= 0; --k){
        zdtemp = std::min(zdtemp, k);
        bdtemp = std::min(bdtemp, k);
        if (z[k].above!=nullptr)
            while (zdtemp>0 && panel[insertIndex][zdtemp-1] == x.panel[z[k].above->id][zdtemp-1])
                zdtemp--;
        else 
            zdtemp = k;
        if (z[k].below->id!=-1)
            while (bdtemp>0 && panel[insertIndex][bdtemp-1] == x.panel[z[k].below->id][bdtemp-1])
                bdtemp--;
        else 
            bdtemp = k;
        z[k].divergence = zdtemp;
        z[k].below->divergence = bdtemp;
    }

    time[1][timeIndex] = get_cpu_time() - oldcputime;
    time[0][timeIndex] = get_wall_time() - oldtime;
    uv[timeIndex] = uvUpdated;
}

void deletedpbwt(double *time[], int *uv, const int & timeIndex, const int &deleteIndex, dpbwt & x, bool debugging){
    double oldtime = 0, oldcputime = 0;
    int uvUpdated = 0;

    oldtime = get_wall_time();
    oldcputime = get_cpu_time();

    int lindex = x.firstcol.size()-1;
    dpbwtnode *botk = &x.bottomfirstcol;
    dpbwtnode *sk = x.firstcol[deleteIndex], *skp; 
    dpbwtnode *lk = x.firstcol[lindex], *lkp;
    for (int k = 0; k<N; ++k){
        //sk++ and lp++ would work in this implementation, however
        //they aren't used to demonstrate the general case.
        if (!x.panel[deleteIndex][k])
            skp = sk->u;
        else 
            skp = sk->v;
        if (!x.panel[lindex][k])
            lkp = lk->u;
        else 
            lkp = lk->v;
        dpbwtnode* temp = sk->above;
        while (temp!=nullptr && x.panel[temp->id][k] != x.panel[deleteIndex][k]){
            if (!x.panel[deleteIndex][k])
                temp->u = sk->below->u;
            else 
                temp->v = sk->below->v;
            temp = temp->above;
            uvUpdated++;
        }
        if (temp == nullptr && x.panel[deleteIndex][k]){
            temp = botk->above;
            botk->u = sk->below->v;
            uvUpdated++;
            while (temp != sk && x.panel[temp->id][k]){
                temp->u = sk->below->v;
                temp = temp->above;
                uvUpdated++;
            }
        }
        lk->id = deleteIndex;
        sk->below->divergence = 
            std::max(sk->divergence, sk->below->divergence);
        x.panel[deleteIndex][k] = x.panel[lindex][k];
        sk->below->above = sk->above;
        if (sk->above != nullptr)
            sk->above->below = sk->below;
        botk = botk->v;
        sk = skp;
        lk = lkp;
    }
    lkp->id = deleteIndex;
    skp->below->divergence = 
        std::max(skp->divergence, skp->below->divergence);
    skp->below->above = sk->above;
    if (skp->above != nullptr)
        skp->above->below = skp->below;

    delete [] x.firstcol[deleteIndex];
    x.firstcol[deleteIndex] = x.firstcol[lindex];
    x.firstcol.pop_back();
    x.panel.pop_back();


    time[1][M+timeIndex] = get_cpu_time() - oldcputime;
    time[0][M+timeIndex] = get_wall_time() - oldtime;
    uv[timeIndex] = uvUpdated;
}

void Randomize(int *order, int num){
    std::random_device seed;
    std::default_random_engine generator(seed());
    for (int i = num-1; i>0; --i){
        std::uniform_int_distribution<int> dis(0, i);
        std::swap(order[i], order[dis(generator)]);
    }
}


//TODO remove qM and integrate
//update help message to updated readme
int main(int argc, char *argv[]){
    std::string in = "panel.vcf", tout = "fullindel_dpbwt_time.txt";
    bool help = false;
    bool debugging = false;
    int delUntil = 1000;
    for (int i = 1; i<argc; ++i){
        switch (argv[i][1]){
            case 'i':
            case 'I':
                in = argv[++i];
                break;
            case 'o':
            case 'O':
                tout = argv[++i];
                break;
            case 'n':
            case 'N':
                delUntil = atoi(argv[++i]);
                break;
            case 'd':
            case 'D':
                debugging = true;
                break;
            case 'h':
            case 'H':
                help = true;
        }
    }
    if (argc < 2 || help){
        std::cout << "This program takes VCF files as input. It builds a PBWT through insertion of all haplotypes in the input VCF, then deletes haplotypes randomly until a user specified amount of haplotypes are left. The run time for all insertions and deletions are outputted. The number of u and v pointer updates per operation is also outputted. It is recommended for multiallelic sites to be removed from the input VCF. If multiallelic sites are not removed, the code will take 0 as 0 and any other allele as 1. Site 0 is the top most site in the vcf and haplotype 0 is the left most haplotype in the vcf. The algorithms used are presented in https://www.biorxiv.org/content/10.1101/2020.01.14.906487v1.\n\nCompile with std=c++17 or higher. The following command may be used:\ng++ -O3 -std=c++17 -o exedpbwt_fullindel.exe dpbwt_fullindel.cpp\n\nSample Usage:\n./exedpbwt_fullindel.exe -i example.vcf -o example_fullindelTime.txt -n 1\n\noptions:\n-i, input VCF file for the panel, default is panel.vcf\n-n, number of haplotypes to keep in the d-PBWT. After insertion of every haplotype in the input panel, haplotypes will be deleted randomly until n haplotypes are left in the d-PBWT. Must be >= 0 and <= number of haplotypes in the input file.\n-o output file for time and u and v pointers updated per operation, default is fullindel_dpbwt_time.txt.\n-h, help\nno parameters, help" << std::endl;
        return 0;
    }
    //even is real, odd is cpu
    double *time[2];
    int *inuv, *deluv;
    bool **panel;
    dpbwt x;
    ReadVCF(in, panel);
    std::cout << "M: " << M << "\nN: " << N << std::endl;
    if (delUntil < 1 || delUntil > M){
        std::cout << "-n must be greater than 0 and less than or equal to M. -n = " << delUntil << ". M = " << M << std::endl;
        return 0;
    }
    if (debugging)
        OutputPanel("fullindel_dpbwt_Panel.txt", panel, M);
    time[0] = new double[M + M-delUntil];
    time[1] = new double[M + M-delUntil];
    inuv = new int[M];
    deluv = new int[M-delUntil];


    int *order = new int[M];
    for (int i = 0; i<M; i++)
        order[i] = i;
    Randomize(order, M);

    initializedpbwt(x);

    std::random_device seed;
    std::default_random_engine generator(seed());

    for (int i = 0; i<M; i++)
        insertdpbwt(time, inuv, i, order[i], panel, x, debugging);
    for (int i = M; i>delUntil; i--){
        std::uniform_int_distribution<int> dis(0, i-1);
        int index = dis(generator);
        deletedpbwt(time, deluv, M-i, index, x, debugging);
    }





    std::ofstream timeout(tout);
    timeout<< "Real(s)\tCPU(s)\tUV\tInsert(1)/Delete(0)\n";
    for (int i = 0; i<M; i++)
        timeout << time[0][i] << '\t' << time[1][i] << '\t' << inuv[i] << "\t1\n";
    for (int i = 0; i<M-delUntil; i++)
        timeout << time[0][M+i] << '\t' << time[1][M+i] << '\t' << deluv[i] << "\t0\n";
    timeout << std::endl;
    timeout.close();
    return 0;
}

