import torch
from diffusers import ModelMixin
from diffusers.modular_pipelines.modular_pipeline_utils import (
    ComponentSpec,
    InputParam,
    OutputParam,
)
from torch.nn import functional as F

from .resolvable import ResolvableBlock


class Pixal3DStructureDecodeStep(ResolvableBlock):
    model_name = "pixal3d"

    @property
    def description(self):
        return "Decode structure latents and extract occupied sparse coordinates."

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

    @property
    def inputs(self):
        return [
            InputParam("latents", required=True, type_hint=torch.Tensor),
            InputParam(
                "structure_resolution",
                default=32,
                type_hint=int,
                description="Occupancy grid size; the decoded 64^3 grid is max-pooled down to it.",
            ),
        ]

    @property
    def intermediate_outputs(self):
        return [OutputParam("coords", type_hint=torch.Tensor)]

    @torch.no_grad()
    def __call__(self, components, state):
        b = self.get_block_state(state)
        model = components.sparse_structure_decoder
        occupancy = model(b.latents.to(model.device)) > 0
        if occupancy.shape[2] != b.structure_resolution:
            if occupancy.shape[2] % b.structure_resolution:
                raise ValueError("Structure resolution must divide the decoded grid size")
            ratio = occupancy.shape[2] // b.structure_resolution
            occupancy = F.max_pool3d(occupancy.float(), ratio, ratio) > 0.5
        b.coords = torch.argwhere(occupancy)[:, [0, 2, 3, 4]].int()
        self.set_block_state(state, b)
        return components, state


class Pixal3DShapeUpsampleStep(ResolvableBlock):
    model_name = "pixal3d"

    @property
    def description(self):
        return "Use the shape decoder to generate the high-resolution sparse coordinate set."

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

    @property
    def inputs(self):
        return [
            InputParam("shape_latents", required=True, type_hint=torch.Tensor),
            InputParam("coords", required=True, type_hint=torch.Tensor),
            InputParam("resolution", default=1024, type_hint=int, description="Output grid resolution: 1024 or 1536."),
            InputParam(
                "max_num_tokens",
                default=49152,
                type_hint=int,
                description="Token budget; 1536 falls back in steps of 128 until the sparse grid fits.",
            ),
        ]

    @property
    def intermediate_outputs(self):
        return [
            OutputParam("coords", type_hint=torch.Tensor),
            OutputParam("actual_resolution", type_hint=int),
        ]

    @torch.no_grad()
    def __call__(self, components, state):
        b = self.get_block_state(state)
        if b.resolution not in (1024, 1536):
            raise ValueError("Supported resolutions are 1024 and 1536")
        model = components.shape_slat_decoder
        coords = model(
            b.shape_latents[0].to(model.device),
            b.coords.to(model.device),
            upsample_times=4,
        )["coords"]
        resolution = b.resolution
        while True:
            grid = resolution // 16
            quantized = torch.cat(
                [
                    coords[:, :1],
                    ((coords[:, 1:] + 0.5) / 512 * (grid - 1)).round().int(),
                ],
                1,
            ).unique(dim=0)
            if len(quantized) < b.max_num_tokens or resolution == 1024:
                break
            resolution -= 128
        b.coords = quantized
        b.actual_resolution = resolution
        self.set_block_state(state, b)
        return components, state


class Pixal3DShapeDecodeStep(ResolvableBlock):
    model_name = "pixal3d"

    @property
    def description(self):
        return "Decode shape latents into a flexible dual grid and texture subdivision guides."

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

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

    @property
    def intermediate_outputs(self):
        return [
            OutputParam("geometry", type_hint=dict),
            OutputParam("subdivisions", type_hint=list),
        ]

    @torch.no_grad()
    def __call__(self, components, state):
        b = self.get_block_state(state)
        model = components.shape_slat_decoder
        decoded = model(
            b.shape_latents[0].to(model.device),
            b.coords.to(model.device),
            return_subs=True,
        )
        features = decoded["features"]
        b.geometry = {
            "coords": decoded["coords"],
            "vertex_offsets": 2 * features[:, :3].sigmoid() - 0.5,
            "intersected": features[:, 3:6] > 0,
            "quad_weights": F.softplus(features[:, 6:7]),
        }
        b.subdivisions = decoded["subdivisions"]
        self.set_block_state(state, b)
        return components, state


class Pixal3DTextureDecodeStep(ResolvableBlock):
    model_name = "pixal3d"

    @property
    def description(self):
        return "Decode material attributes on the shape subdivision grid, without texture baking."

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

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

    @property
    def intermediate_outputs(self):
        return [OutputParam("materials", type_hint=dict)]

    @torch.no_grad()
    def __call__(self, components, state):
        b = self.get_block_state(state)
        model = components.tex_slat_decoder
        guides = [{k: v.to(model.device) for k, v in sub.items()} for sub in b.subdivisions]
        decoded = model(
            b.texture_latents[0].to(model.device),
            b.coords.to(model.device),
            guide_subs=guides,
        )
        b.materials = {
            "coords": decoded["coords"],
            "attributes": decoded["features"] * 0.5 + 0.5,
            "layout": {
                "base_color": [0, 3],
                "metallic": [3, 4],
                "roughness": [4, 5],
                "alpha": [5, 6],
            },
        }
        self.set_block_state(state, b)
        return components, state