import random
import numpy

## Configuration MasterMind

nb_colors = 5
nb_cases = 5


## Black-box
class MasterMind:
    def __init__(self, nb_colors, nb_cases, solution=None):
        self.nb_colors = nb_colors
        self.nb_cases = nb_cases
        self._solution = solution if solution else self.random_candidate()
    
    def random_candidate(self):
        import random
        return [random.randrange(self.nb_colors) for _ in range(self.nb_cases)]
    
    def try_candidate(self, candidate):
        total = 0
        for i, x in enumerate(self._solution):
            if candidate[i] == x:
                total += 1
        return total
    
    def get_max_radius(self):
        return self.nb_cases
    
    def modify_with_radius(self, x, r):
        import random
        x = x[:]
        items = list(range(self.nb_cases))
        random.shuffle(items)
        for i in items[:r]:
            x[i] = random.randrange(self.nb_colors)
        return x
    
    def is_optimal(self, x, v_x=None):
        v_x = v_x if v_x else self.try_candidate(x)
        return (v_x == self.nb_cases)


class SearchHeuristic:
    def set_problem(self, problem):
        self.problem = problem
    
    def run(self):
        raise Exception('Not implemented !')

# Random Local Search
class RLS(SearchHeuristic):
    def run(self):
        pb = self.problem
        x = pb.random_candidate()
        v_x = pb.try_candidate(x)
        nb_calls = 1
        
        while not pb.is_optimal(x, v_x):
            r = 1
            y = pb.modify_with_radius(x, r)
            v_y = pb.try_candidate(y)
            if v_y >= v_x:
                x = y
                v_x = v_y
            nb_calls += 1
        
        return (x, v_x, nb_calls)
        
class EA(SearchHeuristic):
    def __init__(self, proba, lda, mu):
        self.proba = proba
        self.lda = lda
        self.mu = mu
    
    def run(self):
        pb = self.problem
        
        x_list = []
        
        for _ in range(self.mu):
            x = pb.random_candidate()
            vx = pb.try_candidate(x)
            x_list.append((x, vx))
        x_list = sorted(x_list, key=lambda x: x[1], reverse=True)
        x_best = x_list[0]
        nb_calls = 1
        
        while not pb.is_optimal(x_best[0], x_best[1]):
            y_list = []
            
            for _ in range(self.lda):
                r = numpy.random.binomial(pb.get_max_radius(), self.proba)
                y = pb.modify_with_radius(x_best[0], r)
                vy = pb.try_candidate(y)
                y_list.append((y, vy))
            #print(x_list, y_list)
                
            selection = x_list + y_list
            selection = sorted(selection, key=lambda x: x[1], reverse=True)
            x_list = selection[:self.mu]
            x_best = x_list[0]
            nb_calls += 1
                    
        return (x_best[0], x_best[1], nb_calls)

def avg_calls(algo, problem, nb=100):
    nb_calls_list = []
    for _ in range(100):
        algo.set_problem(problem)
        _, _, nb_calls = algo.run()
        nb_calls_list.append(nb_calls)
    return numpy.mean(nb_calls_list)

## Main
avg = avg_calls(RLS(), MasterMind(nb_colors, nb_cases))
print('Average RLS: %s' % avg)
avg = avg_calls(EA(1./nb_cases, 1, 1), MasterMind(nb_colors, nb_cases))
print('Average EA: %s' % avg)


print('=================')
probas = list(numpy.arange(0.05, 0.5, 0.01))
p_best_list = []
for n in range(5, 15):
    avg_list = []
    for p in probas:
        avg = avg_calls(EA(p, 1, 1), MasterMind(n, n))
        avg_list.append(avg)
        #print('{}%'.format(int(p*100)))
    p_best = probas[numpy.argmin(avg_list)]
    p_best *= n
    p_best_list.append(p_best)
    print('{}: {}'.format(n, p_best))

from matplotlib import pyplot as plt
x = list(range(5, 15))
y = p_best_list
plt.plot(x, y)
plt.savefig('foo3.png')