#include "general_dichotomy.hpp"

#include <iostream>
#include <algorithm>

#include "gdconfig.hpp"
#include "gdtypes.hpp"
#include "weight.hpp"
#include "utils.hpp"
#include "gdexception.hpp"

// cdd library header
#ifdef BACKEND_CDD
    #include <cddlib/setoper.h>
    #include <cddlib/cdd.h>
#endif

//-----------------------------------------------------------------------
GeneralDichotomy::GeneralDichotomy(BaseSolver& solver): _solver(solver) {

}



//-----------------------------------------------------------------------
bool GeneralDichotomy::findImprovingSolution(size_t& new_sol_ind, std::set<size_t>& polytope_sol, size_t& origin_weight_ind) {

    dichotomy::CostVec next_sol;
    dichotomy::COST z;

    bool diversify = true;

    // sort the extreme weights to diversify the pareto front by default
    std::vector<size_t> weight_indexes;
    weight_indexes.reserve(_weights.size());
    for(size_t w_ind = 0; w_ind < _weights.size(); w_ind ++) {
        if(!_weights[w_ind].removed && !_weights[w_ind].tested) {

            if(_params.max_depth == 0 || _weights[w_ind].depth <= _params.max_depth) {
                weight_indexes.push_back(w_ind);
            }
        }
    }
    if(diversify) {
        // std::sort(weight_indexes.begin(), weight_indexes.end(), [this](size_t w_left, size_t w_right){ return _weights[w_left].iteration < _weights[w_right].iteration; });
        std::sort(weight_indexes.begin(), weight_indexes.end(), [this](size_t w_left, size_t w_right){ return _weights[w_left].depth < _weights[w_right].depth; });
    }

    auto it = weight_indexes.begin();

    bool found = false;
    while(it != weight_indexes.end() && !found) {

        size_t w_ind = *it;

        // in intensification mode, check only weight adjacent to the initial solution
        if(_params.search_scheme == GeneralDichotomy::INTENSIFICATION) {
            if(_weights[w_ind].adj.find(0) == _weights[w_ind].adj.end()) {
                it ++;
                continue;
            }
        }

        if(_params.verbose >= 1) {
            std::cout << "testing weight w^" << w_ind << " ";
            _weights[w_ind].print(std::cout);
            std::cout << std::endl;
        }

        bool hasSolution = _solver.solve(_weights[w_ind].weight, _weights[w_ind].adj, next_sol, z);
        _nSubSolve ++;
        _weights[w_ind].tested = true;

        if(_params.verbose >= 1) {
            std::cout << "has solution ? " << hasSolution << std::endl;
        }

        // test if there is an improvement
        auto cmp = Utils::compare(z, _weights[w_ind].z);

        if(_params.verbose >= 1) {
            std::cout << "next solution found: ";
            
            Utils::printCostVec(next_sol, std::cout);
            
            std::cout << " cmp=" << cmp << ", z=" << z << std::endl;
            std::cout << std::endl;
        }

        auto temp_size = _solutions.size();

        new_sol_ind = addCheckSolution(next_sol);

        // std::cout << "new sol is y ^" << new_sol_ind << " (temp_size=" << temp_size << ")" << std::endl;

        // notify the user wether the solution is a new one
        _solver.isNew(new_sol_ind == temp_size);

        if(_params.verbose >= 1) {
            _weights[w_ind].print(std::cout);
            std::cout << " " << std::endl;
        }

        // check wether the solution improves the weight
        if(cmp == 1) { // improvement
            
            found = true;
            origin_weight_ind = w_ind;
            if(_params.verbose >= 1) {
                std::cout << "Improvement found!" << std::endl;
            }

            // update the list of solutions to consider to build the polytope
            polytope_sol.insert(_weights[w_ind].adj.begin(), _weights[w_ind].adj.end());

            _weights[w_ind].adj.clear();
            _weights[w_ind].adj.emplace(new_sol_ind); // inside the polytope of the new solution
            if(_weights[w_ind].bnd.size() < _solver.getNObj()-1) {
                _weights[w_ind].removed = true;
                _nRemoved ++;
            }
            break;

        } else { // equality

            if(cmp == 0) {
                _weights[w_ind].adj.emplace(new_sol_ind);
            }

        }

        // to the next weight
        it ++;

    }

    return found;
}

