import matplotlib.pyplot as plt
import tensorflow as tf


def main() -> None:
    print("===| START PROGRAM |===")
    (x_train, y_train), (x_test, y_test) = tf.keras.datasets.mnist.load_data()

    x_train = x_train / 255.0
    y_train = y_train / 255.0
    print("===| CREATE MODEL |===")
    # создаем модель
    model = tf.keras.Sequential(
        [
            tf.keras.layers.Flatten(input_shape=(28, 28)),
            tf.keras.layers.Dense(128, activation="relu"),
            tf.keras.layers.Dense(10, activation="softmax"),
        ]
    )

    print("===| COMPILE MODEL |===")
    # компилируем модель
    model.compile(
        optimizer="adam", loss="sparse_categorical_crossentropy", metrics=["accuracy"]
    )

    history = model.fit(x_train, y_train, epochs=5, validation_split=0.2)
    loss, accuracy = model.evaluate(x_test, y_test)
    print("\nТочность модели: {}".format(accuracy))
    print("===| SAVE MODEL |===")
    model.save("mnist_model.keras")
    prediction = model.predict(x_test[:1])
    print("Предсказанная цифра: ", prediction.argmax())
    print("Правильный ответ: ", y_test[1])

    print("===| Показываем изображение |===")
    plt.imshow(x_test[1], cmap="gray")
    plt.title(f"Предсказание: {prediction.argmax()}")
    plt.axis("off")
    plt.show()


if __name__ == "__main__":
    main()