import torch
import copy
import os
from diffusers import DDPMScheduler
from transformers import AutoTokenizer, CLIPTextModel
from diffusers import AutoencoderKL, UNet2DConditionModel
from peft import LoraConfig
from S2I.modules.utils import sc_vae_encoder_fwd, sc_vae_decoder_fwd, download_models, get_model_path, get_s2i_home


class RelationShipConvolution(torch.nn.Module):
    def __init__(self, conv_in_pretrained, conv_in_curr, r):
        super(RelationShipConvolution, self).__init__()
        self.conv_in_pretrained = copy.deepcopy(conv_in_pretrained)
        self.conv_in_curr = copy.deepcopy(conv_in_curr)
        self.r = r

    def forward(self, x):
        x1 = self.conv_in_pretrained(x).detach()
        x2 = self.conv_in_curr(x)
        return x1 * (1 - self.r) + x2 * self.r


class PrimaryModel:
    def __init__(self, backbone_diffusion_path='stabilityai/sd-turbo'):
        self.backbone_diffusion_path = backbone_diffusion_path
        self.global_unet = None
        self.global_vae = None
        self.global_tokenizer = None
        self.global_text_encoder = None
        self.global_scheduler = None

    @staticmethod
    def _load_model(path, model_class, unet_mode=False):
        model = model_class.from_pretrained(path, subfolder='unet' if unet_mode else 'vae').to('cuda')
        return model


    def one_step_scheduler(self):
        noise_scheduler_1step = DDPMScheduler.from_pretrained(self.backbone_diffusion_path, subfolder="scheduler")
        noise_scheduler_1step.set_timesteps(1, device="cuda")
        noise_scheduler_1step.alphas_cumprod = noise_scheduler_1step.alphas_cumprod.cuda()
        return noise_scheduler_1step

    def skip_connections(self, vae):
        vae.encoder.forward = sc_vae_encoder_fwd.__get__(vae.encoder, vae.encoder.__class__)
        vae.decoder.forward = sc_vae_decoder_fwd.__get__(vae.decoder, vae.decoder.__class__)
        vae.decoder.skip_conv_1 = torch.nn.Conv2d(512, 512, kernel_size=(1, 1), stride=(1, 1), bias=False).cuda()
        vae.decoder.skip_conv_2 = torch.nn.Conv2d(256, 512, kernel_size=(1, 1), stride=(1, 1), bias=False).cuda()
        vae.decoder.skip_conv_3 = torch.nn.Conv2d(128, 512, kernel_size=(1, 1), stride=(1, 1), bias=False).cuda()
        vae.decoder.skip_conv_4 = torch.nn.Conv2d(128, 256, kernel_size=(1, 1), stride=(1, 1), bias=False).cuda()
        vae.decoder.ignore_skip = False
        return vae
    def weights_adapter(self, p_ckpt, model_name):
        if model_name == '350k-adapter':
            home = get_s2i_home()
            sd_sketch = torch.load(os.path.join(home, f"sketch2image_lora_350k.pkl"), map_location="cpu")
            sd = torch.load(p_ckpt, map_location="cpu")
            sd.update(sd_sketch)
            return sd 
        else:
            sd = torch.load(p_ckpt, map_location="cpu")
            return sd
    def from_pretrained(self, model_name, r):
        if self.global_tokenizer is None:
            # self.global_tokenizer = AutoTokenizer.from_pretrained(self.backbone_diffusion_path,
            #                                                       subfolder="tokenizer")
            self.global_tokenizer = AutoTokenizer.from_pretrained("myn0908/stable-diffusion-3", subfolder="tokenizer_2")

        if self.global_text_encoder is None:
            self.global_text_encoder = CLIPTextModel.from_pretrained(self.backbone_diffusion_path,
                                                                     subfolder="text_encoder").to(device='cuda')

        if self.global_scheduler is None:
            self.global_scheduler = self.one_step_scheduler()

        if self.global_vae is None:
            self.global_vae = self._load_model(self.backbone_diffusion_path, AutoencoderKL)
            self.global_vae = self.skip_connections(self.global_vae)

        if self.global_unet is None:
            self.global_unet = self._load_model(self.backbone_diffusion_path, UNet2DConditionModel, unet_mode=True)
            p_ckpt_path = download_models()
            p_ckpt = get_model_path(model_name=model_name, model_paths=p_ckpt_path)
            sd = self.weights_adapter(p_ckpt, model_name)
            conv_in_pretrained = copy.deepcopy(self.global_unet.conv_in)
            self.global_unet.conv_in = RelationShipConvolution(conv_in_pretrained, self.global_unet.conv_in, r)
            unet_lora_config = LoraConfig(r=sd["rank_unet"], init_lora_weights="gaussian",
                                          target_modules=sd["unet_lora_target_modules"])
            vae_lora_config = LoraConfig(r=sd["rank_vae"], init_lora_weights="gaussian",
                                         target_modules=sd["vae_lora_target_modules"])
            self.global_vae.add_adapter(vae_lora_config, adapter_name="vae_skip")
            _sd_vae = self.global_vae.state_dict()
            for k in sd["state_dict_vae"]:
                _sd_vae[k] = sd["state_dict_vae"][k]
            self.global_vae.load_state_dict(_sd_vae)
            self.global_unet.add_adapter(unet_lora_config)
            _sd_unet = self.global_unet.state_dict()
            for k in sd["state_dict_unet"]:
                _sd_unet[k] = sd["state_dict_unet"][k]
            self.global_unet.load_state_dict(_sd_unet, strict=False)