import cv2
import torch
import urllib.request
import matplotlib.pyplot as plt
import numpy as np
from PIL import Image
import queue

from EsoMovieConverter import EsoMovieConverter

cap = cv2.VideoCapture(r'D:\Deep_Learning\MonoDepth2\esophagus\movies\trimed\0.mp4')
# cap = cv2.VideoCapture(1)
eso_movie_converter = EsoMovieConverter()

fps = int(cap.get(cv2.CAP_PROP_FPS))
w = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
h = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
fourcc = cv2.VideoWriter_fourcc('m', 'p', '4', 'v')
video_writer = cv2.VideoWriter('eso0_disp.mp4', fourcc, fps, (480, 352))

model_type = "DPT_Large"  # MiDaS v3 - Large     (highest accuracy, slowest inference speed)
# model_type = "DPT_Hybrid"   # MiDaS v3 - Hybrid    (medium accuracy, medium inference speed)
# model_type = "MiDaS_small"  # MiDaS v2.1 - Small   (lowest accuracy, highest inference speed)
midas = torch.hub.load("intel-isl/MiDaS", model_type)
device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
midas.to(device)
midas.eval()
midas_transforms = torch.hub.load("intel-isl/MiDaS", "transforms")
if model_type == "DPT_Large" or model_type == "DPT_Hybrid":
    transform = midas_transforms.dpt_transform
else:
    transform = midas_transforms.small_transform


def inference(img):
    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
    input_batch = transform(img).to(device)
    with torch.no_grad():
        prediction = midas(input_batch)
        prediction = torch.nn.functional.interpolate(
            prediction.unsqueeze(1),
            size=img.shape[:2],
            mode="bicubic",
            align_corners=False,
        ).squeeze()
    output = prediction.cpu().numpy()
    formatted = (output * 255 / np.max(output)).astype('uint8')
    return formatted


prev_img1 = None
prev_img2 = None
prev_img3 = None
prev_img4 = None
while True:
    ret, frame = cap.read()
    if not ret:
        break

    frame = eso_movie_converter(frame)
    cv2.imshow("input", frame)
    out = inference(frame)
    if (prev_img1 is not None) and (prev_img2 is not None) and (prev_img3 is not None) and (prev_img4 is not None):
        show_img = out / 5 + prev_img1 / 5 + prev_img2 / 5 + prev_img3 / 5 + prev_img4 / 5
        show_img = show_img.astype("uint8")
        cv2.imshow("out", show_img)
        video_writer.write(cv2.cvtColor(show_img[:, :, np.newaxis], cv2.COLOR_GRAY2RGB))
    cv2.waitKey(1)
    if prev_img3 is not None:
        prev_img4 = prev_img3.copy()
    if prev_img2 is not None:
        prev_img3 = prev_img2.copy()
    if prev_img1 is not None:
        prev_img2 = prev_img1.copy()
    prev_img1 = out.copy()

cap.release()
video_writer.release()