# .py file - change file ending

from docplex.mp.model import Model
from random import shuffle, seed
import matplotlib.pylab as plt

class Strategy:
    def get_next_pairs(self, n_pairs, previous_pairs, previous_correct):
        m = Model()
        x = {(letter, envelope): m.binary_var() for letter in range(n_pairs) for envelope in range(n_pairs)}
        # couple condition - one letter per envelope
        for i in range(n_pairs):
            m.add(m.sum(x[(i,j)] for j in range(n_pairs)) == 1)
            m.add(m.sum(x[(j,i)] for j in range(n_pairs)) == 1)
        
        # Use information about number of correct guesses in previous rounds
        for i in range(len(previous_pairs)):
            m.add(m.sum(x[pair] for pair in previous_pairs[i]) == previous_correct[i])
        solution = m.solve()
        if solution:
            pairs = []
            for letter in range(n_pairs):
                for envelope in range(n_pairs):
                    if x[(letter, envelope)].solution_value > 0.5:
                        pairs.append((letter, envelope))
            return pairs
        else:
            return None
        
def generate_random_set_of_pairs(n_pairs):
    letters = list(range(n_pairs))
    envelopes = list(range(n_pairs))
    shuffle(letters)
    shuffle(envelopes)
    matches = [(letters[i], envelopes[i]) for i in range(n_pairs)] 
    return matches
    
def simulate_once(strategy, n_pairs = 10, log_output = True):
    perfect_matches = generate_random_set_of_pairs(n_pairs)
    if log_output:
        print('The matches are as follows')
        print(perfect_matches)
    previous_pairs = []
    previous_correct = []
    for round in range(0, 200):
        if log_output:
            print('Round %i'%round)
        # at the beginning of the round, chose the matches
        pairs = strategy.get_next_pairs(n_pairs, previous_pairs, previous_correct)
        
        if log_output:
            print('Chosen matches:', pairs)
        # at the end of the episode we will know how many of the pairs were perfect matches
        count = 0
        for pair in pairs:
            if pair in perfect_matches:
                count += 1
        if log_output:
            print('Correct pairs in this round: ', count)
        previous_pairs.append(pairs) # update our lists
        previous_correct.append(count)
        if count == n_pairs:
            if log_output:
                print('All pairs were identified correctly in round %i'%round)
            return round
            
def monte_carlo_simulation(n_pairs, strategy, n_simulations = 1000):
    frequency = {i:0 for i in range(0,150)}
    for i in range(n_simulations):
        frequency[simulate_once(strategy, n_pairs, log_output = False)] += 1
    return frequency
    
# simulate_once(Strategy())
    
n_tries = 2000
frequencies = monte_carlo_simulation(10, Strategy(), n_tries)

lists = sorted(frequencies.items()) # sorted by key, return a list of tuples

print(lists)

x, y = zip(*lists) # unpack a list of pairs into two tuples

plt.plot(x, y)
plt.show()

cumulative = []
for i,j in lists:
    if cumulative:
        cumulative.append((i,j/n_tries*100+cumulative[-1][1]))
    else:
        cumulative = [(i,j/n_tries*100)]

print(cumulative)

x, y = zip(*cumulative) # unpack a list of pairs into two tuples

plt.plot(x, y)
plt.show()