import ast

import cv2
import torch
import numpy as np

from diffusers import (
    ControlNetModel,
    DDIMScheduler,
    DPMSolverMultistepScheduler,
    EulerDiscreteScheduler,
    EulerAncestralDiscreteScheduler,
    AutoencoderKL,
    LCMScheduler,
    DDPMScheduler,
    DEISMultistepScheduler,
    PNDMScheduler,
    StableDiffusionControlNetPipeline,
    StableDiffusionXLPipeline
)
from transformers import DPTImageProcessor, DPTForDepthEstimation
from .pipeline_calls import get_depth_map, controlnet_call
from .sa_handler import StyleAlignedArgs, Handler
from .inversion import ddim_inversion, make_inversion_callback
import folder_paths
import comfy.model_management as mm


def process_input(input_string):
    # 第一种格式的处理
    if '\n' in input_string and '"' in input_string:
        # 删除多余的引号和换行符，并分割成列表
        cleaned_string = input_string.replace('"', '').replace(',', '')
        result_list = cleaned_string.split('\n')
    # 第二种格式的处理
    elif '"' in input_string:
        result_list = [input_string.replace('"', '')]
    # 第三种格式的处理
    elif '，\n' in input_string:
        result_list = input_string.split('，\n')
    # 当输入不符合任何已知模式时，引发异常
    else:
        return input_string

    # 除去列表中的空白项并去除每个字符串的前后空格
    result_list = [item.strip() for item in result_list if item.strip()]
    return result_list


def convert_preview_image(images):
    # 转换图像为 torch.Tensor，并调整维度顺序为 NHWC
    images_tensors = []
    for img in images:
        # 将 PIL.Image 转换为 numpy.ndarray
        img_array = np.array(img)
        # 转换 numpy.ndarray 为 torch.Tensor
        img_tensor = torch.from_numpy(img_array).float() / 255.
        # 转换图像格式为 CHW (如果需要)
        if img_tensor.ndim == 3 and img_tensor.shape[-1] == 3:
            img_tensor = img_tensor.permute(2, 0, 1)
        # 添加批次维度并转换为 NHWC
        img_tensor = img_tensor.unsqueeze(0).permute(0, 2, 3, 1)
        images_tensors.append(img_tensor)

    if len(images_tensors) > 1:
        output_image = torch.cat(images_tensors, dim=0)
    else:
        output_image = images_tensors[0]
    return output_image


