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

import numpy as np
from numpy.testing import assert_almost_equal
import matplotlib.pyplot as plt
import scipy

import pickle, random, math

from models import MetricalPositionModel, DiscountedMetricalPositionModel
from stats import beta_kl_divergence

random_seed = 1
#random.seed(random_seed)

bar_length = 16

# Load data
data = pickle.load(open('data/witek.pickle', 'rb'))
essen_data = pickle.load(open('data/german_quadruple_large.pickle', 'rb'))
essen_rhythms = [essen_data[key]['grid'] for key in essen_data]# if not 'Experimenter' in key]

N = len(data)
rhythms = [data[key]['grid'] for key in data]# if not 'Experimenter' in key]
move_ratings = [data[key]['move'] for key in data]# if not 'Experimenter' in key]
pleasure_ratings = [data[key]['pleasure'] for key in data]# if not 'Experimenter' in key]
syncopation_ratings = [data[key]['syncopation'] for key in data]# if not 'Experimenter' in key]

PIRs = []
total_PIRs = []
cross_entropies = []
parameter_information = []

def train_model(model, rhythms):

    for rhythm in rhythms:
        for bar in rhythm.iter_bars():
            for obs in enumerate(bar):
                model.observe(*obs)

ltm_model = MetricalPositionModel(range(bar_length))
ltm_model.set_counts_to_one()
train_model(ltm_model, rhythms)

for rhythm in rhythms:

    cross_entropy = np.zeros(len(rhythm))
    t = 0

    stm_model = MetricalPositionModel(range(bar_length))
    stm_model.set_counts_to_one()
    models = [stm_model, ltm_model]


    for bar in rhythm.iter_bars():

        for position, onset in enumerate(bar):

            obs = position, onset
    
            # Calculate normalised weights using the inverse entropy in the 
            # predictive distributions
            I = [model.entropy(position) ** -1 for model in models]
            W = [i / sum(I) for i in I]

            # Geometric mean
            #p_onset = math.exp(sum([math.log(model.likelihood(position, 1) ** w) for w, model in zip(W, models)]))
            #p_no_onset = math.exp(sum([math.log(model.likelihood(position, 0) ** w) for w, model in zip(W, models)]))

            # Arithmic mean
            p_onset = sum([model.likelihood(position, 1) * w for w, model in zip(W, models)])
            p_no_onset = sum([model.likelihood(position, 0) * w for w, model in zip(W, models)])

            assert_almost_equal(sum(W), 1)
            assert_almost_equal(p_onset + p_no_onset, 1)

            cross_entropy[t] = - math.log(p_onset if onset else p_no_onset)
            stm_model.observe(*obs)

            t += 1

    cross_entropies.append(cross_entropy/len(rhythm))

per_song_ce = [sum(ce) / len(ce) for ce in cross_entropies]

plt.figure()
plt.plot(per_song_ce, move_ratings, 'x')
plt.title('Cross entropy and move')

plt.figure()
plt.plot(per_song_ce, pleasure_ratings, 'x')
plt.title('Cross entropy and pleasure')

plt.figure()
plt.plot(syncopation_ratings, per_song_ce, 'x')
plt.title('Syncopation and cross entropy')

plt.show()
