# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
# ---

# # Recorded data

import pickle
from matplotlib import pyplot as plt
import numpy as np
import tensorflow as tf
from tensorflow.keras import datasets, layers, models

# read data
data = {}
with open('data-test.pkl', 'rb') as f:
    data = pickle.load(f)
    

print(len(data))
print(data[0]["stack"])

plt.imshow(data[0]["stack"][0])

plt.imshow(data[50]["stack"][0])

for d in data[0:20]:
    print(d["value"])
    print(d["stack"][0].shape)

# # Sample model train

# +
model = models.Sequential()
model.add(layers.Conv2D(
    32, (3, 3), activation='relu', input_shape=(40, 120, 3)))
model.add(layers.MaxPooling2D((2, 2)))
model.add(layers.Conv2D(64, (3, 3), activation='relu'))
model.add(layers.MaxPooling2D((2, 2)))
model.add(layers.Conv2D(64, (3, 3), activation='relu'))

model.add(layers.Flatten())
model.add(layers.Dense(64, activation='relu'))
model.add(layers.Dense(64, activation='relu'))
model.add(layers.Dense(2, activation='softmax'))

model.compile(optimizer='adam',
              loss='binary_crossentropy',
              metrics=['accuracy'])

model.summary()

# +
x = []
y = []
for d in data:
    tmp = [d["stack"][0]/255, d["stack"][1]/255, d["stack"][2]/255]
    x.append(np.stack(tmp, axis=2))
    y.append(np.array(d["value"]))
x = np.array(x)
y = np.array(y)

print(len(x))
print(x.shape)
print(len(y))
print(y.shape)

print(y[20])
plt.imshow(x[20])
# -

xx = np.array([x[9], x[0]])
yy = np.array([y[9], y[0]])
print(yy)
history = model.fit(x, y, epochs=20)

# +

plt.plot(history.history['acc'], label='accuracy')
plt.plot(history.history['loss'], label = 'val_loss')
plt.xlabel('Epoch')
plt.ylabel('Accuracy')
plt.ylim([0.5, 1])
plt.legend(loc='lower right')

# -

test_loss, test_acc = model.evaluate(x, y, verbose=2)
print(test_loss)
print(test_acc)


print(model.predict(np.array([x[10]])))
print(model.predict_classes(np.array([x[10]])))

tf.keras.models.save_model(model, 'model.tf')



StarcoderdataJupyterScriptsDedupFiltered