from pathlib import Path
from re import S
from typing import List, Union

from diffusers import EulerDiscreteScheduler, StableDiffusionXLPipeline
from diffusers.loaders.lora import StableDiffusionXLLoraLoaderMixin
from torchvision.datasets.utils import download_url


class LightningMixin:
    LORA_8_STEP_URL = "https://huggingface.co/ByteDance/SDXL-Lightning/resolve/main/sdxl_lightning_8step_lora.safetensors"

    __scheduler_old = None
    __pipe: StableDiffusionXLPipeline = None
    __scheduler = None

    def configure_sdxl_lightning(self, pipe: StableDiffusionXLPipeline):
        lora_path = Path.home() / ".cache" / "lora_8_step.safetensors"

        download_url(self.LORA_8_STEP_URL, str(lora_path.parent), lora_path.name)

        pipe.load_lora_weights(str(lora_path), adapter_name="8step_lora")
        pipe.set_adapters([])

        self.__scheduler = EulerDiscreteScheduler.from_config(
            pipe.scheduler.config, timestep_spacing="trailing"
        )
        self.__scheduler_old = pipe.scheduler
        self.__pipe = pipe

    def enable_sdxl_lightning(self):
        pipe = self.__pipe
        pipe.scheduler = self.__scheduler

        current = pipe.get_active_adapters()
        current.extend(["8step_lora"])

        weights = self.__find_adapter_weights(current)
        pipe.set_adapters(current, adapter_weights=weights)

        return {"guidance_scale": 0, "num_inference_steps": 8}

    def disable_sdxl_lightning(self):
        pipe = self.__pipe
        pipe.scheduler = self.__scheduler_old

        current = pipe.get_active_adapters()
        current = [adapter for adapter in current if adapter != "8step_lora"]

        weights = self.__find_adapter_weights(current)
        pipe.set_adapters(current, adapter_weights=weights)

    def __find_adapter_weights(self, names: List[str]):
        pipe = self.__pipe

        model = pipe.unet

        from peft.tuners.tuners_utils import BaseTunerLayer

        weights = []
        for adapter_name in names:
            weight = 1.0
            for module in model.modules():
                if isinstance(module, BaseTunerLayer):
                    if adapter_name in module.scaling:
                        weight = (
                            module.scaling[adapter_name]
                            * module.r[adapter_name]
                            / module.lora_alpha[adapter_name]
                        )

            weights.append(weight)

        return weights