#pragma once

#include <set>
#include <map>
#include <vector>
#include <bitset>
#include <utility>
#include <algorithm>

#include <random>
#include <chrono>

#include <string>
#include <iostream>

#include <cmath>

#include <boost/math/special_functions/binomial.hpp>

#include <boost/thread.hpp>
#include <boost/bind/bind.hpp>
#include <boost/asio/post.hpp>
#include <boost/asio/thread_pool.hpp>
#include <boost/thread/mutex.hpp>

#include "rtuple.hpp"
#include "network.hpp"
#include "diffusionstate.hpp"
#include "sortmap.hpp"
#include "tools.hpp"
#include "macro.hpp"

// double computeG(DiffusionState_MIC &diffusionState, std::set<int> &S, std::vector<rTuple> &rtup, int n, std::string type, double *result){
//     int count = 0;
//     for(rTuple &rt: rtup){
//         switch(type[0]){
//         case 'u':
//             if(intersection(S, rt.upper)) count++;
//             break;
//         case 'm':
//             if(diffusionState.compute_g(S, rt))  count++;
//             break;
//         case 'l':
//             if(intersection(S, rt.lower)) count++;
//             break;
//         default:
//             std::cout << "invalid type" << std::endl;
//         }
//     }
//     *result = (double)n*count/rtup.size();
//     return *result;
// }

std::bitset<THREAD> scheduler;
boost::mutex mt;

Results HighDegree_computeSeedSet(Network &network, DiffusionState_MIC &diffusionState, int k, int span){
    std::cout << "========== High degree running ==========" << std::endl;
    auto start = std::chrono::high_resolution_clock::now();
    network.sortByDegree();
    int node;
    std::set<int> solution;
    for(int i = 0;i < network.vertexNum;i++){
        node = network.sorted_degree[i];
        if(diffusionState.seed_state[node] == -1 && solution.find(node) == solution.end()){
            solution.insert(node);
            if(solution.size() == k){
                // std::cout << "========== High degree finish ==========" << std::endl << std::endl;
                break;
            }
        }
    }
    Results result;
    for(int i = 0;i < k/span;i++){
        int k = i*span+span;
        std::set<int> temp_solution;
        std::set<int>::iterator iter = solution.begin();
        for(int j = 0;j < k;j++, iter++)
            temp_solution.insert(*iter);
        result.seedset[k] = temp_solution;
        result.supp[k] = 0.;
    }
    auto end = std::chrono::high_resolution_clock::now();
    printTime(start, end);
    std::cout << "========== High degree finish ==========" << std::endl << std::endl;
    return result;
}

std::set<int> NaiveGreedy_computeSeedSet(const Network &network, DiffusionState_MIC &diffusionState, int k, double epsilon_1, double N, int times){
    std::cout << "========== Naive greedy running ==========" << std::endl;
    auto start = std::chrono::high_resolution_clock::now();
    std::set<int> solution;
    int cmaxindex, cindex, *state = new int[network.vertexNum];
    double cmaxvalue, tempvalue;
    memcpy(state, diffusionState.seed_state, network.vertexNum*sizeof(int));

    for(int i = 0;i < k;i++){
        std::cout << "Naive greedy #" << i+1 << std::endl;
        cmaxindex = -1;
        cmaxvalue = -1.;
        for(int j = 0;j < network.vertexNum;j++){
            if(state[j] == -1)  solution.insert(j);
            cindex = diffusionState.seed(solution);
            tempvalue = diffusionState.expInfluenceComplete(network, times, cindex);
            if(tempvalue > cmaxvalue){
                cmaxvalue = tempvalue;
                cmaxindex = j;
            }
            diffusionState.removeSeed(cindex);
            solution.erase(j);
        }
    }
    auto end = std::chrono::high_resolution_clock::now();
    printTime(start, end);
    std::cout << "========== Naive greedy finish ==========" << std::endl << std::endl;
    return solution;
}

void __parallel(DiffusionState_MIC &diffusionState, std::set<int> seed, int v, std::vector<rTuple> &rtup, const std::map<int,int> &coverred, std::set<int> &type1, std::set<int> &type2, std::set<int> &utype1, double *result){
    int ret, temp = 0;
    int cmaxindex = -1;
    int tid = -1;
    double cmaxvalue = -1.;
    mt.lock();
    tid = scheduler._Find_first();
    scheduler.flip(tid);
    mt.unlock();
    type1.clear();
    type2.clear();
    utype1.clear();
    seed.insert(v);
    for(std::pair<int,int> p: coverred){
        ret = diffusionState.compute_g(seed, rtup[p.first], "mid", nullptr, tid);
        switch(p.second){
        case 0:
            switch(ret){
            case 2:
                temp++;
                type2.insert(p.first);
                break;
            case 1:
                temp++;
                type1.insert(p.first);
                break;
            }
            break;
        case 1:
            switch(ret){
            case -1:
                temp--;
                utype1.insert(p.first);
                break;
            }
            break;
        }
    }
    *result = temp;
    mt.lock();
    scheduler.flip(tid);
    mt.unlock();
}

void Sandwich_greedyMid(const Network &network, DiffusionState_MIC &diffusionState, std::vector<rTuple> &rtup, std::set<int> &solution, int k, double *coverred, bool verbose=true){
    std::set<int> candidate;
    std::map<int,int> coverred_state;
    int tid = 0;
    for(rTuple &rt: rtup){
        for(int v: rt.upper)
            if(candidate.find(v) == candidate.end())
                candidate.insert(v);
        coverred_state[tid++] = 0;
    }
    double profit = 0.;
    int candidate_size = candidate.size();
    double *results = new double[candidate_size];
    std::vector<std::set<int>> type1s(candidate_size), type2s(candidate_size), utype1s(candidate_size);
    for(int i = 0;i < k;i++){
        int cmaxindex = -1, max_id = 0;
        double cmaxvalue = -1.;
        tid = 0;
        boost::asio::thread_pool pool(THREAD);
        allSet(scheduler);
        for(int v: candidate){
            auto bind_fn = boost::bind(__parallel, ref(diffusionState), ref(solution), v, ref(rtup), ref(coverred_state), ref(type1s[tid]), ref(type2s[tid]), ref(utype1s[tid]), results+tid);
            boost::asio::post(pool, bind_fn);
            tid++;
        }
        pool.join();
        std::set<int>::iterator it = candidate.begin();
        for(int j = 0;j < candidate_size;j++, it++){
            if(cmaxvalue < results[j]){
                cmaxvalue = results[j];
                cmaxindex = *it;
                max_id = j;
            }
        }
        if(cmaxvalue < 0)
            break;
        profit += cmaxvalue;
        if(coverred)    coverred[i] = profit;
        solution.insert(cmaxindex);
        candidate.erase(cmaxindex);
        for(int j: type2s[max_id])
            coverred_state.erase(j);
        for(int j: type1s[max_id]){
            // if(coverred_state[j] != 0){
            //     std::cout << "ERROR: setting type 1 wrong, != 0" << std::endl;
            //     exit(1);
            // }
            coverred_state[j] = 1;
        }
        for(int j: utype1s[max_id]){
            // if(coverred_state[j] != 1){
            //     std::cout << "ERROR: setting type 1 wrong, != 1" << std::endl;
            //     exit(1);
            // }
            coverred_state[j] = 0;
        }
        if(verbose) std::cout << "greedy mid #" << i+1 << ": " << profit << std::endl;
    }
    delete [] results;
}

Results ReverseGreedy_computeSeedSet(const Network &network, DiffusionState_MIC &diffusionState, int k, int l, int span){
    std::cout << "========== Reverse greedy running ==========" << std::endl;
    auto start = std::chrono::high_resolution_clock::now();
    std::vector<rTuple> rtup;
    std::cout << "l " << l << " " << diffusionState.getRTuples(network, rtup, l) << std::endl;
    std::set<int> mid_solution;
    std::cout << "working on mid-solution..." << std::endl;
    Sandwich_greedyMid(network, diffusionState, rtup, mid_solution, k, nullptr);
    Results result;
    bool flag = true;
    for(int i = 0;i < k/span && flag;i++){
        int nk = i*span+span;
        std::set<int> solution;
        std::set<int>::iterator iter = mid_solution.begin();
        for(int j = 0;j < nk;j++, iter++){
            if(iter == mid_solution.end()){
                flag = false;
                break;
            }
            solution.insert(*iter);
        }
        result.seedset[nk] = solution;
        result.supp[nk] = 0.;
    }
    auto end = std::chrono::high_resolution_clock::now();
    printTime(start, end);
    std::cout << "========== Reverse greedy finish ==========" << std::endl << std::endl;
    return result;
}

Results ReverseGreedy_computeSeedSet_noneoutput(const Network &network, DiffusionState_MIC &diffusionState, int k, int l, int span){
    auto start = std::chrono::high_resolution_clock::now();
    std::vector<rTuple> rtup;
    diffusionState.getRTuples(network, rtup, l);
    std::set<int> mid_solution;
    Sandwich_greedyMid(network, diffusionState, rtup, mid_solution, k, nullptr, false);
    Results result;
    bool flag = true;
    for(int i = 0;i < k/span && flag;i++){
        int nk = i*span+span;
        std::set<int> solution;
        std::set<int>::iterator iter = mid_solution.begin();
        for(int j = 0;j < nk;j++, iter++){
            if(iter == mid_solution.end()){
                flag = false;
                break;
            }
            solution.insert(*iter);
        }
        result.seedset[nk] = solution;
        result.supp[nk] = 0.;
    }
    auto end = std::chrono::high_resolution_clock::now();
    printTime(start, end);
    return result;
}

