import numpy as np
import itertools
import matplotlib.pyplot as plt
import multiprocessing as mp
import time


def save_ascii(population_list, unique_lists):
    data = np.array([population_list, unique_lists])
    data = data.T

    with open("output.txt", 'w+') as datafile_id:
        np.savetxt(datafile_id, data, fmt=['%d', '%.2f'])


def to_trinary(binary_string, twos_index):
    output = 0
    index = 0
    for bit in binary_string:

        if index in twos_index:
            output += 2 * (3 ** index)
        elif bit:
            output += (3 ** index)

        index += 1
    return output


class TaskParams:
    def __str__(self):
        return "Population size: " + str(self.population_size) + " String size: " + str(self.string_size)

    def notifyStart(self):
        self.start_time = time.time()

    def finish(self):
        end_time = time.time()
        elapsed = end_time - self.start_time
        print("Task: " + str(self) + " finished after: " + str(elapsed) + "seconds")

    def __init__(self, n_population, string_size, mean_times):
        self.population_size = n_population
        self.mean_times = mean_times
        self.string_size = string_size


class GeneString:
    def __init__(self, n):
        self.string = np.random.choice([False, True], n)


def calculate_mean_unique_schemas(pool_task):
    print("Calculating for: " + str(pool_task))

    times = pool_task.mean_times
    string_len = pool_task.string_size
    n_population = pool_task.population_size

    pool_task.notifyStart()

    possible_starred_indexes = range(0, string_len)
    possible_starred_indexes_sets = sum([map(list, itertools.combinations(possible_starred_indexes, i)) for i in
                                         range(len(possible_starred_indexes) + 1)], [])

    results = [None] * times
    for i in range(0, times):
        hash_set = set()
        for k in range(n_population):
            g = GeneString(string_len)

            for possible_starred in possible_starred_indexes_sets:
                str_copy = g.string.copy()
                hash_set.add(to_trinary(str_copy, possible_starred))

        results[i] = len(hash_set)

    pool_task.finish()
    return np.mean(results)


def generate_x_axis():
    xs = range(0, 100)
    xs += range(100, 250, 10)
    xs += range(250, 500, 50)
    xs += range(500, 1000, 100)
    return xs


def task_one():
    xs = generate_x_axis()
    pool_tasks = map(lambda x: TaskParams(x, 16, 10), xs)
    pool = mp.Pool(mp.cpu_count())
    out = pool.map(calculate_mean_unique_schemas, (x for x in pool_tasks))
    plt.plot(xs, out)
    plt.xlabel('Population size')
    plt.ylabel('Unique schemas')
    plt.show()

    save_ascii(xs, out)


if __name__ == '__main__':
    mp.freeze_support()
    task_one()