import time

import cv2
import numpy as np
import torch
import torch.nn.functional as F
from torchvision.transforms import transforms


def STDC_process(img0, model, device, new_hw=None):
    if new_hw is None:
        new_hw = [360, 640]
    t_start = time.time()
    img0 = cv2.cvtColor(img0, cv2.COLOR_BGR2RGB)
    t2 = time.time()
    print(f't_bgr2rgb. ({t2 - t_start:.3f}s)')
    # img0 = img0[..., ::-1]
    img = transforms.ToTensor()(img0)
    img = transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225))(img)
    t3 = time.time()
    print(f't_trans. ({t3 - t2:.3f}s)')
    t_togpu = time.time()
    img = img.to(device)
    t_togpu2 = time.time()
    print(f't_togpu. ({t_togpu2 - t_togpu:.3f}s)')
    C, H, W = img.shape
    size = img.size()[-2:]
    # new_hw = [int(H * scale), int(W * scale)]
    # new_hw = [360, 640]
    img = img.unsqueeze(0)
    img = F.interpolate(img, new_hw, mode='bilinear', align_corners=True)
    t_pro = time.time()
    print(f't_interpolate. ({t_pro - t_togpu2:.3f}s)')
    print(f't_pro. ({t_pro - t_start:.3f}s)')
    logits = model(img)[0]
    t_inf = time.time()
    print(f't_inf. ({t_inf - t_pro:.3f}s)')
    logits = F.interpolate(logits, size=size, mode='bilinear', align_corners=True)
    probs = torch.softmax(logits, dim=1)
    preds = torch.argmax(probs, dim=1)
    preds_squeeze = preds.squeeze(0)
    preds_squeeze_predict = np.array(preds_squeeze.cpu())
    t_end = time.time()
    print(f't_post. ({t_end - t_inf:.3f}s)')

    return preds_squeeze_predict