import tensorflow as tf
import random
from tensorflow import keras
from keras import layers, losses, optimizers

seed = random.randint(0, 100000000)

batch_size = 100

dataset = keras.utils.image_dataset_from_directory(
  "{{ .DataDir }}",
  color_mode="rgb",
  validation_split=0.2,
  label_mode='categorical',
  seed=seed,
  shuffle=True,
  subset="training",
  image_size=({{ .Size }}),
  batch_size=batch_size)

dataset_validation = keras.utils.image_dataset_from_directory(
  "{{ .DataDir }}",
  color_mode="rgb",
  validation_split=0.2,
  label_mode='categorical',
  seed=seed,
  shuffle=True,
  subset="validation",
  image_size=({{ .Size }}),
  batch_size=batch_size)

model = keras.Sequential([
    {{- range .Layers }}
    {{- if eq .LayerType 1}}
    layers.Rescaling(1./255),
    {{- else if eq .LayerType 2 }}
    layers.Dense({{ .Shape }}, activation="sigmoid"),
    {{- else if eq .LayerType 3}}
    layers.Flatten(),
    {{- else }}
    ERROR
    {{- end }}
    {{- end }}
])

model.compile(
    loss=losses.CategoricalCrossentropy(),
    optimizer=tf.keras.optimizers.Adam(),
    metrics=['accuracy'])

his = model.fit(dataset, validation_data= dataset_validation, epochs=50)

acc = his.history["accuracy"]

f = open("accuracy.val", "w")
f.write(str(acc[-1]))
f.close()

tf.saved_model.save(model, "model")

# model.save("model.keras", save_format="tf")