import sys
import os.path
import glob
import cv2
import numpy as np
import torch
import architecture as arch

model_path = sys.argv[1]  # models/RRDB_ESRGAN_x4.pth OR models/RRDB_PSNR_x4.pth
device = torch.device('cuda')  # if you want to run on CPU, change 'cuda' -> cpu
# device = torch.device('cpu')

img_source_folder = './LR'
img_output_folder = './results'

model = arch.RRDB_Net(3, 3, 64, 23, gc=32, upscale=4, norm_type=None, act_type='leakyrelu', mode='CNA', res_scale=1,
                      upsample_mode='upconv')
model.load_state_dict(torch.load(model_path), strict=True)
model.eval()
for k, v in model.named_parameters():
    v.requires_grad = False
model = model.to(device)

print('Model path {:s}. \nTesting...'.format(model_path))

idx = 1

for root, dirs, files in os.walk(img_source_folder):
    for dir in dirs:
        dir_to_make = img_output_folder + "/" + root[4:] + "/" + dir
        print("creating subdir:", dir)
        os.makedirs(dir_to_make, exist_ok=True)
    for file in files:
        print("is file, processing...")
        current_image = root + "/" + file
        print(idx, current_image)
        outpath = img_output_folder + "/" + root[4:] + "/" + file
        if os.path.isfile(outpath):
            print("file exists, skipping")
            break
        else:
            print(outpath)
            # read image
            img = cv2.imread(current_image, cv2.IMREAD_COLOR)
            img = img * 1.0 / 255
            img = torch.from_numpy(np.transpose(img[:, :, [2, 1, 0]], (2, 0, 1))).float()
            img_LR = img.unsqueeze(0)
            img_LR = img_LR.to(device)
            output = model(img_LR).data.squeeze().float().cpu().clamp_(0, 1).numpy()
            output = np.transpose(output[[2, 1, 0], :, :], (1, 2, 0))
            output = (output * 255.0).round()
            cv2.imwrite(outpath, output)
            idx += 1