class SchedulerLoader:
    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "scheduler": (
                    [
                        'DPMSolverMultistepScheduler',
                        'DPMSolverMultistepScheduler_SDE_karras',
                        'DDPMScheduler',
                        'DDIMScheduler',
                        'LCMScheduler',
                        'PNDMScheduler',
                        'DEISMultistepScheduler',
                        'EulerDiscreteScheduler',
                        'EulerAncestralDiscreteScheduler'
                    ], {
                        "default": 'DDIMScheduler'
                    }),
                "beta_start": (
                    "FLOAT", {"default": 0.00085, "min": 0.0, "max": 0.001, "step": 0.00005, "display": "slider"}),
                "beta_end": ("FLOAT", {"default": 0.012, "min": 0.0, "max": 0.1, "step": 0.01, "display": "slider"}),
                "beta_schedule": (["scaled_linear"], {"default": "scaled_linear"}),
                "clip_sample": ("BOOLEAN", {"default": False}),
                "set_alpha_to_one": ("BOOLEAN", {"default": False})
            },
        }

    RETURN_TYPES = ("SCHEDULER",)
    RETURN_NAMES = ("scheduler",)
    FUNCTION = "load_scheduler"
    CATEGORY = "StyleAligned"

    def load_scheduler(self,
                       scheduler,
                       beta_start,
                       beta_end,
                       beta_schedule,
                       clip_sample,
                       set_alpha_to_one):
        scheduler_config = {
            'beta_start': beta_start,
            'beta_end': beta_end,
            'beta_schedule': beta_schedule,
            'clip_sample': clip_sample,
            'set_alpha_to_one': set_alpha_to_one
        }
        if scheduler == 'DPMSolverMultistepScheduler':
            noise_scheduler = DPMSolverMultistepScheduler(**scheduler_config)
        elif scheduler == 'DDIMScheduler':
            noise_scheduler = DDIMScheduler(**scheduler_config)
        elif scheduler == 'DPMSolverMultistepScheduler_SDE_karras':
            scheduler_config.update({"algorithm_type": "sde-dpmsolver++"})
            scheduler_config.update({"use_karras_sigmas": True})
            noise_scheduler = DPMSolverMultistepScheduler(**scheduler_config)
        elif scheduler == 'DDPMScheduler':
            noise_scheduler = DDPMScheduler(**scheduler_config)
        elif scheduler == 'LCMScheduler':
            noise_scheduler = LCMScheduler(**scheduler_config)
        elif scheduler == 'PNDMScheduler':
            scheduler_config.update({"set_alpha_to_one": False})
            scheduler_config.update({"trained_betas": None})
            noise_scheduler = PNDMScheduler(**scheduler_config)
        elif scheduler == 'DEISMultistepScheduler':
            noise_scheduler = DEISMultistepScheduler(**scheduler_config)
        elif scheduler == 'EulerDiscreteScheduler':
            noise_scheduler = EulerDiscreteScheduler(**scheduler_config)
        elif scheduler == 'EulerAncestralDiscreteScheduler':
            noise_scheduler = EulerAncestralDiscreteScheduler(**scheduler_config)
        else:
            raise TypeError(f"not support {scheduler}!!!")

        return (noise_scheduler,)


class SASDXL_ModelLoader:
    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "model": (
                    ["stabilityai/stable-diffusion-xl-base-1.0"],
                    {"default": "stabilityai/stable-diffusion-xl-base-1.0"},
                ),
                "scheduler": ("SCHEDULER",)
            },
            "optional": {
                "variant": ("STRING", {"default": "fp16"}),
                "use_safetensors": ("BOOLEAN", {"default": False})
            }
        }

    RETURN_TYPES = ("MODEL",)
    RETURN_NAMES = ("model",)
    FUNCTION = "load_model"
    CATEGORY = "StyleAligned"

    def load_model(self, model, scheduler, variant, use_safetensors):
        device = mm.get_torch_device()
        pipeline = StableDiffusionXLPipeline.from_pretrained(
            model,
            torch_dtype=torch.float16,
            variant=variant,
            use_safetensors=use_safetensors,
            scheduler=scheduler).to(device)
        return (pipeline,)


class SAControlnet_ModelLoader:
    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "model": ([
                              "lllyasviel/sd-controlnet-canny",
                              "diffusers/controlnet-depth-sdxl-1.0"
                          ],
                          {"default": "diffusers/controlnet-depth-sdxl-1.0"}),
                "vae_model": (["madebyollin/sdxl-vae-fp16-fix"], {"default": "madebyollin/sdxl-vae-fp16-fix"}),
                "variant": ("STRING", {"default": "fp16"}),
                "use_safetensors": ("BOOLEAN", {"default": False})
            }
        }

    RETURN_TYPES = ("MODEL",)
    RETURN_NAMES = ("model",)
    FUNCTION = "load_controlnet"
    CATEGORY = "StyleAligned"

    def load_controlnet(self, model, vae_model, variant, use_safetensors):
        device = mm.get_torch_device()
        vae = AutoencoderKL.from_pretrained(vae_model, torch_dtype=torch.float16).to(device)
        controlnet = ControlNetModel.from_pretrained(
            model,
            vae=vae,
            variant=variant,
            use_safetensors=use_safetensors).to(device)
        controlnet.enable_model_cpu_offload()
        return (controlnet,)


