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.interhand import InterhandEstimator
from modules.face_model import FaceKeypointEstimator
from modules.hands_model import HandKeypointEstimator
from modules.inference_utils import ExportModel
from modules.config import CFG


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

    Args:
        model (tf.keras.Model): measured model
        img_size (tuple): input image size in first two dimensions

    Returns:
        int: number of FLOPs
    """
    # random input
    inputs = tf.random.uniform((1, *img_size, 3))
    # freeze computational graph
    real_model = tf.function(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, img_size):
    """Meassure model's latency

    Args:
        model (tf.keras.Model): measured model
        img_size (tuple): input image size in first two dimensions

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


def interhand():
    """Meassure flops, latency and number of parameters of Interhand model
    """
    # load model
    model = InterhandEstimator(CFG.INTERHAND_PARAMS.IMG_SHAPE, CFG.INTERHAND_PARAMS.JOINT_HM_DIMS,
                               CFG.INTERHAND_PARAMS.DEPTH_DIM, CFG.INTERHAND_PARAMS.ROOT_HM_DEPTH)
    # meassure flops, latency and number of parameters
    flops = meassure_flops(model)
    latency = meassure_latency(model)
    params = int(sum([tf.size(w) for w in model.weights]))
    print(f'InterHand - {flops/1e9:.2f} Billion FLOPS\t{params/1e6:.2f} Million params\t{latency*1e3:.2f}ms Latency')


def faces(backbone):
    """Meassure flops, latency and number of parameters of face pose estimation model

    Args:
        backbone (str): model backbone configuration, see model implementation for possible inputs
    """
    img_size = (256, 256)
    # load model
    model = FaceKeypointEstimator(backbone, img_size, training=True)
    # meassure flops, latency and number of parameters
    flops = meassure_flops(model, img_size)
    latency = meassure_latency(model, img_size)
    params = int(sum([tf.size(w) for w in model.weights]))
    print(f'Faces {backbone} - {flops/1e9:.2f} Billion FLOPS\t{params/1e6:.2f} Million params\t{latency*1e3:.2f}ms Latency')


def hands(backbone):
    """Meassure flops, latency and number of parameters of hand pose estimation model

    Args:
        backbone (str): model backbone configuration, see model implementation for possible inputs
    """
    img_size = (256, 256)
    # load model
    model = HandKeypointEstimator(backbone, img_size, training=True)
    # meassure flops, latency and number of parameters
    flops = meassure_flops(model, img_size)
    latency = meassure_latency(model, img_size)
    params = int(sum([tf.size(w) for w in model.weights]))
    print(f'Hands {backbone} - {flops/1e9:.2f} Billion FLOPS\t{params/1e6:.2f} Million params\t{latency*1e3:.2f}ms Latency')


if __name__ == '__main__':
    # interhand()
    faces('efficientnet-b1')
    hands('efficientnet-b1')
