import sys
import os
import warnings

sys.path.append(os.path.abspath(os.path.join(__file__, "..", "..")))
sys.path.append(os.path.abspath(os.path.join(__file__, "..", "..", "FGT")))
sys.path.append(os.path.abspath(os.path.join(__file__, "..", "..", "LAFC")))
warnings.filterwarnings("ignore")

import cvbase
from torchvision.transforms import ToTensor
from tool.get_flowNN_gradient import get_flowNN_gradient
from tool.utils.Poisson_blend_img import Poisson_blend_img
from tool.utils.region_fill import regionfill
from importlib import import_module
import yaml
from RAFT import RAFT
from RAFT import utils
import torch.nn.functional as F2
import torchvision.transforms.functional as F
from skimage.feature import canny
import scipy.ndimage
from PIL import Image
import imageio
import torch
import numpy as np
import copy
import glob
import cv2
import argparse


def to_tensor(img):
    img = Image.fromarray(img)
    img_t = F.to_tensor(img).float()
    return img_t


def diffusion(flows, masks):
    flows_filled = []
    for i in range(flows.shape[0]):
        flow, mask = flows[i], masks[i]
        flow_filled = np.zeros(flow.shape)
        flow_filled[:, :, 0] = regionfill(flow[:, :, 0], mask[:, :, 0])
        flow_filled[:, :, 1] = regionfill(flow[:, :, 1], mask[:, :, 0])
        flows_filled.append(flow_filled)
    return flows_filled


def np2tensor(array, near="c"):
    if isinstance(array, list):
        array = np.stack(array, axis=0)  # [t, h, w, c]
    if near == "c":
        array = (
            torch.from_numpy(np.transpose(array, (3, 0, 1, 2))).unsqueeze(0).float()
        )  # [1, c, t, h, w]
    elif near == "t":
        array = torch.from_numpy(np.transpose(array, (0, 3, 1, 2))).unsqueeze(0).float()
    else:
        raise ValueError(f"Unknown near type: {near}")
    return array


def tensor2np(array):
    array = torch.stack(array, dim=-1).squeeze(0).permute(1, 2, 0, 3).cpu().numpy()
    return array


def gradient_mask(mask):
    gradient_mask = np.logical_or.reduce(
        (
            mask,
            np.concatenate(
                (mask[1:, :], np.zeros((1, mask.shape[1]), dtype=np.bool)), axis=0
            ),
            np.concatenate(
                (mask[:, 1:], np.zeros((mask.shape[0], 1), dtype=np.bool)), axis=1
            ),
        )
    )

    return gradient_mask


def indicesGen(pivot, interval, frames, t):
    singleSide = frames // 2
    results = []
    for i in range(-singleSide, singleSide + 1):
        index = pivot + interval * i
        if index < 0:
            index = abs(index)
        if index > t - 1:
            index = 2 * (t - 1) - index
        results.append(index)
    return results


