# Any4D DiT built on Wan 2.1 Fun 1.3B InP + a canny (or other per-frame) control
# input added at channels [36:52] of the patch_embedding.
#
# Design:
#   - The full InP recipe (gray-padded multi-frame reference latent at [20:36], 4-ch
#     input_mask at [16:20], noisy target at [0:16]) is preserved unchanged from
#     ``Any4DWanDiT``. We inherit and only override the input-reshape helper.
#   - The new [36:52] slot carries the VAE-encoded canny pseudo-RGB latent — same
#     encoding path as the Fun-Control variant, but consumed by the InP DiT's
#     patch_embedding rather than Fun-Control's.
#   - Initialization (handled in ``a4d_pipe_wan_inp_canny.load_checkpoint``): the
#     pretrained Fun-InP checkpoint's ``patch_embedding.weight`` of shape
#     (1536, 36, 1, 2, 2) is zero-padded along the input-channel dim to
#     (1536, 52, 1, 2, 2). At step 0 the canny contribution to the conv output is
#     identically zero, so the model is bit-for-bit equivalent to the pretrained
#     Fun-InP — same ControlNet zero-conv idea applied at the patch_embedding's
#     new input channels.
#
# Per-view stream layout (50 channels — matches the Fun-Control experiment exactly,
# so the dataset side is unchanged):
#     [0:16]   noisy RGB latent
#     [16:17]  input_mask  (1-ch from Any4D)
#     [17:18]  output_mask (1-ch from Any4D, unused at DiT input)
#     [18:34]  Wan reference latent (gray-padded, built by a4d_vae_wan)
#     [34:50]  canny edge_control latent
# Reshaped to the new 52-channel DiT input by the overridden ``_to_wan_36ch``
# (kept the name so the parent ``Any4DWanDiT.forward`` dispatches via MRO):
#     [0:16]   noisy
#     [16:20]  mask (1 -> 4 broadcast)
#     [20:36]  ref
#     [36:52]  canny

from typing import Any

import torch
from custom.wan.a4d_network_wan import Any4DWanDiT


class Any4DWanInPCannyDiT(Any4DWanDiT):
    """Any4DWanDiT with the patch_embedding extended to in_dim=52 for canny control.

    Inherits ``forward()`` from ``Any4DWanDiT`` unchanged. The two deltas:
        * ``WAN_1_3B_INP_KWARGS`` flips ``in_dim`` 36 -> 52 so the patch_embedding
          Conv3d is built with the extra input channels.
        * ``_to_wan_36ch`` is overridden to slice the 50-ch per-view stream into a
          52-ch tensor (noisy + mask + ref + canny). The method name is kept for MRO
          compatibility with the parent's ``forward``; the actual output is 52 ch.
    """

    WAN_1_3B_INP_KWARGS: dict[str, Any] = dict(
        **{**Any4DWanDiT.WAN_1_3B_INP_KWARGS, 'in_dim': 52},
    )

    def _to_wan_36ch(self, x_v_50: torch.Tensor) -> torch.Tensor:
        """Reshape Any4D's 50-ch per-view stream into a 52-ch DiT input.

        Input layout (per-view stream, identical to the Fun-Control experiment):
            x_v_50[:, 0:16]   = noisy rgb VAE latent
            x_v_50[:, 16:17]  = input_mask  (1 at given frames, broadcast read-through)
            x_v_50[:, 17:18]  = output_mask (Any4D bookkeeping, dropped)
            x_v_50[:, 18:34]  = Wan reference latent (gray-padded; built by
                                a4d_vae_wan._build_wan_ref_latent)
            x_v_50[:, 34:50]  = canny VAE latent (VAE-encoded canny pseudo-RGB)

        Output layout (52-ch DiT input — extends Fun InP's 36-ch by appending canny):
            out[:, 0:16]   = noisy
            out[:, 16:20]  = mask (1 -> 4 broadcast)
            out[:, 20:36]  = ref
            out[:, 36:52]  = canny

        Pure tensor slicing + broadcast — no learned parameters here. The pretrained
        Fun-InP patch_embedding weights at [0:36] keep their behavior; the new [36:52]
        slice is zero-initialized at load time (see Any4DWanInPCannyPipeline.
        load_checkpoint) so step 0 reproduces vanilla Fun-InP exactly.
        """
        noisy_slot = x_v_50[:, 0:16]
        imask = x_v_50[:, 16:17]
        ref_slot = x_v_50[:, 18:34]
        canny_slot = x_v_50[:, 34:50]
        mask_slot = imask.expand(-1, 4, -1, -1, -1)
        return torch.cat([noisy_slot, mask_slot, ref_slot, canny_slot], dim=1)


__all__ = ['Any4DWanInPCannyDiT']