"""
Train object detectors
"""
import tensorflow as tf

from modules.config import CFG
from modules.dataset_utils import init_dataset
from modules.yolov5 import YOLOv5
from modules.efficientdet import EfficientDet
from modules.detector_utils import DetectionLoss


def train_hands(model_cfg):
    """Train hand detection model

    Args:
        model_cfg (str): YOLOv5 model configuration, see model implementation for implemented configurations
    """
    # init datasets
    train_data = init_dataset(
        CFG.HANDS_TFRECORDS_PATH['train'],
        CFG.HANDS_PARAMS.BATCH_SIZE,
        CFG.HANDS_PARAMS.ANCHORS,
        CFG.HANDS_PARAMS.IMG_SHAPE,
        CFG.HANDS_PARAMS.MAX_DETECTIONS,
        model_cfg,
        augment=True,
        shuffle=True)
    val_data = init_dataset(
        CFG.HANDS_TFRECORDS_PATH['val'],
        CFG.HANDS_PARAMS.BATCH_SIZE,
        CFG.HANDS_PARAMS.ANCHORS,
        CFG.HANDS_PARAMS.IMG_SHAPE,
        CFG.HANDS_PARAMS.MAX_DETECTIONS,
        model_cfg)
    # load pretrained model
    model_pretrained = YOLOv5(
        CFG.YOLO_PRETRAINED_IMG_SHAPE,
        CFG.YOLO_PRETRAINED_CLASSES,
        CFG.HANDS_PARAMS.ANCHORS,
        model=model_cfg,
        training=True
    )
    model_pretrained.load_weights(CFG.PRETRAINED_TF_WEIGHTS_PATH[model_cfg]).expect_partial()
    # load new model
    model = YOLOv5(
        (CFG.HANDS_PARAMS.IMG_SHAPE[0], CFG.HANDS_PARAMS.IMG_SHAPE[1], 3),
        1,
        CFG.HANDS_PARAMS.ANCHORS,
        model=model_cfg,
        training=True
    )
    # transfer weights from old model to the new model and freeze backbone
    for name in ['yolo_backbone', 'yolo_conv_0', 'yolo_conv_1', 'yolo_conv_2', 'yolo_conv_3']:
        model.get_layer(name).set_weights(model_pretrained.get_layer(name).get_weights())
    model.get_layer('yolo_backbone').trainable = False
    # optimizer and loss function
    optimizer = tf.keras.optimizers.Adam(learning_rate=CFG.HANDS_PARAMS.LEARNING_RATE)
    loss = [DetectionLoss(anchor, classes=1) for anchor in CFG.HANDS_PARAMS.ANCHORS]
    # early stopping and exponential learning rate decay
    callbacks = [
            tf.keras.callbacks.EarlyStopping(monitor='val_loss', patience=5, verbose=True, restore_best_weights=True, min_delta=0.05),
            tf.keras.callbacks.TensorBoard(log_dir='logs/hands'),
            tf.keras.callbacks.LearningRateScheduler(lambda e, lr: lr if e < 5 else lr * tf.math.exp(-0.1))
        ]
    # train model and save weights
    model.compile(optimizer=optimizer, loss=loss, run_eagerly=False)
    model.fit(train_data, epochs=CFG.HANDS_PARAMS.EPOCHS, callbacks=callbacks, validation_data=val_data)
    model.save_weights(CFG.HANDS_WEIGHTS_PATH[model_cfg])


def train_face(model_cfg):
    """Train face detection model

    Args:
        model_cfg (str): YOLOv5 model configuration, see model implementation for implemented configurations
    """
    # init datasets
    train_data = init_dataset(
        CFG.FACES_TFRECORDS_PATH['train'],
        CFG.FACES_PARAMS.BATCH_SIZE,
        CFG.FACES_PARAMS.ANCHORS,
        CFG.FACES_PARAMS.IMG_SHAPE,
        CFG.FACES_PARAMS.MAX_DETECTIONS,
        model_cfg,
        shuffle=True)
    val_data = init_dataset(
        CFG.FACES_TFRECORDS_PATH['val'],
        CFG.FACES_PARAMS.BATCH_SIZE,
        CFG.FACES_PARAMS.ANCHORS,
        CFG.FACES_PARAMS.IMG_SHAPE,
        CFG.FACES_PARAMS.MAX_DETECTIONS,
        model_cfg)
    # load pretrained model
    model_pretrained = YOLOv5(
        CFG.YOLO_PRETRAINED_IMG_SHAPE,
        CFG.YOLO_PRETRAINED_CLASSES,
        CFG.FACES_PARAMS.ANCHORS,
        model=model_cfg,
        training=True
    )
    model_pretrained.load_weights(CFG.PRETRAINED_TF_WEIGHTS_PATH[model_cfg]).expect_partial()
    # load new model
    model = YOLOv5(
        (CFG.FACES_PARAMS.IMG_SHAPE[0], CFG.FACES_PARAMS.IMG_SHAPE[1], 3),
        1,
        CFG.FACES_PARAMS.ANCHORS,
        model=model_cfg,
        training=True
    )
    # transfer weights from old model to the new model and freeze backbone
    for name in ['yolo_backbone', 'yolo_conv_0', 'yolo_conv_1', 'yolo_conv_2', 'yolo_conv_3']:
        model.get_layer(name).set_weights(model_pretrained.get_layer(name).get_weights())
    model.get_layer('yolo_backbone').trainable = False
    # optimizer and loss function
    optimizer = tf.keras.optimizers.Adam(learning_rate=CFG.FACES_PARAMS.LEARNING_RATE)
    loss = [DetectionLoss(anchor, classes=1) for anchor in CFG.FACES_PARAMS.ANCHORS]
    # early stopping and exponential learning rate decay
    callbacks = [
            tf.keras.callbacks.EarlyStopping(monitor='val_loss', patience=5, verbose=True, restore_best_weights=True, min_delta=0.05),
            tf.keras.callbacks.TensorBoard(log_dir='logs/faces'),
            tf.keras.callbacks.LearningRateScheduler(lambda e, lr: lr if e < 5 else lr * tf.math.exp(-0.1))
        ]
    # train model and save weights
    model.compile(optimizer=optimizer, loss=loss, run_eagerly=False)
    model.fit(train_data, epochs=CFG.FACES_PARAMS.EPOCHS, callbacks=callbacks, validation_data=val_data)
    model.save_weights(CFG.FACES_WEIGHTS_PATH[model_cfg])


