import numpy as np
import matplotlib.pyplot as plt
import scipy
import pickle, random, math

from pymuco.representations import MetricalGrid
from pymuco.temperley.models import MetricalPositionModel

from stats import beta_kl_divergence
from shortcuts import load_and_train_model, load_dataset

bar_length = 16

# Load data
model = MetricalPositionModel(bar_length)
load_and_train_model(model, 'essen_german_simple_quadruple_fixed_grid')

witek_data = load_dataset('witek')

test_set = [(d['bass_drum'], d['snare_drum']) for k, d in witek_data.items()]
move_ratings = [witek_data[key]['move'] for key in witek_data]# if not 'Experimenter' in key]
pleasure_ratings = [witek_data[key]['pleasure'] for key in witek_data]# if not 'Experimenter' in key]
syncopation_ratings = [witek_data[key]['syncopation'] for key in witek_data]# if not 'Experimenter' in key]

reference_counts = model.counts.copy()
information_gains = []

for rhythm in test_set:

    observations = model.observations(rhythm[0])
    observations += model.observations(rhythm[0])
    observations += model.observations(rhythm[1])
    model.counts = reference_counts.copy()

    params_before = [(model.counts[position, 0], model.counts[position, 1]) for position in range(bar_length)]

    for obs in observations:

        model.observe(*obs)

    params_after = [(model.counts[position, 0], model.counts[position, 1]) for position in range(bar_length)]

    divergences = [beta_kl_divergence(*before + after) for before, after in zip(params_before, params_after)]
    print(divergences)
    information_gains.append(sum([divergence for divergence in divergences if not np.isnan(divergence)]))

plt.figure()
plt.plot(information_gains, pleasure_ratings, 'o')
plt.figure()
plt.plot(information_gains, move_ratings, 'o')
plt.show()
