import torch
from diffusers import ModelMixin, SchedulerMixin
from diffusers.modular_pipelines.modular_pipeline_utils import ComponentSpec, InputParam, OutputParam
from diffusers.utils.torch_utils import randn_tensor

from .operations import ResolvableBlock


class KimodoLatentsStep(ResolvableBlock):
    model_name = "kimodo"

    @property
    def description(self):
        return "Draw the initial motion-feature noise for an explicit frame count."

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

    @property
    def inputs(self):
        return [
            InputParam(
                "num_frames", type_hint=int, default=120, description="Motion samples at the model frame rate."
            ),
            InputParam(
                "seed",
                type_hint=int,
                default=42,
                description="Seed for a generator on the execution device; ignored when `generator` is given.",
            ),
            InputParam.template("generator"),
        ]

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

    def __call__(self, components, state):
        block = self.get_block_state(state)
        if block.num_frames < 2:
            raise ValueError("Motion requires at least two frames")
        device = components._execution_device
        generator = block.generator
        if generator is None:
            generator = torch.Generator(device=device).manual_seed(block.seed)
        transformer = components.transformer
        block.latents = randn_tensor(
            (1, block.num_frames, transformer.motion_rep_dim),
            generator=generator,
            device=device,
            dtype=transformer.dtype,
        )
        self.set_block_state(state, block)
        return components, state


class KimodoTimestepsStep(ResolvableBlock):
    model_name = "kimodo"

    @property
    def description(self):
        return "Set the Kimodo DDIM schedule for this run."

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

    @property
    def inputs(self):
        return [InputParam.template("num_inference_steps", default=100)]

    @property
    def intermediate_outputs(self):
        return [OutputParam("timesteps", type_hint=torch.Tensor, description="Descending training-step indices.")]

    def __call__(self, components, state):
        block = self.get_block_state(state)
        components.scheduler.set_timesteps(block.num_inference_steps, device=components._execution_device)
        block.timesteps = components.scheduler.timesteps
        self.set_block_state(state, block)
        return components, state