import os
import cv2
import torch
import argparse
import numpy as np
from torch.nn import functional as F
from tqdm import tqdm
import warnings
import _thread
from queue import Queue, SimpleQueue
from moviepy.editor import VideoFileClip
# from moviepy.video.fx.all import crop
from subprocess import Popen, PIPE
# from benchmark.pytorch_msssim import ssim_matlab

warnings.filterwarnings("ignore")

parser = argparse.ArgumentParser(description='Interpolation for a pair of images')
parser.add_argument('--video', dest='video', type=str, default=None)
parser.add_argument('--output', dest='output', type=str, default=None)
parser.add_argument('--drive', dest='drive', type=int, default=1, help='change to 0 to disable flushing/remounting Google Drive in Colab')
parser.add_argument('--model', dest='modelDir', type=str, default='train_log', help='directory with trained model files')
# parser.add_argument('--skip', dest='skip', action='store_true', help='whether to remove static frames before processing')
parser.add_argument('--ext', dest='ext', type=str, default='mkv', help='vid_out video extension')
args = parser.parse_args()
assert (not args.video is None or not args.output is None)
    
device = torch.device("cuda")
torch.set_grad_enabled(False)
torch.backends.cudnn.enabled = True
torch.backends.cudnn.benchmark = True
torch.set_default_tensor_type(torch.cuda.HalfTensor)

from model.RIFE_HDv2 import Model
model = Model()
model.load_model(args.modelDir, -1)
print("Loaded v2.x HD model.")
model.eval()
model.device()

videoCapture = cv2.VideoCapture(args.video)
fps = videoCapture.get(cv2.CAP_PROP_FPS)
tot_frame = videoCapture.get(cv2.CAP_PROP_FRAME_COUNT)
videoCapture.release()
fps *= 2
# videogen = crop(VideoFileClip(f), x1=0, y1=142, width=1920, height=800).iter_frames(dtype="uint8")
videogen = VideoFileClip(args.video).iter_frames(dtype="uint8")
lastframe = next(videogen)

def build_read_buffer(read_buffer, videogen):
    for frame in videogen:
        read_buffer.put(frame)
    read_buffer.put(None)

read_buffer = Queue(maxsize=1000)
_thread.start_new_thread(build_read_buffer, (read_buffer, videogen))
video_path_wo_ext, _ = os.path.splitext(args.video)
print('{}.{}, {} frames in total, {}FPS to {}FPS'.format(video_path_wo_ext, args.ext, tot_frame, fps / 2, fps))
h, w, _ = lastframe.shape
p = Popen(['ffmpeg', '-y', '-fflags', '+genpts', '-loglevel', 'warning', '-nostats', '-hide_banner', '-threads', '1', '-f', 'rawvideo', '-pix_fmt', 'rgb24', '-s', '{}x{}'.format(w, h), '-r', str(fps), '-i', '-', '-an', '-c:v', 'hevc_nvenc', '-profile:v', 'main10', '-level:v', '4.1', '-rc:v', 'vbr_hq', '-cq:v', '21', '-b:v', '0k', '-rc-lookahead:v', '32', '-g', '480', '-bf:v', '4', '-refs:v', '4', '-r', str(fps), '-pix_fmt', 'yuv420p', args.output], stdin=PIPE, stdout=PIPE)

padding = (0, ((w - 1) // 32 + 1) * 32 - w, 0, ((h - 1) // 32 + 1) * 32 - h)
# skip_frame = 1

def clear_write_buffer(write_buffer):
    while True:
        p.stdin.write(write_buffer.get(timeout=4))
        p.stdin.write(write_buffer.get(timeout=4).tobytes())

pbar = tqdm(total=tot_frame)
write_buffer = SimpleQueue()
_thread.start_new_thread(clear_write_buffer, (write_buffer,))

I1 = F.pad((torch.from_numpy(np.transpose(lastframe, (2,0,1))).to(device, non_blocking=True).unsqueeze(0).float() / 255.), padding).half()
while True:
    frame = read_buffer.get(False)
    if frame is not None:
        I0 = I1
        I1 = F.pad((torch.from_numpy(np.transpose(frame, (2,0,1))).to(device, non_blocking=True).unsqueeze(0).float() / 255.), padding).half()
        write_buffer.put(lastframe)
        # ssim = ssim_matlab(I0_small, I1_small)
        # if ssim > 0.995 and args.skip:
            # if skip_frame % 100 == 0:
                # print("Warning: Your video has {} static frames, skipping them may change the duration of the generated video.".format(skip_frame))
            # skip_frame += 1
            # continue
        # if ssim_matlab(F.interpolate(I0, (32, 32), mode='bilinear', align_corners=False), F.interpolate(I1, (32, 32), mode='bilinear', align_corners=False)) >= 0.30:
            # p.stdin.write(((model.inference(I0, I1, 1.0)[0] * 255.).byte().cpu().numpy().transpose(1, 2, 0))[:h, :w].tobytes())
        # else:
            # p.stdin.write((((torch.from_numpy(np.transpose((cv2.addWeighted(frame[:, :, ::-1], 1/2, lastframe[:, :, ::-1], 1/2, 0)[:, :, ::-1].copy()), (2,0,1))).to(device, non_blocking=True).unsqueeze(0).float() / 255.)[0] * 255.).byte().cpu().numpy().transpose(1, 2, 0))[:h, :w].tobytes())
        write_buffer.put(((model.inference(I0, I1, 1.0)[0] * 255.).byte().cpu().numpy().transpose(1, 2, 0))[:h, :w])
        lastframe = frame
        pbar.update(1)
    else:
        break
write_buffer.put(lastframe)
import time
time.sleep(20)
p.stdin.close()
p.wait()
pbar.close()
videogen.close()
if args.drive > 0:
    from google.colab import drive
    time.sleep(1)
    drive.flush_and_unmount()
    time.sleep(5)
    drive.mount('/content/drive', force_remount=True)
    time.sleep(1)
print('\nAll done!')
