#ifndef BEST_PEA_H_INCLUDED
#define BEST_PEA_H_INCLUDED

#include <vector>
#include "heuristics/PEA.h"

class PEAMutationProbaOptimizer : public ParameterOptimizer {
    public:
        PEAMutationProbaOptimizer(unsigned int n_init, std::vector<double> probas, unsigned int flexive_point, std::default_random_engine& randomizer) : ParameterOptimizer(randomizer), n_initial{n_init}, m_probas{probas}, m_flexive_point{flexive_point} {};
        SearchHeuristic* get_context(unsigned int i, double parameter, std::default_random_engine& randomizer) {
            std::vector<double> probas(m_probas);
            for(int k=m_flexive_point; k>=0; k--)
                probas[k] = parameter;
            /*probas[m_flexive_point] = parameter;
            if(m_flexive_point == 0) {
                for(int k=0; k<probas.size(); k++)
                    probas[k] = parameter;
            }*/
            /*std::cout << "[" << m_flexive_point << "] ";
            for(int i=0, s=probas.size(); i<s; i++)
                std::cout << probas[i] << " ";
            std::cout << std::endl;
            */PEA* algo = new PEA(1, 1, PLUS, probas, randomizer);
            MasterMind* mm = new MasterMind(n_initial+i, n_initial+i, randomizer);
            algo->set_problem(mm);
            return algo;
        }
        std::string display_context(unsigned int context) {
           std::stringstream ss;
           ss << "n=" << n_initial+context;
           return ss.str();
        }
        std::string get_context_for_csvfile(unsigned int num_context) {
           return std::to_string(n_initial+num_context);
        }
    private:
        unsigned int n_initial;
        std::vector<double> m_probas;
        unsigned int m_flexive_point;
};

void pea_optimisation(unsigned int n, unsigned int subdivision_size, std::default_random_engine& randomizer) {
    std::vector<double> probas;
    for(unsigned int i=0; i<=subdivision_size; i++)
        probas.push_back(0.);

    for(unsigned int k=subdivision_size; k>=0; k--) {
        PEAMutationProbaOptimizer pea_mpo(n, probas, k, randomizer);
        pea_mpo.set_nb_contexts(1);
        if(k==subdivision_size)
            pea_mpo.set_study_interval(0.01, 0.3, 0.01);
        else if(k > 15)
            pea_mpo.set_study_interval(0.01, 0.5, 0.01);
        else if(k > 10)
            pea_mpo.set_study_interval(0.01, 0.7, 0.01);
        else if(k > 5)
            pea_mpo.set_study_interval(0.01, 0.9, 0.01);
        else
            pea_mpo.set_study_interval(0.01, 0.99, 0.01);

        pea_mpo.add_test_wave(500, 2.2481733216, 0);
        //pea_mpo.add_test_wave(8000, 2.243091696, 1);

        std::vector<ConfidenceInterval> cis = pea_mpo.run(false);
        std::cout << "STEP " << k << ": " << cis[0].down << ", " << cis[0].mean << ", " << cis[0].up << std::endl;

        double p = (cis[0].down + cis[0].up)/2.;
        probas[k] = p;
        /*if(k==0) {
            for(int j=0; j<=subdivision_size; j++)
                probas[j] = p;
        } else {
            probas[k] = p;
            if(k+1<=subdivision_size) probas[k+1] = (2*p+probas[k+1])/3;
            if(k+2<=subdivision_size) probas[k+2] = (p+2*probas[k+2])/3;
        }*/
        std::cout << "==========> " << k << ", " << p << std::endl;
    }
}

#endif // BEST_PEA_H_INCLUDED