# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import copy

from hydra.core.config_store import ConfigStore

from cosmos_predict2.configs.base.config_video2world import (
    PREDICT2_VIDEO2WORLD_PIPELINE_2B,
    PREDICT2_VIDEO2WORLD_PIPELINE_14B,
)
from cosmos_predict2.configs.action_conditioned.config_action_conditioned import (
    ACTION_CONDITIONED_PREDICT2_VIDEO2WORLD_NET_2B,
)
from cosmos_predict2.models.video2world_model import (
    Predict2ModelManagerConfig,
    Predict2Video2WorldModel,
    Predict2Video2WorldModelConfig,
)
from custom.any4d.a4d_config import Any4DConfig
from custom.any4d.a4d_model import VidarModel, Any4DModel
from custom.wan.a4d_config_wan import Any4DWanConfig
from custom.wan.a4d_model_wan import Any4DWanModel
from custom.any4d.a4d_network import Any4DDiT
from imaginaire.lazy_config import LazyCall as L


def make_any4d_dit(*args, any4d_config=None, **kwargs):
    # Dispatch to the legacy network snapshot when any4d_config.legacy_network_behavior == 2,
    if any4d_config is not None and getattr(any4d_config, 'legacy_network_behavior', 0) == 2:
        from custom.legacy.a4d_network_tembv import Any4DDiT as DiTCls
    else:
        DiTCls = Any4DDiT
    return DiTCls(*args, any4d_config=any4d_config, **kwargs)


PREDICT2_VIDEO2WORLD_FSDP_2B = dict(
    trainer=dict(
        distributed_parallelism="fsdp",
    ),
    model=L(Predict2Video2WorldModel)(
        config=Predict2Video2WorldModelConfig(
            pipe_config=PREDICT2_VIDEO2WORLD_PIPELINE_2B,
            model_manager_config=L(Predict2ModelManagerConfig)(
                dit_path="checkpoints/nvidia/Cosmos-Predict2-2B-Video2World/model-720p-16fps.pt",
                text_encoder_path="",  # Do not load text encoder for training.
            ),
            fsdp_shard_size=8,
        ),
        _recursive_=False,
    ),
)

# LFV: Custom model to use with arbitrary networks (e.g., VGGT)
from custom.regression.regression_model import ModelRegression
PREDICT2_VIDEO2WORLD_FSDP_2B_REGRESSION = dict(
    trainer=dict(
        distributed_parallelism="fsdp",
    ),
    model=L(ModelRegression)(
        config=Predict2Video2WorldModelConfig(
            pipe_config=PREDICT2_VIDEO2WORLD_PIPELINE_2B,
            model_manager_config=L(Predict2ModelManagerConfig)(
                dit_path="checkpoints/nvidia/Cosmos-Predict2-2B-Video2World/model-720p-16fps.pt",
                text_encoder_path="",
            ),
            fsdp_shard_size=8,
        ),
        _recursive_=False,
    ),
)

PREDICT2_VIDEO2WORLD_FSDP_14B = dict(
    trainer=dict(
        distributed_parallelism="fsdp",
    ),
    model=L(Predict2Video2WorldModel)(
        config=Predict2Video2WorldModelConfig(
            pipe_config=PREDICT2_VIDEO2WORLD_PIPELINE_14B,
            model_manager_config=L(Predict2ModelManagerConfig)(
                dit_path="checkpoints/nvidia/Cosmos-Predict2-14B-Video2World/model-720p-16fps.pt",
                text_encoder_path="",  # Do not load text encoder for training.
            ),
            fsdp_shard_size=8,
        ),
        _recursive_=False,
    ),
)


# NOTE(bvh): This is the stepping stone from pure vanilla functionality to full Any4D:
# Run the default, untouched Cosmos model, but apply vidar data loading & processing & Any4D visualizations.
VIDAR_FSDP_2B = dict(
    trainer=dict(
        distributed_parallelism="fsdp",
    ),
    model=L(VidarModel)(
        config=Predict2Video2WorldModelConfig(
            vidar_active=True,
            any4d_active=False,
            pipe_config=PREDICT2_VIDEO2WORLD_PIPELINE_2B,
            model_manager_config=L(Predict2ModelManagerConfig)(
                dit_path="checkpoints/nvidia/Cosmos-Predict2-2B-Video2World/model-480p-10fps.pt",
                text_encoder_path="",
            ),
            fsdp_shard_size=8,
        ),
        _recursive_=False,
    ),
)


