import generators.nb_gm_001 as nb_gm_001
import random
import heapq
import csv

def __test_generate_stat_csv(filename):
    alphabetic_choices = ['A', 'B', 'C', 'D', 'E', 'F', 'G', 'H']
    min_trials, max_trials = 24, 100
    n = 2
    with open(filename, mode='w') as stat_dist_file:
        writer = csv.writer(stat_dist_file, delimiter=',', quotechar='"', quoting=csv.QUOTE_MINIMAL)
        writer.writerow(['index'] + alphabetic_choices + ['ralph_skewed'])
        for i in range(10000):
            trials = random.randint(min_trials, max_trials)
            generator = nb_gm_001.SequenceGenerator(alphabetic_choices, n=n, trials=trials)
            seq = generator.generate()
            dist = [float(seq.count(c)) for c in alphabetic_choices]
            ralph_skewed = sum(heapq.nlargest(int(len(alphabetic_choices)/2), dist)) > (trials*2/3)
            writer.writerow([str(i)] + dist + [str(ralph_skewed)])
    __show_skweness_diagram(filename, alphabetic_choices)


def __show_skweness_diagram(filename, choices):
    import pandas as pd
    from matplotlib import pyplot as plt
    print(filename)
    data = pd.read_csv(filename)
    data['trials'] = data[choices].sum(axis=1)
    max_trials = data['trials'].max()
    min_trials = data['trials'].min()
    stats = []
    for t in range(int(min_trials), int(max_trials) + 1):
        dt = data[data.trials == t].trials.count()
        st = data[(data.trials == t) & (data['ralph_skewed']==True)].trials.count() * 100 / data.trials.sum()
        stats.append([t, dt, st])
    stats = pd.DataFrame(stats, columns=['trials', 'num_of_sequences','num_of_skewed_sequences'])
    plt.scatter(stats.trials,stats.num_of_skewed_sequences)
    plt.show()


if __name__ == '__main__':
    __test_generate_stat_csv('../benchmarks/nb_gm_001_2back_24trials.csv')