"""
Evaluate models by performance
"""
import os
os.environ["CUDA_VISIBLE_DEVICES"] = "-1"

import tensorflow as tf

from tensorflow.python.framework.convert_to_constants import convert_variables_to_constants_v2_as_graph
from time import time

from modules.yolov5 import YOLOv5
from modules.efficientdet import EfficientDet
from modules.config import CFG
from modules.inference_utils import ExportModel


def meassure_flops(model):
    """Meassure model's FLOPs

    Args:
        model (ExportModel): exported model for inference

    Returns:
        int: number of FLOPs
    """
    # random input
    inputs = tf.random.uniform((1, *CFG.COMBINED_PARAMS.IMG_SHAPE, 3))
    # freeze computational graph
    real_model = tf.function(model.model).get_concrete_function(inputs)
    frozen_func = convert_variables_to_constants_v2_as_graph(real_model)[0]
    # get model metadata and profiler options
    run_meta = tf.compat.v1.RunMetadata()
    opts = tf.compat.v1.profiler.ProfileOptionBuilder.float_operation()
    # run profiler
    flops = tf.compat.v1.profiler.profile(graph=frozen_func.graph, run_meta=run_meta, cmd='scope', options=opts)
    return flops.total_float_ops


def meassure_latency(model):
    """Meassure model's latency

    Args:
        model (ExportModel): exported model for inference

    Returns:
        float: latency in seconds
    """
    # random testing input
    inputs = tf.random.uniform((1, *CFG.COMBINED_PARAMS.IMG_SHAPE, 3))
    # build model
    model.model.build((1, *CFG.COMBINED_PARAMS.IMG_SHAPE, 3))
    model.f(inputs)
    # average over 50 passes
    latencies = list()
    for _ in range(50):
        s = time()
        model.f(inputs)
        e = time()
        latencies.append(e-s)
    return sum(latencies) / len(latencies)


def main(model_cfgs):
    """Meassure flops, latency and number of parameters of a model

    Args:
        model_cfgs (list): list of model configurations, see model implementation for implemented configurations
    """
    img_shape = (*CFG.COMBINED_PARAMS.IMG_SHAPE, 3)
    flops_results = list()
    params_results = list()
    latency_results = list()
    # for every model config
    for model_cfg in model_cfgs:
        # load model
        model_f = YOLOv5 if 'yolov5' in model_cfg else EfficientDet
        anchors = CFG.COMBINED_PARAMS.YOLO_ANCHORS if 'yolov5' in model_cfg else CFG.COMBINED_PARAMS.EFFDET_ANCHORS
        model = model_f(img_shape, 2, anchors, model=model_cfg, training=False, postprocess=False)
        model = ExportModel(model)
        # meassure flops, latency and number of parameters
        flops = meassure_flops(model)
        latency = meassure_latency(model)
        flops_results.append(flops)
        params_results.append(int(sum([tf.size(w) for w in model.model.weights])))
        latency_results.append(latency)
    # print results to console
    for model_cfg, flops, params, latency in zip(model_cfgs, flops_results, params_results, latency_results):
        print(f'{model_cfg} - {flops/1e9:.2f} Billion FLOPS\t{params/1e6:.2f} Million params\t{latency*1e3:.2f}ms Latency')


if __name__ == '__main__':
    model_cfgs = [
        'yolov5s', 'yolov5m', 'yolov5l',
        'efficientdet-d0', 'efficientdet-d1',
        'efficientdet-d2', ]
    main(model_cfgs)
