%matplotlib notebook

import os
os.environ["CUDA_VISIBLE_DEVICES"] = "2"

from tracker3d import utils
from tracker3d import metrics
from tracker3d import loader

import numpy as np
import pandas as pd

import keras

order = ("phi", "r", "z")

train, target = loader.from_file("datasets/npz/unif25_prz_n10.npz")

#test = train[int(0.2 * len(train)):]
#train = train[:int(0.2 * len(train))]
#print(test.shape)
#print(train.shape)

#testTarget = target[int(0.2 * len(target)):]
#target = target[:int(0.2 * len(target))]
#print(testTarget.shape)
#print(target.shape)

test, testTarget = loader.from_file("datasets/npz/ramp_prz_n10.npz")

event = 1025
utils.display_side_by_side(train[event], target[event], order=order)

from keras.layers import Dense, LSTM, Dropout, GRU, TimeDistributed, Bidirectional
from keras.models import Sequential
from keras.regularizers import l1, l2, l1_l2

input_shape = train[0].shape # Shape of an event.
output_shape = len(target[0][0]) # Number of tracks per event

batch_size = 32
epochs     = 32
valsplit   = .1
opt        = 'rmsprop' # optimizer
model = Sequential()
model.add(Bidirectional(GRU(256, return_sequences=True, recurrent_dropout=0.4, implementation=2, bias_regularizer=l2(0.02)), input_shape=input_shape, merge_mode="mul"))
model.add(Dropout(0.4))
model.add(Bidirectional(GRU(256, return_sequences=True, recurrent_dropout=0.4, implementation=2, bias_regularizer=l2(0.02)), merge_mode="mul"))
model.add(Dropout(0.4))
model.add(Bidirectional(GRU(256, return_sequences=True, recurrent_dropout=0.4, implementation=2, bias_regularizer=l2(0.02)), merge_mode="mul"))
model.add(Dropout(0.4))
model.add(TimeDistributed(Dense(output_shape, activation='softmax')))

model.compile(loss='categorical_crossentropy', optimizer=opt, metrics=['accuracy'])
model.summary()

%%time
modelpath = 'simple.h5'
hist = model.fit(train, target, epochs=epochs, batch_size=batch_size,
                         verbose=1, validation_split=valsplit,
                         #callbacks=[keras.callbacks.ModelCheckpoint(filepath=modelpath, verbose=0)]
                )

predictions = model.predict(test, batch_size=batch_size)

utils.print_scores(model, train, target, batch_size)
utils.graph_losses([("Categorical Cross Entropy", hist)])

tracks, acc = metrics.accuracy_vs_tracks(predictions, testTarget, has_noise=True, has_padding=True)

np.save("testTarget.npy", testTarget)
np.save("predictions.npy", predictions)

print(metrics.discrete_accuracy_all(test, testTarget, predictions))

print(acc)

%matplotlib inline
import matplotlib.pyplot as plt
plt.figure(2)
fig = plt.scatter(tracks, acc)
plt.show(fig)

numTracks = int(np.amax(tracks))
accuracy = np.zeros(numTracks)
for i in range(numTracks):
    tot_acc = 0
    count = 0
    for j, track in enumerate(tracks):
        print(acc[j])
        if track == i + 1:
            tot_acc = tot_acc + acc[j]
            count += 1
    if count > 0:
        accuracy[i] = tot_acc/count

track = np.zeros(numTracks)

for i in range(numTracks):
    track[i] = i + 1
        
print(accuracy)
print(track)

plt.scatter(track, accuracy)
plt.plot(np.unique(track), np.poly1d(np.polyfit(track, accuracy, 1))(np.unique(track)))
plt.title("Average Accuracy vs. Number of Tracks")
plt.xlabel("Number of Tracks")
plt.ylabel("Accuracy")
plt.show(fig)

for i, cell in enumerate(acc):
    if cell < 0.3:
        print(cell)
        print(tracks[i])