Results ReverseGreedy_computeSeedSet(const Network &network, DiffusionState_MIC &diffusionState, int k, std::vector<rTuple> &rtup, int span){
    std::cout << "========== Reverse greedy running ==========" << std::endl;
    auto start = std::chrono::high_resolution_clock::now();
    std::cout << "r_tuple size: " << rtup.size() << std::endl;
    std::set<int> mid_solution;
    std::cout << "working on mid-solution..." << std::endl;
    Sandwich_greedyMid(network, diffusionState, rtup, mid_solution, k, nullptr);
    Results result;
    bool flag = true;
    for(int i = 0;i < k/span && flag;i++){
        int nk = i*span+span;
        std::set<int> solution;
        std::set<int>::iterator iter = mid_solution.begin();
        for(int j = 0;j < nk;j++, iter++){
            if(iter == mid_solution.end()){
                flag = false;
                break;
            }
            solution.insert(*iter);
        }
        result.seedset[nk] = solution;
        result.supp[nk] = 0.;
    }
    auto end = std::chrono::high_resolution_clock::now();
    printTime(start, end);
    std::cout << "========== Reverse greedy finish ==========" << std::endl << std::endl;
    return result;
}

Results ReverseGreedy_computeSeedSet(const Network &network, DiffusionState_MIC &diffusionState, int k, std::vector<rTuple> &rtup){
    std::cout << "========== Reverse greedy running ==========" << std::endl;
    auto start = std::chrono::high_resolution_clock::now();
    std::cout << "r_tuple size: " << rtup.size() << std::endl;
    std::set<int> mid_solution;
    std::cout << "working on mid-solution..." << std::endl;
    Sandwich_greedyMid(network, diffusionState, rtup, mid_solution, k, nullptr);
    Results result;
    result.seedset[k] = mid_solution;
    result.supp[k] = 0.;
    auto end = std::chrono::high_resolution_clock::now();
    printTime(start, end);
    std::cout << "========== Reverse greedy finish ==========" << std::endl << std::endl;
    return result;
}

double Sandwich_greedy(std::vector<rTuple> &rtup, std::set<int> &solution, int k, const std::string &type){
    std::vector<std::set<int>> rrsets;
    for(rTuple &rt: rtup){
        switch(type[0]){
        case 'u':
            rrsets.push_back(rt.upper);
            break;
        case 'l':
            rrsets.push_back(rt.lower);
            break;
        default:
            std::cout << "invalid model" << std::endl;
        }
    }

    double coverred = 0.;
    double profit = 0.;
    std::map<int, std::set<int>> nodes_cover_sets;
    bool coverred_rrsets[rrsets.size()];
    for(int i = 0;i < rrsets.size();i++){
        for(int index: rrsets[i]){
            if(nodes_cover_sets.find(index) == nodes_cover_sets.end())
                nodes_cover_sets[index] = std::set<int>();
            nodes_cover_sets[index].insert(i);
        }
        coverred_rrsets[i] = false;
    }
    std::vector<std::pair<int,int>> sortPair;
    for(auto i = nodes_cover_sets.begin();i != nodes_cover_sets.end();i++)
        sortPair.push_back(std::make_pair(i->first, i->second.size()));
    std::sort(sortPair.begin(), sortPair.end(), [](const std::pair<int,int> a, const std::pair<int,int> b)-> bool {
        return a.second > b.second;
    });
    sortedMap mymap;
    for(std::pair<int,int> p: sortPair)
        mymap.push_back(p.first, p.second);

    int c_bound, t_bound, c_seed;
    bool sign;
    std::set<int> c_seed_cover;
    for(int i = 0;i < k;i++){
        sign = false;
        c_bound = mymap.size();
        while(c_bound > 0){
            c_seed = mymap.get(0);
            c_seed_cover = nodes_cover_sets[c_seed];
            for(auto j = c_seed_cover.begin();j != c_seed_cover.end();){
                std::set<int>::iterator temp_j = ++j;
                j--;
                if(coverred_rrsets[*j]){
                    c_seed_cover.erase(*j);
                    j = temp_j;
                } else  j++;
            }
            // std::cout << "greedy" << std::endl;
            // for(int j: c_seed_cover){
            //     if(coverred_rrsets[j])
            //         c_seed_cover.erase(j);
            // }
            t_bound = mymap.update(c_seed, c_seed_cover.size());
            if(t_bound == 0){
                solution.insert(c_seed);
                sign = true;
                for(int j: c_seed_cover){
                    if(coverred_rrsets[j]){
                        std::cout << "greedy update may wrong" << std::endl;
                        exit(1);
                    }
                    else{
                        coverred_rrsets[j] = true;
                        coverred += 1.;
                    }
                }
                break;
            }
            if(t_bound < c_bound)   c_bound = t_bound;
        }
        if(!sign){
            std::cout << "greedy lazy: no node selected" << std::endl;
            exit(1);
        }
    }
    return coverred;
}

double Sandwich_computeLowerBound(const Network &network, DiffusionState_MIC &diffusionState, int k, double eps1, double N, int p, const std::string &type="lower"){
    int n = network.vertexNum;
    std::vector<rTuple> rtup;
    std::set<int> S;
    double eps0 = 100 * eps1;
    // boost::math::binomial_coefficient<double>(n,k)
    double lambda = (n*(2+eps0)*log(N*k*log(n)*log2(n)))/(eps0*eps0);
    double x, l, g_lower, pw = p;
    for(int i = 1;i < (int)log2(n-1.)/log2(p);i++){
        x = n/pw;
        l = lambda/x;
        pw *= p;
        if(rtup.size() < l)
            diffusionState.getRTuples(network, rtup, (int)(l-rtup.size()));
        S.clear();
        Sandwich_greedy(rtup, S, k, type);
        g_lower = diffusionState.computeG(S, rtup, n, type, &g_lower);
        std::cout << g_lower << " " << (1+eps0)*x << std::endl;
        if(g_lower >= (1+eps0)*x)   return g_lower;
    }
    std::cout << "compute lower bound may wrong, opt too small" << std::endl;
    exit(1);
}

double Sandwich_decideL(int n, int k, double low_bound, double eps1, double eps2, double N, double *l2){
    // double l1 = ((2+eps1)*n*(log(N+boost::math::binomial_coefficient<double>(n,k))))/(eps1*eps1);
    double l1 = ((2+eps1)*n*log(N*k*log(n)))/(eps1*eps1);
    *l2 = 2*n*log(N)/(eps2*eps2);
    return (l1>*l2?l1:*l2)/low_bound;
}

int Sandwich_computeSeedSet(const Network &network, DiffusionState_MIC &diffusionState, int k, double eps1, double N, std::set<int> &solution, int p, double *l2){
    // std::cout << "========== Sandwich running ==========" << std::endl;
    auto start = std::chrono::high_resolution_clock::now();
    double eps0 = eps1, eps2 = (eps1*log(N))/(log(network.vertexNum)+log(N));
    eps2 = eps1;
    double low_bound = Sandwich_computeLowerBound(network, diffusionState, k, eps1, N, p);
    int l = (int)Sandwich_decideL(network.vertexNum, k, low_bound, eps1, eps2, N, l2);
    *l2 /= low_bound;
    // std::cout << "l2  " << l2 << std::endl;
    std::vector<rTuple> rtup;
    // std::cout << "l " << l << " " << diffusionState.getRTuples(network, rtup, l) << std::endl;
    diffusionState.getRTuples(network, rtup, l);
    std::set<int> upper_solution, lower_solution;
    // std::cout << "working on upper solution..." << std::endl;
    Sandwich_greedy(rtup, upper_solution, k, "upper");
    // std::cout << "working on lower solution..." << std::endl;
    Sandwich_greedy(rtup, lower_solution, k, "lower");
    // std::cout << "calculating upper G... " << std::endl;
    double upper_g = diffusionState.computeG(upper_solution, rtup, network.vertexNum, "mid", &upper_g);
    // std::cout << "  upper G = " << upper_g << std::endl;
    // std::cout << "calculating lower G... " << std::endl;
    double lower_g = diffusionState.computeG(lower_solution, rtup, network.vertexNum, "mid", &lower_g);
    // std::cout << "  lower G = " << lower_g << std::endl;
    if(upper_g > lower_g)
        for(int upper_v: upper_solution)
            solution.insert(upper_v);
    else
        for(int lower_v: lower_solution)
            solution.insert(lower_v);
    auto end = std::chrono::high_resolution_clock::now();
    printTime(start, end);
    // std::cout << "========== Sandwich finish ==========" << std::endl << std::endl;
    return l;
}

Results Sandwich_computeSeedSet(const Network &network, DiffusionState_MIC &diffusionState, int k, double eps1, double N, std::vector<rTuple> &rtup, int p, int span, double *l2){
    std::cout << "========== Sandwich running ==========" << std::endl;
    auto start = std::chrono::high_resolution_clock::now();
    double eps0 = eps1, eps2 = (eps1*log(N))/(log(network.vertexNum)+log(N));
    eps2 = eps1;
    double low_bound = Sandwich_computeLowerBound(network, diffusionState, k, eps1, N, p);
    int l = (int)Sandwich_decideL(network.vertexNum, k, low_bound, eps1, eps2, N, l2);
    *l2 /= low_bound;
    // *l2 = l/log2(network.vertexNum);
    std::cout << "l " << l;
    std::cout << " " << diffusionState.getRTuples(network, rtup, l) << std::endl;
    // diffusionState.getRTuples(network, rtup, l);
    std::set<int> upper_solution, lower_solution;
    std::cout << "working on upper solution..." << std::endl;
    Sandwich_greedy(rtup, upper_solution, k, "upper");
    std::cout << "working on lower solution..." << std::endl;
    Sandwich_greedy(rtup, lower_solution, k, "lower");
    Results result;
    for(int i = 0;i < k/span;i++){
        int nk = i*span+span;
        std::set<int> upper_solution_k, lower_solution_k;
        std::set<int>::iterator upper_iter = upper_solution.begin();
        std::set<int>::iterator lower_iter = lower_solution.begin();
        for(int j = 0;j < nk;j++, upper_iter++, lower_iter++){
            upper_solution_k.insert(*upper_iter);
            lower_solution_k.insert(*lower_iter);
        }
        double upper_G = diffusionState.computeG(upper_solution_k, rtup, network.vertexNum, "mid", nullptr);
        double lower_G = diffusionState.computeG(lower_solution_k, rtup, network.vertexNum, "mid", nullptr);
        double ratio = upper_G / diffusionState.computeG(upper_solution_k, rtup, network.vertexNum, "upper", nullptr);

        result.supp[nk] = ratio;
        if(upper_G > lower_G)
            result.seedset[nk] = upper_solution_k;
        else
            result.seedset[nk] = lower_solution_k;
    }
    auto end = std::chrono::high_resolution_clock::now();
    printTime(start, end);
    std::cout << "========== Sandwich finish ==========" << std::endl << std::endl;
    return result;
}

