import os
import random
import sys
from glob import glob
os.environ["CUDA_VISIBLE_DEVICES"] = "1"
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
from keras.metrics import Precision, Recall
from keras.optimizer_v2.adam import Adam
import keras.backend as K
from click_captcha.metrics import precision, recall, f1
from click_captcha.loss import focal_loss, contrastive_loss, l2_focal_loss, l1_focal_loss
from click_captcha.model import siamese_net, mobile_net, cnn_net, cnn_net_drag, u_net_drag
from click_captcha.pre_process import gen_mobile, gen_drag

PRETRAINED = False
random.seed(42)
image_shape = (128, 256, 3)
project_root = os.path.dirname(os.path.abspath(__file__)) + "/../"


if __name__ == "__main__":
    model = u_net_drag(input_shape=image_shape)

    if PRETRAINED:
        _path = "./models/e14-f10.71.h5"
        model.load_weights(_path, skip_mismatch=True, by_name=True)
        print("read pretrained model", _path)
    else:
        print("no pretrained")

    paths = glob("../data/drag/*.jpg")

    # data path split into train,test
    testP = random.sample(paths, int(len(paths)*0.1))
    trainP = []
    for p in paths:
        if p not in testP:
            trainP.append(p)
    random.shuffle(trainP)
    print('total:', len(paths), 'train:', len(trainP), 'test:', len(testP))

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

    # 模型权重存放位置
    filepath = 'models/e{epoch:02d}-f1{val_f1:.2f}-drag.h5'
    check_pointer = ModelCheckpoint(filepath=filepath, monitor='val_acc', verbose=0,
                                    save_weights_only=True, save_best_only=True,
                                    mode="max", save_freq='epoch')
    rlu = ReduceLROnPlateau(monitor='val_f1', factor=0.5, patience=10,
                            verbose=1, mode='max', cooldown=0, min_lr=0)
    model.compile(optimizer=Adam(lr=0.003), loss=focal_loss(2, 0.99),
                  metrics=['acc', f1])

    # data loader
    train_loader = gen_drag(trainP, batch_size=batch_size, shape=image_shape)
    test_loader = gen_drag(testP, batch_size=batch_size, shape=image_shape)
    # train
    model.fit_generator(train_loader,
                        steps_per_epoch=steps_per_epoch,
                        callbacks=[check_pointer],
                        validation_data=test_loader,
                        validation_steps=validation_steps,
                        epochs=1000)