import os
import os.path as osp

import cv2
import numpy as np
import torch
from basicsr.utils import img2tensor, tensor2img
from pytorch_lightning import seed_everything
from ldm.models.diffusion.plms import PLMSSampler
from ldm.modules.encoders.adapter import Adapter
from ldm.util import instantiate_from_config
from model_edge import pidinet
import gradio as gr
from omegaconf import OmegaConf

import pathlib
import random
import shlex
import subprocess
import sys

sys.path.append('T2I-Adapter')

config_path =  'https://github.com/TencentARC/T2I-Adapter/raw/main/configs/stable-diffusion/'
model_path = 'https://github.com/TencentARC/T2I-Adapter/raw/main/models/'

def load_model_from_config(config, ckpt, verbose=False):
    print(f"Loading model from {ckpt}")
    pl_sd = torch.load(ckpt, map_location="cpu")
    if "global_step" in pl_sd:
        print(f"Global Step: {pl_sd['global_step']}")
    if "state_dict" in pl_sd:
        sd = pl_sd["state_dict"]
    else:
        sd = pl_sd
    model = instantiate_from_config(config.model)
    m, u = model.load_state_dict(sd, strict=False)
    # if len(m) > 0 and verbose:
    #     print("missing keys:")
    #     print(m)
    # if len(u) > 0 and verbose:
    #     print("unexpected keys:")
    #     print(u)

    model.cuda()
    model.eval()
    return model

class Model:
    def __init__(self,
                 model_config_path: str = 'ControlNet/models/cldm_v15.yaml',
                 model_dir: str = 'models',
                 use_lightweight: bool = True):
        self.device = torch.device(
            'cuda:0' if torch.cuda.is_available() else 'cpu')    
        self.model_dir = model_dir

        self.download_models()



    def download_models(self) -> None:
        self.model_dir.mkdir(exist_ok=True, parents=True)
        device = 'cuda'
        
        config_file = config_path+'test_sketch.yaml'
        subprocess.run(shlex.split(f'wget {config_file} -O config_sketch.yaml'))    
        config = OmegaConf.load("config_sketch.yaml")
        config.model.params.cond_stage_config.params.device = device

        base_model_file = model_path+"sd-v1-4.ckpt"
        sketch_adapter_file = model_path+"t2iadapter_sketch_sd14v1.pth"
        pidinet_file = model_path+"table5_pidinet.pth"
        
        subprocess.run(shlex.split(f'wget {base_model_file} -O models/sd-v1-4.ckpt'))
        subprocess.run(shlex.split(f'wget {sketch_adapter_file} -O models/t2iadapter_sketch_sd14v1.pth'))
        subprocess.run(shlex.split(f'wget {pidinet_file} -O models/table5_pidinet.pth'))
        
        model = load_model_from_config(config, "models/sd-v1-4.ckpt").to(device)
        current_base = 'sd-v1-4.ckpt'
        model_ad = Adapter(channels=[320, 640, 1280, 1280][:4], nums_rb=2, ksize=1, sk=True, use_conv=False).to(device)
        model_ad.load_state_dict(torch.load("models/t2iadapter_sketch_sd14v1.pth"))
        net_G = pidinet()
        ckp = torch.load('models/table5_pidinet.pth', map_location='cpu')['state_dict']
        net_G.load_state_dict({k.replace('module.',''):v for k, v in ckp.items()})
        net_G.to(device)
        sampler = PLMSSampler(model)
        save_memory=True

    @torch.inference_mode()
    def process_sketch(self, input_img, type_in, color_back, prompt, neg_prompt, fix_sample, scale, con_strength, base_model):    
        global current_base
        if current_base != base_model:
            ckpt = os.path.join("models", base_model)
            pl_sd = torch.load(ckpt, map_location="cpu")
            if "state_dict" in pl_sd:
                sd = pl_sd["state_dict"]
            else:
                sd = pl_sd
            model.load_state_dict(sd, strict=False) #load_model_from_config(config, os.path.join("models", base_model)).to(device)
            current_base = base_model
        con_strength = int((1-con_strength)*50)
        if fix_sample == 'True':
            seed_everything(42)
        
        im = cv2.resize(input_img,(512,512))
    
        if type_in == 'Sketch':
            # net_G = net_G.cpu()
            if color_back == 'White':
                im = 255-im
            im_edge = im.copy()
            im = img2tensor(im)[0].unsqueeze(0).unsqueeze(0)/255.
            # edge = 1-edge # for white background
            im = im>0.5
            im = im.float()
        elif type_in == 'Image':
            im = img2tensor(im).unsqueeze(0)/255.
            im = net_G(im.to(device))[-1]
            im = im>0.5
            im = im.float()
            im_edge = tensor2img(im)
    
        c = model.get_learned_conditioning([prompt])
        nc = model.get_learned_conditioning([neg_prompt])
        
        with torch.no_grad():
            # extract condition features
            features_adapter = model_ad(im.to(device))
    
        shape = [4, 64, 64]
    
        # sampling
        samples_ddim, _ = sampler.sample(S=50,
                                        conditioning=c,
                                        batch_size=1,
                                        shape=shape,
                                        verbose=False,
                                        unconditional_guidance_scale=scale,
                                        unconditional_conditioning=nc,
                                        eta=0.0,
                                        x_T=None,
                                        features_adapter1=features_adapter,
                                        mode = 'sketch',
                                        con_strength = con_strength)
    
        x_samples_ddim = model.decode_first_stage(samples_ddim)
        x_samples_ddim = torch.clamp((x_samples_ddim + 1.0) / 2.0, min=0.0, max=1.0)
        x_samples_ddim = x_samples_ddim.permute(0, 2, 3, 1).numpy()[0]
        x_samples_ddim = 255.*x_samples_ddim
        x_samples_ddim = x_samples_ddim.astype(np.uint8)
    
        return [im_edge, x_samples_ddim]