Results Sandwich_computeSeedSet(const Network &network, DiffusionState_MIC &diffusionState, int k, int l, int p, int span){
    std::cout << "========== Sandwich running ==========" << std::endl;
    auto start = std::chrono::high_resolution_clock::now();
    std::vector<rTuple> rtup;
    std::cout << "l " << l << " " << diffusionState.getRTuples(network, rtup, l) << std::endl;
    std::set<int> upper_solution, lower_solution;
    std::cout << "working on upper solution..." << std::endl;
    Sandwich_greedy(rtup, upper_solution, k, "upper");
    std::cout << "working on lower solution..." << std::endl;
    Sandwich_greedy(rtup, lower_solution, k, "lower");
    Results result;
    for(int i = 0;i < k/span;i++){
        int nk = i*span+span;
        std::set<int> upper_solution_k, lower_solution_k;
        std::set<int>::iterator upper_iter = upper_solution.begin();
        std::set<int>::iterator lower_iter = lower_solution.begin();
        for(int j = 0;j < nk;j++, upper_iter++, lower_iter++){
            upper_solution_k.insert(*upper_iter);
            lower_solution_k.insert(*lower_iter);
        }
        double upper_G = diffusionState.computeG(upper_solution_k, rtup, network.vertexNum, "mid", nullptr);
        double lower_G = diffusionState.computeG(lower_solution_k, rtup, network.vertexNum, "mid", nullptr);
        double ratio = upper_G / diffusionState.computeG(upper_solution_k, rtup, network.vertexNum, "upper", nullptr);

        result.supp[nk] = ratio;
        if(upper_G > lower_G)
            result.seedset[nk] = upper_solution_k;
        else
            result.seedset[nk] = lower_solution_k;
    }
    auto end = std::chrono::high_resolution_clock::now();
    printTime(start, end);
    std::cout << "========== Sandwich finish ==========" << std::endl << std::endl;
    return result;
}

Results Sandwich_computeSeedSet(const Network &network, DiffusionState_MIC &diffusionState, int k, int l, std::vector<rTuple> &rtup){
    // std::cout << "========== Sandwich running ==========" << std::endl;
    auto start = std::chrono::high_resolution_clock::now();
    // std::cout << "l " << l << " " << diffusionState.getRTuples(network, rtup, l) << std::endl;
    diffusionState.getRTuples(network, rtup, l);
    std::set<int> upper_solution, lower_solution;
    // std::cout << "working on upper solution..." << std::endl;
    Sandwich_greedy(rtup, upper_solution, k, "upper");
    // std::cout << "working on lower solution..." << std::endl;
    Sandwich_greedy(rtup, lower_solution, k, "lower");
    Results result;
    double upper_G = diffusionState.computeG(upper_solution, rtup, network.vertexNum, "mid", nullptr);
    double lower_G = diffusionState.computeG(lower_solution, rtup, network.vertexNum, "mid", nullptr);
    double ratio = upper_G / diffusionState.computeG(upper_solution, rtup, network.vertexNum, "upper", nullptr);

    result.supp[k] = ratio;
    if(upper_G > lower_G)
        result.seedset[k] = upper_solution;
    else
        result.seedset[k] = lower_solution;
    auto end = std::chrono::high_resolution_clock::now();
    // printTime(start, end);
    // std::cout << "========== Sandwich finish ==========" << std::endl << std::endl;
    return result;
}

Results Sandwich_computeSeedSet_upper(const Network &network, DiffusionState_MIC &diffusionState, int k, double eps1, double N, std::vector<rTuple> &rtup, int p, int span, double *l2){
    diffusionState.setUpper();
    std::cout << "========== Sandwich running ==========" << std::endl;
    auto start = std::chrono::high_resolution_clock::now();
    double eps0 = eps1, eps2 = (eps1*log(N))/(log(network.vertexNum)+log(N));
    eps2 = eps1;
    double low_bound = Sandwich_computeLowerBound(network, diffusionState, k, eps1, N, p, "upper");
    int l = (int)Sandwich_decideL(network.vertexNum, k, low_bound, eps1, eps2, N, l2);
    *l2 /= low_bound;
    // *l2 = l/log2(network.vertexNum);
    std::cout << "l " << l << " " << diffusionState.getRTuples(network, rtup, l) << std::endl;
    // diffusionState.getRTuples(network, rtup, l);
    std::set<int> upper_solution;
    std::cout << "working on upper solution..." << std::endl;
    Sandwich_greedy(rtup, upper_solution, k, "upper");
    // std::cout << "working on lower solution..." << std::endl;
    // Sandwich_greedy(rtup, lower_solution, k, "lower");
    Results result;
    for(int i = 0;i < k/span;i++){
        int nk = i*span+span;
        std::set<int> upper_solution_k;
        std::set<int>::iterator upper_iter = upper_solution.begin();
        for(int j = 0;j < nk;j++, upper_iter++)
            upper_solution_k.insert(*upper_iter);
        double upper_G = diffusionState.computeG(upper_solution_k, rtup, network.vertexNum, "mid", nullptr);
        result.seedset[nk] = upper_solution_k;
        // double lower_G = diffusionState.computeG(lower_solution_k, rtup, network.vertexNum, "mid", nullptr);
        // double ratio = upper_G / diffusionState.computeG(upper_solution_k, rtup, network.vertexNum, "upper", nullptr);

        // result.supp[nk] = ratio;
        // if(upper_G > lower_G)
        //     result.seedset[nk] = upper_solution_k;
        // else
        //     result.seedset[nk] = lower_solution_k;
    }
    auto end = std::chrono::high_resolution_clock::now();
    printTime(start, end);
    std::cout << "========== Sandwich finish ==========" << std::endl << std::endl;
    diffusionState.setBack();
    return result;
}

Results Sandwich_computeSeedSet_lower(const Network &network, DiffusionState_MIC &diffusionState, int k, double eps1, double N, std::vector<rTuple> &rtup, int p, int span, double *l2){
    diffusionState.setLower();
    std::cout << "========== Sandwich running ==========" << std::endl;
    auto start = std::chrono::high_resolution_clock::now();
    double eps0 = eps1, eps2 = (eps1*log(N))/(log(network.vertexNum)+log(N));
    eps2 = eps1;
    double low_bound = Sandwich_computeLowerBound(network, diffusionState, k, eps1, N, p);
    int l = (int)Sandwich_decideL(network.vertexNum, k, low_bound, eps1, eps2, N, l2);
    *l2 /= low_bound;
    // *l2 = l/log2(network.vertexNum);
    std::cout << "l " << l << " " << diffusionState.getRTuples(network, rtup, l) << std::endl;
    // diffusionState.getRTuples(network, rtup, l);
    std::set<int> lower_solution;
    // std::cout << "working on upper solution..." << std::endl;
    // Sandwich_greedy(rtup, upper_solution, k, "upper");
    std::cout << "working on lower solution..." << std::endl;
    Sandwich_greedy(rtup, lower_solution, k, "lower");
    Results result;
    for(int i = 0;i < k/span;i++){
        int nk = i*span+span;
        std::set<int> lower_solution_k;
        std::set<int>::iterator lower_iter = lower_solution.begin();
        for(int j = 0;j < nk;j++, lower_iter++)
            lower_solution_k.insert(*lower_iter);
        double lower_G = diffusionState.computeG(lower_solution_k, rtup, network.vertexNum, "mid", nullptr);
        result.seedset[nk] = lower_solution_k;
    }
    auto end = std::chrono::high_resolution_clock::now();
    printTime(start, end);
    std::cout << "========== Sandwich finish ==========" << std::endl << std::endl;
    diffusionState.setBack();
    return result;
}


//
//=======================================================================
// Copyright 2012 Fernando Vilas
//           2010 Daniel Trebbien
//
// Distributed under the Boost Software License, Version 1.0. (See
// accompanying file LICENSE_1_0.txt or copy at
// http://www.boost.org/LICENSE_1_0.txt)
//=======================================================================
//

// The maximum adjacency search algorithm was originally part of the
// Stoer-Wagner min cut implementation by Daniel Trebbien. It has been
// broken out into its own file to be a public search algorithm, with
// visitor concepts.
#ifndef BOOST_GRAPH_MAXIMUM_ADJACENCY_SEARCH_H
#define BOOST_GRAPH_MAXIMUM_ADJACENCY_SEARCH_H

/**
 * This is an implementation of the maximum adjacency search on an
 * undirected graph. It allows a visitor object to perform some
 * operation on each vertex as that vertex is visited.
 *
 * The algorithm runs as follows:
 *
 * Initialize all nodes to be unvisited (reach count = 0)
 *   and call vis.initialize_vertex
 * For i = number of nodes in graph downto 1
 *   Select the unvisited node with the highest reach count
 *     The user provides the starting node to break the first tie,
 *     but future ties are broken arbitrarily
 *   Visit the node by calling vis.start_vertex
 *   Increment the reach count for all unvisited neighbors
 *     and call vis.examine_edge for each of these edges
 *   Mark the node as visited and call vis.finish_vertex
 *
 */

#include <boost/concept_check.hpp>
#include <boost/concept/assert.hpp>
#include <boost/graph/buffer_concepts.hpp>
#include <boost/graph/exception.hpp>
#include <boost/graph/graph_concepts.hpp>
#include <boost/graph/iteration_macros.hpp>
#include <boost/graph/named_function_params.hpp>
#include <boost/graph/visitors.hpp>
#include <boost/tuple/tuple.hpp>

#include <set>

