import argparse
import os
from runpy import run_path
from collections import OrderedDict
import cv2
import torchvision.transforms.functional as TF

import torch

os.environ["CUDA_VISIBLE_DEVICES"] = "0"

parser = argparse.ArgumentParser(description='Demo MPRNet')
parser.add_argument('--input_dir', default='./samples/input/', type=str, help='Input images')
parser.add_argument('--result_dir', default='./samples/output/', type=str, help='Directory for results')
parser.add_argument('--task', default='Denoising',  type=str, help='Task to run', choices=['Deblurring', 'Denoising', 'Deraining'])

args = parser.parse_args()

def load_checkpoint(model, weights):
    checkpoint = torch.load(weights)
    try:
        model.load_state_dict(checkpoint["state_dict"])
    except:
        state_dict = checkpoint["state_dict"]
        new_state_dict = OrderedDict()
        for k, v in state_dict.items():
            name = k[7:] # remove `module.`
            new_state_dict[name] = v
        model.load_state_dict(new_state_dict)

# Load corresponding model architecture and weights
task = args.task
load_file = run_path(os.path.join(task, "MPRNet.py"))
model = load_file['MPRNet']()

weights = os.path.join(task, "pretrained_models", "model_"+task.lower()+".pth")
load_checkpoint(model, weights)
model.eval()
img = cv2.imread('samples/input/4.jpg')
input_ = TF.to_tensor(img).unsqueeze(0)

with torch.no_grad():
    example = torch.ones((1, 3, 128, 128))
    trace_model = torch.jit.trace(model, example)
    r1 = model(input_)
    r2 = trace_model(input_)
    print(r1)
    print(r2)
    trace_model.save('denoise.pt')