from collections.abc import Callable

import torch
from diffusers import ModelMixin, SchedulerMixin
from diffusers.configuration_utils import FrozenDict
from diffusers.modular_pipelines import LoopSequentialPipelineBlocks
from diffusers.modular_pipelines.modular_pipeline_utils import ComponentSpec, InputParam, OutputParam

from .guidance import KimodoGuidance
from .operations import ResolvableBlock, checkpoint


class KimodoLoopDenoiser(ResolvableBlock):
    model_name = "kimodo"
    # Text-conditioned and unconditional batches; the unconditional prompt features are zeros.
    _guider_input_fields = {"encoder_hidden_states": ("prompt_embeds", "negative_prompt_embeds")}

    @property
    def description(self):
        return "Step within the denoising loop: run the transformer for each guider batch, then apply guidance."

    @property
    def expected_components(self):
        return self.resolve(
            [
                ComponentSpec("transformer", ModelMixin),
                ComponentSpec(
                    "guider", KimodoGuidance, config=FrozenDict(text_scale=2.0), default_creation_method="from_config"
                ),
            ]
        )

    @property
    def inputs(self):
        return [
            InputParam("prompt_embeds", required=True, type_hint=torch.Tensor, description="LLM2Vec prompt features.")
        ]

    @property
    def intermediate_outputs(self):
        return [OutputParam("noise_pred", type_hint=torch.Tensor, description="Guided clean-motion prediction.")]

    @torch.no_grad()
    def __call__(self, components, block, i, t):
        latents = block.latents
        guider = components.guider
        guider.set_state(step=i, num_inference_steps=len(block.timesteps), timestep=t)
        batches = guider.prepare_inputs_from_block_state(block, self._guider_input_fields)
        timestep = t.reshape(1).repeat(latents.shape[0])
        heading = latents.new_zeros(latents.shape[0])
        for batch in batches:
            guider.prepare_models(components.transformer)
            batch.noise_pred = components.transformer(
                hidden_states=latents,
                timestep=timestep,
                encoder_hidden_states=batch.encoder_hidden_states,
                first_heading_angle=heading,
            ).sample
            guider.cleanup_models(components.transformer)
        block.noise_pred = guider(batches).pred
        return components, block


class KimodoLoopSchedulerStep(ResolvableBlock):
    model_name = "kimodo"

    @property
    def description(self):
        return "Step within the denoising loop: apply one DDIM update to the motion features."

    @property
    def expected_components(self):
        return self.resolve([ComponentSpec("scheduler", SchedulerMixin)])

    @property
    def inputs(self):
        return [InputParam("noise_pred", required=True, type_hint=torch.Tensor)]

    @property
    def intermediate_outputs(self):
        return [OutputParam("latents", type_hint=torch.Tensor, description="Updated normalized motion features.")]

    def __call__(self, components, block, i, t):
        block.latents = components.scheduler.step(block.noise_pred, t, block.latents).prev_sample
        return components, block


class KimodoDenoiseStep(LoopSequentialPipelineBlocks):
    model_name = "kimodo"
    block_classes = [KimodoLoopDenoiser, KimodoLoopSchedulerStep]
    block_names = ["denoiser", "scheduler_step"]

    @property
    def description(self):
        return (
            "Iteratively denoise the motion features over `timesteps`. Each iteration runs:\n"
            " - `KimodoLoopDenoiser`: transformer per guider batch, then guidance\n"
            " - `KimodoLoopSchedulerStep`: DDIM update"
        )

    @property
    def loop_inputs(self):
        return [
            InputParam("latents", type_hint=torch.Tensor, required=True, description="From the latents step."),
            InputParam("timesteps", type_hint=torch.Tensor, required=True, description="From the timesteps step."),
            InputParam(
                "progress_callback",
                type_hint=Callable,
                default=None,
                description="Called after each step with (step, total, stage).",
            ),
            InputParam("cancel_event", default=None, description="Optional threading.Event that cancels the run."),
        ]

    @property
    def loop_intermediate_outputs(self):
        return [OutputParam("latents", type_hint=torch.Tensor, description="Denoised normalized motion features.")]

    @torch.no_grad()
    def __call__(self, components, state):
        block = self.get_block_state(state)
        block.negative_prompt_embeds = torch.zeros_like(block.prompt_embeds)
        for i, t in enumerate(block.timesteps):
            checkpoint(block.cancel_event)
            components, block = self.loop_step(components, block, i=i, t=t)
            if block.progress_callback is not None:
                block.progress_callback(i + 1, len(block.timesteps), "Denoising")
        checkpoint(block.cancel_event)
        self.set_block_state(state, block)
        return components, state