def get_ref_index(f, neighbor_ids, length, ref_length, num_ref):
    ref_index = []
    if num_ref == -1:
        for i in range(0, length, ref_length):
            if i not in neighbor_ids:
                ref_index.append(i)
    else:
        start_idx = max(0, f - ref_length * (num_ref // 2))
        end_idx = min(length, f + ref_length * (num_ref // 2))
        for i in range(start_idx, end_idx + 1, ref_length):
            if i not in neighbor_ids:
                if len(ref_index) > num_ref:
                    break
                ref_index.append(i)
    return ref_index


def save_flows(output, videoFlowF, videoFlowB):
    create_dir(os.path.join(output, "completed_flow", "forward_flo"))
    create_dir(os.path.join(output, "completed_flow", "backward_flo"))
    create_dir(os.path.join(output, "completed_flow", "forward_png"))
    create_dir(os.path.join(output, "completed_flow", "backward_png"))
    N = videoFlowF.shape[-1]
    for i in range(N):
        forward_flow = videoFlowF[..., i]
        backward_flow = videoFlowB[..., i]
        forward_flow_vis = cvbase.flow2rgb(forward_flow)
        backward_flow_vis = cvbase.flow2rgb(backward_flow)
        cvbase.write_flow(
            forward_flow,
            os.path.join(
                output, "completed_flow", "forward_flo", "{:05d}.flo".format(i)
            ),
        )
        cvbase.write_flow(
            backward_flow,
            os.path.join(
                output, "completed_flow", "backward_flo", "{:05d}.flo".format(i)
            ),
        )
        imageio.imwrite(
            os.path.join(
                output, "completed_flow", "forward_png", "{:05d}.png".format(i)
            ),
            forward_flow_vis,
        )
        imageio.imwrite(
            os.path.join(
                output, "completed_flow", "backward_png", "{:05d}.png".format(i)
            ),
            backward_flow_vis,
        )


def save_fgcp(output, frames, masks):
    create_dir(os.path.join(output, "prop_frames"))
    create_dir(os.path.join(output, "masks_left"))
    create_dir(os.path.join(output, "prop_frames_npy"))
    create_dir(os.path.join(output, "masks_left_npy"))

    assert len(frames) == masks.shape[2]
    for i in range(len(frames)):
        cv2.imwrite(
            os.path.join(output, "prop_frames", "%05d.png" % i), frames[i] * 255.0
        )
        cv2.imwrite(
            os.path.join(output, "masks_left", "%05d.png" % i), masks[:, :, i] * 255.0
        )
        np.save(
            os.path.join(output, "prop_frames_npy", "%05d.npy" % i), frames[i] * 255.0
        )
        np.save(
            os.path.join(output, "masks_left_npy", "%05d.npy" % i),
            masks[:, :, i] * 255.0,
        )


def create_dir(dir):
    """Creates a directory if not exist."""
    if not os.path.exists(dir):
        os.makedirs(dir)


def initialize_RAFT(args, device):
    """Initializes the RAFT model."""
    model = torch.nn.DataParallel(RAFT(args))
    if not torch.cuda.is_available():
        model.load_state_dict(torch.load(args.raft_model, map_location=lambda storage, loc: storage))
    else:
        model.load_state_dict(torch.load(args.raft_model))

    model = model.module
    model.to(device)
    model.eval()

    return model


def initialize_LAFC(args, device):
    print(args.lafc_ckpts)
    assert len(os.listdir(args.lafc_ckpts)) == 2
    checkpoint, config_file = (
        glob.glob(os.path.join(args.lafc_ckpts, "*.tar"))[0],
        glob.glob(os.path.join(args.lafc_ckpts, "*.yaml"))[0],
    )
    with open(config_file, "r") as f:
        configs = yaml.full_load(f)
    model = configs["model"]
    pkg = import_module("LAFC.models.{}".format(model))
    model = pkg.Model(configs)
    if not torch.cuda.is_available():
        state = torch.load(
        checkpoint, map_location=lambda storage, loc: storage
        )
    else:
        state = torch.load(
            checkpoint, map_location=lambda storage, loc: storage.cuda(device)
        )
    model.load_state_dict(state["model_state_dict"])
    model = model.to(device)
    return model, configs


def initialize_FGT(args, device):
    assert len(os.listdir(args.fgt_ckpts)) == 2
    checkpoint, config_file = (
        glob.glob(os.path.join(args.fgt_ckpts, "*.tar"))[0],
        glob.glob(os.path.join(args.fgt_ckpts, "*.yaml"))[0],
    )
    with open(config_file, "r") as f:
        configs = yaml.full_load(f)
    model = configs["model"]
    net = import_module("FGT.models.{}".format(model))
    model = net.Model(configs).to(device)
    if not torch.cuda.is_available():
        state = torch.load(
        checkpoint, map_location=lambda storage, loc: storage
        )
    else:
        state = torch.load(
            checkpoint, map_location=lambda storage, loc: storage.cuda(device)
        )
    model.load_state_dict(state["model_state_dict"])
    return model, configs


def calculate_flow(args, model, video, mode):
    """Calculates optical flow."""
    if mode not in ["forward", "backward"]:
        raise NotImplementedError

    imgH, imgW = args.imgH, args.imgW
    Flow = np.empty(((imgH, imgW, 2, 0)), dtype=np.float32)

    if args.vis_flows:
        create_dir(os.path.join(args.outroot, "flow", mode + "_flo"))
        create_dir(os.path.join(args.outroot, "flow", mode + "_png"))

    with torch.no_grad():
        for i in range(video.shape[0] - 1):
            print(
                "Calculating {0} flow {1:2d} <---> {2:2d}".format(mode, i, i + 1),
                "\r",
                end="",
            )
            if mode == "forward":
                # Flow i -> i + 1
                image1 = video[i, None]
                image2 = video[i + 1, None]
            elif mode == "backward":
                # Flow i + 1 -> i
                image1 = video[i + 1, None]
                image2 = video[i, None]
            else:
                raise NotImplementedError

            _, flow = model(image1, image2, iters=20, test_mode=True)
            flow = flow[0].permute(1, 2, 0).cpu().numpy()
            # resize optical flows
            h, w = flow.shape[:2]
            if h != imgH or w != imgW:
                flow = cv2.resize(flow, (imgW, imgH), cv2.INTER_LINEAR)
                flow[:, :, 0] *= float(imgW) / float(w)
                flow[:, :, 1] *= float(imgH) / float(h)

            Flow = np.concatenate((Flow, flow[..., None]), axis=-1)

            if args.vis_flows:
                # Flow visualization.
                flow_img = utils.flow_viz.flow_to_image(flow)
                flow_img = Image.fromarray(flow_img)

                # Saves the flow and flow_img.
                flow_img.save(
                    os.path.join(args.outroot, "flow", mode + "_png", "%05d.png" % i)
                )
                utils.frame_utils.writeFlow(
                    os.path.join(args.outroot, "flow", mode + "_flo", "%05d.flo" % i),
                    flow,
                )

    return Flow


def extrapolation(args, video_ori, corrFlowF_ori, corrFlowB_ori):
    """Prepares the data for video extrapolation."""
    imgH, imgW, _, nFrame = video_ori.shape

    # Defines new FOV.
    imgH_extr = int(args.H_scale * imgH)
    imgW_extr = int(args.W_scale * imgW)
    imgH_extr = imgH_extr - imgH_extr % 4
    imgW_extr = imgW_extr - imgW_extr % 4
    H_start = int((imgH_extr - imgH) / 2)
    W_start = int((imgW_extr - imgW) / 2)

    # Generates the mask for missing region.
    flow_mask = np.ones(((imgH_extr, imgW_extr)), dtype=np.bool)
    flow_mask[H_start : H_start + imgH, W_start : W_start + imgW] = 0

    mask_dilated = gradient_mask(flow_mask)

    # Extrapolates the FOV for video.
    video = np.zeros(((imgH_extr, imgW_extr, 3, nFrame)), dtype=np.float32)
    video[H_start : H_start + imgH, W_start : W_start + imgW, :, :] = video_ori

    for i in range(nFrame):
        print("Preparing frame {0}".format(i), "\r", end="")
        video[:, :, :, i] = (
            cv2.inpaint(
                (video[:, :, :, i] * 255).astype(np.uint8),
                flow_mask.astype(np.uint8),
                3,
                cv2.INPAINT_TELEA,
            ).astype(np.float32)
            / 255.0
        )

    # Extrapolates the FOV for flow.
    corrFlowF = np.zeros(((imgH_extr, imgW_extr, 2, nFrame - 1)), dtype=np.float32)
    corrFlowB = np.zeros(((imgH_extr, imgW_extr, 2, nFrame - 1)), dtype=np.float32)
    corrFlowF[H_start : H_start + imgH, W_start : W_start + imgW, :] = corrFlowF_ori
    corrFlowB[H_start : H_start + imgH, W_start : W_start + imgW, :] = corrFlowB_ori

    return (
        video,
        corrFlowF,
        corrFlowB,
        flow_mask,
        mask_dilated,
        (W_start, H_start),
        (W_start + imgW, H_start + imgH),
    )


def complete_flow(config, flow_model, flows, flow_masks, mode, device):
    if mode not in ["forward", "backward"]:
        raise NotImplementedError(f"Error flow mode {mode}")
    flow_masks = np.moveaxis(flow_masks, -1, 0)  # [N, H, W]
    flows = np.moveaxis(flows, -1, 0)  # [N, H, W, 2]
    if len(flow_masks.shape) == 3:
        flow_masks = flow_masks[:, :, :, np.newaxis]
    if mode == "forward":
        flow_masks = flow_masks[0:-1]
    else:
        flow_masks = flow_masks[1:]

    num_flows, flow_interval = config["num_flows"], config["flow_interval"]

    diffused_flows = diffusion(flows, flow_masks)

    flows = np2tensor(flows)
    flow_masks = np2tensor(flow_masks)
    diffused_flows = np2tensor(diffused_flows)

    flows = flows.to(device)
    flow_masks = flow_masks.to(device)
    diffused_flows = diffused_flows.to(device)

    t = diffused_flows.shape[2]
    filled_flows = [None] * t
    pivot = num_flows // 2
    for i in range(t):
        indices = indicesGen(i, flow_interval, num_flows, t)
        print("Indices: ", indices, "\r", end="")
        cand_flows = flows[:, :, indices]
        cand_masks = flow_masks[:, :, indices]
        inputs = diffused_flows[:, :, indices]
        pivot_mask = cand_masks[:, :, pivot]
        pivot_flow = cand_flows[:, :, pivot]
        with torch.no_grad():
            output_flow = flow_model(inputs, cand_masks)
        if isinstance(output_flow, tuple) or isinstance(output_flow, list):
            output_flow = output_flow[0]
        comp = output_flow * pivot_mask + pivot_flow * (1 - pivot_mask)
        if filled_flows[i] is None:
            filled_flows[i] = comp
    assert None not in filled_flows
    return filled_flows


def read_flow(flow_dir, video):
    nFrame, _, imgH, imgW = video.shape
    Flow = np.empty(((imgH, imgW, 2, 0)), dtype=np.float32)
    flows = sorted(glob.glob(os.path.join(flow_dir, "*.flo")))
    for flow in flows:
        flow_data = cvbase.read_flow(flow)
        h, w = flow_data.shape[:2]
        flow_data = cv2.resize(flow_data, (imgW, imgH), cv2.INTER_LINEAR)
        flow_data[:, :, 0] *= float(imgW) / float(w)
        flow_data[:, :, 1] *= float(imgH) / float(h)
        Flow = np.concatenate((Flow, flow_data[..., None]), axis=-1)
    return Flow


def norm_flows(flows):
    assert len(flows.shape) == 5, "FLow shape: {}".format(flows.shape)
    flattened_flows = flows.flatten(3)
    flow_max = torch.max(flattened_flows, dim=-1, keepdim=True)[0]
    flows = flows / flow_max.unsqueeze(-1)
    return flows


def save_results(outdir, comp_frames):
    out_dir = os.path.join(outdir, "frames")
    if not os.path.exists(out_dir):
        os.makedirs(out_dir)
    for i in range(len(comp_frames)):
        out_path = os.path.join(out_dir, "{:05d}.png".format(i))
        cv2.imwrite(out_path, comp_frames[i][:, :, ::-1])


def video_inpainting(args, imgArr, imgMaskArr):
    #device = torch.device("cuda:{}".format(args.gpu))
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    print(args)
    if args.opt is not None:
        with open(args.opt, "r") as f:
            opts = yaml.full_load(f)

    for k in opts.keys():
        if k in args:
            setattr(args, k, opts[k])

    print(args)
    # Flow model.
    RAFT_model = initialize_RAFT(args, device)
    # LAFC (flow completion)
    LAFC_model, LAFC_config = initialize_LAFC(args, device)
    # FGT
    FGT_model, FGT_config = initialize_FGT(args, device)

    # Loads frames.
    # filename_list = glob.glob(os.path.join(args.path, '*.png')) + \
    #                glob.glob(os.path.join(args.path, '*.jpg'))

    # Obtains imgH, imgW and nFrame.
    imgH, imgW = args.imgH, args.imgW
    # nFrame = len(filename_list)
    nFrame = len(imgArr)

    if imgH < 350:
        flowH, flowW = imgH * 2, imgW * 2
    else:
        flowH, flowW = imgH, imgW

    # Load video.
    video, video_flow = [], []
    if args.mode == "watermark_removal":
        maskname_list = glob.glob(os.path.join(args.path_mask, "*.png")) + glob.glob(
            os.path.join(args.path_mask, "*.jpg")
        )
        assert len(filename_list) == len(maskname_list)
        for filename, maskname in zip(sorted(filename_list), sorted(maskname_list)):
            frame = (
                torch.from_numpy(np.array(Image.open(filename)).astype(np.uint8))
                .permute(2, 0, 1)
                .float()
                .unsqueeze(0)
            )
            mask = (
                torch.from_numpy(np.array(Image.open(maskname)).astype(np.uint8))
                .permute(2, 0, 1)
                .float()
                .unsqueeze(0)
            )
            mask[mask > 0] = 1
            frame = frame * (1 - mask)
            frame = F2.upsample(
                frame, size=(imgH, imgW), mode="bilinear", align_corners=False
            )
            frame_flow = F2.upsample(
                frame, size=(flowH, flowW), mode="bilinear", align_corners=False
            )
            video.append(frame)
            video_flow.append(frame_flow)
    else:
        """for filename in sorted(filename_list):
        frame = torch.from_numpy(np.array(Image.open(filename)).astype(np.uint8)).permute(2, 0, 1).float().unsqueeze(0)
        frame = F2.upsample(frame, size=(imgH, imgW), mode='bilinear', align_corners=False)
        frame_flow = F2.upsample(frame, size=(flowH, flowW), mode='bilinear', align_corners=False)
        video.append(frame)
        video_flow.append(frame_flow)"""
        for im in imgArr:
            frame = (
                torch.from_numpy(np.array(im).astype(np.uint8))
                .permute(2, 0, 1)
                .float()
                .unsqueeze(0)
            )
            frame = F2.upsample(
                frame, size=(imgH, imgW), mode="bilinear", align_corners=False
            )
            frame_flow = F2.upsample(
                frame, size=(flowH, flowW), mode="bilinear", align_corners=False
            )
            video.append(frame)
            video_flow.append(frame_flow)

    video = torch.cat(video, dim=0)  # [n, c, h, w]
    video_flow = torch.cat(video_flow, dim=0)
    gts = video.clone()
    video = video.to(device)
    video_flow = video_flow.to(device)

    # Calcutes the corrupted flow.
    forward_flows = calculate_flow(
        args, RAFT_model, video_flow, "forward"
    )  # [B, C, 2, N]
    backward_flows = calculate_flow(args, RAFT_model, video_flow, "backward")

    # Makes sure video is in BGR (opencv) format.
    video = (
        video.permute(2, 3, 1, 0).cpu().numpy()[:, :, ::-1, :] / 255.0
    )  # np array -> [h, w, c, N] (0~1)

    if args.mode == "video_extrapolation":
        # Creates video and flow where the extrapolated region are missing.
        (
            video,
            forward_flows,
            backward_flows,
            flow_mask,
            mask_dilated,
            start_point,
            end_point,
        ) = extrapolation(args, video, forward_flows, backward_flows)
        imgH, imgW = video.shape[:2]

        # mask indicating the missing region in the video.
        mask = np.tile(flow_mask[..., None], (1, 1, nFrame))
        flow_mask = np.tile(flow_mask[..., None], (1, 1, nFrame))
        mask_dilated = np.tile(mask_dilated[..., None], (1, 1, nFrame))

    else:
        # Loads masks.
        filename_list = glob.glob(os.path.join(args.path_mask, "*.png")) + glob.glob(
            os.path.join(args.path_mask, "*.jpg")
        )

        mask = []
        mask_dilated = []
        flow_mask = []
        """for filename in sorted(filename_list):
            mask_img = np.array(Image.open(filename).convert('L'))
            mask_img = cv2.resize(mask_img, dsize=(imgW, imgH), interpolation=cv2.INTER_NEAREST)

            if args.flow_mask_dilates > 0:
                flow_mask_img = scipy.ndimage.binary_dilation(mask_img, iterations=args.flow_mask_dilates)
            else:
                flow_mask_img = mask_img
            flow_mask.append(flow_mask_img)

            if args.frame_dilates > 0:
                mask_img = scipy.ndimage.binary_dilation(mask_img, iterations=args.frame_dilates)
            mask.append(mask_img)
            mask_dilated.append(gradient_mask(mask_img))"""

        for f_mask in imgMaskArr:
            mask_img = np.array(f_mask)
            mask_img = cv2.resize(
                mask_img, dsize=(imgW, imgH), interpolation=cv2.INTER_NEAREST
            )

            if args.flow_mask_dilates > 0:
                flow_mask_img = scipy.ndimage.binary_dilation(
                    mask_img, iterations=args.flow_mask_dilates
                )
            else:
                flow_mask_img = mask_img
            flow_mask.append(flow_mask_img)

            if args.frame_dilates > 0:
                mask_img = scipy.ndimage.binary_dilation(
                    mask_img, iterations=args.frame_dilates
                )
            mask.append(mask_img)
            mask_dilated.append(gradient_mask(mask_img))

        # mask indicating the missing region in the video.
        mask = np.stack(mask, -1).astype(np.bool)  # [H, W, C, N]
        mask_dilated = np.stack(mask_dilated, -1).astype(np.bool)
        flow_mask = np.stack(flow_mask, -1).astype(np.bool)

    # Completes the flow.
    videoFlowF = complete_flow(
        LAFC_config, LAFC_model, forward_flows, flow_mask, "forward", device
    )
    videoFlowB = complete_flow(
        LAFC_config, LAFC_model, backward_flows, flow_mask, "backward", device
    )
    videoFlowF = tensor2np(videoFlowF)
    videoFlowB = tensor2np(videoFlowB)
    print("\nFinish flow completion.")

    if args.vis_completed_flows:
        save_flows(args.outroot, videoFlowF, videoFlowB)

    # Prepare gradients
    gradient_x = np.empty(((imgH, imgW, 3, 0)), dtype=np.float32)
    gradient_y = np.empty(((imgH, imgW, 3, 0)), dtype=np.float32)

    for indFrame in range(nFrame):
        img = video[:, :, :, indFrame]
        img[mask[:, :, indFrame], :] = 0
        img = (
            cv2.inpaint(
                (img * 255).astype(np.uint8),
                mask[:, :, indFrame].astype(np.uint8),
                3,
                cv2.INPAINT_TELEA,
            ).astype(np.float32)
            / 255.0
        )

        gradient_x_ = np.concatenate(
            (np.diff(img, axis=1), np.zeros((imgH, 1, 3), dtype=np.float32)), axis=1
        )
        gradient_y_ = np.concatenate(
            (np.diff(img, axis=0), np.zeros((1, imgW, 3), dtype=np.float32)), axis=0
        )
        gradient_x = np.concatenate(
            (gradient_x, gradient_x_.reshape(imgH, imgW, 3, 1)), axis=-1
        )
        gradient_y = np.concatenate(
            (gradient_y, gradient_y_.reshape(imgH, imgW, 3, 1)), axis=-1
        )

        gradient_x[mask_dilated[:, :, indFrame], :, indFrame] = 0
        gradient_y[mask_dilated[:, :, indFrame], :, indFrame] = 0

    gradient_x_filled = gradient_x
    gradient_y_filled = gradient_y
    mask_gradient = mask_dilated
    video_comp = video

    # Gradient propagation.
    gradient_x_filled, gradient_y_filled, mask_gradient = get_flowNN_gradient(
        args,
        gradient_x_filled,
        gradient_y_filled,
        mask,
        mask_gradient,
        videoFlowF,
        videoFlowB,
        None,
        None,
    )

    # if there exist holes in mask, Poisson blending will fail. So I did this trick. I sacrifice some value. Another solution is to modify Poisson blending.
    for indFrame in range(nFrame):
        mask_gradient[:, :, indFrame] = scipy.ndimage.binary_fill_holes(
            mask_gradient[:, :, indFrame]
        ).astype(np.bool)

    # After one gradient propagation iteration
    # gradient --> RGB
    frameBlends = []
    for indFrame in range(nFrame):
        print("Poisson blending frame {0:3d}".format(indFrame))

        if mask[:, :, indFrame].sum() > 0:
            try:
                frameBlend, UnfilledMask = Poisson_blend_img(
                    video_comp[:, :, :, indFrame],
                    gradient_x_filled[:, 0 : imgW - 1, :, indFrame],
                    gradient_y_filled[0 : imgH - 1, :, :, indFrame],
                    mask[:, :, indFrame],
                    mask_gradient[:, :, indFrame],
                )
            except:
                frameBlend, UnfilledMask = (
                    video_comp[:, :, :, indFrame],
                    mask[:, :, indFrame],
                )

            frameBlend = np.clip(frameBlend, 0, 1.0)
            tmp = (
                cv2.inpaint(
                    (frameBlend * 255).astype(np.uint8),
                    UnfilledMask.astype(np.uint8),
                    3,
                    cv2.INPAINT_TELEA,
                ).astype(np.float32)
                / 255.0
            )
            frameBlend[UnfilledMask, :] = tmp[UnfilledMask, :]

            video_comp[:, :, :, indFrame] = frameBlend
            mask[:, :, indFrame] = UnfilledMask

            frameBlend_ = copy.deepcopy(frameBlend)
            # Green indicates the regions that are not filled yet.
            frameBlend_[mask[:, :, indFrame], :] = [0, 1.0, 0]
        else:
            frameBlend_ = video_comp[:, :, :, indFrame]
        frameBlends.append(frameBlend_)

    if args.vis_prop:
        save_fgcp(args.outroot, frameBlends, mask)

    video_length = len(frameBlends)

    for i in range(len(frameBlends)):
        frameBlends[i] = frameBlends[i][:, :, ::-1]

    frames_first = np2tensor(frameBlends, near="t").to(device)
    mask = np.moveaxis(mask, -1, 0)
    mask = mask[:, :, :, np.newaxis]
    masks = np2tensor(mask, near="t").to(device)
    normed_frames = frames_first * 2 - 1
    comp_frames = [None] * video_length

    ref_length = args.step
    num_ref = args.num_ref
    neighbor_stride = args.neighbor_stride

    videoFlowF = np.moveaxis(videoFlowF, -1, 0)

    videoFlowF = np.concatenate([videoFlowF, videoFlowF[-1:, ...]], axis=0)

    flows = np2tensor(videoFlowF, near="t")
    flows = norm_flows(flows).to(device)

    for f in range(0, video_length, neighbor_stride):
        neighbor_ids = [
            i
            for i in range(
                max(0, f - neighbor_stride), min(video_length, f + neighbor_stride + 1)
            )
        ]
        ref_ids = get_ref_index(f, neighbor_ids, video_length, ref_length, num_ref)
        print(f, len(neighbor_ids), len(ref_ids))
        selected_frames = normed_frames[:, neighbor_ids + ref_ids]
        selected_masks = masks[:, neighbor_ids + ref_ids]
        masked_frames = selected_frames * (1 - selected_masks)
        selected_flows = flows[:, neighbor_ids + ref_ids]
        with torch.no_grad():
            filled_frames = FGT_model(masked_frames, selected_flows, selected_masks)
        filled_frames = (filled_frames + 1) / 2
        filled_frames = filled_frames.cpu().permute(0, 2, 3, 1).numpy() * 255
        for i in range(len(neighbor_ids)):
            idx = neighbor_ids[i]
            valid_frame = frames_first[0, idx].cpu().permute(1, 2, 0).numpy() * 255.0
            valid_mask = masks[0, idx].cpu().permute(1, 2, 0).numpy()
            comp = np.array(filled_frames[i]).astype(np.uint8) * valid_mask + np.array(
                valid_frame
            ).astype(np.uint8) * (1 - valid_mask)
            if comp_frames[idx] is None:
                comp_frames[idx] = comp
            else:
                comp_frames[idx] = (
                    comp_frames[idx].astype(np.float32) * 0.5
                    + comp.astype(np.float32) * 0.5
                )
    if args.vis_frame:
        save_results(args.outroot, comp_frames)
    create_dir(args.outroot)
    for i in range(len(comp_frames)):
        comp_frames[i] = comp_frames[i].astype(np.uint8)
    imageio.mimwrite(
        os.path.join(args.outroot, args.outfilename), comp_frames, fps=args.out_fps, quality=8
    )
    print(f"Done, please check your result in {args.outroot} ")


def main(args):
    assert args.mode in (
        "object_removal",
        "video_extrapolation",
        "watermark_removal",
    ), (
        "Accepted modes: 'object_removal', 'video_extrapolation', and 'watermark_removal', but input is %s"
    ) % args.mode
    video_inpainting(args)


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--opt",
        default="configs/object_removal.yaml",
        help="Please select your config file for inference",
    )
    # video completion
    parser.add_argument(
        "--mode",
        default="object_removal",
        choices=["object_removal", "watermark_removal", "video_extrapolation"],
        help="modes: object_removal / video_extrapolation",
    )
    parser.add_argument(
        "--path", default="/myData/davis_resized/walking", help="dataset for evaluation"
    )
    parser.add_argument(
        "--path_mask",
        default="/myData/dilateAnnotations_4/walking",
        help="mask for object removal",
    )
    parser.add_argument(
        "--outroot", default="quick_start/walking3", help="output directory"
    )
    parser.add_argument(
        "--consistencyThres",
        dest="consistencyThres",
        default=5,
        type=float,
        help="flow consistency error threshold",
    )
    parser.add_argument("--alpha", dest="alpha", default=0.1, type=float)
    parser.add_argument("--Nonlocal", dest="Nonlocal", default=False, type=bool)

    # RAFT
    parser.add_argument(
        "--raft_model",
        default="../LAFC/flowCheckPoint/raft-things.pth",
        help="restore checkpoint",
    )
    parser.add_argument("--small", action="store_true", help="use small model")
    parser.add_argument(
        "--mixed_precision", action="store_true", help="use mixed precision"
    )
    parser.add_argument(
        "--alternate_corr",
        action="store_true",
        help="use efficent correlation implementation",
    )

    # LAFC
    parser.add_argument("--lafc_ckpts", type=str, default="../LAFC/checkpoint")

    # FGT
    parser.add_argument("--fgt_ckpts", type=str, default="../FGT/checkpoint")

    # extrapolation
    parser.add_argument(
        "--H_scale", dest="H_scale", default=2, type=float, help="H extrapolation scale"
    )
    parser.add_argument(
        "--W_scale", dest="W_scale", default=2, type=float, help="W extrapolation scale"
    )

    # Image basic information
    parser.add_argument("--imgH", type=int, default=256)
    parser.add_argument("--imgW", type=int, default=432)
    parser.add_argument("--flow_mask_dilates", type=int, default=8)
    parser.add_argument("--frame_dilates", type=int, default=0)

    parser.add_argument("--gpu", type=int, default=0)

    # FGT inference parameters
    parser.add_argument("--step", type=int, default=10)
    parser.add_argument("--num_ref", type=int, default=-1)
    parser.add_argument("--neighbor_stride", type=int, default=5)

    # visualization
    parser.add_argument(
        "--vis_flows", action="store_true", help="Visualize the initialized flows"
    )
    parser.add_argument(
        "--vis_completed_flows",
        action="store_true",
        help="Visualize the completed flows",
    )
    parser.add_argument(
        "--vis_prop",
        action="store_true",
        help="Visualize the frames after stage-I filling (flow guided content propagation)",
    )
    parser.add_argument("--vis_frame", action="store_true", help="Visualize frames")

    args = parser.parse_args()

    main(args)