import os
import random
import sys
from glob import glob
os.environ["CUDA_VISIBLE_DEVICES"] = "0"
import tensorflow as tf
import keras.backend as K
# tf.compat.v1.disable_eager_execution()
sys.path.append(os.path.dirname(os.path.abspath(__file__)))
sys.path.append(os.path.dirname(os.path.abspath(__file__)) + "/../")
from keras.callbacks import ReduceLROnPlateau, ModelCheckpoint
from keras.losses import BinaryCrossentropy, mse, CategoricalCrossentropy, MSE
from keras.optimizer_v2.adam import Adam
from click_captcha.metrics import precision, recall, f1
from click_captcha.loss import focal_loss, contrastive_loss, l2_focal_loss, l1_focal_loss, ctc_accuracy
from click_captcha.model import crnn_ctc_equation, crnn_ctc_equation_large, crnn_ctc_equation_less, \
    crnn_ctc_equation_loss
from click_captcha.pre_process import gen_equation, gen_equation2

PRETRAINED = True
random.seed(42)
image_shape = (32, 192, 1)
project_root = os.path.dirname(os.path.abspath(__file__)) + "/../"
class_num = 35 + 2
data_path = 'equation2'


if __name__ == "__main__":
    model = crnn_ctc_equation_loss(input_shape=image_shape, class_num=class_num)

    if PRETRAINED:
        _path = "./models/e83-loss0.06-equation.h5"
        model.load_weights(_path, skip_mismatch=True, by_name=True)
        print("read pretrained model", _path)
    else:
        print("no pretrained")

    # with open(project_root + "data/click/map.txt", "r") as f:
    #     paths = f.readlines()
    # print("len(paths)", len(paths))
    paths = glob("../data/" + data_path + "/*.jpg")

    # data path split into train,test
    random.shuffle(paths)
    # paths = paths[:100000]
    trainP = paths[:int(len(paths)*0.9)]
    testP = paths[int(len(paths)*0.9):]
    print('total:', len(paths), 'train:', len(trainP), 'test:', len(testP))

    # batch num
    batch_size = 32
    steps_per_epoch = max(1, len(trainP) // batch_size)
    validation_steps = max(1, len(testP) // batch_size)

    # 模型权重存放位置
    filepath = 'models/e{epoch:02d}-loss{val_loss:.2f}-equation.h5'
    check_pointer = ModelCheckpoint(filepath=filepath, monitor='val_loss', verbose=0,
                                    save_weights_only=True, save_best_only=True,
                                    mode="min", save_freq='epoch')
    rlu = ReduceLROnPlateau(monitor='val_loss', factor=0.5, patience=10,
                            verbose=1, mode='min', cooldown=0, min_lr=0)
    model.compile(optimizer=Adam(lr=0.0003), loss={'ctc': lambda y_true, y_pred: y_pred})

    # 使用ctc center loss 所需
    # sess = K.get_session()
    # sess.run(tf.compat.v1.global_variables_initializer())

    # data loader
    train_loader = gen_equation2(trainP, batch_size=batch_size, shape=image_shape, cls_num=class_num, data_path=data_path)
    test_loader = gen_equation2(testP, batch_size=batch_size, shape=image_shape, cls_num=class_num, data_path=data_path)
    # train
    steps_per_epoch = 500
    validation_steps = int(steps_per_epoch * 0.1)
    model.fit_generator(train_loader,
                        steps_per_epoch=steps_per_epoch,
                        callbacks=[check_pointer, rlu],
                        validation_data=test_loader,
                        validation_steps=validation_steps,
                        epochs=1000,
                        max_queue_size=100,
                        use_multiprocessing=False)