"""
Train face and hand pose estimation, search for optimal compensation delta for heatmap decoding
"""
import os
import tensorflow as tf

from modules.config import CFG
from modules.wflw_utils import init_dataset as wflw_init_dataset
from modules.hands_utils import init_dataset as hands_init_dataset
from modules.face_model import FaceKeypointEstimator, load_model
from modules.hands_model import HandKeypointEstimator, HandPoseLoss
from modules.metrics import FaceKeypointMetric
from modules.estimation_utils import decode_keypoints_compensated_2d, decode_keypoints_argmax_2d


def train_hand_estimator(backbone):
    """Train hand pose estimator

    Args:
        backbone (str): model backbone configuration, see model implementation for possible inputs
    """
    # load train and val subsets
    train_data = hands_init_dataset(
        CFG.HANDS_TFRECORD_PATHS['train'],
        CFG.HANDS_PARAMS.BATCH_SIZE,
        CFG.HANDS_PARAMS.IMG_SHAPE,
        CFG.HANDS_PARAMS.HM_DIMS,
        CFG.HANDS_PARAMS.SIGMA
    )
    val_data = hands_init_dataset(
        CFG.HANDS_TFRECORD_PATHS['val'],
        CFG.HANDS_PARAMS.BATCH_SIZE,
        CFG.HANDS_PARAMS.IMG_SHAPE,
        CFG.HANDS_PARAMS.HM_DIMS,
        CFG.HANDS_PARAMS.SIGMA
    )
    # load model
    model = HandKeypointEstimator(backbone, CFG.HANDS_PARAMS.IMG_SHAPE, CFG.HANDS_PARAMS.HM_DIMS, imagenet=True, training=True, loss_weights=CFG.HANDS_PARAMS.LOSS_WEIGHTS)
    model.summary()
    # loss function
    loss = HandPoseLoss(*CFG.HANDS_PARAMS.LOSS_WEIGHTS)
    # train model with frozen backbone
    # callbacks = [
    #         tf.keras.callbacks.TensorBoard(log_dir=f'logs/hand_keypoints/init')
    #     ]
    # optimizer = tf.keras.optimizers.SGD(learning_rate=CFG.HANDS_PARAMS.INIT_LEARNING_RATE)
    # model.layers[1].trainable = False
    # model.compile(loss=loss, optimizer=optimizer)
    # model.fit(train_data, epochs=CFG.HANDS_PARAMS.INIT_EPOCHS, callbacks=callbacks, validation_data=val_data)
    
    # finetune model with trainable backbone
    callbacks = [
            tf.keras.callbacks.TensorBoard(log_dir=f'logs/hand_keypoints'),
            tf.keras.callbacks.EarlyStopping(monitor='val_loss', patience=5, verbose=True, restore_best_weights=False, min_delta=0.1),
            tf.keras.callbacks.LearningRateScheduler(lambda e, lr: lr if e < 5 else lr * tf.math.exp(-0.1)),
            tf.keras.callbacks.ModelCheckpoint(os.path.join(CFG.CHECKPOINTS_PATH, model.layers[1].name), save_weights_only=True),
        ]
    optimizer = tf.keras.optimizers.SGD(learning_rate=CFG.HANDS_PARAMS.FINETUNE_LEARNING_RATE, momentum=0.9, nesterov=True)
    model.layers[1].trainable = True
    model.compile(loss=loss, optimizer=optimizer)
    model.fit(train_data, epochs=CFG.HANDS_PARAMS.FINETUNE_EPOCHS, callbacks=callbacks, validation_data=val_data)
    # save weights
    model.save_weights(CFG.HANDS_WEIGHTS[backbone])


