# %%
# !pip install pyreadr tensorflow -U

import pyreadr
import pandas as pd
import numpy as np

dataset_file = 'data/CL2015.RData'

# read n-back dataset
datasets = pyreadr.read_r(dataset_file, use_objects="NB")
dataset = datasets["NB"]

X = dataset.stimulus.cat.codes.values
X = np.reshape(X, (X.shape[0], 1, 1))

y = pd.get_dummies(dataset.correct).values
y = np.reshape(y, (y.shape[0], 1, y.shape[1]))

# %%
from tensorflow import keras
from tensorflow.keras import layers

model = keras.Sequential()

model.add(layers.SimpleRNN(units=32, input_dim=1, activation="relu"))
model.add(layers.Dense(8, activation="relu"))
model.add(layers.Dense(1))
model.compile(loss='mean_squared_error', optimizer='rmsprop')
print(model.summary())

model.fit(X, y, batch_size=32)

# model.add(layers.Embedding(input_dim=28, output_dim=64))
# model.add(layers.LSTM(128))
# model.add(layers.Dense(10))

# model.summary()

# stimulus -> correct

# %%
model.summary()