import matplotlib.pyplot as plt
import tqdm

from utils import *

num_epoch = 5
step = 1000
learning_rate = 1e-5

eps = 1e-5

all_loss = []

model = tf.keras.applications.VGG16(input_shape=(224, 224, 3))

new_model = tf.keras.Sequential(model)
new_model.add(tf.keras.layers.Dropout(0.5, input_shape=(None, 1024), name="dropout"))
new_model.add(tf.keras.layers.Dense(128, activation=tf.nn.elu, name="output"))

loader = FaceLoader("face\\")

new_model.build(input_shape=(None, 224, 224, 3))

new_model.summary()

tf.keras.utils.plot_model(new_model, "model.png", show_shapes=True)

optimizer = tf.keras.optimizers.Adam(learning_rate=learning_rate)
for e in range(num_epoch):
    with tqdm.tqdm(total=step) as pbar:
        for i in range(step):
            with tf.GradientTape() as tape:
                src, same, diff = loader.getTrainData()
                src_pred = new_model(src)
                same_pred = new_model(same)
                diff_pred = new_model(diff)

                src_pred = 100 * tf.reshape(src_pred, -1)
                same_pred = 100 * tf.reshape(same_pred, -1)
                diff_pred = 100 * tf.reshape(diff_pred, -1)

                same_distance = get_cos_distance(src_pred, same_pred)
                diff_distance = get_cos_distance(src_pred, diff_pred)

                P_s = same_distance / 2 + 0.5
                P_d = diff_distance / 2 + 0.5

                d_loss_real = -tf.math.log(P_s + eps)
                d_loss_fake = -tf.math.log(1 - P_d + eps)

                loss = d_loss_fake + d_loss_real
                all_loss.append(loss)

                pbar.set_description("total loss = %s,"
                                     "d_loss_real = %s,"
                                     "d_loss_fake = %s" % (loss.numpy(),
                                                           d_loss_real.numpy(),
                                                           d_loss_fake.numpy()))
            grads = tape.gradient(loss, new_model.trainable_variables)
            optimizer.apply_gradients(grads_and_vars=zip(grads, new_model.trainable_variables))
            pbar.update(1)
# tf.saved_model.save(new_model, "saved/1")
xaxis = [x for x in range(num_epoch * step)]
plt.plot(xaxis, all_loss)
plt.show()

# test
test_loss = []
with tqdm.tqdm(total=step) as pbar:
    for i in range(step):
        src, same, diff = loader.getTrainData()
        src_pred = new_model(src)
        same_pred = new_model(same)
        diff_pred = new_model(diff)

        src_pred = 100 * tf.reshape(src_pred, -1)
        same_pred = 100 * tf.reshape(same_pred, -1)
        diff_pred = 100 * tf.reshape(diff_pred, -1)

        same_distance = get_cos_distance(src_pred, same_pred)
        diff_distance = get_cos_distance(src_pred, diff_pred)

        P_s = same_distance / 2 + 0.5
        P_d = diff_distance / 2 + 0.5

        d_loss_real = -tf.math.log(P_s + eps)
        d_loss_fake = -tf.math.log(1 - P_d + eps)

        loss = d_loss_fake + d_loss_real
        test_loss.append(loss)

        pbar.set_description("total loss = %s,"
                             "d_loss_real = %s,"
                             "d_loss_fake = %s" % (loss.numpy(),
                                                   d_loss_real.numpy(),
                                                   d_loss_fake.numpy()))
        pbar.update(1)
xaxis = [x for x in range(step)]
plt.plot(xaxis, test_loss)
plt.show()
