import tensorflow as tf
import numpy as np
from matplotlib import pyplot as plt
from PIL import Image, ImageTk


def __prep_data():
    (x_train, y_train), (x_test, y_test) = tf.keras.datasets.mnist.load_data()

    x_train = x_train.astype("float32") / 255.
    x_test = x_test.astype("float32") / 255.

    x_train = x_train.reshape(-1, 784)
    x_test = x_test.reshape(-1, 784)

    y_train = tf.keras.utils.to_categorical(y_train, 10)
    y_test = tf.keras.utils.to_categorical(y_test, 10)

    return (x_train, y_train), (x_test, y_test)


def __prep_conf_matr(m):
    output_matrix = np.zeros([10, 10])

    (_, _), (x_test, y_test) = __prep_data()
    pred = m.predict(x_test)

    for i, v in enumerate(pred):
        output_matrix[np.argmax(v)][np.argmax(y_test[i])] += 1

    return output_matrix


def __plot_conf_matr(m):
    matr = __prep_conf_matr(m)

    _, ax = plt.subplots()

    ax.matshow(matr, cmap = plt.cm.Blues)

    for i, x in enumerate(matr):
        for j, y in enumerate(x):
            ax.text(i,
                    j,
                    str(round(y)),
                    va = "center",
                    ha = "center")

    plt.show()


def __plot_acc_rate(h):
    plt.plot(h.history['accuracy'], label = 'train_acc')
    plt.plot(h.history['val_accuracy'], label = 'valid_acc')

    plt.legend()

    plt.show()


def train(m, label):
    (x_train, y_train), (x_test, y_test) = __prep_data()

    m.compile(optimizer = "adam",
              loss = "categorical_crossentropy",
              metrics = ["accuracy"])

    h = m.fit(x_train,
              y_train,
              epochs = 30,
              batch_size = 512,
              validation_data = (x_test, y_test))

    m.save_weights(f"save-{label}.weights.h5")

    __plot_acc_rate(h)
    __plot_conf_matr(m)


def classify(m, label, imgfn):
    m.compile(optimizer = "adam",
              loss = "categorical_crossentropy",
              metrics = ["accuracy"])

    m.load_weights(f"save-{label}.weights.h5")

    img = Image.open(imgfn).convert("L")
    flat_img = np.array(img).reshape(-1, 784)

    res = m.predict(flat_img)

    plt.imshow(flat_img.reshape(28, 28),
               cmap = "gray")

    plt.title(np.argmax(res))
    plt.show()

put_active = 0
take_active = 0

def classify_live(m, label):
    import lv
    lv.classify_live(m, label)