import numpy as np
import sys

samples = np.loadtxt(sys.argv[1] + '.csv', delimiter=',')
np.random.default_rng().shuffle(samples)
num_val = int(0.01*len(samples))
val = samples[:num_val]
train = samples[num_val:]

import tensorflow as tf

nr_patterns = 9
pattern_size = 8
nr_features = 2

patterns = tf.keras.Input(shape=(nr_patterns,), dtype='int32', name='patterns')
x = tf.one_hot(patterns, depth=3**pattern_size)
x = tf.keras.layers.Flatten()(x)
features = tf.keras.Input(shape=(nr_features,), name='features')
x = tf.keras.layers.Concatenate()([x, features])
output = tf.keras.layers.Dense(1, activation='sigmoid',
                                  kernel_initializer='zeros',
                                  name='dense')(x)
model = tf.keras.Model([patterns, features], output)

model.summary()
tf.keras.utils.plot_model(model, "model_shapes.png", show_shapes=True)

model.compile(optimizer='adam', loss='mse')

batch = 1000
epochs = 5

model.fit(
    [train[:, :-1-nr_features], train[:, -1-nr_features:-1]],
    train[:, -1], 
    batch_size=batch,
    epochs=epochs,
    validation_data=([val[:, :-1-nr_features], val[:, -1-nr_features:-1]],
                     val[:, -1])
)

np.set_printoptions(suppress=True)
np.set_printoptions(threshold=np.inf)
f = open(sys.argv[1], 'w') 
f.write(np.array2string(model.get_layer('dense').get_weights()[0]))
f.close()
