import os
import argparse
import numpy as np
from skimage import color, io
import torch
import torch.nn.functional as F
from PIL import Image
from models import ColorEncoder, ColorUNet
from extractor.manga_panel_extractor import PanelExtractor

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

def mkdirs(path):
    if not os.path.exists(path):
        os.makedirs(path)

def Lab2RGB_out(img_lab):
    img_lab = img_lab.detach().cpu()
    img_l = img_lab[:,:1,:,:]
    img_ab = img_lab[:,1:,:,:]
    img_l = img_l + 50
    pred_lab = torch.cat((img_l, img_ab), 1)[0,...].numpy()
    out = (np.clip(color.lab2rgb(pred_lab.transpose(1, 2, 0)), 0, 1) * 255).astype("uint8")
    return out

def RGB2Lab(inputs):
    return color.rgb2lab(inputs)

def Normalize(inputs):
    l = inputs[:, :, 0:1]
    ab = inputs[:, :, 1:3]
    l = l - 50
    lab = np.concatenate((l, ab), 2)
    return lab.astype('float32')

def numpy2tensor(inputs):
    out = torch.from_numpy(inputs.transpose(2, 0, 1))
    return out

def tensor2numpy(inputs):
    out = inputs[0, ...].detach().cpu().numpy().transpose(1, 2, 0)
    return out

def preprocessing(inputs):
    img_lab = Normalize(RGB2Lab(inputs))
    img = np.array(inputs, 'float32')
    img = numpy2tensor(img)
    img_lab = numpy2tensor(img_lab)
    return img.unsqueeze(0), img_lab.unsqueeze(0)

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("-r", "--reference", type=str, help="ruta de la imagen de referencia")
    parser.add_argument("-o", "--output", type=str, help="carpeta de salida para las imágenes coloreadas")
    parser.add_argument("-ckpt", "--model_checkpoint", type=str, help="ruta del modelo de checkpoint")
    args = parser.parse_args()

    device = "cuda"

    ckpt_path = args.model_checkpoint or 'experiments/Color2Manga_gray/074000_gray.pt'
    test_dir_path = 'test_datasets/gray_test'
    no_extractor = False

    # ... (resto del código)

    while True:
        # ... (resto del código)

        with torch.no_grad():
            img2_resize = F.interpolate(img2 / 255., size=(256, 256), mode='bilinear', recompute_scale_factor=False, align_corners=False)
            img1_L_resize = F.interpolate(img1_lab[:,:1,:,:] / 50., size=(256, 256), mode='bilinear', recompute_scale_factor=False, align_corners=False)

            color_vector = colorEncoder(img2_resize)

            fake_ab = colorUNet((img1_L_resize, color_vector))
            fake_ab = F.interpolate(fake_ab * 110, size=(height, width), mode='bilinear', recompute_scale_factor=False, align_corners=False)

            fake_img = torch.cat((img1_lab[:,:1,:,:], fake_ab), 1)
            fake_img = Lab2RGB_out(fake_img)

            out_folder = os.path.join(output_folder, 'color')
            if not os.path.exists(out_folder):
                os.makedirs(out_folder)
            out_img_path = os.path.join(out_folder, f'{img_name}_color.png')

            # show image
            Image.fromarray(fake_img).show()
            # save image
            io.imsave(out_img_path, fake_img)