namespace boost
{
template < class Visitor, class Graph > struct MASVisitorConcept
{
    void constraints()
    {
        boost::function_requires<
            boost::CopyConstructibleConcept< Visitor > >();
        vis.initialize_vertex(u, g);
        vis.start_vertex(u, g);
        vis.examine_edge(e, g);
        vis.finish_vertex(u, g);
    }
    Visitor vis;
    Graph g;
    typename boost::graph_traits< Graph >::vertex_descriptor u;
    typename boost::graph_traits< Graph >::edge_descriptor e;
};

template < class Visitors = null_visitor > class mas_visitor
{
public:
    mas_visitor() {}
    mas_visitor(Visitors vis) : m_vis(vis) {}

    template < class Vertex, class Graph >
    void initialize_vertex(Vertex u, Graph& g)
    {
        invoke_visitors(m_vis, u, g, ::boost::on_initialize_vertex());
    }

    template < class Vertex, class Graph > void start_vertex(Vertex u, Graph& g)
    {
        invoke_visitors(m_vis, u, g, ::boost::on_start_vertex());
    }

    template < class Edge, class Graph > void examine_edge(Edge e, Graph& g)
    {
        invoke_visitors(m_vis, e, g, ::boost::on_examine_edge());
    }

    template < class Vertex, class Graph >
    void finish_vertex(Vertex u, Graph& g)
    {
        invoke_visitors(m_vis, u, g, ::boost::on_finish_vertex());
    }

    BOOST_GRAPH_EVENT_STUB(on_initialize_vertex, mas)
    BOOST_GRAPH_EVENT_STUB(on_start_vertex, mas)
    BOOST_GRAPH_EVENT_STUB(on_examine_edge, mas)
    BOOST_GRAPH_EVENT_STUB(on_finish_vertex, mas)

protected:
    Visitors m_vis;
};
template < class Visitors >
mas_visitor< Visitors > make_mas_visitor(Visitors vis)
{
    return mas_visitor< Visitors >(vis);
}
typedef mas_visitor<> default_mas_visitor;

namespace detail
{
    template < class Graph, class WeightMap, class MASVisitor,
        class VertexAssignmentMap, class KeyedUpdatablePriorityQueue >
    void maximum_adjacency_search(const Graph& g, WeightMap weights,
        MASVisitor vis,
        const typename boost::graph_traits< Graph >::vertex_descriptor start,
        VertexAssignmentMap assignments, KeyedUpdatablePriorityQueue pq)
    {
        typedef typename boost::graph_traits< Graph >::vertex_descriptor
            vertex_descriptor;
        typedef typename boost::property_traits< WeightMap >::value_type
            weight_type;

        std::set< vertex_descriptor > assignedVertices;

        // initialize `assignments` (all vertices are initially
        // assigned to themselves)
        BGL_FORALL_VERTICES_T(v, g, Graph) { put(assignments, v, v); }

        typename KeyedUpdatablePriorityQueue::key_map keys = pq.keys();

        // set number of visited neighbors for all vertices to 0
        BGL_FORALL_VERTICES_T(v, g, Graph)
        {
            if (v == get(assignments, v))
            { // foreach u \in V do
                put(keys, v, weight_type(0));
                vis.initialize_vertex(v, g);

                pq.push(v);
            }
        }
        BOOST_ASSERT(pq.size() >= 2);

        // Give the starting vertex high priority
        put(keys, start, get(keys, start) + num_vertices(g) + 1);
        pq.update(start);

        // start traversing the graph
        // vertex_descriptor s, t;
        // weight_type w;
        while (!pq.empty())
        { // while PQ \neq {} do
            const vertex_descriptor u = pq.top(); // u = extractmax(PQ)
            /* weight_type w = */ get(keys, u);
            vis.start_vertex(u, g);
            pq.pop(); //            vis.start_vertex(u, g);

            BGL_FORALL_OUTEDGES_T(u, e, g, Graph)
            { // foreach (u, v) \in E do
                vis.examine_edge(e, g);

                const vertex_descriptor v = get(assignments, target(e, g));

                if (pq.contains(v))
                { // if v \in PQ then
                    put(keys, v,
                        get(keys, v)
                            + get(weights,
                                e)); // increasekey(PQ, v, wA(v) + w(u, v))
                    pq.update(v);
                }
            }

            typename std::set< vertex_descriptor >::const_iterator
                assignedVertexIt,
                assignedVertexEnd = assignedVertices.end();
            for (assignedVertexIt = assignedVertices.begin();
                 assignedVertexIt != assignedVertexEnd; ++assignedVertexIt)
            {
                const vertex_descriptor uPrime = *assignedVertexIt;

                if (get(assignments, uPrime) == u)
                {
                    BGL_FORALL_OUTEDGES_T(uPrime, e, g, Graph)
                    { // foreach (u, v) \in E do
                        vis.examine_edge(e, g);

                        const vertex_descriptor v
                            = get(assignments, target(e, g));

                        if (pq.contains(v))
                        { // if v \in PQ then
                            put(keys, v,
                                get(keys, v)
                                    + get(weights, e)); // increasekey(PQ, v,
                                                        // wA(v) + w(u, v))
                            pq.update(v);
                        }
                    }
                }
            }
            vis.finish_vertex(u, g);
        }
    }
} // end namespace detail

template < class Graph, class WeightMap, class MASVisitor,
    class VertexAssignmentMap, class KeyedUpdatablePriorityQueue >
void maximum_adjacency_search(const Graph& g, WeightMap weights, MASVisitor vis,
    const typename boost::graph_traits< Graph >::vertex_descriptor start,
    VertexAssignmentMap assignments, KeyedUpdatablePriorityQueue pq)
{
    BOOST_CONCEPT_ASSERT((boost::IncidenceGraphConcept< Graph >));
    BOOST_CONCEPT_ASSERT((boost::VertexListGraphConcept< Graph >));
    typedef typename boost::graph_traits< Graph >::vertex_descriptor
        vertex_descriptor;
    typedef typename boost::graph_traits< Graph >::vertices_size_type
        vertices_size_type;
    typedef
        typename boost::graph_traits< Graph >::edge_descriptor edge_descriptor;
    BOOST_CONCEPT_ASSERT((boost::Convertible<
        typename boost::graph_traits< Graph >::directed_category,
        boost::undirected_tag >));
    BOOST_CONCEPT_ASSERT(
        (boost::ReadablePropertyMapConcept< WeightMap, edge_descriptor >));
    // typedef typename boost::property_traits<WeightMap>::value_type
    // weight_type;
    boost::function_requires< MASVisitorConcept< MASVisitor, Graph > >();
    BOOST_CONCEPT_ASSERT(
        (boost::ReadWritePropertyMapConcept< VertexAssignmentMap,
            vertex_descriptor >));
    BOOST_CONCEPT_ASSERT((boost::Convertible< vertex_descriptor,
        typename boost::property_traits< VertexAssignmentMap >::value_type >));
    BOOST_CONCEPT_ASSERT(
        (boost::KeyedUpdatableQueueConcept< KeyedUpdatablePriorityQueue >));

    vertices_size_type n = num_vertices(g);
    if (n < 2)
        throw boost::bad_graph(
            "the input graph must have at least two vertices.");
    else if (!pq.empty())
        throw std::invalid_argument(
            "the max-priority queue must be empty initially.");

    detail::maximum_adjacency_search(g, weights, vis, start, assignments, pq);
}

namespace graph
{
    namespace detail
    {
        template < typename WeightMap > struct mas_dispatch
        {
            typedef void result_type;
            template < typename Graph, typename ArgPack >
            static result_type apply(const Graph& g,
                // const bgl_named_params<P,T,R>& params,
                const ArgPack& params, WeightMap w)
            {

                using namespace boost::graph::keywords;
                typedef typename boost::graph_traits< Graph >::vertex_descriptor
                    vertex_descriptor;
                typedef typename WeightMap::value_type weight_type;

                typedef boost::detail::make_priority_queue_from_arg_pack_gen<
                    boost::graph::keywords::tag::max_priority_queue,
                    weight_type, vertex_descriptor,
                    std::greater< weight_type > >
                    default_pq_gen_type;

                default_pq_gen_type pq_gen(
                    choose_param(get_param(params, boost::distance_zero_t()),
                        weight_type(0)));

                typename boost::result_of< default_pq_gen_type(
                    const Graph&, const ArgPack&) >::type pq
                    = pq_gen(g, params);

                boost::null_visitor null_vis;
                boost::mas_visitor< boost::null_visitor > default_visitor(
                    null_vis);
                vertex_descriptor v = vertex_descriptor();
                boost::detail::make_property_map_from_arg_pack_gen<
                    boost::graph::keywords::tag::vertex_assignment_map,
                    vertex_descriptor >
                    map_gen(v);
                typename boost::detail::map_maker< Graph, ArgPack,
                    boost::graph::keywords::tag::vertex_assignment_map,
                    vertex_descriptor >::map_type default_map
                    = map_gen(g, params);
                boost::maximum_adjacency_search(g, w,
                    params[_visitor | default_visitor],
                    params[_root_vertex | *vertices(g).first],
                    params[_vertex_assignment_map | default_map], pq);
            }
        };

