import os
from typing import Type, Callable, TypeVar, Dict, Any
import torch
import diffusers
from transformers.models.clip.modeling_clip import CLIPTextModel, CLIPTextModelWithProjection


class ENVStore:
    __DESERIALIZER: Dict[Type, Callable[[str,], Any]] = {
        bool: lambda x: bool(int(x)),
        int: int,
        str: lambda x: x,
    }
    __SERIALIZER: Dict[Type, Callable[[Any,], str]] = {
        bool: lambda x: str(int(x)),
        int: str,
        str: lambda x: x,
    }

    def __getattr__(self, name: str):
        value = os.environ.get(f"SDNEXT_OLIVE_{name}", None)
        if value is None:
            return
        ty = self.__class__.__annotations__[name]
        deserialize = self.__DESERIALIZER[ty]
        return deserialize(value)

    def __setattr__(self, name: str, value) -> None:
        if name not in self.__class__.__annotations__:
            return
        ty = self.__class__.__annotations__[name]
        serialize = self.__SERIALIZER[ty]
        os.environ[f"SDNEXT_OLIVE_{name}"] = serialize(value)

    def __delattr__(self, name: str) -> None:
        if name not in self.__class__.__annotations__:
            return
        key = f"SDNEXT_OLIVE_{name}"
        if key not in os.environ:
            return
        os.environ.pop(key)


class OliveOptimizerConfig(ENVStore):
    from_diffusers_cache: bool

    is_sdxl: bool

    vae: str
    vae_sdxl_fp16_fix: bool

    width: int
    height: int
    batch_size: int

    cross_attention_dim: int
    time_ids_size: int


config = OliveOptimizerConfig()


def get_variant():
    from modules.shared import opts

    if opts.diffusers_model_load_variant == 'default':
        from modules import devices

        if devices.dtype == torch.float16:
            return 'fp16'

        return None
    elif opts.diffusers_model_load_variant == 'fp32':
        return None
    else:
        return opts.diffusers_model_load_variant


def get_loader_arguments(no_variant: bool = False):
    kwargs = {}

    if config.from_diffusers_cache:
        from modules.shared import opts
        kwargs["cache_dir"] = opts.diffusers_dir
        if not no_variant:
            kwargs["variant"] = get_variant()

    return kwargs


T = TypeVar("T")
def from_pretrained(cls: Type[T], pretrained_model_name_or_path: os.PathLike, *args, no_variant: bool = False, **kwargs) -> T:
    pretrained_model_name_or_path = str(pretrained_model_name_or_path)
    if pretrained_model_name_or_path.endswith(".onnx"):
        cls = diffusers.OnnxRuntimeModel
        pretrained_model_name_or_path = os.path.dirname(pretrained_model_name_or_path)
    return cls.from_pretrained(pretrained_model_name_or_path, *args, **kwargs, **get_loader_arguments(no_variant))


# -------------------------------------------------------------------------
# Copyright (c) Microsoft Corporation. All rights reserved.
# Licensed under the MIT License.
# --------------------------------------------------------------------------


# Helper latency-only dataloader that creates random tensors with no label
class RandomDataLoader:
    def __init__(self, create_inputs_func, batchsize, torch_dtype):
        self.create_input_func = create_inputs_func
        self.batchsize = batchsize
        self.torch_dtype = torch_dtype

    def __getitem__(self, idx):
        label = None
        return self.create_input_func(self.batchsize, self.torch_dtype), label

# -----------------------------------------------------------------------------
# TEXT ENCODER
# -----------------------------------------------------------------------------


def text_encoder_inputs(batchsize, torch_dtype):
    input_ids = torch.zeros((config.batch_size, 77), dtype=torch_dtype)
    return {
        "input_ids": input_ids,
        "output_hidden_states": True,
    } if config.is_sdxl else input_ids


def text_encoder_load(model_name):
    model = from_pretrained(CLIPTextModel, model_name, subfolder="text_encoder")
    return model


def text_encoder_conversion_inputs(model):
    return text_encoder_inputs(1, torch.int32)


def text_encoder_data_loader(data_dir, batchsize, *_, **__):
    return RandomDataLoader(text_encoder_inputs, config.batch_size, torch.int32)


# -----------------------------------------------------------------------------
# TEXT ENCODER 2
# -----------------------------------------------------------------------------


def text_encoder_2_inputs(batchsize, torch_dtype):
    return {
        "input_ids": torch.zeros((config.batch_size, 77), dtype=torch_dtype),
        "output_hidden_states": True,
    }


def text_encoder_2_load(model_name):
    model = from_pretrained(CLIPTextModelWithProjection, model_name, subfolder="text_encoder_2")
    return model


def text_encoder_2_conversion_inputs(model):
    return text_encoder_2_inputs(1, torch.int64)


def text_encoder_2_data_loader(data_dir, batchsize, *_, **__):
    return RandomDataLoader(text_encoder_2_inputs, config.batch_size, torch.int64)


# -----------------------------------------------------------------------------
# UNET
# -----------------------------------------------------------------------------


