from dataclasses import dataclass, field

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

from .guidance import Pixal3DGuidance
from .resolvable import ResolvableBlock


@dataclass(frozen=True)
class Pixal3DStage:
    """Components and guidance settings of one flow stage."""

    transformer: str
    scheduler: str
    guider: str
    guider_config: dict = field(default_factory=dict)
    sparse: bool = True
    concat_cond: bool = False


STRUCTURE_STAGE = Pixal3DStage(
    "sparse_structure_flow_model",
    "structure_scheduler",
    "structure_guider",
    {"guidance_scale": 7.5, "guidance_rescale": 0.7, "guidance_interval": (0.6, 1.0)},
    sparse=False,
)
SHAPE_GUIDANCE = {"guidance_scale": 7.5, "guidance_rescale": 0.5, "guidance_interval": (0.6, 1.0)}
SHAPE_LOW_STAGE = Pixal3DStage("shape_slat_flow_model_512", "shape_scheduler", "shape_guider", SHAPE_GUIDANCE)
SHAPE_HIGH_STAGE = Pixal3DStage("shape_slat_flow_model_1024", "shape_scheduler", "shape_guider", SHAPE_GUIDANCE)
TEXTURE_STAGE = Pixal3DStage(
    "tex_slat_flow_model_1024",
    "texture_scheduler",
    "texture_guider",
    {"guidance_scale": 1.0, "guidance_rescale": 0.0, "guidance_interval": (0.6, 0.9)},
    concat_cond=True,
)


class Pixal3DPrepareLatentsStep(ResolvableBlock):
    model_name = "pixal3d"

    def __init__(self, stage: Pixal3DStage):
        self.stage = stage
        super().__init__()

    @property
    def description(self):
        return "Draw the initial noise for a dense structure volume or for the sparse shape/texture tokens."

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

    @property
    def inputs(self):
        inputs = [InputParam.template("generator")]
        if self.stage.sparse:
            inputs.append(
                InputParam("coords", required=True, type_hint=torch.Tensor, description="Sparse (N, 4) coordinates.")
            )
        if self.stage.concat_cond:
            inputs.append(
                InputParam("concat_cond", required=True, type_hint=torch.Tensor, description="Shape latents.")
            )
        return inputs

    @property
    def intermediate_outputs(self):
        return [OutputParam("latents", type_hint=torch.Tensor, description="Initial noise for this stage.")]

    @torch.no_grad()
    def __call__(self, components, state):
        b = self.get_block_state(state)
        config = getattr(components, self.stage.transformer).config
        if not self.stage.sparse:
            shape = (1, config.in_channels, *([config.resolution] * 3))
        else:
            if len(b.coords) == 0:
                raise ValueError("Nonempty sparse coordinates are required")
            channels = config.in_channels - (b.concat_cond.shape[-1] if self.stage.concat_cond else 0)
            shape = (1, len(b.coords), channels)
        b.latents = randn_tensor(
            shape, generator=b.generator, device=components._execution_device, dtype=torch.float32
        )
        self.set_block_state(state, b)
        return components, state


class Pixal3DSetTimestepsStep(ResolvableBlock):
    model_name = "pixal3d"

    def __init__(self, stage: Pixal3DStage):
        self.stage = stage
        super().__init__()

    @property
    def description(self):
        return "Set the shifted flow-matching sigmas for this stage."

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

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

    @property
    def intermediate_outputs(self):
        return [OutputParam("timesteps", type_hint=list, description="Sigmas in [0, 1], one per step.")]

    def __call__(self, components, state):
        b = self.get_block_state(state)
        scheduler = getattr(components, self.stage.scheduler)
        scheduler.set_timesteps(b.num_inference_steps)
        b.timesteps = scheduler.timesteps
        self.set_block_state(state, b)
        return components, state


