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.model import SignRecognitionModel
from modules.inference_utils import ExportModel


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

    Args:
        model (tf.keras.Model): measured model
        input_shape (tuple): input image shape

    Returns:
        int: number of FLOPs
    """
    inputs = tf.random.uniform((1, *input_shape))
    real_model = tf.function(model).get_concrete_function(inputs)
    frozen_func = convert_variables_to_constants_v2_as_graph(real_model)[0]
    run_meta = tf.compat.v1.RunMetadata()
    opts = tf.compat.v1.profiler.ProfileOptionBuilder.float_operation()
    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, input_shape):
    """Meassure model's latency

    Args:
        model (tf.keras.Model): measured model
        input_shape (tuple): input image shape

    Returns:
        float: latency in seconds
    """
    inputs = tf.random.uniform((1, *input_shape))
    model.build((1, *input_shape))
    model = ExportModel(model)
    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 wlasl_models():
    """Meassure flops latency and number of parameters of trained classification models
    """
    # list of class splits, input tensor shape
    classes = (100, 300, 1000, 2000)
    input_shape = (50, 2*21*2+1*98*2)
    # results lists
    flops_list, latency_list, params_list = list(), list(), list()
    # for every model
    for n_classes in classes:
        # get model
        model = SignRecognitionModel(input_shape, n_classes)
        # measure flops, latency and number of parameters
        flops = meassure_flops(model, input_shape)
        flops_list.append(flops)
        latency = meassure_latency(model, input_shape)
        latency_list.append(latency)
        params = int(sum([tf.size(w) for w in model.weights]))
        params_list.append(params)
    # write to console
    for flops, latency, params, n_classes in zip(flops_list, latency_list, params_list, classes):
        print(f'Translation@{n_classes} - {flops/1e6:.2f} Million FLOPS\t{params/1e3:.2f} Thousand params\t{latency*1e3:.2f}ms Latency')


if __name__ == '__main__':
    wlasl_models()
