import tensorflow as tf

print("# GPUs Available: ", len(tf.config.experimental.list_physical_devices('GPU')))


def create_modulation_model(nary):
    model = tf.keras.models.Sequential()
    model.add(tf.keras.layers.Dense(units=32, activation='relu', input_shape=(nary,), dtype='bool'))
    model.add(tf.keras.layers.Dropout(0.2))
    model.add(tf.keras.layers.Dense(units=2, activation='softmax'))
    model.compile(optimizer='adam', loss='sparse_categorical_crossentropy', metrics=['sparse_categorical_accuracy'])
    return model


def create_demodulation_model(nary):
    model = tf.keras.models.Sequential()
    model.add(tf.keras.layers.Dense(units=32, activation='relu', input_shape=(2,)))
    model.add(tf.keras.layers.Dropout(0.2))
    model.add(tf.keras.layers.Dense(units=nary, activation='softmax'))
    model.compile(optimizer='adam', loss='sparse_categorical_crossentropy', metrics=['sparse_categorical_accuracy'])
    return model