# NOTE(bvh): These fields mirror / are modified versions of:
# - cosmos_predict2/configs/base/defaults/model.py:PREDICT2_VIDEO2WORLD_FSDP_2B
# - cosmos_predict2/configs/base/config_video2world.py:PREDICT2_VIDEO2WORLD_PIPELINE_2B
# - cosmos_predict2/configs/action_conditioned/config_action_conditioned.py:ACTION_CONDITIONED_PREDICT2_VIDEO2WORLD_NET_2B
# and are instantiated in / controlled by:
# - register_model() below
# - custom/experiment/basile.py:override /model

ANY4D_NET_2B = copy.copy(ACTION_CONDITIONED_PREDICT2_VIDEO2WORLD_NET_2B)
ANY4D_NET_2B._target_ = make_any4d_dit
ANY4D_NET_2B.action_dim = -1  # must be overwritten later, this default value means invalid

# Use stateless checkpointing policy for Any4D - the default "predict2_2b_720" policy counts
# operations with hardcoded thresholds that don't work with Any4DBlock's variable / different
# op count whenever extra_sattn is active.
# ANY4D_NET_2B.sac_config.mode = "mm_only"

ANY4D_PIPELINE_2B = copy.copy(PREDICT2_VIDEO2WORLD_PIPELINE_2B)
ANY4D_PIPELINE_2B.net = ANY4D_NET_2B
ANY4D_PIPELINE_2B.any4d_active = True
# Any4D bypasses text2image_dit.prepare_embedded_sequence (its own net handles embedding), and
# in_channels_override skips the +1 padding_mask channel at patch-embed construction. The
# padding_mask embedder in the conditioner is therefore unused; drop it so data batches don't
# need to carry a dummy tensor.
ANY4D_PIPELINE_2B.conditioner = copy.deepcopy(PREDICT2_VIDEO2WORLD_PIPELINE_2B.conditioner)
del ANY4D_PIPELINE_2B.conditioner.padding_mask

S3_PRETRAINED_PREFIX = r's3://tri-ml-sandbox-16011-us-west-2-datasets/cosmos-predict-2/checkpoints'

ANY4D_FSDP_2B = dict(
    trainer=dict(
        distributed_parallelism="fsdp",
    ),
    model=L(Any4DModel)(  # formerly Predict2Video2WorldModel
        config=Any4DConfig(  # formerly Predict2Video2WorldModelConfig
            vidar_active=True,
            any4d_active=True,
            pipe_config=ANY4D_PIPELINE_2B,  # formerly PREDICT2_VIDEO2WORLD_PIPELINE_2B
            model_manager_config=L(Predict2ModelManagerConfig)(
                dit_path=f"{S3_PRETRAINED_PREFIX}/nvidia/Cosmos-Predict2-2B-Video2World/model-480p-10fps.pt",
                text_encoder_path="",
            ),
            fsdp_shard_size=8,
        ),
        _recursive_=False,
    ),
)


# ============================================================
# Any4D on top of Wan 2.1 Fun 1.3B InP (image-to-video base)
# ============================================================
# Mirrors ANY4D_FSDP_2B but swaps:
#   - DiT:       Any4DWanDiT  (custom/wan/a4d_network_wan.py)
#   - Pipeline:  Any4DWanPipeline (custom/wan/a4d_pipe_wan.py)
#   - Tokenizer: WanVAEAdapter (custom/wan/wan_vae_adapter.py)
#   - VAE wrapper: custom/wan/a4d_vae_wan.py (via vae='a4d_vae_wan' in experiment)
# The existing Any4DConfig (masks, streams, multi-view) is kept — only the backbone changes.
from copy import deepcopy as _deepcopy
from custom.wan.a4d_network_wan import Any4DWanDiT
from custom.wan.wan_vae_adapter import WanVAEAdapter

