import numpy as np
import matplotlib.pyplot as plt
import os
from os.path import join
import cv2
from skimage.transform import resize
from tqdm import tqdm
from datetime import datetime
import functools

os.environ['TF_CPP_MIN_VLOG_LEVEL'] = '3'
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3'

from keras import optimizers
from keras.models import Model, load_model
from keras.layers import Flatten, Dense, Conv2D, Dropout, MaxPooling2D, BatchNormalization, Input
from keras.callbacks import ModelCheckpoint, Callback
import json
import tensorflow as tf

tf.reset_default_graph()
from keras import backend as K

K.image_dim_ordering()

###################### Hyperparameters ######################
#                    Mode paramaters
INPUT_SHAPE = (256, 256, 3)
IMAGE_NUMBER = 30000
WEIGHT_FILENAME = os.path.join('models', 'tag_prediction.hdf5')
ROOT = 'dress'
VALIDATION_SPLIT = 0.9
TAG_END = "_dress"
FILE_END = 'S'
MIN_TAG_USE = 500
#                    Training parameters
BATCH_SIZE = 8
EPOCHS = 100
LEARNING_RATE = 1 * 10 ** -3
DROPOUT = 0.5
MOMENTUM = 0.5
WEIGHT_DECAY = 4 * 10 ** -5  # weight decay
ACTIVATION = 'selu'
NEURON_BASIS = 32


class PlotLearning(Callback):
    def __init__(self, examples=False):
        super().__init__()
        self.examples = examples
        self.x = []
        self.losses, self.val_losses = [], []
        self.logs = []

    def on_epoch_end(self, epoch, logs={}):
        self.logs.append(logs)
        self.x.append(epoch)
        self.losses.append(logs.get('loss'))
        self.val_losses.append(logs.get('val_loss'))

        plt.yscale('log')
        plt.plot(self.x, self.losses)
        plt.plot(self.x, self.val_losses)
        plt.xlabel('Epochs')
        plt.ylabel('Crossentropy')
        plt.legend(['Training', 'Validation'])

        plt.tight_layout()
        plt.savefig('plot.png')
        plt.close()

        if self.examples:
            z = self.model.predict(self.model.example[0][:6])
            plot_example(self.model.example[0][:6], self.model.example[1][:6], z, self.model.labels)
            plt.savefig(f"epochs/epoch{self.x[-1]}.png")
            plt.close()


def plot_example(xs, ys, zs, labels):
    n = xs.shape[0]
    plt.figure(figsize=(12, 8))
    plt.tight_layout()
    for i, (x, y, z) in enumerate(zip(xs, ys, zs)):
        if i != 0:
            tick_label = [' ' for label in labels]
        else:
            tick_label = labels
        plt.subplot(3, n, i + 1)
        plt.imshow(x)
        plt.subplot(3, n, i + n + 1)
        plt.barh(labels, y, tick_label=tick_label)
        plt.xlim(0, 1)
        plt.subplot(3, n, i + 2 * n + 1)
        plt.barh(labels, z, tick_label=tick_label)
        plt.xlim(0, 1)


def import_model(tag_number, input_shape=INPUT_SHAPE):
    inputs = Input(input_shape)

    layer = Conv2D(NEURON_BASIS, (3, 3), activation=ACTIVATION, padding='same')(inputs)
    layer = Conv2D(NEURON_BASIS, (3, 3), activation=ACTIVATION, padding='same')(layer)
    layer = Conv2D(NEURON_BASIS, (3, 3), activation=ACTIVATION, padding='same')(layer)
    layer = MaxPooling2D(pool_size=(2, 2))(layer)

    layer = Conv2D(NEURON_BASIS * 2, (3, 3), activation=ACTIVATION, padding='same')(layer)
    layer = Conv2D(NEURON_BASIS * 2, (3, 3), activation=ACTIVATION, padding='same')(layer)
    layer = Conv2D(NEURON_BASIS * 2, (3, 3), activation=ACTIVATION, padding='same')(layer)
    layer = MaxPooling2D(pool_size=(2, 2))(layer)

    layer = Conv2D(NEURON_BASIS * 4, (3, 3), activation=ACTIVATION, padding='same')(layer)
    layer = Conv2D(NEURON_BASIS * 4, (3, 3), activation=ACTIVATION, padding='same')(layer)
    layer = Conv2D(NEURON_BASIS * 4, (3, 3), activation=ACTIVATION, padding='same')(layer)
    layer = MaxPooling2D(pool_size=(2, 2))(layer)

    layer = Conv2D(NEURON_BASIS * 8, (3, 3), activation=ACTIVATION, padding='same')(layer)
    layer = Conv2D(NEURON_BASIS * 8, (3, 3), activation=ACTIVATION, padding='same')(layer)
    layer = Conv2D(NEURON_BASIS * 8, (3, 3), activation=ACTIVATION, padding='same')(layer)
    layer = Conv2D(NEURON_BASIS * 4, (1, 1), activation=ACTIVATION, padding='same')(layer)
    layer = MaxPooling2D(pool_size=(2, 2))(layer)

    layer = Flatten()(layer)
    layer = BatchNormalization()(layer)
    layer = Dense(512, activation=ACTIVATION)(layer)
    layer = Dropout(DROPOUT)(layer)
    layer = Dense(2048, activation=ACTIVATION)(layer)
    layer = Dropout(DROPOUT)(layer)
    layer = Dense(tag_number, activation='sigmoid')(layer)

    model = Model(inputs=[inputs], outputs=[layer])
    sgd = optimizers.SGD(lr=LEARNING_RATE, momentum=MOMENTUM, nesterov=True)
    model.compile(optimizer='adam', loss='binary_crossentropy')
    model.summary()

    return model