        template <> struct mas_dispatch< boost::param_not_found >
        {
            typedef void result_type;

            template < typename Graph, typename ArgPack >
            static result_type apply(
                const Graph& g, const ArgPack& params, param_not_found)
            {

                using namespace boost::graph::keywords;
                typedef typename boost::graph_traits< Graph >::vertex_descriptor
                    vertex_descriptor;

                // get edge_weight_t as the weight type
                typedef typename boost::property_map< Graph, edge_weight_t >
                    WeightMap;
                typedef typename WeightMap::value_type weight_type;

                typedef boost::detail::make_priority_queue_from_arg_pack_gen<
                    boost::graph::keywords::tag::max_priority_queue,
                    weight_type, vertex_descriptor,
                    std::greater< weight_type > >
                    default_pq_gen_type;

                default_pq_gen_type pq_gen(
                    choose_param(get_param(params, boost::distance_zero_t()),
                        weight_type(0)));

                typename boost::result_of< default_pq_gen_type(
                    const Graph&, const ArgPack&) >::type pq
                    = pq_gen(g, params);

                boost::null_visitor null_vis;
                boost::mas_visitor< boost::null_visitor > default_visitor(
                    null_vis);
                vertex_descriptor v = vertex_descriptor();
                boost::detail::make_property_map_from_arg_pack_gen<
                    boost::graph::keywords::tag::vertex_assignment_map,
                    vertex_descriptor >
                    map_gen(v);
                typename boost::detail::map_maker< Graph, ArgPack,
                    boost::graph::keywords::tag::vertex_assignment_map,
                    vertex_descriptor >::map_type default_map
                    = map_gen(g, params);
                boost::maximum_adjacency_search(g, get(edge_weight, g),
                    params[_visitor | default_visitor],
                    params[_root_vertex | *vertices(g).first],
                    params[_vertex_assignment_map | default_map], pq);
            }
        };
    } // end namespace detail
} // end namespace graph

// Named parameter interface
// BOOST_GRAPH_MAKE_OLD_STYLE_PARAMETER_FUNCTION(maximum_adjacency_search, 1)
template < typename Graph, typename P, typename T, typename R >
void maximum_adjacency_search(
    const Graph& g, const bgl_named_params< P, T, R >& params)
{

    typedef bgl_named_params< P, T, R > params_type;
    BOOST_GRAPH_DECLARE_CONVERTED_PARAMETERS(params_type, params)

    // do the dispatch based on WeightMap
    typedef typename get_param_type< edge_weight_t,
        bgl_named_params< P, T, R > >::type W;
    graph::detail::mas_dispatch< W >::apply(
        g, arg_pack, get_param(params, edge_weight));
}

namespace graph
{
    namespace detail
    {
        template < typename Graph > struct maximum_adjacency_search_impl
        {
            typedef void result_type;

            template < typename ArgPack >
            void operator()(const Graph& g, const ArgPack& arg_pack) const
            {
                // call the function that does the dispatching
                typedef
                    typename get_param_type< edge_weight_t, ArgPack >::type W;
                graph::detail::mas_dispatch< W >::apply(
                    g, arg_pack, get_param(arg_pack, edge_weight));
            }
        };
    } // end namespace detail
    BOOST_GRAPH_MAKE_FORWARDING_FUNCTION(maximum_adjacency_search, 1, 5)
} // end namespace graph

} // end namespace boost

#include <boost/graph/iteration_macros_undef.hpp>

#endif // BOOST_GRAPH_MAXIMUM_ADJACENCY_SEARCH_H


/**
 * @Author: Maxime Agor (4rzael)
 * @Date:   Sat Aug 18 2018
 * @Email:  maxime.agor23@gmail.com
 * @Project: CUDA-Based Simulator of Quantum Systems
 * @Filename: BasicMeasurementResultsTree.cpp
 * @Last modified by:   4rzael
 * @Last modified time: Sat Aug 18 2018, 23:36:59
 * @License: MIT License
 */

#include <exception>
#include <optional>
#include <limits>
#include <boost/optional.hpp>
#include "TaskScheduling/BasicMeasurementResultsTree.hpp"
#include "utils.hpp"
#include "Logger.hpp"

using namespace MeasurementResultsTree;

BasicMeasurementResultsTree::BasicMeasurementResultsTree(uint samples): greatestId(MEASUREMENT_NODE_NONE) {
    root = std::make_shared<MeasurementResultsNode>(this, ++greatestId);
    root->samples = samples;
}

std::shared_ptr<MeasurementResultsNode> BasicMeasurementResultsTree::getRoot() const {
    return root;
}
std::shared_ptr<MeasurementResultsNode> BasicMeasurementResultsTree::getNodeWithId(NodeId id) const {
    // Making a recursive helper function
    const std::function<std::shared_ptr<MeasurementResultsNode>(std::shared_ptr<MeasurementResultsNode>)> helper =
    [&](std::shared_ptr<MeasurementResultsNode> currentNode) {
        if (!currentNode) return std::shared_ptr<MeasurementResultsNode>(); // if not a node, abort
        if (currentNode->id == id) return currentNode; // if found, return self
        std::shared_ptr<MeasurementResultsNode> res;
        // try child 0
        res = helper(currentNode->results[0].node);
        if (res) return res;
        // if didn't work, try child 1
        return helper(currentNode->results[1].node);
    };
    // Run the recursive helper from the root node
    const auto res = helper(root);
    if (!res) throw std::logic_error("Measurement Node not found");
    return res;
}

uint BasicMeasurementResultsTree::getCregValueAtNode(std::string cregName, NodeId id) const {
    // Making a recursive helper function
    const std::function<boost::optional<uint>(uint, std::shared_ptr<MeasurementResultsNode>)> helper =
    [&](uint currentValue, std::shared_ptr<MeasurementResultsNode> node) -> boost::optional<uint> {
        if (!node) return boost::none; // if not a node, abort
        if (node->id == id) return currentValue; // if node is self, return the currentValue
        //try child 0
        boost::optional<uint> res;
        res = helper(
            (node->results[0].measuredCBit.registerName != cregName) // if the current node modifies the good creg, apply modification
            ? currentValue
            : (currentValue & (std::numeric_limits<int>::max() - (1 << node->results[0].measuredCBit.element))), // set bit to 0
            node->results[0].node
        );
        if (res) return res;
        // if didn't work, try child 1
        return helper(
            (node->results[1].measuredCBit.registerName != cregName) // if the current node modifies the good creg, apply modification
            ? currentValue
            : (currentValue | ((1 << node->results[1].measuredCBit.element))),  // set bit to 1
            node->results[1].node
        );
    };
    // Run the recursive helper from the root node
    const auto res = helper(0, root);
    if (!res) throw std::logic_error("Measurement Node not found");
    return res.value();
}

std::vector<std::shared_ptr<MeasurementResultsNode>> BasicMeasurementResultsTree::makeChildrens(NodeId parentId, Circuit::Qubit creg) {
    auto parent = getNodeWithId(parentId);
    parent->results[0].node = std::make_shared<MeasurementResultsNode>(this, ++greatestId);
    parent->results[0].measuredCBit = creg;
    parent->results[0].value = 0;
    parent->results[0].probability = 0;
    parent->results[0].node->samples = 0;
    parent->results[1].node = std::make_shared<MeasurementResultsNode>(this, ++greatestId);
    parent->results[1].measuredCBit = creg;
    parent->results[1].value = 1;
    parent->results[1].probability = 0;
    parent->results[1].node->samples = 0;
    return {parent->results[0].node, parent->results[1].node};
}

std::vector<NodeId> BasicMeasurementResultsTree::addMeasurement(NodeId nodeId, double zeroProbability) {
    auto node = getNodeWithId(nodeId);
    auto zeroSamples = sampleBinomialDistribution(node->samples, zeroProbability);
    node->results[0].probability = zeroProbability;
    node->results[0].node->samples = zeroSamples;

    node->results[1].probability = 1 - zeroProbability;
    node->results[1].node->samples = node->samples - zeroSamples;

    node->status = MeasurementResultsNodeStatus::COMPLETE;
    return {node->results[0].node->id, node->results[1].node->id};
}

void BasicMeasurementResultsTree::printResults(std::vector<Circuit::Register> const &cregs) const {
    LOG(Logger::INFO, "Sampled " << getRoot()->samples << " executions");

    const std::function<void(std::shared_ptr<MeasurementResultsNode> const &)> helper = [&](std::shared_ptr<MeasurementResultsNode> const &node) {
        // no childrens => we can print the results for this node
        if (!(node->results[0].node)) {
            if (node->samples == 0) return;
            LOG(Logger::INFO, "On " << node->samples << " samples:");
            for (auto const &reg : cregs) {
                const auto value = getCregValueAtNode(reg.name, node->id);
                LOG(Logger::INFO, reg.name << " = " << value);
            }
        } else {
            helper(node->results[0].node);
            helper(node->results[1].node);
        }
    };

    helper(getRoot());
}


/*
 * Copyright (c) 2011-2022, The DART development contributors
 * All rights reserved.
 *
 * The list of contributors can be found at:
 *   https://github.com/dartsim/dart/blob/master/LICENSE
 *
 * This file is provided under the following "BSD-style" License:
 *   Redistribution and use in source and binary forms, with or
 *   without modification, are permitted provided that the following
 *   conditions are met:
 *   * Redistributions of source code must retain the above copyright
 *     notice, this list of conditions and the following disclaimer.
 *   * Redistributions in binary form must reproduce the above
 *     copyright notice, this list of conditions and the following
 *     disclaimer in the documentation and/or other materials provided
 *     with the distribution.
 *   THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND
 *   CONTRIBUTORS "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES,
 *   INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF
 *   MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
 *   DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR
 *   CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
 *   SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
 *   LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF
 *   USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED
 *   AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT
 *   LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN
 *   ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
 *   POSSIBILITY OF SUCH DAMAGE.
 */

#ifndef DART_CONSTRAINT_BOXEDLCPSOLVER_HPP_
#define DART_CONSTRAINT_BOXEDLCPSOLVER_HPP_

#include <string>

#include <Eigen/Core>

#include "dart/common/Castable.hpp"

namespace dart {
namespace constraint {

class BoxedLcpSolver : public common::Castable<BoxedLcpSolver>
{
public:
  /// Destructor
  virtual ~BoxedLcpSolver() = default;

  /// Returns the type
  virtual const std::string& getType() const = 0;

  /// Solves constriant impulses for a constrained group. The LCP formulation
  /// setting that this function solve is A*x = b + w where each x[i], w[i]
  /// satisfies one of
  ///   (1) x = lo, w >= 0
  ///   (2) x = hi, w <= 0
  ///   (3) lo < x < hi, w = 0
  ///
  /// \param[in] n Dimension of constraints.
  /// \param[in] A A term of the LCP formulation.
  /// \param[in] x x term of the LCP formulation.
  /// \param[in] b b term of the LCP formulation.
  /// \param[in] nub Number of the first unbounded constraints.
  /// \param[in] lo Lower bound of x where it's restricted to be lo <= 0.
  /// \param[in] hi Upper bound of x where it's enforced to be hi >= 0.
  /// \param[in] findex Indices to corresponding normal contact constraint. Set
  /// the index to itself (e.g., findex[k] = k) for normal contacts or
  /// non-contact constraints. For friction constraint, set the cooresponding
  /// normal contact constraint.
  /// \param[in] earlyTermination Set true to return false as soon as the solver
  /// find the solution doesn't exist. Otherwise, the solver will continue to
  /// push hard to solve the problem using some hacks.
  ///
  /// \return Success.
  // Note: The function signature is ODE specific for now. Consider changing
  // this to Eigen friendly version once own Dantzig LCP solver is available.
  virtual bool solve(
      int n,
      double* A,
      double* x,
      double* b,
      int nub,
      double* lo,
      double* hi,
      int* findex,
      bool earlyTermination = false)
      = 0;

#ifndef NDEBUG
  virtual bool canSolve(int n, const double* A) = 0;
#endif
};

} // namespace constraint
} // namespace dart