ANY4D_WAN_NET_1_3B = L(Any4DWanDiT)(
    any4d_config=None,  # will be set by pipeline at instantiation time
)

ANY4D_WAN_PIPELINE_1_3B = _deepcopy(PREDICT2_VIDEO2WORLD_PIPELINE_2B)
ANY4D_WAN_PIPELINE_1_3B.net = ANY4D_WAN_NET_1_3B
ANY4D_WAN_PIPELINE_1_3B.any4d_active = True
# Wan VAE tokenizer (16 latent channels, temporal compression 4, spatial compression 8).
ANY4D_WAN_PIPELINE_1_3B.tokenizer = L(WanVAEAdapter)(
    vae_pth="",            # populated at runtime from vae_path
    chunk_duration=41,     # raw frames per VAE chunk
    name="wan_vae",
    load_mean_std=False,
    dtype="bfloat16",
)
ANY4D_WAN_PIPELINE_1_3B.state_ch = 16   # Wan VAE latent channels
ANY4D_WAN_PIPELINE_1_3B.state_t = 11    # 41 pixel frames -> 11 latent frames

# Disable the conditioner's TextAttr dropout (which would multiply embeddings by 0).
# `Any4DVAEWan.encode_text_as_needed` handles training-time text dropout itself by
# replacing dropped samples with the cached UMT5("") embedding from
# default_text_wan.pkl — matching the inference uncond branch
# (`is_negative_prompt=True`). Keeping the conditioner dropout on would multiply
# those UMT5("") tensors by 0 for ~4% of samples, producing zero-tensor uncond
# again and undoing the fix.
ANY4D_WAN_PIPELINE_1_3B.conditioner.text.dropout_rate = 0.0

ANY4D_WAN_FSDP_1_3B = dict(
    trainer=dict(
        distributed_parallelism="fsdp",
    ),
    # Any4DWanModel (not Any4DModel) so training goes through Wan-native flow-match:
    #   - σ sampled via WanFlowMatchScheduler.sample_sigma (shift-transformed uniform s∈[0,1])
    #   - forward: yt = (1-s)*x0 + s*ε  (flow matching, not EDM's yt = x0 + σ*ε)
    #   - unit loss weighting
    #   - video_noise_multiplier disabled; high_sigma_ratio disabled
    model=L(Any4DWanModel)(
        config=Any4DWanConfig(
            vidar_active=True,
            any4d_active=True,
            use_wan_backbone=True,       # routes Predict2Video2WorldModel.__init__ to Wan
            pipe_config=ANY4D_WAN_PIPELINE_1_3B,
            model_manager_config=L(Predict2ModelManagerConfig)(
                # Expected to point to a directory containing
                #   diffusion_pytorch_model*.safetensors (the Wan 2.1 Fun 1.3B InP DiT weights)
                dit_path="",
                text_encoder_path="",    # disabled -> use default_text_wan.pkl
                # Expected to point to a path or directory containing Wan2.1_VAE.pth
                vae_path="",
                tokenizer_chunk_duration=41,
            ),
            fsdp_shard_size=8,
            vae="a4d_vae_wan",
            # Cosmos-era σ tweaks (`adjust_video_noise`, `high_sigma_ratio`) don't apply:
            # Any4DWanModel forces `video_noise_multiplier=1.0` and overrides
            # `draw_training_sigma_and_epsilon`, so those config values are never read.
        ),
        _recursive_=False,
    ),
)


# ============================================================
# Any4D on top of Wan 2.1 Fun-V1.1-1.3B-Control (text + canny -> video)
# ============================================================
# Same shape as ANY4D_WAN_*_1_3B but with the Control DiT (in_dim=48, ref_conv)
# and wan_variant='control' so the model dispatch routes through
# build_wan_control_pipeline. Stream-level `*_edge_control` (and future
# `*_depth_control`, `*_seg_control`, ...) entries land at a per-view channel
# slot defined by the experiment config.
from custom.wan.a4d_network_wan_control import Any4DWanControlDiT

