import os
import cv2
import argparse
import warnings
import megengine as mge
import megengine.functional as F
warnings.filterwarnings("ignore")

parser = argparse.ArgumentParser(description='Interpolation for a pair of images')
parser.add_argument('--img', dest='img', nargs=2, required=True)
parser.add_argument('--exp', default=4, type=int)
parser.add_argument('--ratio', default=0, type=float, help='inference ratio between two images with 0 - 1 range')
parser.add_argument('--rthreshold', default=0.02, type=float, help='returns image when actual ratio falls in given range threshold')
parser.add_argument('--rmaxcycles', default=8, type=int, help='limit max number of bisectional cycles')
parser.add_argument('--model', dest='modelDir', type=str, default='train_log', help='directory with trained model files')

args = parser.parse_args()


from model.RIFE import Model
model = Model()
model.load_model(args.modelDir, -1)
print("Loaded model")
model.eval()

if args.img[0].endswith('.exr') and args.img[1].endswith('.exr'):
    img0 = cv2.imread(args.img[0], cv2.IMREAD_COLOR | cv2.IMREAD_ANYDEPTH)
    img1 = cv2.imread(args.img[1], cv2.IMREAD_COLOR | cv2.IMREAD_ANYDEPTH)
    img0 = F.expand_dims(mge.Tensor(img0.transpose(2, 0, 1)), 0)
    img1 = F.expand_dims(mge.Tensor(img1.transpose(2, 0, 1)), 0)

else:
    img0 = cv2.imread(args.img[0], cv2.IMREAD_UNCHANGED)
    img1 = cv2.imread(args.img[1], cv2.IMREAD_UNCHANGED)
    img0 = F.expand_dims(mge.Tensor(img0.transpose(2, 0, 1)) / 255. , 0)
    img1 = F.expand_dims(mge.Tensor(img1.transpose(2, 0, 1)) / 255. , 0)

n, c, h, w = img0.shape
ph = ((h - 1) // 32 + 1) * 32
pw = ((w - 1) // 32 + 1) * 32
padding = ((0, 0), (0, 0), (0, ph - h), (0, pw - w))
img0 = F.nn.pad(img0, padding)
img1 = 
F.nn.pad(img1, padding)