def unet_inputs(batchsize, torch_dtype, is_conversion_inputs=False):
    if config.is_sdxl:
        inputs = {
            "sample": torch.rand((2 * config.batch_size, 4, config.height // 8, config.width // 8), dtype=torch_dtype),
            "timestep": torch.rand((1,), dtype=torch_dtype),
            "encoder_hidden_states": torch.rand((2 * config.batch_size, 77, config.cross_attention_dim), dtype=torch_dtype),
        }

        if is_conversion_inputs:
            inputs["additional_inputs"] = {
                "added_cond_kwargs": {
                    "text_embeds": torch.rand((2 * config.batch_size, 1280), dtype=torch_dtype),
                    "time_ids": torch.rand((2 * config.batch_size, config.time_ids_size), dtype=torch_dtype),
                }
            }
        else:
            inputs["text_embeds"] = torch.rand((2 * config.batch_size, 1280), dtype=torch_dtype)
            inputs["time_ids"] = torch.rand((2 * config.batch_size, config.time_ids_size), dtype=torch_dtype)
    else:
        inputs = {
            "sample": torch.rand((config.batch_size, 4, config.height // 8, config.width // 8), dtype=torch_dtype),
            "timestep": torch.rand((config.batch_size,), dtype=torch_dtype),
            "encoder_hidden_states": torch.rand((config.batch_size, 77, config.cross_attention_dim), dtype=torch_dtype),
        }

        # use as kwargs since they won't be in the correct position if passed along with the tuple of inputs
        kwargs = {
            "return_dict": False,
        }
        if is_conversion_inputs:
            inputs["additional_inputs"] = {
                **kwargs,
                "added_cond_kwargs": {
                    "text_embeds": torch.rand((1, 1280), dtype=torch_dtype),
                    "time_ids": torch.rand((1, 5), dtype=torch_dtype),
                },
            }
        else:
            inputs.update(kwargs)
            inputs["onnx::Concat_4"] = torch.rand((1, 1280), dtype=torch_dtype)
            inputs["onnx::Shape_5"] = torch.rand((1, 5), dtype=torch_dtype)

    return inputs


def unet_load(model_name):
    model = from_pretrained(diffusers.UNet2DConditionModel, model_name, subfolder="unet")
    return model


def unet_conversion_inputs(model):
    return tuple(unet_inputs(1, torch.float32, True).values())


def unet_data_loader(data_dir, batchsize, *_, **__):
    return RandomDataLoader(unet_inputs, config.batch_size, torch.float16)


# -----------------------------------------------------------------------------
# VAE ENCODER
# -----------------------------------------------------------------------------


def vae_encoder_inputs(batchsize, torch_dtype):
    return {
        "sample": torch.rand((config.batch_size, 3, config.height, config.width), dtype=torch_dtype),
        "return_dict": False,
    }


def vae_encoder_load(model_name):
    subfolder = "vae_encoder" if os.path.isdir(os.path.join(model_name, "vae_encoder")) else "vae"

    if config.vae_sdxl_fp16_fix:
        model_name = "madebyollin/sdxl-vae-fp16-fix"
        subfolder = ""

    if config.vae is None:
        model = from_pretrained(diffusers.AutoencoderKL, model_name, subfolder=subfolder, no_variant=config.vae_sdxl_fp16_fix)
    else:
        model = diffusers.AutoencoderKL.from_single_file(config.vae)

    model.forward = lambda sample, return_dict: model.encode(sample, return_dict)[0].sample()

    return model


def vae_encoder_conversion_inputs(model):
    return tuple(vae_encoder_inputs(1, torch.float32).values())


def vae_encoder_data_loader(data_dir, batchsize, *_, **__):
    return RandomDataLoader(vae_encoder_inputs, config.batch_size, torch.float16)


# -----------------------------------------------------------------------------
# VAE DECODER
# -----------------------------------------------------------------------------


def vae_decoder_inputs(batchsize, torch_dtype):
    return {
        "latent_sample": torch.rand((config.batch_size, 4, config.height // 8, config.width // 8), dtype=torch_dtype),
        "return_dict": False,
    }


def vae_decoder_load(model_name):
    subfolder = "vae_decoder" if os.path.isdir(os.path.join(model_name, "vae_decoder")) else "vae"

    if config.vae_sdxl_fp16_fix:
        model_name = "madebyollin/sdxl-vae-fp16-fix"
        subfolder = ""

    if config.vae is None:
        model = from_pretrained(diffusers.AutoencoderKL, model_name, subfolder=subfolder, no_variant=config.vae_sdxl_fp16_fix)
    else:
        model = diffusers.AutoencoderKL.from_single_file(config.vae)

    model.forward = model.decode

    return model


def vae_decoder_conversion_inputs(model):
    return tuple(vae_decoder_inputs(1, torch.float32).values())


def vae_decoder_data_loader(data_dir, batchsize, *_, **__):
    return RandomDataLoader(vae_decoder_inputs, config.batch_size, torch.float16)