"""
Evaluate trained models with mean average precision metric
"""
import os
# os.environ["CUDA_VISIBLE_DEVICES"] = "-1"

from functools import partial

from modules.config import CFG
from modules.mean_average_precision import mAP

from modules.dataset_utils import init_dataset
from modules.yolov5 import YOLOv5, boxes_from_output, non_maximum_suppression
from modules.efficientdet import EfficientDet


def evaluate_model(model, metric_f, dataset, classes, iou_thresholds, score_thresholds, anchors, max_detections):
    """Evaluate general model with multiple iou and score thresholds

    Args:
        model (tf.keras.Model): evaluated model
        metric_f (callable): metric construction callable
        dataset (tf.data.Dataset): evaluation dataset
        classes (int): number of classes
        iou_thresholds (list): list of mean average precision iou thresholds to evaluate the model with
        score_thresholds (list): list of score thresholds to evaluate the model with
        anchors (tf.Tensor): anchor boxes
        max_detections (int): maximum detections outputted by the model

    Returns:
        dictionary: mean average precision metric results
    """
    # construct a metric for each iou and score threshold combination
    metrics = list()
    for iou_thold in iou_thresholds:
        metrics_at_iou = list()
        for _ in score_thresholds:
            metrics_at_iou.append(metric_f(iou_threshold=iou_thold))
        metrics.append(metrics_at_iou)
    # iterate over dataset
    for x, y_true in dataset:
        # get raw outputs
        outputs = model(x)
        # postprocess outputs with specific iou and score thresholds
        for iou_thold, metric_at_iou in zip(iou_thresholds, metrics):
            for score_thold, metric in zip(score_thresholds, metric_at_iou):
                boxes = [boxes_from_output(output, single_anchors, classes) for output, single_anchors in zip(outputs, anchors)]
                y_pred = non_maximum_suppression(boxes, classes, max_detections, 0.1, score_thold)
                metric.update_state(y_true, y_pred)
    # calculate results
    results = {}
    for iou_thold, metric_at_iou in zip(iou_thresholds, metrics):
        for score_thold, metric in zip(score_thresholds, metric_at_iou):
            results[f'{iou_thold} {score_thold}'] = metric.result()
            # reset metric after evaluation
            metric.reset_state()
    return results


def evaluate_faces(model_cfg, iou_threshold=0.5, score_threshold=CFG.FACES_PARAMS.SCORE_THRESHOLD):
    """Evaluate face detection

    Args:
        model_cfg (str): model configuration, see model implementation for implemented configurations
        iou_threshold (float, optional): mAP iou threshold or list of thresholds. Defaults to 0.5.
        score_threshold (float, optional): non-max suppression threshold or list of these. Defaults to CFG.FACES_PARAMS.SCORE_THRESHOLD.
    """
    # convert passed thresholds to list
    iou_thold = iou_threshold if isinstance(iou_threshold, list) else [iou_threshold]
    score_threshold = score_threshold if isinstance(score_threshold, list) else [score_threshold]
    # get anchors
    anchors = CFG.FACES_PARAMS.ANCHORS if 'yolov5' in model_cfg else CFG.FACES_PARAMS.EFFDET_ANCHORS
    # init evaluation dataset
    dataset = init_dataset(
        CFG.FACES_TFRECORDS_PATH['val'],
        CFG.FACES_PARAMS.BATCH_SIZE,
        anchors,
        CFG.FACES_PARAMS.IMG_SHAPE,
        100,
        model_cfg,
        testing=True)
    # load model and its weights
    model_f = YOLOv5 if 'yolov5' in model_cfg else EfficientDet
    model = model_f(
            (*CFG.FACES_PARAMS.IMG_SHAPE, 3),
            1,
            CFG.FACES_PARAMS.ANCHORS,
            model=model_cfg,
            training=False,
            postprocess=False
        )
    model.load_weights(CFG.FACES_WEIGHTS_PATH[model_cfg]).expect_partial()
    # metric callable - partial function
    metric = partial(mAP, classes=1, single_class_labels=True)
    # evaluate model
    results = evaluate_model(model, metric, dataset, 1, iou_thold, score_threshold, CFG.FACES_PARAMS.ANCHORS, 100)
    # print results to console
    print('--- FACES ---')
    for key, value in results.items():
        iou, score = key.split(' ')
        mean_ap, ap = value
        ap = [float(f'{val:0.5f}') for val in ap.numpy().tolist()]
        print(f'mAP@{iou} Score threshold: {score} - Faces: {model_cfg}: {mean_ap:.5f} ~ {ap}')