class Pixal3DLoopDenoiser(ResolvableBlock):
    model_name = "pixal3d"
    # Conditional batches use the image features; the unconditional batch uses zeros of the same shape.
    _guider_input_fields = {
        "global_features": ("global_features", "negative_global_features"),
        "projected_features": ("projected_features", "negative_projected_features"),
        "latents": "latents",
    }

    def __init__(self, stage: Pixal3DStage):
        self.stage = stage
        super().__init__()

    @property
    def description(self):
        return "Step within the denoising loop: run the flow model for each guider batch and combine them."

    @property
    def expected_components(self):
        return self.resolve(
            [
                ComponentSpec(self.stage.transformer, ModelMixin),
                ComponentSpec(
                    self.stage.guider,
                    Pixal3DGuidance,
                    default_creation_method="from_config",
                    config=self.stage.guider_config,
                ),
            ]
        )

    @property
    def inputs(self):
        return [
            InputParam("global_features", required=True, type_hint=torch.Tensor, description="DINOv3 global tokens."),
            InputParam(
                "projected_features", required=True, type_hint=torch.Tensor, description="Camera-projected features."
            ),
        ]

    @torch.no_grad()
    def __call__(self, components, block_state, i, t):
        model = getattr(components, self.stage.transformer)
        guider = getattr(components, self.stage.guider)
        guider.set_state(step=i, num_inference_steps=len(block_state.timesteps), timestep=t)
        batches = guider.prepare_inputs_from_block_state(block_state, self._guider_input_fields)
        latents = block_state.latents
        timestep = torch.full((latents.shape[0],), 1000 * t, device=latents.device, dtype=torch.float32)
        for batch in batches:
            guider.prepare_models(model)
            batch.noise_pred = model(
                latents,
                timestep,
                batch.global_features,
                batch.projected_features,
                coords=block_state.model_coords,
                concat_cond=block_state.model_concat_cond,
            )
            guider.cleanup_models(model)
        block_state.noise_pred = guider(batches).pred
        return components, block_state


class Pixal3DLoopAfterDenoiser(ResolvableBlock):
    model_name = "pixal3d"

    def __init__(self, stage: Pixal3DStage):
        self.stage = stage
        super().__init__()

    @property
    def description(self):
        return "Step within the denoising loop: take one Euler step with the guided velocity."

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

    @torch.no_grad()
    def __call__(self, components, block_state, i, t):
        scheduler = getattr(components, self.stage.scheduler)
        prediction = block_state.noise_pred
        block_state.latents = scheduler.step(prediction, t, block_state.latents.to(prediction.device)).prev_sample
        return components, block_state


class Pixal3DDenoiseLoop(LoopSequentialPipelineBlocks):
    model_name = "pixal3d"

    def __init__(self, stage: Pixal3DStage):
        self.stage = stage
        self.block_classes = [Pixal3DLoopDenoiser(stage), Pixal3DLoopAfterDenoiser(stage)]
        self.block_names = ["denoiser", "after_denoiser"]
        super().__init__()

    @property
    def description(self):
        return (
            "Iteratively denoise one stage's latents over `timesteps`. Each iteration runs:\n"
            " - `Pixal3DLoopDenoiser`: flow model per guider batch, then guidance\n"
            " - `Pixal3DLoopAfterDenoiser`: Euler step"
        )

    @property
    def loop_inputs(self):
        inputs = [
            InputParam("timesteps", required=True, type_hint=list, description="From the set-timesteps step."),
            InputParam("latents", required=True, type_hint=torch.Tensor, description="From the prepare-latents step."),
        ]
        if self.stage.sparse:
            inputs.append(InputParam("coords", required=True, type_hint=torch.Tensor))
        if self.stage.concat_cond:
            inputs.append(InputParam("concat_cond", required=True, type_hint=torch.Tensor))
        return inputs

    @property
    def loop_intermediate_outputs(self):
        return [OutputParam("latents", type_hint=torch.Tensor, description="Denoised latents of this stage.")]

    @torch.no_grad()
    def __call__(self, components, state):
        b = self.get_block_state(state)
        device = components._execution_device
        b.latents = b.latents.to(device)
        b.global_features = b.global_features.to(device)
        b.projected_features = b.projected_features.to(device)
        b.negative_global_features = torch.zeros_like(b.global_features)
        b.negative_projected_features = torch.zeros_like(b.projected_features)
        b.model_coords = b.coords[None, :, 1:].to(device) if self.stage.sparse else None
        b.model_concat_cond = b.concat_cond.to(device) if self.stage.concat_cond else None
        with self.progress_bar(total=len(b.timesteps)) as progress_bar:
            for i, t in enumerate(b.timesteps):
                components, b = self.loop_step(components, b, i=i, t=t)
                progress_bar.update()
        self.set_block_state(state, b)
        return components, state