def train_combined(model_cfg):
    """Train combined hand and face detection model

    Args:
        model_cfg (str): model configuration, see model implementation for implemented configurations

    Raises:
        ValueError: when model_cfg argument is not recognized
    """
    if 'yolov5' in model_cfg:
        # load pretrained model
        anchors = CFG.COMBINED_PARAMS.YOLO_ANCHORS
        model_pretrained = YOLOv5(
            CFG.YOLO_PRETRAINED_IMG_SHAPE,
            CFG.YOLO_PRETRAINED_CLASSES,
            anchors,
            model=model_cfg,
            training=True
        )
        model_pretrained.load_weights(CFG.PRETRAINED_TF_WEIGHTS_PATH[model_cfg]).expect_partial()
        # load new model
        model = YOLOv5(
            (*CFG.COMBINED_PARAMS.IMG_SHAPE, 3),
            2,
            anchors,
            model=model_cfg,
            training=True
        )
        # transfer weights from old model to the new model and freeze backbone
        for name in ['yolo_backbone', 'yolo_conv_0', 'yolo_conv_1', 'yolo_conv_2', 'yolo_conv_3']:
            model.get_layer(name).set_weights(model_pretrained.get_layer(name).get_weights())
        model.get_layer('yolo_backbone').trainable = False
    elif 'efficientdet' in model_cfg:
        # load model and its weights and freeze backbone
        anchors = CFG.COMBINED_PARAMS.EFFDET_ANCHORS
        model = EfficientDet(
            (*CFG.COMBINED_PARAMS.IMG_SHAPE, 3),
            2,
            anchors,
            model=model_cfg,
            training=True
        )
        model.load_weights(CFG.PRETRAINED_TF_WEIGHTS_PATH[model_cfg]).expect_partial()
        model.get_layer(model_cfg.replace('det-d', 'net-b')).trainable = False
    else:
        raise ValueError(f'Unrecognized model configuration "{model_cfg}"')
    # init datasets
    train_data = init_dataset(
        CFG.COMBINED_TFRECORDS_PATH['train'],
        CFG.COMBINED_PARAMS.BATCH_SIZE,
        anchors,
        CFG.COMBINED_PARAMS.IMG_SHAPE,
        100,
        model_cfg,
        shuffle=True)
    val_data = init_dataset(
        CFG.COMBINED_TFRECORDS_PATH['val'],
        CFG.COMBINED_PARAMS.BATCH_SIZE,
        anchors,
        CFG.COMBINED_PARAMS.IMG_SHAPE,
        100,
        model_cfg)
    # optimizer and loss function
    loss = [DetectionLoss(a, classes=2) for a in anchors]
    optimizer = tf.keras.optimizers.Adam(learning_rate=CFG.COMBINED_PARAMS.LEARNING_RATE)
    # early stopping and exponential learning rate decay
    callbacks = [
            tf.keras.callbacks.EarlyStopping(monitor='val_loss', patience=5, verbose=True, restore_best_weights=True, min_delta=0.05),
            tf.keras.callbacks.TensorBoard(log_dir=f'logs/combined/{model_cfg}'),
            tf.keras.callbacks.LearningRateScheduler(lambda e, lr: lr if e < 5 else lr * tf.math.exp(-0.1))
        ]
    # train model and save weights
    model.compile(optimizer=optimizer, loss=loss, run_eagerly=False)
    model.fit(train_data, epochs=CFG.COMBINED_PARAMS.EPOCHS, callbacks=callbacks, validation_data=val_data)
    model.save_weights(CFG.COMBINED_WEIGHTS_PATH[model_cfg])


if __name__ == '__main__':
    # train_face('yolov5s')
    # train_hands('yolov5s')
    train_combined('yolov5s')