//-----------------------------------------------------------------------
dichotomy::Result GeneralDichotomy::run(Parameters parameters) {

    _params = parameters;

    auto compareWeights = [](Weight* left, Weight* right) { 
        if( abs(left->weight[0] - right->weight[0]) <= dichotomy::epsilon1) {
            return true;
        } else {
            return left->weight[0] < right->weight[0];
        } };
    
    #ifdef BACKEND_CDD
    if(_params.backend == Backend::CDD) {
        dd_set_global_constants();
    }
    #endif

    dichotomy::Result result(_weights, _weightGraph, _solutions, _nSubSolve);

    // brute force version: each compute the polytope of a solution against every other solution

    // test with the initial weights
    _nRemoved = 0;

    _nSubSolve = 0; // number of sub-problems solved

    _nTotalWeights = 0;

    // Weight set decomposition initialization (W^0)
    initWeights();

    // initialize the weight for the first solution
    // can be anything
    // the adjacency for the first solution computed is always be empty
    dichotomy::WeightVec init_weight;
    
    if(_params.initial_weight.size() > 0) {
        init_weight = _params.initial_weight;        
    } else {
        init_weight = dichotomy::WeightVec(_solver.getNObj(), static_cast<double>(parameters.wnorm)/_solver.getNObj());
        if(_params.directions.size() > 0) {
            for(size_t obj_ind = 0; obj_ind < init_weight.size(); obj_ind ++) {
                if(_params.directions[obj_ind] == dichotomy::MAX) {
                    init_weight[obj_ind] = -init_weight[obj_ind]; 
                }
            }
        }
    }
    
    std::set<size_t> empty_adj;

    // dichotomy::WeightVec init_weight(_solver.getNObj(), 0.);
    // init_weight[0] = parameters.wnorm;

    // pick one solution as starter
    size_t obj_ind = 0;
    dichotomy::CostVec solution;
    dichotomy::COST z;

    bool hasSolution = _solver.solve(init_weight, empty_adj, solution, z);
    _solutions.push_back(solution);
    _nSubSolve ++;

    // temporary
    if(!hasSolution) {
        std::cerr << "Error, the problem has no solution!" << std::endl;
        throw Dichotomy::GDException();
    }

    // the first solution is always new
    _solver.isNew(true);

    if(_params.verbose >= 1) {
        std::cout << "initial solution found: ";
        Utils::printCostVec(solution, std::cout);
        std::cout << std::endl; 
    }

    // update the extreme weights accordingly
    for(std::size_t wind = 0; wind < _weights.size(); wind ++) {
        _weights[wind].adj.emplace(0);
        _weights[wind].z = Utils::dot_product(_weights[wind].weight, solution);
    }
    // this should be true only of the initial weight tested is one of the initial extreme weight
    // _weights[obj_ind].tested = true;

    bool stop = false;
    size_t iter = 0;

    std::set<size_t> polytope_sol; // list of solutions to consider to build the polytope

    while(!stop) {    

        iter += 1;
        if(_params.max_iter > 0 && iter >= _params.max_iter) {
            break;
        }

        if(_params.verbose >= 1) {
            std::cout << "--------------------------------------------" << std::endl;
            std::cout << "           iteration " << iter << std::endl;
            std::cout << "--------------------------------------------" << std::endl;
        }

        // look for an untested weight, optimize it, if the solution improves some weights, the weights are removed and the solution's polytope is computed

        // std::cout << "n weights: " << _weights.size() << std::endl;
        // std::cout << "n sol: " << _solutions.size() << std::endl;

        size_t new_sol_ind;

        size_t origin_weight_ind; // if a new solution is found, this is the index of the weight vector from which it was found

        polytope_sol.clear();

        bool found = findImprovingSolution(new_sol_ind, polytope_sol, origin_weight_ind); 

        if(!found) {
            stop = true;
            break;
        }

        std::set<dichotomy::WeightVec, weight_compare> equal_weights;

        // remove weights that are not extreme anymore
        for(size_t w_ind = 0; w_ind < _weights.size(); w_ind ++) {

            // if(_weights[w_ind].removed || _weights[w_ind].tested) {
            if(_weights[w_ind].removed) {
                continue;
            }

            dichotomy::WEIGHT z_cmp = Utils::dot_product(_weights[w_ind].weight, _solutions[new_sol_ind]);

            auto cmp = Utils::compare(z_cmp, _weights[w_ind].z);

            if(cmp == 1) { // improvement
                if(_weights[w_ind].bnd.size() < _solver.getNObj()-1) {
                    _weights[w_ind].removed = true;
                    _nRemoved ++;
                }
                polytope_sol.insert(_weights[w_ind].adj.begin(), _weights[w_ind].adj.end());
            } else if(cmp == 0) { // equality
                equal_weights.insert(_weights[w_ind].weight);
                polytope_sol.insert(_weights[w_ind].adj.begin(), _weights[w_ind].adj.end());
                _weights[w_ind].adj.emplace(new_sol_ind);
            }
        }

        // std::cout << "equal weights: " << std::endl;
        // for(auto weight: equal_weights) {
        //     for(auto val: weight) {
        //         std::cout << val << ", ";
        //     }
        //     std::cout << std::endl;
        // }

        // add the new weights
        std::vector<Weight> new_weights;

        std::vector<size_t> other_sol;
        for(auto& elt: polytope_sol) {
            if(elt != new_sol_ind) {
                other_sol.push_back(elt);
            }
        }
        
        if(_params.backend == Backend::CDD) {
            computePolytopeCDD(new_sol_ind, other_sol, new_weights, _weights[origin_weight_ind].depth);
        } else if(_params.backend == Backend::PPL) {
            computePolytopePPL(new_sol_ind, other_sol, new_weights, _weights[origin_weight_ind].depth);
        }

        for(auto weight: new_weights) {
            if(weight.bnd.size() < _solver.getNObj()-1) {
                auto it = equal_weights.find(weight.weight);
                if(it == equal_weights.end()) {
                    _weights.emplace_back(weight);
                }
            }
        }

        if(_params.verbose >= 1) {
            std::cout << std::endl << "Output extreme weights after polytope enumeration:" << std::endl;
            size_t wnew_ind = 0;
            for(auto weight: new_weights) {
                std::cout << "w^" << wnew_ind << ": ";
                weight.print(std::cout);
                std::cout << std::endl;
                wnew_ind += 1;
            }
        }

        if(_params.verbose >= 1) {
            std::cout << std::endl;
            std::cout << "Current extreme weights: " << _weights.size() << std::endl;
            for(size_t ind = 0; ind < _weights.size(); ind ++) {
                std::cout << "w^" << ind << " ";
                _weights[ind].print(std::cout);
                std::cout << std::endl;
            }
            std::cout << std::endl;
        }

        if(_params.verbose >= 1) {
            std::cout << "Current solutions: " << _solutions.size() << std::endl;
            for(size_t sol_ind = 0; sol_ind < _solutions.size(); sol_ind ++) {
                std::cout << "y^" << sol_ind << ": ";
                Utils::printCostVec(_solutions[sol_ind], std::cout);
                std::cout << std::endl;
            }
            std::cout << std::endl;
        }

        // clean weight set
        if(3*_nRemoved > _weights.size()) {
            _nTotalWeights += _nRemoved;
            _nRemoved = 0;
            std::vector<Weight> wtemp;
            for(auto& weight: _weights) {
                if(!weight.removed) {
                    wtemp.emplace_back(weight);
                }
            }
            _weights = wtemp;
        }

        
    }

    _nTotalWeights += _weights.size();

    // remove discarded weights
    std::vector<Weight> wtemp;
    for(auto& weight: _weights) {
        if(!weight.removed) {
            wtemp.emplace_back(weight);
        }
    }
    _weights = wtemp;

    // optional computation of the weight graph
    if(_params.export_weight_graph) {
        computeGraph();
    }

    #ifdef BACKEND_CDD
    if(_params.backend == Backend::CDD) {
        dd_free_global_constants();
    }
    #endif

    return result;

}