ANY4D_WAN_CONTROL_NET_1_3B = L(Any4DWanControlDiT)(
    any4d_config=None,
    view_emb_std=0.06,
)

ANY4D_WAN_CONTROL_PIPELINE_1_3B = _deepcopy(PREDICT2_VIDEO2WORLD_PIPELINE_2B)
ANY4D_WAN_CONTROL_PIPELINE_1_3B.net = ANY4D_WAN_CONTROL_NET_1_3B
ANY4D_WAN_CONTROL_PIPELINE_1_3B.any4d_active = True
ANY4D_WAN_CONTROL_PIPELINE_1_3B.tokenizer = L(WanVAEAdapter)(
    vae_pth="", chunk_duration=41, name="wan_vae",
    load_mean_std=False, dtype="bfloat16",
)
ANY4D_WAN_CONTROL_PIPELINE_1_3B.state_ch = 16
ANY4D_WAN_CONTROL_PIPELINE_1_3B.state_t = 11
# Same TextAttr-dropout-off rule as the InP variant — a4d_vae_wan handles dropout.
ANY4D_WAN_CONTROL_PIPELINE_1_3B.conditioner.text.dropout_rate = 0.0

ANY4D_WAN_CONTROL_FSDP_1_3B = dict(
    trainer=dict(
        distributed_parallelism="fsdp",
    ),
    model=L(Any4DWanModel)(
        config=Any4DWanConfig(
            vidar_active=True,
            any4d_active=True,
            use_wan_backbone=True,
            wan_variant='control',
            pipe_config=ANY4D_WAN_CONTROL_PIPELINE_1_3B,
            model_manager_config=L(Predict2ModelManagerConfig)(
                dit_path="",
                text_encoder_path="",
                vae_path="",
                tokenizer_chunk_duration=41,
            ),
            fsdp_shard_size=8,
            vae="a4d_vae_wan",
        ),
        _recursive_=False,
    ),
)


# ============================================================
# Any4D on top of Wan 2.1 Fun-InP base + canny control extension (in_dim=52)
# ============================================================
# Same shape as ANY4D_WAN_*_1_3B (Fun-InP) but with the InP DiT's patch_embedding
# extended from in_dim=36 to in_dim=52. The new [36:52] input channels carry the
# VAE-encoded canny pseudo-RGB latent (same encoding path as the Fun-Control
# variant). At checkpoint load time the pretrained Fun-InP patch_embedding weights
# are zero-padded along the input-channel dim, so at step 0 the model is
# bit-for-bit equivalent to vanilla Fun-InP (ControlNet zero-conv idea applied at
# the patch_embedding's new input channels). Dispatched by wan_variant='inp_canny'.
from custom.wan.a4d_network_wan_inp_canny import Any4DWanInPCannyDiT

ANY4D_WAN_INP_CANNY_NET_1_3B = L(Any4DWanInPCannyDiT)(
    any4d_config=None,
)

ANY4D_WAN_INP_CANNY_PIPELINE_1_3B = _deepcopy(PREDICT2_VIDEO2WORLD_PIPELINE_2B)
ANY4D_WAN_INP_CANNY_PIPELINE_1_3B.net = ANY4D_WAN_INP_CANNY_NET_1_3B
ANY4D_WAN_INP_CANNY_PIPELINE_1_3B.any4d_active = True
ANY4D_WAN_INP_CANNY_PIPELINE_1_3B.tokenizer = L(WanVAEAdapter)(
    vae_pth="", chunk_duration=41, name="wan_vae",
    load_mean_std=False, dtype="bfloat16",
)
ANY4D_WAN_INP_CANNY_PIPELINE_1_3B.state_ch = 16
ANY4D_WAN_INP_CANNY_PIPELINE_1_3B.state_t = 11
# Same TextAttr-dropout-off rule as the InP and Control variants — a4d_vae_wan
# handles dropout (replacing dropped samples with the cached UMT5("") embedding).
ANY4D_WAN_INP_CANNY_PIPELINE_1_3B.conditioner.text.dropout_rate = 0.0