#endif // DART_CONSTRAINT_BOXEDLCPSOLVER_HPP_


/*
 *            Copyright 2009-2017 The VOTCA Development Team
 *                       (http://www.votca.org)
 *
 *      Licensed under the Apache License, Version 2.0 (the "License")
 *
 * You may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *              http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 *
 */

#include "votca/xtp/orbitals.h"
#include "votca/tools/globals.h"
#include "votca/xtp/elements.h"
#include <stdio.h>
#include <iostream>
#include <iomanip>
#include <fstream>
#include <map>
#include <iterator>
#include <numeric>
#include <boost/algorithm/string.hpp>
#include <boost/lexical_cast.hpp>
#include <boost/numeric/ublas/io.hpp>
#include <boost/numeric/ublas/vector.hpp>

namespace votca {
    namespace xtp {

        Orbitals::Orbitals() {

            _basis_set_size = 0;
            _occupied_levels = 0;
            _unoccupied_levels = 0;
            _number_of_electrons = 0;
            _self_energy = 0.0;
            _qm_energy = 0.0;
            _couplingsA = 0;
            _couplingsB = 0;
            _ECP = "";
            _bsetype="";
            //_has_atoms = false;

            // GW-BSE
            _qpmin = 0;
            _qpmax = 0;
            _qptotal = 0;

            _rpamin = 0;
            _rpamax = 0;

            _ScaHFX = 0.0;

            _bse_cmin = 0;
            _bse_cmax = 0;
            _bse_vmin = 0;
            _bse_vmax = 0;
            _bse_vtotal = 0;
            _bse_ctotal = 0;
            _bse_size = 0;
            _bse_nmax = 0;

        };

        Orbitals::~Orbitals() {
            _mo_energies.clear();
            _mo_coefficients.clear();
            _overlap.clear();

            std::vector< ctp::QMAtom* >::iterator it;
            for (it = _atoms.begin(); it != _atoms.end(); ++it) delete *it;

        };

        void Orbitals::setNumberOfLevels(const int &occupied_levels, const int &unoccupied_levels) {
            // _has_occupied_levels = true;
            // _has_unoccupied_levels = true;
            _occupied_levels = occupied_levels;
            _unoccupied_levels = unoccupied_levels;
        }

        /**
         *
         * @param _energy_difference [ev] Two levels are degenerate if their energy is smaller than this value
         * @return A map with key as a level and a vector which is a list of close lying orbitals
         */
        bool Orbitals::CheckDegeneracy(double _energy_difference) {

            ub::vector<double>::iterator it1 = _mo_energies.begin();
            bool _degenerate = false;

            if (tools::globals::verbose) {
                cout << endl << "... ... Checking level degeneracy " << endl;
            }

            _level_degeneracy.clear();

            while (it1 != _mo_energies.end()) {

                // in all containers counters start with 0; real life - with 1
                int _level1 = std::distance(_mo_energies.begin(), it1) + 1;

                // add the level itself - it is easier to loo over all levels later
                _level_degeneracy[_level1].push_back(_level1);

                ub::vector<double>::iterator it2 = it1;
                it2++;

                while (it2 != _mo_energies.end()) {
                    //cout << _level1 << ":" << *it1 << ":" << *it2 << endl;
                    double energy1 = *it1;
                    double energy2 = *it2;

                    // in all containers counters start with 0; real life - with 1
                    int _level2 = std::distance(_mo_energies.begin(), it2) + 1;

                    if (std::abs(energy1 - energy2) * tools::conv::hrt2ev < _energy_difference) {
                        _level_degeneracy[_level1].push_back(_level2);
                        _level_degeneracy[_level2].push_back(_level1);
                        _degenerate = true;
                    }
                    it2++;
                }
                it1++;
            }

            if (tools::globals::verbose) {

                if (_degenerate) {
                    cout << "... ... Some levels are degenerate" << endl;
                    for (std::map<int, std::vector<int> >::iterator it = _level_degeneracy.begin();
                            it != _level_degeneracy.end();
                            ++it) {
                        // output only degenerate levels
                        if ((it->second).size() > 1) {
                            std::cout << "... ... level  " << it->first << " : ";
                            for (vector<int>::iterator lev = (it->second).begin(); lev != (it->second).end(); lev++)
                                cout << *lev << " ";
                            cout << endl;
                        }
                    }
                } else {
                    cout << "... ... No degeneracy found" << endl;
                }
                cout << "... ... Done checking level degeneracy" << endl;
            }
            return _degenerate;
        }

        std::vector<int>* Orbitals::getDegeneracy(int level, double _energy_difference) {
            if (!hasDegeneracy()) {
                CheckDegeneracy(_energy_difference);
            }

            return &_level_degeneracy.at(level);
        }

        std::vector<int> Orbitals::SortEnergies() {
            if (tools::globals::verbose) cout << "... ... Sorting energies" << endl;
            std::vector<int>index=std::vector<int>(_mo_energies.size());
            std::iota(index.begin(), index.end(), 0);
            std::stable_sort(index.begin(), index.end(),[this](int i1, int i2) {return this->MOEnergies()[i1] > this->MOEnergies()[i2];});
            return index;
        }

        /// Writes a PDB file

        void Orbitals::WritePDB(FILE *out, string tag) {
            string tag_extended = "HEADER ! GENERATED BY VOTCA::XTP::" + tag + "\n";
            fprintf(out, "%s", tag_extended.c_str());
            vector < ctp::QMAtom* > ::iterator atom;
            int id = 0;

            for (atom = _atoms.begin(); atom < _atoms.end(); ++atom) {
                id++;
                string resname = ((*atom)->from_environment) ? "MM" : "QM";
                int resnr = 1;

                fprintf(out, "ATOM  %5d %4s%1s%3s %1s%4d%1s   %8.3f%8.3f%8.3f%6.2f%6.2f      %4s%2s  %8.3f\n",
                        id, // Atom serial number           %5d
                        (*atom)->type.c_str(), // Atom name                    %4s
                        " ", // alternate location indicator.%1s
                        resname.c_str(), // Residue name.                %3s
                        "A", // Chain identifier             %1s
                        resnr, // Residue sequence number      %4d
                        " ", // Insertion of residues.       %1s
                        (*atom)->x, // X in Angstroms               %8.3f
                        (*atom)->y, // Y in Angstroms               %8.3f
                        (*atom)->z, // Z in Angstroms               %8.3f
                        1.0, // Occupancy                    %6.2f
                        0.0, // Temperature factor           %6.2f
                        " ", // Segment identifier           %4s
                        (*atom)->type.c_str(), // Element symbol               %2s
                        (*atom)->charge // Charge on the atom.          %2s
                        );
            }
            return;
        }

        // reduces the number of virtual orbitals to factor*number_of_occupied_orbitals

        void Orbitals::Trim(int factor) {

            if (hasMOCoefficients()) {
                _mo_coefficients.resize(factor * _occupied_levels, _basis_set_size, true);
                _unoccupied_levels = (factor - 1) * _occupied_levels;
            }

            if (hasMOEnergies()) {
                _mo_energies.resize(factor * _occupied_levels, true);
                _unoccupied_levels = (factor - 1) * _occupied_levels;
            }
            return;
        }

        // reduces the number of orbitals to [HOMO-degH:LUMO+degL]

        void Orbitals::Trim(int degH, int degL) {

            if (hasMOCoefficients()) {
                _mo_coefficients = ub::project(_mo_coefficients, ub::range(_occupied_levels - degH, _occupied_levels + degL), ub::range(0, _basis_set_size));
            }

            if (hasMOEnergies()) {
                ub::vector<double> _temp(degH + degL);
                for (int i = 0; i < degH + degL; i++) {
                    _temp(i) = _mo_energies(_occupied_levels - degH + i);
                }
                _mo_energies = _temp;
            }
            _occupied_levels = degH;
            _unoccupied_levels = degL;

            return;
        }

        bool Orbitals::Load(string file_name) {
            try {
                std::ifstream ifs(file_name.c_str());
                boost::archive::binary_iarchive ia(ifs);
                ia >> *this;
                ifs.close();
            } catch (std::exception &err) {
                std::cerr << "Could not load orbitals from " << file_name << flush;
                std::cerr << "An error occurred:\n" << err.what() << endl;
                return false;
            }
            return true;
        }

        /* Save to archive */
        bool Orbitals::Save(std::string file_name) {

            try {
                std::ofstream ofs((file_name).c_str());
                boost::archive::binary_oarchive oa(ofs);
                oa << *this;
                ofs.close();
            } catch (std::exception &err) {
                std::cerr << "Could not save orbitals to " << file_name << flush;
                std::cerr << "An error occurred:\n" << err.what() << endl;
                return false;
            }
            return true;
        }

        // Determine ground state density matrix

        ub::matrix<double> Orbitals::DensityMatrixGroundState() {
            ub::matrix<double> dmatGS = ub::zero_matrix<double>(_basis_set_size, _basis_set_size);
#pragma omp parallel for
            for (int _i = 0; _i < _basis_set_size; _i++) {
                for (int _j = 0; _j < _basis_set_size; _j++) {
                    for (int _level = 0; _level < _occupied_levels; _level++) {

                        dmatGS(_i, _j) += 2.0 * _mo_coefficients(_level, _i) * _mo_coefficients(_level, _j);

                    }
                }
            }
            return dmatGS;
        }