class SADepth:
    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "image": ("IMAGE",),
                "model": (
                    ["Intel/dpt-hybrid-midas"],
                    {"default": "Intel/dpt-hybrid-midas"}),
            },
        }

    RETURN_TYPES = ("IMAGE",)
    RETURN_NAMES = ("image",)
    FUNCTION = "depth"
    CATEGORY = "StyleAligned"

    def depth(self, image, model):
        device = mm.get_torch_device()
        depth_estimator = DPTForDepthEstimation.from_pretrained(model).to(device)
        feature_processor = DPTImageProcessor.from_pretrained(model)
        depth_image = get_depth_map(image, feature_processor, depth_estimator)
        return (depth_image,)


class SAInversion:
    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "image": ("IMAGE",),
                "model": ("MODEL"),
                "prompt": ("STRING", {"default": "a photo of an astronaut riding a horse on mars"}),
                "size": ("INT", {"default": 1024, "min": 512, "max": 1024}),
                "num_inference_steps": ("INT", {"default": 50, "min": 20, "max": 50}),
            }
        }

    RETURN_TYPES = ("LATENTS",)
    RETURN_NAMES = ("latents",)
    FUNCTION = "inversion"
    CATEGORY = "StyleAligned"

    def inversion(self, image, model, prompt, size, num_inference_steps):
        x0 = cv2.resize(image, (size, size), interpolation=cv2.INTER_AREA)
        zts = ddim_inversion(
            model,
            x0,
            prompt,
            num_inference_steps,
            2
        )
        return (zts,)


class SAHandler:
    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "share_group_norm": ("BOOLEAN", {"default": False, }),
                "share_layer_norm": ("BOOLEAN", {"default": False, }),
                "share_attention": ("BOOLEAN", {"default": True, }),
                "adain_queries": ("BOOLEAN", {"default": True, }),
                "adain_keys": ("BOOLEAN", {"default": True, }),
                "adain_values": ("BOOLEAN", {"default": False, }),
                "model": ("MODEL",)
            },
            "optional": {
                "shared_score_scale": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 5}),
                "shared_score_shift": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 5}),
            }
        }

    RETURN_TYPES = ("MODEL",)
    RETURN_NAMES = ("model",)
    FUNCTION = "handler"
    CATEGORY = "StyleAligned"

    def handler(self,
                share_group_norm,
                share_layer_norm,
                share_attention,
                adain_queries,
                adain_keys,
                adain_values,
                shared_score_scale,
                shared_score_shift,
                model):
        sa_args = StyleAlignedArgs(
            share_group_norm=share_group_norm,
            share_layer_norm=share_layer_norm,
            share_attention=share_attention,
            adain_queries=adain_queries,
            adain_keys=adain_keys,
            adain_values=adain_values,
            shared_score_scale=shared_score_scale,
            shared_score_shift=shared_score_shift,
        )
        handler = Handler(model)
        handler.register(sa_args, )
        return (model,)


class SASDXLKampler:
    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "model": ("MODEL",),
                "prompt": ("STRING", {"multiline": True, "dynamicPrompts": True, }),
                "num_images_per_prompt": ("INT", {"default": 2, "min": 2, "max": 8, "step": 1, "display": "slider"}),
            },
        }

    RETURN_TYPES = ("IMAGE",)
    RETURN_NAMES = ("image",)
    FUNCTION = "kampler"
    CATEGORY = "StyleAligned"

    def kampler(self,
                model,
                prompt,
                num_images_per_prompt, ):
        prompt = process_input(prompt)

        images = model(
            prompt=prompt,
            num_images_per_prompt=num_images_per_prompt,
        ).images

        output_image = convert_preview_image(images)
        return (output_image,)


