from pathlib import Path
import sys
from common.color import PIL_CONVERT_COLOR
from common.utils import pil2numpy
from models import LoadCheckpoint
import torch
import numpy as np
import cv2
from PIL import Image
from datetime import datetime
import argparse 

class ImageDemoOptions():
    def __init__(self):
        self.parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)
        default_paths = ["../experiments/last_trained_net.pth","../experiments/last_transfered_lut.pth"]
        default_paths = [f for f in default_paths if Path(f).exists()]
        self.parser.add_argument('--model_paths', '-n', nargs='+', type=str, default=default_paths, help="Model paths for comparison. Example --model_paths ./A/ ./B/")
        self.parser.add_argument('--hr_image_path', '-a', type=str, default="../data/Set14/HR/monarch.png", help="HR image path")
        self.parser.add_argument('--lr_image_path', '-b', type=str, default="../data/Set14/LR/X4/monarch.png", help="LR image path")
        self.parser.add_argument('--output_path', type=str, default="../experiments/", help="Output path.")
        self.parser.add_argument('--batch_size', type=int, default=2**10, help="Size of the batch for the input domain values.")
        self.parser.add_argument('--mirror', action='store_true', default=False)
        self.parser.add_argument('--device', default='cuda', help='Device of the model')
        self.parser.add_argument('--color_model', type=str, default="RGB", help="Color model for input image.") 

    def parse_args(self): 
        args = self.parser.parse_args()       
        args.output_path = Path(args.output_path).resolve() 
        args.hr_image_path = Path(args.hr_image_path).resolve()
        args.lr_image_path = Path(args.lr_image_path).resolve()
        args.model_paths = [Path(x).resolve() for x in args.model_paths]
        args.output_name = args.hr_image_path.stem + "_" + "_".join([Path(f).resolve().stem for f in args.model_paths]) + ".png"
        return args

    def __repr__(self):
        config = self.parse_args() 
        message = ''
        message += '----------------- Options ---------------\n'
        for k, v in sorted(vars(config).items()):
            comment = ''
            default = self.parser.get_default(k)
            if v != default:
                comment = '\t[default: %s]' % str(default)
            message += '{:>25}: {:<30}{}\n'.format(str(k), str(v), comment)
        message += '----------------- End -------------------'
        return message

config_inst = ImageDemoOptions()
config = config_inst.parse_args()

start_script_time = datetime.now()
print(config_inst)
models = [LoadCheckpoint(x).to(config.device) for x in config.model_paths]
for m in models:
    print(m)

lr_image = pil2numpy(PIL_CONVERT_COLOR[config.color_model](Image.open(str(config.lr_image_path))))
image_gt = pil2numpy(PIL_CONVERT_COLOR[config.color_model](Image.open(str(config.hr_image_path))))
image_gt = np.concatenate([image_gt]*3, axis=-1) if image_gt.shape[-1] == 1 else image_gt

if config.mirror:
    lr_image = lr_image[:,::-1,:].copy()
    image_gt = image_gt[:,::-1,:].copy()

input_image = torch.tensor(lr_image).type(torch.float32).permute(2,0,1)[None,...].to(config.device)

predictions = []
for model in models:
    with torch.inference_mode():
        prediction = model(input_image).cpu().type(torch.uint8).squeeze(0).permute(1,2,0).numpy().copy()
        predictions.append(prediction)

image_gt = cv2.putText(image_gt, 'GT', org=(20, 50) , fontFace=cv2.FONT_HERSHEY_SIMPLEX, fontScale=1, color=(255,255,255), thickness=2, lineType= cv2.LINE_AA) 
images_predicted = []
for model_path, model, prediction in zip(config.model_paths, models, predictions):
    prediction = np.concatenate([prediction]*3, axis=-1) if prediction.shape[-1] == 1 else prediction
    prediction = cv2.putText(prediction, model_path.stem, org=(20, 50) , fontFace=cv2.FONT_HERSHEY_SIMPLEX, fontScale=1, color=(255,255,255), thickness=2, lineType= cv2.LINE_AA) 
    images_predicted.append(prediction)

image_count = 1 + len(images_predicted)
t = np.sqrt(image_count).astype(np.int32)
residual = image_count % t
if residual != 0:
    column_count = image_count
    row_count = 1
else:
    column_count = image_count // t
    row_count = t 
images = [image_gt] + images_predicted

columns = []
for i in range(row_count):
    row = []
    for j in range(column_count):
        row.append(images[i*column_count + j])
    columns.append(np.concatenate(row, axis=1))
canvas = np.concatenate(columns, axis=0).astype(np.uint8)
Image.fromarray(canvas).save(config.output_path / config.output_name)

print(f"Saved to {config.output_path / config.output_name}")
print(datetime.now() - start_script_time )