# Any4D pipeline for Wan 2.1 Fun-InP base + canny control input.
#
# Differences from Any4DWanPipeline (vanilla Fun-InP):
#   1. Uses Any4DWanInPCannyDiT instead of Any4DWanDiT. Per-view stream layout is
#      50 ch (rgb + 1ch in_mask + 1ch out_mask + 16ch ref + 16ch edge_control);
#      reshaped to the 52 ch DiT input by Any4DWanInPCannyDiT._to_wan_36ch.
#   2. _rewrite_masks_for_wan also zeroes the supervise mask for rgb{N}_edge_control
#      (canny is read-through conditioning, not a denoising target).
#   3. load_checkpoint pads patch_embedding.weight from in_dim=36 -> 52 with zeros
#      on the new [36:52] input channels. At step 0 this makes the new DiT
#      bit-for-bit equivalent to vanilla Fun-InP (ControlNet zero-conv idea applied
#      at the patch_embedding's new input channels).


from typing import Any

import torch
from custom.wan.a4d_pipe_wan import (
    Any4DWanPipeline,
    _remap_wan_dit_state_dict,
    _strip_common_prefixes,
)
from imaginaire.utils import log

_PATCH_EMBED_KEY_CANDIDATES = (
    'wan.patch_embedding.weight',
    'patch_embedding.weight',
)


def _zero_pad_patch_embedding(state_dict: dict[str, torch.Tensor],
                              target_in_dim: int) -> dict[str, torch.Tensor]:
    """Zero-pad ``patch_embedding.weight`` along the input-channel dim to target_in_dim.

    The pretrained Fun-InP checkpoint ships a Conv3d weight of shape
    ``(out=1536, in=36, kT=1, kH=2, kW=2)``. The InP+canny DiT has in_dim=52, so
    we append 16 zero channels at the end of the input-channel axis. This is the
    "ControlNet zero conv" trick: the new conv input channels contribute exactly 0
    to the output at step 0, so the model behaves identically to vanilla Fun-InP
    on InP-recipe inputs; gradients flow into the new channels and they learn the
    canny pathway from scratch.

    Bias is not touched — it is per-output-channel and the output dim is unchanged.
    """
    for key in _PATCH_EMBED_KEY_CANDIDATES:
        if key not in state_dict:
            continue
        weight = state_dict[key]
        if weight.ndim != 5:
            log.warning(f'[InPCanny] {key} has unexpected ndim={weight.ndim}; skipping pad')
            continue
        out_ch, in_ch, kT, kH, kW = weight.shape
        if in_ch == target_in_dim:
            log.info(f'[InPCanny] {key} already at in_dim={target_in_dim}; no pad needed')
            return state_dict
        if in_ch != target_in_dim - 16:
            log.warning(
                f'[InPCanny] {key} has in_ch={in_ch}, expected {target_in_dim - 16} '
                f'(target {target_in_dim}); skipping pad')
            continue
        pad = torch.zeros(out_ch, target_in_dim - in_ch, kT, kH, kW, dtype=weight.dtype)
        state_dict[key] = torch.cat([weight, pad], dim=1)
        log.success(
            f'[InPCanny] zero-padded {key} from in_ch={in_ch} to in_ch={target_in_dim} '
            f'(new {target_in_dim - in_ch} channels init=0)')
        return state_dict
    log.warning(
        f'[InPCanny] no patch_embedding.weight key found in state_dict; '
        f'tried {_PATCH_EMBED_KEY_CANDIDATES}')
    return state_dict


