import torch
from diffusers.modular_pipelines.modular_pipeline_utils import InputParam, OutputParam

from .normalization import SHAPE_MEAN, SHAPE_STD, TEXTURE_MEAN, TEXTURE_STD
from .resolvable import ResolvableBlock


class Pixal3DShapeNormalizeStep(ResolvableBlock):
    model_name = "pixal3d"

    @property
    def description(self):
        return "Restore shape latents to the decoder training distribution."

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

    @property
    def intermediate_outputs(self):
        return [OutputParam("shape_latents", type_hint=torch.Tensor, description="Shape latents for the decoder.")]

    def __call__(self, components, state):
        b = self.get_block_state(state)
        b.shape_latents = b.latents * b.latents.new_tensor(SHAPE_STD) + b.latents.new_tensor(SHAPE_MEAN)
        self.set_block_state(state, b)
        return components, state


class Pixal3DTextureConditionStep(ResolvableBlock):
    model_name = "pixal3d"

    @property
    def description(self):
        return "Normalize the shape latents into the texture flow model's concatenated condition."

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

    @property
    def intermediate_outputs(self):
        return [OutputParam("concat_cond", type_hint=torch.Tensor, description="Normalized shape latents.")]

    def __call__(self, components, state):
        b = self.get_block_state(state)
        b.concat_cond = (b.shape_latents - b.shape_latents.new_tensor(SHAPE_MEAN)) / b.shape_latents.new_tensor(
            SHAPE_STD
        )
        self.set_block_state(state, b)
        return components, state


class Pixal3DTextureNormalizeStep(ResolvableBlock):
    model_name = "pixal3d"

    @property
    def description(self):
        return "Restore material latents to the decoder training distribution."

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

    @property
    def intermediate_outputs(self):
        return [
            OutputParam("texture_latents", type_hint=torch.Tensor, description="Material latents for the decoder.")
        ]

    def __call__(self, components, state):
        b = self.get_block_state(state)
        b.texture_latents = b.latents * b.latents.new_tensor(TEXTURE_STD) + b.latents.new_tensor(TEXTURE_MEAN)
        self.set_block_state(state, b)
        return components, state