# Copyright 2026 Sperid Labs and The HuggingFace Team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import torch
from diffusers.configuration_utils import FrozenDict
from diffusers.guiders import ClassifierFreeGuidance
from diffusers.modular_pipelines import (
    BlockState,
    LoopSequentialPipelineBlocks,
    ModularPipeline,
    ModularPipelineBlocks,
    PipelineState,
)
from diffusers.modular_pipelines.modular_pipeline_utils import (
    ComponentSpec,
    InputParam,
    OutputParam,
)

from .components import scheduler_spec, transformer_spec


def pack_latents(latents: torch.Tensor, camera_latents: torch.Tensor) -> torch.Tensor:
    """Object latent [B, N, C] and camera latent [B, 1, K] -> one [B, N * C + K] tensor for the guider and scheduler."""
    return torch.cat([latents.flatten(1), camera_latents.flatten(1)], dim=1)


def unpack_latents(packed: torch.Tensor, latents_shape: torch.Size, camera_shape: torch.Size):
    num_latent_values = latents_shape[1] * latents_shape[2]
    latents = packed[:, :num_latent_values].reshape(latents_shape)
    camera_latents = packed[:, num_latent_values:].reshape(camera_shape)
    return latents, camera_latents


class IdolLoopDenoiser(ModularPipelineBlocks):
    model_name = "idol"

    # transformer argument -> (conditional, unconditional) block_state fields
    _guider_input_fields = {
        "encoder_hidden_states": ("image_embeds", "negative_image_embeds"),
        "encoder_hidden_states_vae": ("image_vae_embeds", "negative_image_vae_embeds"),
    }

    @property
    def description(self) -> str:
        return "Step within the denoising loop that predicts the guided velocity of the object and camera latents."

    @property
    def expected_components(self) -> list[ComponentSpec]:
        return [
            ComponentSpec(
                "guider",
                ClassifierFreeGuidance,
                config=FrozenDict({"guidance_scale": 4.0}),
                default_creation_method="from_config",
            ),
            transformer_spec(),
        ]

    @property
    def inputs(self) -> list[InputParam]:
        return [
            InputParam("latents", required=True, type_hint=torch.Tensor, description="The object latent."),
            InputParam("camera_latents", required=True, type_hint=torch.Tensor, description="The camera latent."),
            InputParam("num_inference_steps", required=True, type_hint=int),
            InputParam("image_embeds", required=True, type_hint=torch.Tensor),
            InputParam("image_vae_embeds", required=True, type_hint=torch.Tensor),
            InputParam("negative_image_embeds", type_hint=torch.Tensor),
            InputParam("negative_image_vae_embeds", type_hint=torch.Tensor),
        ]

    @torch.no_grad()
    def __call__(self, components: ModularPipeline, block_state: BlockState, i: int, t: torch.Tensor):
        components.guider.set_state(step=i, num_inference_steps=block_state.num_inference_steps, timestep=t)
        guider_state = components.guider.prepare_inputs_from_block_state(block_state, self._guider_input_fields)

        timestep = t.expand(block_state.latents.shape[0]).to(torch.float32)
        for guider_state_batch in guider_state:
            components.guider.prepare_models(components.transformer)
            cond_kwargs = {name: getattr(guider_state_batch, name) for name in self._guider_input_fields}
            sample, camera = components.transformer(
                hidden_states=block_state.latents,
                camera=block_state.camera_latents,
                timestep=timestep,
                return_dict=False,
                **cond_kwargs,
            )
            guider_state_batch.noise_pred = pack_latents(sample, camera)
            components.guider.cleanup_models(components.transformer)

        block_state.noise_pred = components.guider(guider_state)[0]
        return components, block_state


class IdolLoopAfterDenoiser(ModularPipelineBlocks):
    model_name = "idol"

    @property
    def description(self) -> str:
        return "Step within the denoising loop that takes one Euler step on the object and camera latents."

    @property
    def expected_components(self) -> list[ComponentSpec]:
        return [scheduler_spec()]

    @property
    def intermediate_outputs(self) -> list[OutputParam]:
        return [
            OutputParam("latents", type_hint=torch.Tensor, description="The denoised object latent."),
            OutputParam("camera_latents", type_hint=torch.Tensor, description="The denoised camera latent."),
        ]

    @torch.no_grad()
    def __call__(self, components: ModularPipeline, block_state: BlockState, i: int, t: torch.Tensor):
        packed = pack_latents(block_state.latents, block_state.camera_latents)
        packed = components.scheduler.step(block_state.noise_pred, t, packed, return_dict=False)[0]
        block_state.latents, block_state.camera_latents = unpack_latents(
            packed, block_state.latents.shape, block_state.camera_latents.shape
        )
        return components, block_state


class IdolDenoiseLoopWrapper(LoopSequentialPipelineBlocks):
    model_name = "idol"

    @property
    def description(self) -> str:
        return "Iteratively denoises the object and camera latents over `timesteps`."

    @property
    def loop_expected_components(self) -> list[ComponentSpec]:
        return [scheduler_spec()]

    @property
    def loop_inputs(self) -> list[InputParam]:
        return [
            InputParam("timesteps", required=True, type_hint=torch.Tensor),
            InputParam("num_inference_steps", required=True, type_hint=int),
        ]

    @torch.no_grad()
    def __call__(self, components: ModularPipeline, state: PipelineState) -> PipelineState:
        block_state = self.get_block_state(state)
        with self.progress_bar(total=block_state.num_inference_steps) as progress_bar:
            for i, t in enumerate(block_state.timesteps):
                components, block_state = self.loop_step(components, block_state, i=i, t=t)
                progress_bar.update()
        self.set_block_state(state, block_state)
        return components, state


class IdolDenoiseStep(IdolDenoiseLoopWrapper):
    block_classes = [IdolLoopDenoiser, IdolLoopAfterDenoiser]
    block_names = ["denoiser", "after_denoiser"]

    @property
    def description(self) -> str:
        return (
            "Denoise step that iteratively denoises the object and camera latents. At each step it runs "
            "`IdolLoopDenoiser` then `IdolLoopAfterDenoiser`."
        )