def train_face_estimator(backbone):
    """Train face pose estimator

    Args:
        backbone (str): model backbone configuration, see model implementation for possible inputs
    """
    # load train and val subsets
    train_data = wflw_init_dataset(
        CFG.WFLW_TFRECORD_PATHS['train'],
        CFG.WFLW_PARAMS.BATCH_SIZE,
        CFG.WFLW_PARAMS.IMG_SHAPE,
        CFG.WFLW_PARAMS.HM_DIMS,
        CFG.WFLW_PARAMS.SIGMA
    )
    val_data = wflw_init_dataset(
        CFG.WFLW_TFRECORD_PATHS['val'],
        CFG.WFLW_PARAMS.BATCH_SIZE,
        CFG.WFLW_PARAMS.IMG_SHAPE,
        CFG.WFLW_PARAMS.HM_DIMS,
        CFG.WFLW_PARAMS.SIGMA
    )
    # load model
    model = FaceKeypointEstimator(backbone, CFG.WFLW_PARAMS.IMG_SHAPE, CFG.WFLW_PARAMS.HM_DIMS, imagenet=True, training=True)
    model.summary()
    # loss function
    loss = tf.keras.losses.MeanSquaredError()
    # train model with frozen backbone
    # callbacks = [
    #         tf.keras.callbacks.TensorBoard(log_dir=f'logs/face_keypoints/init')
    #     ]
    # optimizer = tf.keras.optimizers.SGD(learning_rate=CFG.WFLW_PARAMS.INIT_LEARNING_RATE)
    # model.layers[1].trainable = False
    # model.compile(loss=loss, optimizer=optimizer)
    # model.fit(train_data, epochs=CFG.WFLW_PARAMS.INIT_EPOCHS, callbacks=callbacks, validation_data=val_data)

    # finetune model with trainable backbone
    callbacks = [
            tf.keras.callbacks.TensorBoard(log_dir=f'logs/face_keypoints'),
            tf.keras.callbacks.EarlyStopping(monitor='val_loss', patience=10, verbose=True, restore_best_weights=True, min_delta=0.1),
            tf.keras.callbacks.LearningRateScheduler(lambda e, lr: lr if e < 5 else lr * tf.math.exp(-0.1)),
        ]
    optimizer = tf.keras.optimizers.SGD(learning_rate=CFG.WFLW_PARAMS.FINETUNE_LEARNING_RATE)
    model.layers[1].trainable = True
    model.compile(loss=loss, optimizer=optimizer)
    model.fit(train_data, epochs=CFG.WFLW_PARAMS.FINETUNE_EPOCHS, callbacks=callbacks, validation_data=val_data)
    # save weights
    model.save_weights(CFG.WFLW_WEIGHTS[backbone])


def find_optimal_delta(backbone, deltas):
    """Find optimal delta for face pose estimation model

    Args:
        backbone (str): model backbone configuration, see model implementation for possible inputs
        deltas (list): list of deltas to search
    """
    # load training dataset
    train_data = wflw_init_dataset(
        CFG.WFLW_TFRECORD_PATHS['train'],
        16,
        CFG.WFLW_PARAMS.IMG_SHAPE,
        CFG.WFLW_PARAMS.HM_DIMS,
        CFG.WFLW_PARAMS.SIGMA,
        testing=True
    )
    # load trained model
    model = load_model(CFG.WFLW_WEIGHTS[backbone], backbone, CFG.WFLW_PARAMS.IMG_SHAPE, CFG.WFLW_PARAMS.HM_DIMS, imagenet=False, decode=None)
    # define metric for argmax and for all passed deltas
    argmax_metric = FaceKeypointMetric()
    compensation_metrics = [FaceKeypointMetric() for _ in deltas]
    # iterate over training dataset
    for x, y_true in train_data:
        # forward pass
        heatmaps = model(x)
        # decode keypoints with argmax and update argmax metric
        y_pred = decode_keypoints_argmax_2d(heatmaps) / CFG.WFLW_PARAMS.HM_DIMS
        argmax_metric.update_state(y_true, y_pred)
        # decode keypoints with all deltas and update metrics
        for delta, metric in zip(deltas, compensation_metrics):
            y_pred = decode_keypoints_compensated_2d(heatmaps, delta) / CFG.WFLW_PARAMS.HM_DIMS
            metric.update_state(y_true, y_pred)
    # scale results and print them
    overall_results, attrib_results = argmax_metric.result()
    scale = CFG.WFLW_PARAMS.IMG_SHAPE[0]
    attrib_results = ''.join(f'{name}: {scale*res:.2f}  ' for name, res in attrib_results.items())
    print(f'Argmax: \tOverall: {scale*overall_results:.2f}  {attrib_results}')
    for delta, metric in zip(deltas, compensation_metrics):
        overall_results, attrib_results = metric.result()
        attrib_results = ''.join(f'{name}: {scale*res:.2f}  ' for name, res in attrib_results.items())
        print(f'Delta: {delta}\tOverall: {scale*overall_results:.2f}  {attrib_results}')


if __name__ == '__main__':
    # train_hand_estimator('efficientnet-b2')
    # train_face_estimator('efficientnet-b1')
    find_optimal_delta('efficientnet-b1', list(range(17)))
