#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 insertInitializedpbwt(const int  & insertIndex, bool **panel, dpbwt & x, bool debugging){
    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;
        }
        if (temp == nullptr && panel[insertIndex][k]){
            temp = botk->above;
            botk->u = &z[k+1];
            while (temp != &z[k] && x.panel[temp->id][k]){
                temp->u = &z[k+1];
                temp = temp->above;
            }
        }
        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;
    }
}

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][timeIndex] = get_cpu_time() - oldcputime;
    time[0][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)]);
    }
}

void sortN(int *order, int num){
    int min, minpos;
    for (int i = 0; i < num; i++){
        min = order[i];
        minpos = i;
        for (int j = i+1; j<num; j++)
            if (order[j] < min){
                min = order[j];
                minpos = j;
            }
        for (int j = minpos; j>i; j--)
            order[j] = order[j-1];
        order[i] = min;
    }
}


//TODO remove qM and integrate
//update help message to updated readme
int main(int argc, char *argv[]){
    std::string in = "panel.vcf", tout = "update_dpbwt_time.txt";
    bool help = false;
    bool debugging = false;
    int generateIndel = 10;
    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':
                generateIndel = 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 d-PBWT through insertion of haplotypes in the VCF file, then outputs the time taken for some insertions and deletions. 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 initial d-PBWT is built with the number of haplotypes in the input panel minus n randomly selected haplotypes. Then a random sequence of insertions and deletions is generated with n insertions and n deletions. The deletions are chosen randomly from the d-PBWT while the insertions must be new haplotypes to the d-PBWT. 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_update.exe dpbwt_update.cpp\n\nSample Usage:\n./exedpbwt_update.exe -i example.vcf -o example_updateTime.txt -n 1\n\noptions:\n-i, input VCF file for the panel, default is panel.vcf\n-o, output file for time, default is update_dpbwt_time.txt.\n-n, number of haplotypes to test insertion and deletion, must be less than half the number of haplotypes in the input VCF.\n-h, help\nno parameters, help\n" << std::endl;
        return 0;
    }
    //even is real, odd is cpu
    double *time[2];
    int *uv;
    bool **panel;
    dpbwt x;
    ReadVCF(in, panel);
    std::cout << "M: " << M << "\nN: " << N << std::endl;
    if (generateIndel > M/2){
        std::cout << "-n must be less than half of M. -n = " << generateIndel << ". M = " << M << std::endl;
        return 0;
    }
    if (debugging)
        OutputPanel("indel_dpbwt_Panel.txt", panel, M);
    time[0] = new double[generateIndel*2];
    time[1] = new double[generateIndel*2];
    uv = new int[generateIndel*2];


    int *order = new int[M];
    for (int i = 0; i<M; i++)
        order[i] = i;
    Randomize(order, M);
    sortN(order, generateIndel);
    
    initializedpbwt(x);
    int current = 0;
    for (int i = 0; i < generateIndel; i++){
        for (int j = current; j < order[i]; j++)
            insertInitializedpbwt(j, panel, x, debugging);
        current = order[i]+1;
    }
            

    
    bool *indelOrder = new bool[generateIndel*2];
    std::random_device seed;
    std::default_random_engine generator(seed());
    {
        int inSum = 0, delSum = 0;
        std::uniform_int_distribution<int> dis(0, 1);
        for (int i = 0; i < generateIndel*2; i++){
            indelOrder[i] = dis(generator);
            if (indelOrder[i]){
                inSum++;
                if (inSum == generateIndel){
                    for (int j = i+1; j<generateIndel*2; j++){
                        indelOrder[j] = false;
                    }
                    break;
                }
            }
            else{
                delSum++;
                if (delSum == generateIndel){
                    for (int j = i+1; j<generateIndel*2; j++){
                        indelOrder[j] = true;
                    }
                    break;
                }

            }
        }
    }

    int inNum = 0;
    for (int i = 0; i < generateIndel*2; i++){
        if (indelOrder[i]){
            insertdpbwt(time, uv, i, order[inNum], panel, x, debugging);
            inNum++;
        }
        else{
            std::uniform_int_distribution<int> dis(0, x.panel.size()-1);
            int index = dis(generator);
            deletedpbwt(time, uv, i, index, x, debugging); 
        }
    }
        



    std::ofstream timeout(tout);
    timeout<< "Real(s)\tCPU(s)\tUV\tInsert(1)/Delete(0)\n";
    for (int s = 0; s<generateIndel*2; ++s)
        timeout << time[0][s] << '\t' << time[1][s] <<  '\t' << uv[s] << '\t' << indelOrder[s] <<  '\n';
    timeout.close();
    return 0;
}

