import numpy as np
import keras


def createModel(loss="mse", optimizer="adam"):
  return createModel1(loss, optimizer)


def createModel1(loss="mse", optimizer="adam"):
  inputs = keras.Input(shape=(512*512,))
  x = keras.layers.Dense(128, activation="softsign")(inputs)
  x = keras.layers.Dense(32, activation="softsign")(x)
  outputs = keras.layers.Dense(4, activation="relu")(x)
  model = keras.Model(inputs=inputs, outputs=outputs)
  model.compile(optimizer=optimizer, loss=loss)
  model.summary()
  return model

def createModel2(loss="mse", optimizer="adam"):
  inputs = keras.Input(shape=(512*512,))
  x = keras.layers.Dense(1024, activation="softsign")(inputs)
  x = keras.layers.Dense(128, activation="softsign")(x)
  x = keras.layers.Dense(32, activation="softsign")(x)
  outputs = keras.layers.Dense(4, activation="relu")(x)
  model = keras.Model(inputs=inputs, outputs=outputs)
  model.compile(optimizer=optimizer, loss=loss)
  model.summary()
  return model


def save(model, fileName="model.keras"):
  model.save(fileName)


def load(fileName="model.keras"):
  return keras.saving.load_model(fileName)