#include <iostream>
#include <cstdlib>
#include <ctime>
#include <numeric>
#include <algorithm>
#include <vector>
#include <random>

using namespace std;

class MasterMind {
    public:
        MasterMind(int nb_colors, int nb_cases) {
            m_nb_colors = nb_colors;
            m_nb_cases = nb_cases;
            m_solution = random_candidate();
        }
        
        MasterMind(int nb_colors, int nb_cases, int* solution) {
            m_nb_colors = nb_colors;
            m_nb_cases = nb_cases;
            m_solution = solution;
        }
        
        int* random_candidate() {
            int* cand = new int[m_nb_cases];
            for(int i=0; i<m_nb_cases; i++)
                cand[i] = (rand() % m_nb_colors);
            return cand;
        }
        
        int try_candidate(int* candidate){
            int total = 0;
            for(int i=0; i<m_nb_cases; i++)
                if(candidate[i] == m_solution[i])
                    total++;
            return total;
        }
        
        int get_max_radius() {
            return m_nb_cases;
        }

        int* modify_with_radius(int* x, int r, bool copy=false) {
            vector<int> items(m_nb_cases);
            iota(items.begin(), items.end(), 0);
            mt19937 rng(std::random_device{}());
            shuffle(begin(items), end(items), rng);
            
            int* x2 = x;
            if(copy) {
                x2 = new int[m_nb_cases];
                for(int i=0; i<m_nb_cases; i++)
                    x2[i] = x[i];
            }
            
            for(int i=0; i<r; i++)
                x2[items[i]] = rand() % m_nb_colors;
            return x2;
        }

        bool is_optimal(int* x) {
            int v_x = try_candidate(x);
            return (v_x == m_nb_cases);
        }
        
        bool is_optimal(int* x, int v_x) {
            return (v_x == m_nb_cases);
        }

    private:
        int m_nb_colors;
        int m_nb_cases;
        int* m_solution;
};

struct Result {
    Result(int* s, int v, int c) {
        solution = s;
        v_solution = v;
        nb_calls = c;
    }
    
    int* solution;
    int v_solution;
    int nb_calls;
};

class SearchHeuristic {
    public:
        virtual Result run() = 0;
        void set_problem(MasterMind* problem) {
            m_problem = problem;
        }
    
    protected:
        MasterMind* m_problem;
};

struct Candidate {
    Candidate(int* s, int v) {
        sol = s;
        value = v;
    }
    
    int* sol;
    int value;
};

class RLS : public SearchHeuristic {
    public:
        Result run() {
            MasterMind* pb = m_problem;
            int* x = pb->random_candidate();
            int v_x = pb->try_candidate(x);
            int nb_calls = 1;
            
            while(! pb->is_optimal(x, v_x)) {
                int r = 1;
                int* y = pb->modify_with_radius(x, r, true);
                int v_y = pb->try_candidate(y);
                if(v_y >= v_x) {
                    delete[] x;
                    x = y;
                    v_x = v_y;
                } else {
                    delete[] y;
                }
                nb_calls++;
            }
            
            return Result(x, v_x, nb_calls);
        }
};

class EA : public SearchHeuristic {
    public:
        EA(double proba, int lda, int mu) {
            m_proba = proba;
            m_lambda = lda;
            m_mu = mu;
        }
        
        Result run() {
            MasterMind* pb = m_problem;
            
            vector<Candidate> x_list;
            for(int i=0; i<m_mu; i++) {
                int* x = pb->random_candidate();
                int vx = pb->try_candidate(x);
                x_list.push_back(Candidate(x, vx));
            }
            
            struct {
                bool operator()(Candidate a, Candidate b) const
                {
                    return a.value > b.value;
                }
            } customLess;
            sort(x_list.begin(), x_list.end(), customLess);
            
            Candidate x_best = x_list[0];
            int nb_calls = 1;
            
            default_random_engine generator;
            binomial_distribution<int> distribution(pb->get_max_radius(), m_proba);
            
            while(! pb->is_optimal(x_best.sol, x_best.value)) {
                for(int i=0; i<m_lambda; i++) {
                    int r = distribution(generator);
                    int* y = pb->modify_with_radius(x_best.sol, r, true);
                    int vy = pb->try_candidate(y);
                    x_list.push_back(Candidate(y, vy));
                }
                
                sort(x_list.begin(), x_list.end(), customLess);
                for(int i=0; i<x_list.size()-m_mu; i++)
                    x_list.pop_back();
                x_best = x_list[0];
                nb_calls++;
            }
            
            return Result(x_best.sol, x_best.value, nb_calls);
        }
    
    private:
        double m_proba;
        int m_lambda;
        int m_mu;
};

double avg_calls(SearchHeuristic& algo, MasterMind problem, int nb=100) {
    int sum = 0;
    for(int i=0; i<nb; i++) {
        algo.set_problem(&problem);
        Result res = algo.run();
        sum += res.nb_calls;
        //cout << res.nb_calls << "; ";
    }
    return (double) sum / nb;
}

int main() {
    int nb_colors = 5;
    int nb_cases = 5;
    
    srand(time(NULL));
    
    RLS algo1;
    double avg = avg_calls(algo1, MasterMind(nb_colors, nb_cases));
    cout << "Average RLS: " << avg << endl;
    EA algo2(1./nb_cases, 1, 1);
    avg = avg_calls(algo2, MasterMind(nb_colors, nb_cases));
    cout << "Average EA: " << avg << endl;

    cout << "=========================" << endl;
    for(int n=5; n<15; n++) {
        
        vector<double> avg_list;
        vector<double> probas;
        
        for(double p=0.05; p<0.5; p+=0.01) {
            EA algo(p, 1, 1);
            double avg = avg_calls(algo, MasterMind(n, n));
            avg_list.push_back(avg);
            probas.push_back(p);
        }
        
        int min_ind = 0;
        for(int i=1; i<avg_list.size(); i++)
            if(avg_list[min_ind] > avg_list[i])
                min_ind = i;
        double p_best = probas[min_ind];
        cout << n << ": " << p_best << endl;
        
    }
    
    return 0;
}