class Any4DWanInPCannyPipeline(Any4DWanPipeline):
    """Pipeline for Any4D + Wan 2.1 Fun-InP base with a canny control extension.

    Inherits the entire Any4DWanPipeline flow (denoise / generate / per-stream noise
    seeding / Wan-native flow-match math). Two overrides:
        * ``_rewrite_masks_for_wan`` also zeros the supervise mask on the new
          ``rgb{N}_edge_control`` slot so canny stays read-through (clean,
          un-supervised) just like the Fun-Control pipeline does.
        * ``load_checkpoint`` pads ``patch_embedding.weight`` from 36 -> 52 input
          channels with zeros before delegating to the parent loader, so the
          pretrained Fun-InP DiT slots straight into our 52-ch patch_embedding.
    """

    # Target in_dim for the InP+canny DiT (must match Any4DWanInPCannyDiT.WAN_1_3B_INP_KWARGS).
    _TARGET_IN_DIM = 52

    def _rewrite_masks_for_wan(
        self,
        x0_streams: dict[str, torch.Tensor],
        masks: dict[str, dict[str, torch.Tensor]],
    ) -> None:
        # Run the InP base rewrite first (rgb noisy slot at all frames, ref read-through).
        super()._rewrite_masks_for_wan(x0_streams, masks)
        # Then zero the supervise mask for the canny slot — same as the Fun-Control pipeline.
        cfg = self.any4d_config
        for v_idx in range(cfg.num_views):
            v_key = f'v{v_idx}'
            if v_key not in x0_streams:
                continue
            entries = cfg.video_entries[v_idx]
            control_k = f'rgb{v_idx}_edge_control'
            if control_k not in entries:
                raise AssertionError(
                    f'Any4DWanInPCannyPipeline view {v_idx} missing required entry '
                    f'{control_k!r} in video_entries.')
            control_s, control_e = entries[control_k][0], entries[control_k][1]
            masks['supervise'][v_key][:, control_s:control_e] = 0.0

    def load_checkpoint(self, ckpt: dict, any4d_config: Any = None) -> None:
        """Pad patch_embedding.weight to in_dim=52, then load via the parent.

        Two checkpoint shapes need handling:
          (a) Raw Fun-InP release with in_dim=36 weights — pad to 52 (zero-init
              the new channels). This is the common case for warm starts.
          (b) Any4D-saved checkpoint already at in_dim=52 — pass through unchanged.

        Detection is shape-based (see _zero_pad_patch_embedding).
        """
        stripped = {_strip_common_prefixes(k): v for k, v in ckpt.items()}
        already_any4d_format = any(k.startswith('wan.') for k in stripped)
        remapped = stripped if already_any4d_format else _remap_wan_dit_state_dict(ckpt)
        remapped = _zero_pad_patch_embedding(remapped, target_in_dim=self._TARGET_IN_DIM)
        missing, surplus = self.dit.load_state_dict(remapped, strict=False, assign=True)
        new_parameter_prefixes = (
            'view_embs_newviews', 'view_emb_existing',
            'sattn_inproj', 'sattn_outproj', 'xattn_inproj', 'action_embedder',
        )
        unexpected_missing = [k for k in missing
                              if not any(k.startswith(p) for p in new_parameter_prefixes)
                              and not k.startswith('wan.')]
        log.warning(f'[Any4DWanInPCannyPipeline] missing (sample): {missing[:8]}')
        log.warning(f'[Any4DWanInPCannyPipeline] surplus (sample): {surplus[:8]}')
        log.info(f'[Any4DWanInPCannyPipeline] missing: {len(missing)}, '
                 f'surplus: {len(surplus)}, unexpected_missing: {len(unexpected_missing)}')