def evaluate_hands(model_cfg, iou_threshold=0.5, score_threshold=CFG.HANDS_PARAMS.SCORE_THRESHOLD):
    """Evaluate hand detection

    Args:
        model_cfg (str): model configuration, see model implementation for implemented configurations
        iou_threshold (float, optional): mAP iou threshold or list of thresholds. Defaults to 0.5.
        score_threshold (float, optional): non-max suppression threshold or list of these. Defaults to CFG.HANDS_PARAMS.SCORE_THRESHOLD.
    """
    # convert passed thresholds to list
    iou_thold = iou_threshold if isinstance(iou_threshold, list) else [iou_threshold]
    score_threshold = score_threshold if isinstance(score_threshold, list) else [score_threshold]
    # get anchors
    anchors = CFG.HANDS_PARAMS.ANCHORS if 'yolov5' in model_cfg else CFG.HANDS_PARAMS.EFFDET_ANCHORS
    # init evaluation dataset
    dataset = init_dataset(
        CFG.HANDS_TFRECORDS_PATH['test'],
        CFG.HANDS_PARAMS.BATCH_SIZE,
        anchors,
        CFG.HANDS_PARAMS.IMG_SHAPE,
        100,
        model_cfg,
        testing=True)
    # load model and its weights
    model_f = YOLOv5 if 'yolov5' in model_cfg else EfficientDet
    model = model_f(
            (*CFG.HANDS_PARAMS.IMG_SHAPE, 3),
            1,
            CFG.HANDS_PARAMS.ANCHORS,
            model=model_cfg,
            training=False,
            postprocess=False
        )
    model.load_weights(CFG.HANDS_WEIGHTS_PATH[model_cfg]).expect_partial()
    # metric callable - partial function
    metric = partial(mAP, classes=1, single_class_labels=True)
    # evaluate model
    results = evaluate_model(model, metric, dataset, 1, iou_thold, score_threshold, CFG.HANDS_PARAMS.ANCHORS, 100)
    # print results to console
    print('--- HANDS ---')
    for key, value in results.items():
        iou, score = key.split(' ')
        mean_ap, ap = value
        ap = [float(f'{val:0.5f}') for val in ap.numpy().tolist()]
        print(f'mAP@{iou} Score threshold: {score} - Hands: {model_cfg}: {mean_ap:.5f} ~ {ap}')


def evaluate_combined(model_cfg, iou_threshold=0.5, score_threshold=CFG.COMBINED_PARAMS.SCORE_THRESHOLD):
    """Evaluate combined hand and face detection

    Args:
        model_cfg (str): model configuration, see model implementation for implemented configurations
        iou_threshold (float, optional): mAP iou threshold or list of thresholds. Defaults to 0.5.
        score_threshold (float, optional): non-max suppression threshold or list of these. Defaults to CFG.COMBINED_PARAMS.SCORE_THRESHOLD.
    """
    # convert passed thresholds to list
    iou_thold = iou_threshold if isinstance(iou_threshold, list) else [iou_threshold]
    score_threshold = score_threshold if isinstance(score_threshold, list) else [score_threshold]
    # get anchors
    anchors = CFG.COMBINED_PARAMS.YOLO_ANCHORS if 'yolov5' in model_cfg else CFG.COMBINED_PARAMS.EFFDET_ANCHORS
    # init evaluation dataset
    dataset = init_dataset(
        CFG.COMBINED_TFRECORDS_PATH['test'],
        CFG.COMBINED_PARAMS.BATCH_SIZE,
        anchors,
        CFG.COMBINED_PARAMS.IMG_SHAPE,
        100,
        model_cfg,
        testing=True)
    # load model and its weights
    model_f = YOLOv5 if 'yolov5' in model_cfg else EfficientDet
    model = model_f(
            (*CFG.COMBINED_PARAMS.IMG_SHAPE, 3),
            2,
            anchors,
            model=model_cfg,
            training=False,
            postprocess=False
        )
    model.load_weights(CFG.COMBINED_WEIGHTS_PATH[model_cfg]).expect_partial()
    # metric callable - partial function
    metric = partial(mAP, classes=2, single_class_labels=True)
    # evaluate model
    results = evaluate_model(model, metric, dataset, 2, iou_thold, score_threshold, anchors, 100)
    # print results to console
    print('--- COMB ---')
    for key, value in results.items():
        iou, score = key.split(' ')
        mean_ap, ap = value
        ap = [float(f'{val:0.5f}') for val in ap.numpy().tolist()]
        print(f'mAP@{iou} Score threshold: {score} - Combined: {model_cfg}: {mean_ap:.5f} ~ {ap}')


if __name__ == '__main__':
    iou_tholds = [0.3, 0.5, 0.7]
    score_tholds = [0.2, 0.3, 0.4, 0.5]
    evaluate_faces('yolov5s', iou_threshold=iou_tholds, score_threshold=score_tholds)
    evaluate_hands('yolov5s', iou_threshold=iou_tholds, score_threshold=score_tholds)
    evaluate_combined('yolov5s', iou_threshold=iou_tholds, score_threshold=score_tholds)
    evaluate_combined('yolov5m', iou_threshold=iou_tholds, score_threshold=score_tholds)
    evaluate_combined('yolov5l', iou_threshold=iou_tholds, score_threshold=score_tholds)
    evaluate_combined('efficientdet-d0', iou_threshold=iou_tholds, score_threshold=score_tholds)
    evaluate_combined('efficientdet-d1', iou_threshold=iou_tholds, score_threshold=score_tholds)
    evaluate_combined('efficientdet-d2', iou_threshold=iou_tholds, score_threshold=score_tholds)