//---------------------------------------------------------------
void GeneralDichotomy::initWeights() {

    for(size_t obj_ind = 0; obj_ind < _solver.getNObj(); obj_ind ++) {
        _weights.emplace_back(Weight());
        _weights.back().weight = std::vector<dichotomy::WEIGHT>(_solver.getNObj(), 0.);

        if(!_params.directions.empty() && _params.directions[obj_ind] == dichotomy::MAX) {
            _weights.back().weight[obj_ind] = -_params.wnorm;
        } else {
            _weights.back().weight[obj_ind] = _params.wnorm;
        }
        for(int i = static_cast<int>(_solver.getNObj())-1; i >= 0; i --) {
            if(i == static_cast<int>(obj_ind)) {
                continue;
            }
            _weights.back().bnd.emplace(-i-1);
        }
        _weights.back().initial = true;
        _weights.back().iteration = 0;
        _weights.back().depth = 1;

    }

}

//---------------------------------------------------------------
size_t GeneralDichotomy::addCheckSolution(dichotomy::CostVec& sol) {

    // TODO more efficient data structure

    size_t result;
    bool found = false;

    for(size_t sol_ind = 0; sol_ind < _solutions.size(); sol_ind ++) {
        if(Utils::compareCostVectors(_solutions[sol_ind], sol)) {
            found = true;
            result = sol_ind;
            break;
        }
    }

    if(!found) {
        result = _solutions.size();
        _solutions.push_back(sol);
    }

    return result;

}

//---------------------------------------------------------------
void GeneralDichotomy::computeGraph() {
    
    _weightGraph.resize(_weights.size(), std::set<size_t>());

    for(size_t wind1 = 0; wind1 < _weights.size(); wind1 ++) {
        for(size_t wind2 = wind1+1; wind2 < _weights.size(); wind2 ++) {
        
            // performing set intersection
            std::set<size_t> adj_common;
            std::set<int> bnd_common;

            for(auto elt: _weights[wind1].bnd) {
                if(_weights[wind2].bnd.find(elt) != _weights[wind2].bnd.end()) {
                    bnd_common.emplace(elt);
                }
            }
            for(auto elt: _weights[wind1].adj) {
                if(_weights[wind2].adj.find(elt) != _weights[wind2].adj.end()) {
                    adj_common.emplace(elt);
                }
            }

            if(adj_common.size()+bnd_common.size() >= _solver.getNObj()-1)  {
                _weightGraph[wind1].emplace(wind2);
                _weightGraph[wind2].emplace(wind1);
            }

        }
    }

}

//---------------------------------------------------------------
size_t GeneralDichotomy::getNSubSolve() {
    return _nSubSolve;
}

//---------------------------------------------------------------
size_t GeneralDichotomy::getNComputedWeights() {
    return _nTotalWeights;
}