def get_tags(root, files, min_tag_use=MIN_TAG_USE, suffix=TAG_END):
    with open(join('..', 'imgs', root + '.json'), 'r') as file:
        tags = json.load(file)
    tag_count = {}
    files = [os.path.split(file)[-1] for file in files]
    for key, value in tags.items():
        if key + f'{FILE_END}.png' not in files:
            continue
        for tag in value.split():
            if tag not in tag_count:
                tag_count[tag] = 1
            else:
                tag_count[tag] += 1
    with open(join('..', 'imgs', root + '_count.json'), 'w') as file:
        json.dump(tag_count, file, sort_keys=True, indent=4)
    print(f'Have {len(list(tag_count.keys()))} tags')
    tags_count = {tag: count for tag, count in tag_count.items() if count > min_tag_use and tag.endswith(suffix)}
    print(f'Keep tags with >{min_tag_use} use: {len(tag_count)} tags')
    for tag, count in tags_count.items():
        print(f'{tag}: {count}')
    input('Continue?')
    return tags, tags_count


def make_output(files, tags, tags_count):
    output = {}
    for file in tqdm(files):
        i = os.path.splitext(os.path.split(file)[-1])[0]
        if FILE_END:
            i = i[:-1]
        truth = tags[i].split()
        output[file] = []
        for tag in tags_count.keys():
            if tag in truth:
                output[file].append(1)
            else:
                output[file].append(0)
    return output


def metrics(model, files, output, tags_count):
    true_positives = np.zeros(len(output))
    positives = np.zeros(len(output))
    truth = np.zeros(len(output))
    for file in tqdm(files):
        img = image_process(file)
        img = np.expand_dims(img, axis=0)
        prediction = model.predict(img)[0]
        for i, coef in enumerate(prediction):
            f = tags_count.values()[i] / len(files)
            if coef > f:
                positives[i] += 1
            if output[file][i] > f:
                truth[i] += 1
            if output[file][i] > f and coef > f:
                true_positives[i] += 1

    print('Tag\tPrecision\tRecall')
    for i, k, l, key in zip(true_positives, positives, truth, tags_count.keys()):
        if k != 0:
            precision = int(i / k * 1000) / 100
        else:
            precision = 0
        if l != 0:
            recall = int(i / l * 1000) / 100
        else:
            recall = 0
        print(f'{key}\t{precision}%\t{recall}%\t')


# @functools.lru_cache(maxsize=IMAGE_NUMBER)
def image_process(file):
    img = cv2.imread(file)
    img = img[:, :, [2, 1, 0]]
    # img = resize(img, INPUT_SHAPE, mode='reflect', preserve_range=True, anti_aliasing=True)
    return img


def generator(files, output, batch_size=BATCH_SIZE):
    while True:
        batch_files = np.random.choice(files, size=batch_size)
        # j += 1
        # print(index, j, [(k + j) % n for k in index], [(k + j) for k in index], index+j)
        batch_output = np.array([output[file] for file in batch_files])

        batch_input = np.zeros([batch_size] + [shape for shape in INPUT_SHAPE])
        for i, file in enumerate(batch_files):
            batch_input[i] = image_process(file)

        yield batch_input / 255, batch_output


def train(model, files, output, tags_count, weight_filename=WEIGHT_FILENAME,
          validation_split=VALIDATION_SPLIT, epochs=EPOCHS, batch_size=BATCH_SIZE):
    class_weights = {i: len(files) / count for i, count in enumerate(tags_count.values())}
    index = int(len(files) * validation_split)
    training_generator = generator(files[:index], output)
    validation_generator = generator(files[index:], output)

    calls = [ModelCheckpoint(weight_filename, save_best_only=True),
             PlotLearning(examples=True)]
    model.example = next(validation_generator)
    model.labels = list(tags_count.keys())

    model.fit_generator(generator=training_generator,
                        validation_data=validation_generator,
                        verbose=1,
                        steps_per_epoch=int(len(files) * validation_split) // batch_size,
                        validation_steps=int(len(files) * (1 - validation_split)) // batch_size,
                        epochs=epochs,
                        callbacks=calls,
                        class_weight=class_weights
                        )


def test(files, output, tags_count, weight_filename=WEIGHT_FILENAME):
    model = load_model(weight_filename)
    metrics(model, files, output, tags_count)
    image_generator = generator(files, output, batch_size=1)
    fs = [count / len(files) for count in tags_count.values()]
    fs = [0.5 for i in fs]
    while True:
        print('---')
        im, truth = next(image_generator)
        truth_string = ' '.join([tags_count.keys()[j] for j, v in enumerate(truth[0]) if v > fs[j]])
        print('TRUTH:', truth_string)
        print(im.shape)

        prediction = model.predict(im)[0]
        prediction_string = ' '.join([tags_count.keys()[j] for j, v in enumerate(prediction) if v > fs[j]])
        print('PREDICTION:', prediction_string)
        plt.imshow(im[0])
        plt.show()


def main():
    root = join('..', 'imgs', ROOT)
    files = [join(root, folder, file) for folder in os.listdir(root) for file in os.listdir(join(root, folder))][
            :IMAGE_NUMBER]
    tags, tags_count = get_tags(ROOT, files)
    output = make_output(files, tags, tags_count)
    # test(files, output, tags_count)
    model = import_model(len(tags_count))
    train(model, files, output, tags_count)
    print('DONE')


if __name__ == '__main__':
    main()
