import sys
import os
sys.path.append(os.path.abspath("../../.."))
os.environ['KERAS_BACKEND'] = 'tensorflow'
import keras
from keras.metrics import categorical_accuracy
from BiddingKG.dl.table_head.metrics import precision, recall, f1
from keras import optimizers, Model
from BiddingKG.dl.table_head.models.model import get_model
from BiddingKG.dl.table_head.loss import focal_loss
from keras.callbacks import ModelCheckpoint, ReduceLROnPlateau
from BiddingKG.dl.table_head.pre_process import get_data_from_file, get_data_from_sql, my_data_loader
import numpy as np
from keras import backend as K


input_shape = (3, 10, 60)
output_shape = (3,)
batch_size = 1024
epochs = 1000
pretrained_path = "checkpoints/best.hdf5"
checkpoint_path = "checkpoints/"
PRETRAINED = False
CHECKPOINT = False


def train():
    # GPU available
    print("gpus", K.tensorflow_backend._get_available_gpus())

    # Data
    data_x, data_y = get_data_from_file('txt')
    print("finish read data", len(data_x))

    # Split -> Train, Test
    split_size = int(len(data_x)*0.1)
    test_x, test_y = data_x[:split_size], data_y[:split_size]
    train_x, train_y = data_x[split_size:], data_y[split_size:]

    # Data Loader
    train_data_loader = my_data_loader(train_x, train_y, batch_size=batch_size)
    test_data_loader = my_data_loader(test_x, test_y, batch_size=batch_size)

    # Model
    model = get_model(input_shape, output_shape)
    if PRETRAINED:
        model.load_weights(pretrained_path)
        print("read pretrained model", pretrained_path)
    else:
        print("no pretrained")
    if CHECKPOINT:
        model.load_weights(checkpoint_path)
        print("read checkpoint model", checkpoint_path)
    else:
        print("no checkpoint")

    filepath = 'e{epoch:02d}-loss{val_loss:.2f}'
    checkpoint = ModelCheckpoint(checkpoint_path+filepath+".hdf5", monitor='val_f1',
                                 verbose=1, save_best_only=True, mode='max')

    model.compile(optimizer=optimizers.Adam(lr=0.005), loss=focal_loss(),
                  metrics=[categorical_accuracy,
                           precision, recall, f1])

    rlu = ReduceLROnPlateau(monitor='val_f1', factor=0.1, patience=5,
                            verbose=1, mode='max', cooldown=0, min_lr=0)

    model.fit_generator(train_data_loader,
                        steps_per_epoch=max(1, len(train_x) // batch_size),
                        callbacks=[checkpoint, rlu],
                        validation_data=test_data_loader,
                        validation_steps=max(1, len(test_x) // batch_size),
                        epochs=epochs)

    # model.fit(x=[train_x[0], train_x[1], train_x[2]], y=train_y,
    #           validation_data=([test_x[0], test_x[1], test_x[2]], test_y),
    #           epochs=epochs, batch_size=256, shuffle=True,
    #           callbacks=[checkpoint, rlu])

    return model, test_x


def print_layer_output(model, data):
    middle_layer = Model(inputs=model.inputs,
                         outputs=model.get_layer('input_2').output)

    middle_layer_output = middle_layer.predict([data[0], data[1]])
    print(middle_layer_output)
    return


if __name__ == '__main__':
    model, data = train()