import heapq
import csv
import random

from benchmarks.common import *


def to_csv_row(sample_index, seq):
    trials = len(seq)
    freqs = [float(seq.count(c)) for c in alphabetic_choices]
    ralph_skewed = sum(heapq.nlargest(int(len(alphabetic_choices) / 2), freqs)) > (trials * 2 / 3)
    return [str(sample_index)] + freqs + [''.join(seq) , str(ralph_skewed)]


def benchmark(filename):
    import generators.nb_gm_001 as nb_gm_001

    generator = nb_gm_001.SequenceGenerator(alphabetic_choices, n, trials_range[0])

    # TODO store timing details
    with open(filename, mode='w') as benchmark_results_file:
        writer = csv.writer(benchmark_results_file, delimiter=',', quotechar='"', quoting=csv.QUOTE_MINIMAL)
        writer.writerow(['index'] + alphabetic_choices + ['seq', 'ralph_skewed'])
        for i in range(sample_size):
            trials = random.randint(trials_range[0], trials_range[1])
            seq = generator.generate(trials)
            print(f"sequence {i}/{sample_size}: {trials} trials")
            writer.writerow(to_csv_row(i, seq))


def skewness_diagram(csv_filename, figure_title):
    import pandas as pd
    import numpy as np
    from matplotlib import pyplot as plt
    data = pd.read_csv(csv_filename)
    data['trials'] = data[choices].sum(axis=1)
    max_trials = data['trials'].max()
    min_trials = data['trials'].min()
    stats = []
    for sequence_length in range(int(min_trials), int(max_trials) + 1):
        num_of_sequences = np.sum(data[data.trials == sequence_length].trials)
        skewed_sequences = np.sum(data[np.logical_and(data.trials == sequence_length, data.ralph_skewed)].trials)
        if num_of_sequences == 0:
            continue
        skewness = skewed_sequences * 100.0 / num_of_sequences
        stats.append([sequence_length, skewness])
    stats = pd.DataFrame(stats, columns=['trials', 'skewness'])  # .dropna(subset=['skewness'])
    # print(stats.trials)
    plt.ylim([0, 110])
    plt.scatter(stats.trials, stats.skewness, alpha=0.1)
    p = np.poly1d(np.polyfit(stats.trials, stats.skewness, 3))
    plt.plot(stats.trials, p(stats.trials), color='red')
    plt.title(figure_title)
    plt.ylabel('skewed blocks (%)')
    plt.xlabel('# of trials')
    plt.savefig(f'results/{figure_title}.png', bbox_inches='tight')
    plt.show()


def targets_ratio_diagram(csv_filename, title):
    import pandas as pd
    from matplotlib import pyplot as plt
    data = pd.read_csv(csv_filename)
    data['trials'] = data[choices].sum(axis=1)
    data['targets'] = data.apply(lambda s: count_targets_and_lures(s.seq, n)[0], axis=1)
    data['targets_ratio'] = data.apply(lambda s: s.targets / s.trials, axis=1)
    plt.scatter(data.trials, data.targets_ratio, alpha=0.2)
    plt.show()

if __name__ == '__main__':
    alg = 'nb_gm_001'
    fn = f'results/{alg}_{n}back.csv'
    benchmark(fn)
    # skewness_diagram(fn, alg)
    targets_ratio_diagram(fn, alg)