class SASDXLTransferKsampler:
    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "model": ("MODEL",),
                "latents": ("LATENTS",),
                "prompt": ("STRING", {"multiline": True, "dynamicPrompts": True, }),
                "num_inference_steps": ("INT", {"default": 50, "min": 20, "max": 50, "disaply": "slider"}),
                "num_images_per_prompt": ("INT", {"default": 2, "min": 2, "max": 8, "step": 1, "display": "slider"}),
            }
        }

    RETURN_TYPES = ("IMAGE",)
    RETURN_NAMES = ("image",)
    FUNCTION = "transfer_kampler"
    CATEGORY = "StyleAligned"

    def transfer_kampler(self, model, latents, prompt, num_inference_steps, num_images_per_prompt):
        zts = latents
        zT, inversion_callback = make_inversion_callback(zts, offset=5)
        g_cpu = torch.Generator(device='cpu')
        g_cpu.manual_seed(10)

        prompt = process_input(prompt)
        latents = torch.randn(len(prompt), 4, 128, 128, device='cpu', generator=g_cpu,
                              dtype=model.unet.dtype, ).to('cuda')
        latents[0] = zT

        images = model(prompt,
                       latents=latents,
                       callback_on_step_end=inversion_callback,
                       num_inference_steps=num_inference_steps,
                       guidance_scale=10.0).images
        output_image = convert_preview_image(images)
        return (output_image,)


class SASDXLControlnetKsampler:
    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "model": ("MODEL",),
                "image": ("IMAGE",),
                "prompt": ("STRING", {"multiline": True, "dynamicPrompts": True, }),
                "reference_prompt": ("STRING", {"multiline": True, "dynamicPrompts": True, }),
                "num_inference_steps": ("INT", {"default": 50, "min": 20, "max": 50, "disaply": "slider"}),
                "num_images_per_prompt": ("INT", {"default": 2, "min": 2, "max": 8, "step": 1, "display": "slider"}),
                "controlnet_conditioning_scale": ("FLOAT", {"default": 0.8, "min": 0.0, "max": 5.0, "step": 0.1})
            }
        }

    RETURN_TYPES = ("IMAGE",)
    RETURN_NAMES = ("image",)
    FUNCTION = "controlnet_kampler"
    CATEGORY = "StyleAligned"

    def controlnet_kampler(self,
                           model,
                           image,
                           prompt,
                           reference_prompt,
                           num_inference_steps,
                           num_images_per_prompt,
                           controlnet_conditioning_scale):
        prompt = process_input(prompt)
        latents = torch.randn(1 + num_images_per_prompt, 4, 128, 128).to(model.unet.dtype)
        images = controlnet_call(
            model,
            [reference_prompt, prompt],
            image=image,
            num_inference_steps=num_inference_steps,
            controlnet_conditioning_scale=controlnet_conditioning_scale,
            num_images_per_prompt=num_images_per_prompt,
            latents=latents
        )

        output_image = convert_preview_image(images)
        return (output_image,)


NODE_CLASS_MAPPINGS = {
    "SASDXL_ModelLoader": SASDXL_ModelLoader,
    "SchedulerLoader": SchedulerLoader,
    "SAControlnet_ModelLoader": SAControlnet_ModelLoader,
    "SADepth": SADepth,
    "SAInversion": SAInversion,
    "SAHandler": SAHandler,
    "SASDXLKampler": SASDXLKampler,
    "SASDXLTransferKsampler": SASDXLTransferKsampler,
    "SASDXLControlnetKsampler": SASDXLControlnetKsampler
}

NODE_DISPLAY_NAME_MAPPINGS = {
    "SASDXL_ModelLoader": "SA SDXL Model Loader",
    "SchedulerLoader": "SA Scheduler Loader",
    "SAControlnet_ModelLoader": "SA Controlnet ModelLoader",
    "SADepth": "SA Depth",
    "SAInversion": "SA Inversion",
    "SAHandler": "SA Handler",
    "SASDXLKampler": "SA SDXL Kampler",
    "SASDXLTransferKsampler": "SA SDXL Transfer Kampler",
    "SASDXLControlnetKsampler": "SA SDXL Controlnet Kampler"
}