ANY4D_WAN_INP_CANNY_FSDP_1_3B = dict(
    trainer=dict(
        distributed_parallelism="fsdp",
    ),
    model=L(Any4DWanModel)(
        config=Any4DWanConfig(
            vidar_active=True,
            any4d_active=True,
            use_wan_backbone=True,
            wan_variant='inp_canny',
            pipe_config=ANY4D_WAN_INP_CANNY_PIPELINE_1_3B,
            model_manager_config=L(Predict2ModelManagerConfig)(
                dit_path="",
                text_encoder_path="",
                vae_path="",
                tokenizer_chunk_duration=41,
            ),
            fsdp_shard_size=8,
            vae="a4d_vae_wan",
        ),
        _recursive_=False,
    ),
)


# ============================================================
# Any4D on top of the pure Wan 2.1 T2V-1.3B base (text -> image/video)
# ============================================================
# Same shape as ANY4D_WAN_*_1_3B but with the pure T2V DiT (in_dim=16,
# has_image_input=False — no InP mask/ref slot, no CLIP image branch) and
# wan_variant='t2v'. Per-view streams carry only the noisy rgb latent (+ Any4D mask
# bookkeeping); there is no reference/control slot. Routes through the standard
# build_wan_pipeline (Any4DWanPipeline) — the Fun-Control / InP+canny dispatch
# branches in video2world_model.py are skipped because wan_variant is not
# 'control'/'inp_canny'. Reference checkpoint: Wan-AI/Wan2.1-T2V-1.3B.
from custom.wan.a4d_network_wan_t2v import Any4DWanT2VDiT

ANY4D_WAN_T2V_NET_1_3B = L(Any4DWanT2VDiT)(
    any4d_config=None,
)

ANY4D_WAN_T2V_PIPELINE_1_3B = _deepcopy(PREDICT2_VIDEO2WORLD_PIPELINE_2B)
ANY4D_WAN_T2V_PIPELINE_1_3B.net = ANY4D_WAN_T2V_NET_1_3B
ANY4D_WAN_T2V_PIPELINE_1_3B.any4d_active = True
ANY4D_WAN_T2V_PIPELINE_1_3B.tokenizer = L(WanVAEAdapter)(
    vae_pth="", chunk_duration=41, name="wan_vae",
    load_mean_std=False, dtype="bfloat16",
)
ANY4D_WAN_T2V_PIPELINE_1_3B.state_ch = 16
ANY4D_WAN_T2V_PIPELINE_1_3B.state_t = 11
# Same TextAttr-dropout-off rule as the other Wan variants — a4d_vae_wan handles dropout.
ANY4D_WAN_T2V_PIPELINE_1_3B.conditioner.text.dropout_rate = 0.0

ANY4D_WAN_T2V_FSDP_1_3B = dict(
    trainer=dict(
        distributed_parallelism="fsdp",
    ),
    model=L(Any4DWanModel)(
        config=Any4DWanConfig(
            vidar_active=True,
            any4d_active=True,
            use_wan_backbone=True,
            wan_variant='t2v',
            pipe_config=ANY4D_WAN_T2V_PIPELINE_1_3B,
            model_manager_config=L(Predict2ModelManagerConfig)(
                dit_path="",
                text_encoder_path="",
                vae_path="",
                tokenizer_chunk_duration=41,
            ),
            fsdp_shard_size=8,
            vae="a4d_vae_wan",
        ),
        _recursive_=False,
    ),
)


# ============================================================
# Any4D on top of Wan 2.2 Fun A14B InP (image-to-video base)
# ============================================================
# Same structure as 1.3B but with 14B DiT dims (dim=5120, 40 heads, 40 layers).
ANY4D_WAN_NET_14B = L(Any4DWanDiT)(
    any4d_config=None,
    wan_model_kwargs=dict(
        dim=5120,
        in_dim=36,
        ffn_dim=13824,
        out_dim=16,
        text_dim=4096,
        freq_dim=256,
        eps=1e-6,
        patch_size=(1, 2, 2),
        num_heads=40,
        num_layers=40,
        has_image_input=True,
        has_image_pos_emb=False,
    ),
)

