# Copyright 2021 Huawei Technologies Co., Ltd
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

# Thanks for the contribution of KopiSoftware https://github.com/KopiSoftware

import torch
import time
import numpy as np
from model.model import parsingNet
import torchvision.transforms as transforms
import cv2
from matplotlib import pyplot as plt
from PIL import Image

img_transforms = transforms.Compose([
    transforms.Resize((288, 800)),
    transforms.ToTensor(),
    transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)),
])


def resize(x, y):
    global cap
    cap.set(3, x)
    cap.set(4, y)


def test_practical_without_readtime():
    global cap
    for i in range(10):
        _, img = cap.read()
        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
        img2 = Image.fromarray(img)
        x = img_transforms(img2)
        x = x.unsqueeze(0).cuda() + 1
        y = net(x)

    print("pracrical image input size:", img.shape)
    print("pracrical tensor input size:", x.shape)
    t_all = []
    for i in range(100):
        _, img = cap.read()

        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
        img2 = Image.fromarray(img)
        x = img_transforms(img2)
        x = x.unsqueeze(0).cuda() + 1

        t1 = time.time()
        y = net(x)
        t2 = time.time()
        t_all.append(t2 - t1)

    print("practical with out read time:")
    print('\taverage time:', np.mean(t_all) / 1)
    print('\taverage fps:', 1 / np.mean(t_all))

    # print('fastest time:', min(t_all) / 1)
    # print('fastest fps:',1 / min(t_all))

    # print('slowest time:', max(t_all) / 1)
    # print('slowest fps:',1 / max(t_all))


def test_practical():
    global cap
    for i in range(10):
        _, img = cap.read()
        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
        img2 = Image.fromarray(img)
        x = img_transforms(img2)
        x = x.unsqueeze(0).cuda() + 1
        y = net(x)

    print("pracrical image input size:", img.shape)
    print("pracrical tensor input size:", x.shape)
    t_all = []
    t_capture = []
    t_preprocessing = []
    t_net = []
    for i in range(100):
        t1 = time.time()

        t3 = time.time()
        _, img = cap.read()
        t4 = time.time()

        t5 = time.time()
        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
        img2 = Image.fromarray(img)
        x = img_transforms(img2)
        x = x.unsqueeze(0).cuda() + 1
        t6 = time.time()

        y = net(x)
        t2 = time.time()
        t_all.append(t2 - t1)
        t_capture.append(t4 - t3)
        t_preprocessing.append(t6 - t5)
        t_net.append(t2 - t6)

    print("practical with read time:")
    print('\taverage time:', np.mean(t_all) / 1)
    print('\taverage fps:', 1 / np.mean(t_all))
    print('\tcapture time:', np.mean(t_capture) / 1)
    print('\tpre-processing time:', np.mean(t_preprocessing) / 1)
    print('\tdetect time:', np.mean(t_net) / 1)

    # print('fastest time:', min(t_all) / 1)
    # print('fastest fps:',1 / min(t_all))

    # print('slowest time:', max(t_all) / 1)
    # print('slowest fps:',1 / max(t_all))


###x = torch.zeros((1,3,288,800)).cuda() + 1
def test_theoretical():
    x = torch.zeros((1, 3, 288, 800)).cuda() + 1
    for i in range(10):
        y = net(x)

    t_all = []
    for i in range(100):
        t1 = time.time()
        y = net(x)
        t2 = time.time()
        t_all.append(t2 - t1)
    print("theortical")
    print('\taverage time:', np.mean(t_all) / 1)
    print('\taverage fps:', 1 / np.mean(t_all))

    # print('fastest time:', min(t_all) / 1)
    # print('fastest fps:',1 / min(t_all))

    # print('slowest time:', max(t_all) / 1)
    # print('slowest fps:',1 / max(t_all))


if __name__ == "__main__":
    # captrue data from camera or video
    # cap = cv2.VideoCapture("video.mp4") #uncommen to activate a video input
    cap = cv2.VideoCapture(0)  # uncommen to activate a camera imput
    # resize(480, 640) #ucommen to change input size

    # torch.backends.cudnn.deterministic = False
    torch.backends.cudnn.benchmark = True
    net = parsingNet(pretrained=False, backbone='18', cls_dim=(100 + 1, 56, 4), use_aux=False).cuda()
    # net = parsingNet(pretrained = False, backbone='18',cls_dim = (200+1,18,4),use_aux=False).cuda()
    net.eval()

    test_practical_without_readtime()
    test_practical()
    cap.release()
    test_theoretical()