        // Determine QuasiParticle Density Matrix
        ub::matrix<double> Orbitals::DensityMatrixQuasiParticle( int state){
            ub::matrix<double> dmatQP = ub::zero_matrix<double>(_basis_set_size, _basis_set_size);
            //ub::matrix<double> _mos = ub::project(_mo_coefficients, ub::range(_qpmin, _qpmax + 1), ub::range(0, _basis_set_size));

            //ub::matrix<real_gwbse> lambda = ub::prod(_QPdiag_coefficients,_mos);

            ub::matrix<double> lambda = LambdaMatrixQuasiParticle();

            #pragma omp parallel for
            for (int _i = 0; _i < _basis_set_size; _i++) {
                for (int _j = 0; _j < _basis_set_size; _j++) {
                    dmatQP(_i,_j) = lambda(state,_i)*lambda(state,_j);
                }
            }

            return dmatQP;
        }

        // Determine QuasiParticle Lambda Matrix
        ub::matrix<double> Orbitals::LambdaMatrixQuasiParticle(){
            ub::matrix<double> _mos = ub::project(_mo_coefficients, ub::range(_qpmin, _qpmax + 1), ub::range(0, _basis_set_size));
            ub::matrix<double> lambda = ub::prod(_QPdiag_coefficients,_mos);
            return lambda;
        }

        ub::matrix<double> Orbitals::TransitionDensityMatrix(const string& spin, int state) {
            if(!(spin=="singlet" || spin=="triplet")){
                throw runtime_error("Spin type not known for density matrix. Available are singlet and triplet");
            }
            ub::matrix<real_gwbse>& _BSECoefs = (spin=="singlet") ? _BSE_singlet_coefficients : _BSE_triplet_coefficients;

            ub::matrix<double> dmatTS = ub::zero_matrix<double>(_basis_set_size);
            // The Transition dipole is sqrt2 bigger because of the spin, the excited state is a linear combination of 2 slater determinants, where either alpha or beta spin electron is excited
            double sqrt2 = sqrt(2.0);
            /*Trying to implement D_{alpha,beta}= sqrt2*sum_{i}^{occ}sum_{j}^{virt}{BSEcoef(i,j)*MOcoef(alpha,i)*MOcoef(beta,j)} */
            // c stands for conduction band and thus virtual orbitals
            // v stand for valence band and thus occupied orbitals

            if (_bse_size == 0) {
                _bse_vtotal = _bse_vmax - _bse_vmin + 1;
                _bse_ctotal = _bse_cmax - _bse_cmin + 1;
                _bse_size = _bse_vtotal * _bse_ctotal;
                // indexing info BSE vector index to occupied/virtual orbital
                for (unsigned _v = 0; _v < _bse_vtotal; _v++) {
                    for (unsigned _c = 0; _c < _bse_ctotal; _c++) {
                        _index2v.push_back(_bse_vmin + _v);
                        _index2c.push_back(_bse_cmin + _c);
                    }
                }
            }
            if (_bsetype == "full" && spin == "singlet") {
                const ub::matrix<real_gwbse>& _BSECoefs_AR=_BSE_singlet_coefficients_AR;
#pragma omp parallel for
                for (unsigned a = 0; a < dmatTS.size1(); a++) {
                    for (unsigned b = 0; b < dmatTS.size2(); b++) {
                        for (unsigned i = 0; i < _bse_size; i++) {
                            int occ = _index2v[i];
                            int virt = _index2c[i];
                            dmatTS(a, b) += sqrt2 * (_BSECoefs(i, state) + _BSECoefs_AR(i, state)) * _mo_coefficients(occ, a) * _mo_coefficients(virt, b); //check factor 2??
                        }
                    }
                }
            } else {

#pragma omp parallel for
                for (unsigned a = 0; a < dmatTS.size1(); a++) {
                    for (unsigned b = 0; b < dmatTS.size2(); b++) {
                        for (unsigned i = 0; i < _bse_size; i++) {
                            int occ = _index2v[i];
                            int virt = _index2c[i];
                            dmatTS(a, b) += sqrt2 * _BSECoefs(i, state) * _mo_coefficients(occ, a) * _mo_coefficients(virt, b); //check factor 2??
                        }
                    }
                }

            }
            return dmatTS;
        }

        std::vector<ub::matrix<double> > Orbitals::DensityMatrixExcitedState(const string& spin,int state) {

            std::vector<ub::matrix<double> > dmat = DensityMatrixExcitedState_R(spin,state);
            if(_bsetype=="full" && spin=="singlet"){
                std::vector<ub::matrix<double> > dmat_AR = DensityMatrixExcitedState_AR(spin, state);
                dmat[0] -= dmat_AR[0];
                dmat[1] -= dmat_AR[1];
            }
            return dmat;
        }

        // Excited state density matrix

        std::vector<ub::matrix<double> > Orbitals::DensityMatrixExcitedState_R(const string& spin, int state) {
            if(!(spin=="singlet" || spin=="triplet")){
                throw runtime_error("Spin type not known for density matrix. Available are singlet and triplet");
            }

            ub::matrix<real_gwbse>& _BSECoefs = (spin=="singlet") ? _BSE_singlet_coefficients : _BSE_triplet_coefficients;

            /******
             *
             *    Density matrix for GW-BSE based excitations
             *
             *    - electron contribution
             *      D_ab = \sum{vc} \sum{c'} A_{vc}A_{vc'} mo_a(c)mo_b(c')
             *
             *    - hole contribution
             *      D_ab = \sum{vc} \sum{v'} A_{vc}A_{v'c} mo_a(v)mo_b(v')
             *
             *
             *   more efficient:
             *
             *   - electron contribution
             *      D_ab = \sum{c} \sum{c'} mo_a(c)mo_b(c') [ \sum{v} A_{vc}A_{vc'} ]
             *           = \sum{c} \sum{c'} mo_a(c)mo_b(c') A_{cc'}
             *
             *   - hole contribution
             *      D_ab = \sum{v} \sum{v'} mo_a(v)mo_b(v') [ \sum{c} A_{vc}A_{v'c} ]
             *           = \sum{v} \sum{v'} mo_a(v)mo_b(v') A_{vv'}
             *
             */
            std::vector<ub::matrix<double> > dmatEX;
            dmatEX.resize(2);
            dmatEX[0] = ub::zero_matrix<double>(_basis_set_size, _basis_set_size);
            dmatEX[1] = ub::zero_matrix<double>(_basis_set_size, _basis_set_size);

            int _vmin = this->_bse_vmin;
            int _vmax = this->_bse_vmax;
            int _cmin = this->_bse_cmin;
            int _cmax = this->_bse_cmax;

            if (_bse_size == 0) {
                _bse_vtotal = _bse_vmax - _bse_vmin + 1;
                _bse_ctotal = _bse_cmax - _bse_cmin + 1;
                _bse_size = _bse_vtotal * _bse_ctotal;
                // indexing info BSE vector index to occupied/virtual orbital
                for (unsigned _v = 0; _v < _bse_vtotal; _v++) {
                    for (unsigned _c = 0; _c < _bse_ctotal; _c++) {
                        _index2v.push_back(_bse_vmin + _v);
                        _index2c.push_back(_bse_cmin + _c);
                    }
                }
            }

            // electron assist matrix A_{cc'}
            ub::matrix<real_gwbse> _Acc = ub::zero_matrix<real_gwbse>(_bse_ctotal, _bse_ctotal);
            ub::matrix<real_gwbse> _Avv = ub::zero_matrix<real_gwbse>(_bse_vtotal, _bse_vtotal);

            for (unsigned _idx1 = 0; _idx1 < _bse_size; _idx1++) {

                int _v = this->_index2v[_idx1];
                int _c = this->_index2c[_idx1];

                // electron assist matrix A_{cc'}
#pragma omp parallel for
                for (int _c2 = _cmin; _c2 <= _cmax; _c2++) {
                    int _idx2 = (_cmax - _cmin + 1)*(_v - _vmin)+(_c2 - _cmin);

                    _Acc(_c - _cmin, _c2 - _cmin) += _BSECoefs(_idx1, state) * _BSECoefs(_idx2, state);
                }

                // hole assist matrix A_{vv'}
#pragma omp parallel for
                for (int _v2 = _vmin; _v2 <= _vmax; _v2++) {
                    int _idx2 = (_cmax - _cmin + 1)*(_v2 - _vmin)+(_c - _cmin);

                    _Avv(_v - _vmin, _v2 - _vmin) += _BSECoefs(_idx1, state) * _BSECoefs(_idx2, state);

                }

            }

            // hole part as matrix products
            // get slice of MOs of occs only
            ub::matrix<double> _occlevels = ub::project(_mo_coefficients, ub::range(_vmin, _vmax + 1), ub::range(0, _basis_set_size));
            ub::matrix<double> _temp = ub::prod(_Avv, _occlevels);
            dmatEX[0] = ub::prod(ub::trans(_occlevels), _temp);

            // electron part as matrix products
            // get slice of MOs of virts only
            ub::matrix<double> _virtlevels = ub::project(_mo_coefficients, ub::range(_cmin, _cmax + 1), ub::range(0, _basis_set_size));
            _temp = ub::prod(_Acc, _virtlevels);
            dmatEX[1] = ub::prod(ub::trans(_virtlevels), _temp);

            return dmatEX;
        }

        // Excited state density matrix

