# D:/workplace/python
# -*- coding: utf-8 -*-
# @File  :train.py
# @Author:Guido    LuXiaohao
# @Date  :2020/4/8
# @Software:PyCharm

import numpy as np
import tensorflow as tf

from keras.callbacks import ReduceLROnPlateau
from keras.optimizers import Adam, SGD, RMSprop
from keras.callbacks import ModelCheckpoint
from keras.models import Model
from keras.callbacks import LearningRateScheduler
from keras.layers import Dense, GlobalAveragePooling2D, Dropout
from dataset.class_dataloader import load_sample, DataGenerator, myDataloderV2, MydatasetV2
from dataset.augmention import get_training_augmentation, get_valing_augmentation
from nets.temp_mobilenetv3 import MobileNetV3_Small
from nets.temp_ghost import GhostModel
from nets.temp_multinetV1 import mutinet
from nets.densenet.densenet import DenseNet, DenseNetImageNet121, DenseNetImageNet161, DenseNetImageNet169, \
    DenseNetImageNet201, DenseNetImageNet264
from nets.efficientnets.efficientnet import EfficientNetB0, EfficientNetB1, EfficientNetB2, EfficientNetB3, \
    EfficientNetB4, EfficientNetB5, EfficientNetB6, EfficientNetB7

config = tf.ConfigProto()
config.gpu_options.allow_growth = True
sess = tf.Session(config=config)

def lr_schedule(epoch):#定义一个函数，接受轮数为参数, 返回学习率
    """Learning Rate Schedule

    Learning rate is scheduled to be reduced after 80, 120, 160, 180 epochs.
    Called automatically every epoch as part of callbacks during training.

    # Arguments
        epoch (int): The number of epochs

    # Returns
        lr (float32): learning rate
    """
    lr = 1e-3
    if epoch > 1800:  # 0.5e-6
        lr *= 0.5e-3
    elif epoch > 100:  # 1e-6
        lr *= 1e-3
    elif epoch > 50:  # 1e-5
        lr *= 1e-1
    elif epoch > 15:  # 1e-4
        lr *= 5e-1
    print('Learning rate: ', lr)
    return lr

def add_new_last_layer(base_model, nb_classes):
    x = base_model.output
    x = GlobalAveragePooling2D()(x)
    x = Dense(1024, activation='relu')(x)
    x = Dense(512, activation='relu')(x)
    x = Dense(128, activation='relu')(x)
    x = Dropout(0.5)(x)
    predictions = Dense(nb_classes, activation='softmax')(x)
    model = Model(input=base_model.input, output=predictions)
    return model

if __name__ == '__main__':

    # Training parameters超参数
    batch_size = 2
    epochs = 3000
    num_classes = 4
    model_save_path = r""
    model_name = 'model.{epoch:04d}-{val_acc:.5f}.h5'
    train_file_path = r'C:\Users\VINNO\Desktop\1111111\train'
    test_file_path = r'C:\Users\VINNO\Desktop\1111111\val'

    # 按文件夹生成标签
    trainX, trainY = load_sample(train_file_path)
    testX, testY = load_sample(test_file_path)

    traindata = MydatasetV2(trainX, trainY, augmentation=get_training_augmentation(), size=(160, 160))
    valdata = MydatasetV2(testX, testY, augmentation=get_valing_augmentation(), size=(160, 160))

    # training_generator = DataGenerator(trainX, trainY, Is_train=True, shuffle=True, **params)
    # testing_generator = DataGenerator(testX, testY, Is_train=False, shuffle=False, **params)
    training_generator = myDataloderV2(traindata, n_classes=num_classes, batch_size=batch_size, shuffle=True, train=False)
    testing_generator = myDataloderV2(valdata, n_classes=num_classes, batch_size=batch_size, shuffle=False, train=False)

    A = MobileNetV3_Small((160, 160, 3), 4, alpha=1.0, include_top=True)
    model = A.build()
    # 备选模型
    # model = DenseNet(input_shape=(256, 256, 3), depth=40, classes=10, activation='softmax', attention_module='se_block')
    # model = DenseNetImageNet121(input_shape=(256, 256, 3), classes=10, activation='softmax', attention_module='cbam_block')
    # model = DenseNetImageNet161(input_shape=(256, 256, 3), classes=10, activation='softmax', attention_module=None)
    # model = DenseNetImageNet169(input_shape=(256, 256, 3), classes=10, activation='softmax', attention_module=None)
    # model = DenseNetImageNet201(input_shape=(256, 256, 3), classes=10, activation='softmax', attention_module=None)
    # model = DenseNetImageNet264(input_shape=(256, 256, 3), classes=10, activation='softmax', attention_module=None)
    # model = EfficientNetB0(input_shape=(256, 256, 3),  weights=None, classes=10)
    # model = EfficientNetB1(input_shape=(256, 256, 3), weights=None, classes=10)
    # model = EfficientNetB2(input_shape=(256, 256, 3), weights=None, classes=10)
    # model = EfficientNetB3(input_shape=(256, 256, 3), weights=None, classes=10)
    # model = EfficientNetB4(input_shape=(256, 256, 3), weights=None, classes=10)
    # model = EfficientNetB5(input_shape=(256, 256, 3), weights=None, classes=10)
    # model = EfficientNetB6(input_shape=(256, 256, 3), weights=None, classes=10)
    # model = EfficientNetB7(input_shape=(256, 256, 3), weights=None, classes=10)

    lr_scheduler = LearningRateScheduler(lr_schedule)
    lr_reducer = ReduceLROnPlateau(factor=np.sqrt(0.1),  # 学习率减少器， 学不动的时候降低学习率
                                   cooldown=0,
                                   patience=5,
                                   min_lr=0.5e-6)
    opt = Adam(lr=lr_schedule(0), decay=lr_schedule(0) / epochs)
    model.compile(loss='categorical_crossentropy', optimizer=Adam(lr=lr_schedule(0)),
                  metrics=["accuracy"])
    filepath = model_save_path + model_name
    checkpoint = ModelCheckpoint(filepath, monitor='val_acc', verbose=1, save_best_only=True,
                                 mode='max')
    # 回调学习率
    callbacks_list = [checkpoint, lr_reducer, lr_scheduler]
    # train the network
    print("[INFO] training network...")
    H = model.fit_generator(generator=training_generator,
                            validation_data=testing_generator, steps_per_epoch=len(trainX)//batch_size,
                            epochs=epochs, callbacks=callbacks_list, verbose=1)