import argparse
from genericpath import isfile
import cv2
import glob
import os
from basicsr.archs.rrdbnet_arch import RRDBNet
from basicsr.utils.download_util import load_file_from_url

from ctsgan import SCANer
from ctsgan.archs.srvgg_arch import SRVGGNetCompact

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('-i', '--input', type=str, default='inputs', help='Input image or folder')
    parser.add_argument('-n', '--model_name', type=str, default='CTSGAN_x4', help=('Model names: CTSGAN_x4 | CTSGAN_x2'))
    parser.add_argument('-o', '--output', type=str, default='output', help='Output folder')
    parser.add_argument('-s', '--outscale', type=float, default=4, help='The final upsampling scale of the image')
    parser.add_argument('--suffix', type=str, default='out', help='Suffix of the restored image')
    parser.add_argument('-t', '--tile', type=int, default=0, help='Tile size, 0 for no tile during testing')
    parser.add_argument('--tile_pad', type=int, default=10, help='Tile padding')
    parser.add_argument('--pre_pad', type=int, default=0, help='Pre padding size at each border')
    parser.add_argument('--face_enhance', action='store_true', help='Use GFPGAN to enhance face')
    parser.add_argument('--half', action='store_true', help='Use fp32 precision during inference. Default: fp16 (half precision).')
    parser.add_argument('--alpha_upsampler',type=str, default='ctsgan',help='The upsampler for the alpha channels. Options: realesrgan | bicubic')
    parser.add_argument('--ext',type=str,default='auto',help='Image extension. Options: auto | jpg | png, auto means using the same extension as inputs')

    args = parser.parse_args()

    #determine models according to model names
    args.model_name = args.model_name.split('.')[0]
    if args.model_name == 'CTSGAN_x4':
        model = RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64, num_block=23, num_grow_ch=32, scale=4)
        netscale = 4
        file_url = ['https://github.com/Sidhved/CTSGAN_Repo/releases/download/v1/CTSGAN_x4.pth']
    elif args.model_name == 'CTSGAN_x2':
        model = RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64, num_block=23, num_grow_ch=32, scale=2)
        netscale = 2
        file_url = ['https://github.com/Sidhved/CTSGAN_Repo/releases/download/v1/CTSGAN_x4.pth']
    elif args.model_name == 'CTSGAN_x1':
        model = RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64, num_block=23, num_grow_ch=32, scale=2)
        netscale = 1
        file_url = ['https://github.com/Sidhved/CTSGAN_Repo/releases/download/v1/CTSGAN_x4.pth']
    
    #determine model paths
    # if args.model_path is None:
    #     model_path = args.model_path
    # else:
    #     model_path = os.path.join('weights', args.model_name + '.pth')
    #     if not os.path.isfile(model_path):
    #         ROOT_DIR = os.path.dirname(os.path.abspath(__file__))
    #         for url in file_url:
    #             model_path = load_file_from_url(url = url, model_dir=os.path.join(ROOT_DIR, 'weights'), progress=True, file_name=None)


    model_path = os.path.join('weights', args.model_name + '.pth')
    # if not os.path.isfile(model_path):
    #     ROOT_DIR = os.path.dirname(os.path.abspath(__file__))
    #     for url in file_url:
    #         # model_path will be updated
    #         model_path = load_file_from_url(
    #             url=url, model_dir=os.path.join(ROOT_DIR, 'weights'), progress=True, file_name=None)

    # restorer
    upsampler = SCANer(
        scale=netscale,
        model_path=model_path,
        model=model,
        tile=args.tile,
        tile_pad=args.tile_pad,
        pre_pad=args.pre_pad,
        half=args.half)

    if args.face_enhance:  # Use GFPGAN for face enhancement
        from gfpgan import GFPGANer
        face_enhancer = GFPGANer(
            model_path='https://github.com/TencentARC/GFPGAN/releases/download/v1.3.0/GFPGANv1.3.pth',
            upscale=args.outscale,
            arch='clean',
            channel_multiplier=2,
            bg_upsampler=upsampler)
    os.makedirs(args.output, exist_ok=True)

    if os.path.isfile(args.input):
        paths = [args.input]
    else:
        paths = sorted(glob.glob(os.path.join(args.input, '*')))

    for idx, path in enumerate(paths):
        imgname, extension = os.path.splitext(os.path.basename(path))
        print('Testing', idx, imgname)

        img = cv2.imread(path, cv2.IMREAD_UNCHANGED)
        if len(img.shape) == 3 and img.shape[2] == 4:
            img_mode = 'RGBA'
        else:
            img_mode = None

        try:
            if args.face_enhance:
                _, _, output = face_enhancer.enhance(img, has_aligned=False, only_center_face=False, paste_back=True)
            else:
                output, _ = upsampler.enhance(img, outscale=args.outscale)
        except RuntimeError as error:
            print('Error', error)
            print('If you encounter CUDA out of memory, try to set --tile with a smaller number.')
        else:
            if args.ext == 'auto':
                extension = extension[1:]
            else:
                extension = args.ext
            if img_mode == 'RGBA':  # RGBA images should be saved in png format
                extension = 'png'
            if args.suffix == '':
                save_path = os.path.join(args.output, f'{imgname}.{extension}')
            else:
                save_path = os.path.join(args.output, f'{imgname}_{args.suffix}.{extension}')
            cv2.imwrite(save_path, output)


if __name__ == '__main__':
    main()