from collections import Counter
from datetime import datetime
from math import log, sqrt

import pandas as pd
from scipy.stats import uniform, beta

from funcs import smallest_odd_larger, run_experiment

batch_size = 10000
iterations = 100

total_runs = batch_size * iterations

output_file = f'data/experiment-{total_runs}-runs-{datetime.now().strftime("%Y%m%d%H%M%S")}.csv'

n_vals = [3, 10, 50, 200, 500]

k_funcs = {
    'const1': lambda n: 1,
    'const3': lambda n: 3,
    'logn': lambda n: smallest_odd_larger(log(n)),
    'n^.4': lambda n: smallest_odd_larger(n ** .4),
    'sqrt': lambda n: smallest_odd_larger(sqrt(3 / 4 * n)),
    'n/2': lambda n: smallest_odd_larger(n / 2),
}

dists = {
    'uniform[0,1]': uniform(0, 1),
    'uniform[.1,.9]': uniform(.1, .8),  # scipy has weird syntax, first param is lower end and second is length
    'beta[2,2]': beta(2, 2),
}

if __name__ == '__main__':
    rows = []
    for n in n_vals:
        # Compute actual k values for this n.
        k_vals = {k_name: k_func(n) for k_name, k_func in k_funcs.items()}
        for dist_name, dist in dists.items():
            # Keep track of total number of correct for each k value.
            counts = Counter()
            for i in range(iterations):
                if i % 10 == 0:
                    print(f'{dist_name}, n={n}: Iteration {i}')
                counts += run_experiment(n, k_vals, dist, batch_size)
            row = {'dist': dist_name, 'n': n, **counts}
            rows.append(row)

    df = pd.DataFrame(rows, columns=['dist', 'n', *k_funcs.keys()])
    df.to_csv(output_file, index=False)