def build_wan_inp_canny_pipeline(
    config: Any,
    dit_path: str = '',
    text_encoder_path: str = '',
    vae_path: str = '',
    device: str = 'cuda',
    torch_dtype: torch.dtype = torch.bfloat16,
    any4d_config: Any = None,
) -> Any4DWanInPCannyPipeline:
    """Drop-in replacement for build_wan_pipeline, building an InP+canny pipeline.

    Mirrors ``build_wan_pipeline`` (Fun-InP) with two deltas:
      * Instantiates Any4DWanInPCannyPipeline (and via config.net, Any4DWanInPCannyDiT).
      * Pads pretrained Fun-InP patch_embedding from in_dim=36 -> 52 at load time.
    """
    from cosmos_predict2.models.utils import init_weights_on_device
    from cosmos_predict2.schedulers.wan_flow_match_scheduler import WanFlowMatchScheduler
    from custom.any4d import a4d_surgery
    from imaginaire.lazy_config import instantiate
    from imaginaire.utils import misc

    any4d_config = a4d_surgery.validate_config(any4d_config)

    pipe = Any4DWanInPCannyPipeline(device=device, torch_dtype=torch_dtype)
    pipe.any4d_config = any4d_config
    pipe.config = config
    pipe.precision = {
        'float32': torch.float32, 'float16': torch.float16, 'bfloat16': torch.bfloat16,
    }[config.precision]
    pipe.tensor_kwargs = {'device': 'cuda', 'dtype': pipe.precision}

    pipe.sigma_data = config.sigma_data
    pipe.setup_data_key()
    pipe.scheduler = WanFlowMatchScheduler(
        shift=5.0,
        sigma_min=0.003 / 1.002,
        sigma_max=1.0,
        num_train_timesteps=1000,
        extra_one_step=True,
    )
    pipe.scaling = None

    # Wan VAE tokenizer (shared with InP / Control — same VAE weights).
    if not config.remove_dit:
        pipe.tokenizer = instantiate(config.tokenizer)
        assert pipe.tokenizer.latent_ch == pipe.config.state_ch, \
            f'Wan VAE latent_ch {pipe.tokenizer.latent_ch} != state_ch {pipe.config.state_ch}'
        pipe.tokenizer.reset_dtype()
        pipe.tokenizer.to(device=device)

    # Text encoder: not loaded — fall back to default_text_wan.pkl (same pattern as
    # vanilla Fun-InP build_wan_pipeline).
    if text_encoder_path and len(text_encoder_path) > 10:
        raise NotImplementedError(
            'UMT5-XXL text encoder is not wired up in this integration. '
            'Set text_encoder_path="" to use default_text_wan.pkl.')
    log.warning('[Wan-InP-Canny] Skipping UMT5 text encoder; will use default_text_wan.pkl')
    pipe.text_encoder = None

    if not config.remove_dit:
        pipe.conditioner = instantiate(config.conditioner)

    pipe.prompt_refiner = None
    pipe.text_guardrail_runner = None
    pipe.video_guardrail_runner = None

    # DiT: instantiate Any4DWanInPCannyDiT via LazyCall `config.net`.
    if not config.remove_dit:
        with init_weights_on_device(device=torch.device('cpu')):
            dit_config = config.net
            dit_config.any4d_config = any4d_config
            with misc.timer('instantiate Any4DWanInPCannyDiT'):
                pipe.dit = instantiate(dit_config).eval()

    if dit_path:
        state_dict = Any4DWanPipeline.gather_wan_checkpoint(dit_path)
        with misc.timer('Any4DWanInPCannyDiT load_checkpoint'):
            pipe.load_checkpoint(state_dict, any4d_config=any4d_config)
        del state_dict
        log.success(f'Loaded Wan 2.1 Fun-InP DiT (+ canny zero-pad) from {dit_path}')

    if not config.remove_dit:
        pipe.dit = pipe.dit.to(device=device, dtype=pipe.precision)
        pipe.dit.eval()

    # EMA: build a frozen copy of the DiT for the training loop to exponentially average
    # into (Karras EDM2 power EMA), mirroring build_wan_control_pipeline. Must exist before
    # the inherited apply_fsdp so it gets sharded + the DTensor EMA worker installed; the
    # inherited on_before_zero_grad then updates it each step, ema_scope() swaps it in for
    # validation, and state_dict() persists it under net_ema.*. Seeded from the base DiT
    # here; the checkpointer reseeds from loaded weights on resume/init.
    if config.ema.enabled and not config.remove_dit:
        import numpy as np

        from imaginaire.utils.ema import FastEmaModelUpdater

        pipe.dit_ema = instantiate(config.net).eval()
        pipe.dit_ema.requires_grad_(False)
        pipe.dit_ema = pipe.dit_ema.to(device=device, dtype=pipe.precision)
        pipe.dit_ema_worker = FastEmaModelUpdater()  # swapped to DTensor variant by apply_fsdp
        s = config.ema.rate
        pipe.ema_exp_coefficient = np.roots([1, 7, 16 - s**-2, 12 - s**-2]).real.max()
        pipe.dit_ema_worker.copy_to(src_model=pipe.dit, tgt_model=pipe.dit_ema)
        log.success(f'Built EMA DiT (rate={s}, ema_exp_coefficient={pipe.ema_exp_coefficient:.4f})')
    else:
        pipe.dit_ema = None

    from megatron.core import parallel_state
    if parallel_state.is_initialized():
        pipe.data_parallel_size = parallel_state.get_data_parallel_world_size()
    else:
        pipe.data_parallel_size = 1

    return pipe