import warnings

warnings.simplefilter("ignore", FutureWarning)
warnings.simplefilter("ignore", UserWarning)

import torch
from torch.profiler import profile, ProfilerActivity
from model.model import Thundernet
from models_repo.model_attention import Thundernet as Thundernet_attention
from models_repo.model_attention_2 import Thundernet as Thundernet_attention2
from models_repo.model_ppm_factors import Thundernet as Thundernet_ppm

import time
import cv2
import numpy as np
import tensorflow as tf

device: torch.device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
# Define input shape
input_shape = (480, 640, 3)


def execute_profiler(model: Thundernet) -> None:
    """
    Function to measure de CPU and CUDA times.
    It prints the results on the console
    Args:
    - model: loaded model to profile
    Returns:
    - None
    """
    image = torch.randn(1, 480, 640, 3).cpu().numpy()

    # model = torch.jit.trace(model, (image, depth))
    with profile(
        activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA],
        record_shapes=True,
        profile_memory=True,
    ) as prof:
        with torch.no_grad():
            _ = model.predict(image)

    print(prof.key_averages().table(sort_by="cuda_time_total", row_limit=20))

    return


def measure_inference_time(model: Thundernet) -> None:
    """
    Function to measure the average inference time
    and the FPS of a given AsymFormer model.
    It prints the results on the console
    Args:
    - model: loaded model to profile
    Returns:
    - None
    """
    image = torch.randn(1, 480, 640, 3).cpu().numpy()

    for _ in range(5):
        _ = model.predict(image)

    times = []
    for _ in range(20):
        tf.constant(0).numpy()
        start = time.time()
        _ = model.predict(image)
        tf.constant(0).numpy()
        times.append((time.time() - start) * 1000)

    avg_time = sum(times) / len(times)
    print(f"Average inference time: {avg_time:.2f} ms")

    fps = 1000 / avg_time
    print(f"FPS: {fps:.2f}")
    return


def main() -> None:

    # PATH TO THE BEST MODEL SO FAR (.hdf5)
    # weights_path = "D:/RealTimeSemanticSegmentation_Sofia/keras.hdf5"
    # weights_path = "C:/Users/user/Documents/Thundernet/pruebas_modelos/32_ppm/BS4_lossBCE_weights_lr_0.00013713842558297858_reg-1.1743577101671763e-05-ep-13-val_loss0.11463435739278793-train_loss0.053004469722509384-val_iou0.8959722518920898-train_iou0.9606077075004578.hdf5"
    weights_path = "keras.hdf5"
    # Load the model. Change it depending on where it was trained
    # ThunderNet = Thundernet_ppm(input_shape=input_shape, resnet_trainable=False, n_classes = 2)
    ThunderNet = Thundernet(
        input_shape=input_shape, resnet_trainable=False, n_classes=2
    )
    model = ThunderNet.model
    ThunderNet.model.load_weights(weights_path)
    execute_profiler(model)
    measure_inference_time(model)
    return


if __name__ == "__main__":
    main()