ANY4D_WAN_PIPELINE_14B = _deepcopy(PREDICT2_VIDEO2WORLD_PIPELINE_2B)
ANY4D_WAN_PIPELINE_14B.net = ANY4D_WAN_NET_14B
ANY4D_WAN_PIPELINE_14B.any4d_active = True
ANY4D_WAN_PIPELINE_14B.tokenizer = L(WanVAEAdapter)(
    vae_pth="",
    chunk_duration=41,
    name="wan_vae",
    load_mean_std=False,
    dtype="bfloat16",
)
ANY4D_WAN_PIPELINE_14B.state_ch = 16
ANY4D_WAN_PIPELINE_14B.state_t = 11

ANY4D_WAN_FSDP_14B = dict(
    trainer=dict(
        distributed_parallelism="fsdp",
    ),
    model=L(Any4DWanModel)(
        config=Any4DWanConfig(
            vidar_active=True,
            any4d_active=True,
            use_wan_backbone=True,
            pipe_config=ANY4D_WAN_PIPELINE_14B,
            model_manager_config=L(Predict2ModelManagerConfig)(
                dit_path="",
                text_encoder_path="",
                vae_path="",
                tokenizer_chunk_duration=41,
            ),
            fsdp_shard_size=8,
            vae="a4d_vae_wan",
            adjust_video_noise=False,
            high_sigma_ratio=0.0,
            view_emb_std=0.02,
        ),
        _recursive_=False,
    ),
)


def register_model() -> None:
    cs = ConfigStore.instance()
    # predict2 v2w 2b model
    cs.store(
        group="model", package="_global_", name="predict2_video2world_fsdp_2b", node=PREDICT2_VIDEO2WORLD_FSDP_2B
    )
    # predict2 v2w 14b model
    cs.store(
        group="model", package="_global_", name="predict2_video2world_fsdp_14b", node=PREDICT2_VIDEO2WORLD_FSDP_14B
    )

    # Custom Vidar 2B model
    cs.store(
        group="model", package="_global_", name="vidar_2b", node=VIDAR_FSDP_2B
    )

    # Custom Any4D 2B model
    cs.store(
        group="model", package="_global_", name="any4d_2b", node=ANY4D_FSDP_2B
    )

    # Custom Any4D + Wan 2.1 Fun 1.3B InP model
    cs.store(
        group="model", package="_global_", name="any4d_wan_1_3b", node=ANY4D_WAN_FSDP_1_3B
    )

    # Custom Any4D + Wan 2.1 Fun-V1.1-1.3B-Control model (text + canny -> video)
    cs.store(
        group="model", package="_global_", name="any4d_wan_control_1_3b",
        node=ANY4D_WAN_CONTROL_FSDP_1_3B
    )

    # Custom Any4D + Wan 2.1 Fun-InP base with canny control extension (in_dim=52,
    # zero-init on the new channels). Identical to Fun-InP at step 0.
    cs.store(
        group="model", package="_global_", name="any4d_wan_inp_canny_1_3b",
        node=ANY4D_WAN_INP_CANNY_FSDP_1_3B
    )

    # Custom Any4D + pure Wan 2.1 T2V-1.3B base (text -> image/video, in_dim=16,
    # has_image_input=False; no InP mask/ref slot, no CLIP image branch).
    cs.store(
        group="model", package="_global_", name="any4d_wan_t2v_1_3b",
        node=ANY4D_WAN_T2V_FSDP_1_3B
    )

    # Custom Any4D + Wan 2.2 Fun A14B InP model
    cs.store(
        group="model", package="_global_", name="any4d_wan_14b", node=ANY4D_WAN_FSDP_14B
    )

    # Custom Regression model
    cs.store(
        group="model", package="_global_", name="predict2_video2world_fsdp_2b_regression", node=PREDICT2_VIDEO2WORLD_FSDP_2B_REGRESSION
    )

    # NOTE(bvh): see also custom/any4d/a4d_config.py