        std::vector<ub::matrix<double> > Orbitals::DensityMatrixExcitedState_AR(const string& spin, int state) {
             if(!(spin=="singlet" )){
                throw runtime_error("Spin type not known for density matrix. Available is singlet");
            }

            ub::matrix<real_gwbse>& _BSECoefs_AR = _BSE_singlet_coefficients_AR;

            /******
             *
             *    Density matrix for GW-BSE based excitations
             *
             *    - electron contribution
             *      D_ab = \sum{vc} \sum{v'} B_{vc}B_{v'c} mo_a(v)mo_b(v')
             *
             *    - hole contribution
             *      D_ab = \sum{vc} \sum{c'} B_{vc}B_{vc'} mo_a(c)mo_b(c')
             *
             *
             *   more efficient:
             *
             *   - electron contribution
             *      D_ab = \sum{v} \sum{v'} mo_a(v)mo_b(v') [ \sum{c} B_{vc}B_{v'c} ]
             *           = \sum{v} \sum{v'} mo_a(v)mo_b(v') B_{vv'}
             *
             *   - hole contribution
             *      D_ab = \sum{c} \sum{c'} mo_a(c)mo_b(c') [ \sum{v} B_{vc}B_{vc'} ]
             *           = \sum{c} \sum{c'} mo_a(c)mo_b(c') B_{cc'}
             *
             */
            std::vector<ub::matrix<double> > dmatAR;
            dmatAR.resize(2);
            dmatAR[0] = ub::zero_matrix<double>(_basis_set_size, _basis_set_size);
            dmatAR[1] = ub::zero_matrix<double>(_basis_set_size, _basis_set_size);

            int _vmin = this->_bse_vmin;
            int _vmax = this->_bse_vmax;
            int _cmin = this->_bse_cmin;
            int _cmax = this->_bse_cmax;

            if (_bse_size == 0) {
                _bse_vtotal = _bse_vmax - _bse_vmin + 1;
                _bse_ctotal = _bse_cmax - _bse_cmin + 1;
                _bse_size = _bse_vtotal * _bse_ctotal;
                // indexing info BSE vector index to occupied/virtual orbital
                for (unsigned _v = 0; _v < _bse_vtotal; _v++) {
                    for (unsigned _c = 0; _c < _bse_ctotal; _c++) {
                        _index2v.push_back(_bse_vmin + _v);
                        _index2c.push_back(_bse_cmin + _c);
                    }
                }
            }

            // hole assist matrix B_{cc'}
            ub::matrix<real_gwbse> _Bcc = ub::zero_matrix<real_gwbse>(_bse_ctotal, _bse_ctotal);
            ub::matrix<real_gwbse> _Bvv = ub::zero_matrix<real_gwbse>(_bse_vtotal, _bse_vtotal);

            for (unsigned _idx1 = 0; _idx1 < _bse_size; _idx1++) {

                int _v = this->_index2v[_idx1];
                int _c = this->_index2c[_idx1];

                // hole assist matrix B_{cc'}
#pragma omp parallel for
                for (int _c2 = _cmin; _c2 <= _cmax; _c2++) {
                    int _idx2 = (_cmax - _cmin + 1)*(_v - _vmin)+(_c2 - _cmin);

                    _Bcc(_c - _cmin, _c2 - _cmin) += _BSECoefs_AR(_idx1, state) * _BSECoefs_AR(_idx2, state);
                }

                // electron assist matrix B_{vv'}
#pragma omp parallel for
                for (int _v2 = _vmin; _v2 <= _vmax; _v2++) {
                    int _idx2 = (_cmax - _cmin + 1)*(_v2 - _vmin)+(_c - _cmin);

                    _Bvv(_v - _vmin, _v2 - _vmin) += _BSECoefs_AR(_idx1, state) * _BSECoefs_AR(_idx2, state);

                }
            }

            //hole part as matrix products
            // get slice of MOs of virts only
            ub::matrix<double> _virtlevels = ub::project(_mo_coefficients, ub::range(_cmin, _cmax + 1), ub::range(0, _basis_set_size));
            ub::matrix<double> _temp = ub::prod(_Bcc, _virtlevels);
            dmatAR[0] = ub::prod(ub::trans(_virtlevels), _temp);

            // electron part as matrix products
            // get slice of MOs of occs only
            ub::matrix<double> _occlevels = ub::project(_mo_coefficients, ub::range(_vmin, _vmax + 1), ub::range(0, _basis_set_size));
            _temp = ub::prod(_Bvv, _occlevels);
            dmatAR[1] = ub::prod(ub::trans(_occlevels), _temp);

            return dmatAR;
        }

        ub::vector<double> Orbitals::LoewdinPopulation(const ub::matrix<double>& _densitymatrix, const ub::matrix<double>& _overlapmatrix, int _frag) {

            ub::vector<double> fragmentCharges = ub::vector<double>(2, 0.0);
            ub::matrix<double> overlap = _overlapmatrix;

            linalg_matrixsqrt(overlap);
            ub::matrix<double> temp = ub::prod(_densitymatrix, overlap);
            ub::matrix<double> _prodmat = ub::prod(overlap, temp);
            for (int _i = 0; _i < _frag; _i++) {
                fragmentCharges(0) += _prodmat(_i, _i);
            }
            for (unsigned _i = _frag; _i < _overlapmatrix.size1(); _i++) {
                fragmentCharges(1) += _prodmat(_i, _i);
            }

            return fragmentCharges;
        }

        std::vector<double> Orbitals::Oscillatorstrengths() {
            std::vector<double> oscs;
            unsigned size = _transition_dipoles.size();
            if (size > _BSE_singlet_energies.size()) {
                size = _BSE_singlet_energies.size();
            }
            for (unsigned i = 0; i < size; ++i) {
                double osc = (_transition_dipoles[i] * _transition_dipoles[i]) * 2.0 / 3.0 * (_BSE_singlet_energies(i));
                oscs.push_back(osc);
            }
            return oscs;
        }

        double Orbitals::GetTotalEnergy(string _spintype, int _opt_state) {

            // total energy of the excited state
            double _total_energy;
            double _omega = 0.0;

            double _dft_energy = getQMEnergy();

            if (_spintype == "singlet") {
                _omega = BSESingletEnergies()[_opt_state - 1];
            } else if (_spintype == "triplet") {
                _omega = BSETripletEnergies()[_opt_state - 1];
            } else {
                throw std::runtime_error("GetTotalEnergy only knows spintypes:singlet,triplet");
            }

            // DFT total energy is stored in eV
            // singlet energies are stored in Ryd...

            return _total_energy = _dft_energy * tools::conv::ev2hrt + _omega; //  e.g. hartree
        }

        ub::vector<double> Orbitals::FragmentNuclearCharges(int _frag) {
            Elements _elements;

            // go through atoms and count
            vector < ctp::QMAtom* > ::iterator atom;
            int id = 0;

            if (_frag < 0) {
                throw runtime_error("Orbitals::FragmentNuclearCharges Fragment index is smaller than zero");
            }

            ub::vector<double> fragmentNuclearCharges = ub::vector<double>(2, 0.0);

            for (atom = _atoms.begin(); atom < _atoms.end(); ++atom) {
                id++;
                // get element type and determine its nuclear charge
                double crg;
                if ( _ECP == ""){
                    crg = _elements.getNucCrg((*atom)->type);
                }else{
                    crg = _elements.getNucCrgECP((*atom)->type);
                }
                // add to either fragment
                if (id <= _frag) {
                    fragmentNuclearCharges(0) += crg;
                } else {
                    fragmentNuclearCharges(1) += crg;
                }
            }
            return fragmentNuclearCharges;
        }

/**
         * \brief Guess for a dimer based on monomer orbitals
         *
         * Given two monomer orbitals (A and B) constructs a guess for dimer
         * orbitals: | A 0 | and energies: [EA, EB]
         *           | 0 B |
         */
        void Orbitals::PrepareGuess(Orbitals* _orbitalsA, Orbitals* _orbitalsB, Orbitals* _orbitalsAB) {

            // constructing the direct product orbA x orbB
            int _basisA = _orbitalsA->getBasisSetSize();
            int _basisB = _orbitalsB->getBasisSetSize();

            int _levelsA = _orbitalsA->getNumberOfLevels();
            int _levelsB = _orbitalsB->getNumberOfLevels();

            int _electronsA = _orbitalsA->getNumberOfElectrons();
            int _electronsB = _orbitalsB->getNumberOfElectrons();

            ub::zero_matrix<double> zeroB(_levelsA, _basisB);
            ub::zero_matrix<double> zeroA(_levelsB, _basisA);

            ub::matrix<double>& _mo_coefficients = _orbitalsAB->MOCoefficients();

            // AxB = | A 0 |  //   A = [EA, EB]  //
            //       | 0 B |  //                 //
            _mo_coefficients.resize(_levelsA + _levelsB, _basisA + _basisB);
            _orbitalsAB->setBasisSetSize(_basisA + _basisB);
            _orbitalsAB->setNumberOfLevels(_electronsA - _electronsB,
                    _levelsA + _levelsB - _electronsA - _electronsB);
            _orbitalsAB->setNumberOfElectrons(_electronsA + _electronsB);

            ub::project(_mo_coefficients, ub::range(0, _levelsA), ub::range(_basisA, _basisA + _basisB)) = zeroB;
            ub::project(_mo_coefficients, ub::range(_levelsA, _levelsA + _levelsB), ub::range(0, _basisA)) = zeroA;
            ub::project(_mo_coefficients, ub::range(0, _levelsA), ub::range(0, _basisA)) = _orbitalsA->MOCoefficients();
            ub::project(_mo_coefficients, ub::range(_levelsA, _levelsA + _levelsB), ub::range(_basisA, _basisA + _basisB)) = _orbitalsB->MOCoefficients();

            ub::vector<double>& _energies = _orbitalsAB->MOEnergies();
            _energies.resize(_levelsA + _levelsB);

            ub::project(_energies, ub::range(0, _levelsA)) = _orbitalsA->MOEnergies();
            ub::project(_energies, ub::range(_levelsA, _levelsA + _levelsB)) = _orbitalsB->MOEnergies();

            return;
        }

        void Orbitals::LoadFromXYZ(std::string filename) {

            string line;
            std::ifstream in;

            string type;

            in.open(filename.c_str(), std::ios::in);
            if (!in) throw runtime_error(string("Error reading coordinates from: ")
                    + filename);

            int atomCount = 1;

            if (in.is_open()) {
                while (in.good()) {
                    std::getline(in, line);

                    vector< string > split;
                    Tokenizer toker(line, " \t");
                    toker.ToVector(split);
                    if (!split.size() ||
                            split.size() != 4 ||
                            split[0] == "#" ||
                            split[0].substr(0, 1) == "#") {
                        continue;
                    }

                    // Interesting information written here: e.g. 'C 0.000 0.000 0.000'
                    atomCount++;
                    string element = split[0];
                    double x = boost::lexical_cast<double>(split[1]); //°A to NM
                    double y = boost::lexical_cast<double>(split[2]);
                    double z = boost::lexical_cast<double>(split[3]);
                    AddAtom(element, x, y, z);

                }
            } else {
                throw std::runtime_error("No such file: '" + filename + "'.");
            }
            return;
        }

    }
}