# 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 attrs

from cosmos_predict2.conditioner import BooleanFlag, ReMapkey, TextAttr
from cosmos_predict2.configs.base.defaults.ema import EMAConfig
from cosmos_predict2.configs.action_conditioned.defaults.conditioner import ActionConditionedConditioner
from cosmos_predict2.models.text2image_dit import SACConfig
from cosmos_predict2.models.action_video2world_dit import ActionConditionedMinimalV1LVGDiT
from cosmos_predict2.tokenizers.tokenizer import TokenizerInterface
from cosmos_predict2.configs.base.config_video2world import Video2WorldPipelineConfig
from imaginaire.config import make_freezable
from imaginaire.lazy_config import LazyCall as L
from imaginaire.lazy_config import LazyDict
from cosmos_predict2.configs.base.config_video2world import ConditioningStrategy, CosmosReason1Config, CosmosGuardrailConfig, SolverTimestampConfig



# Cosmos Predict2 Video2World 2B
# NOTE(bvh): we are borrowing this one for Any4D (see cosmos_predict2/configs/base/defaults/model.py)
ACTION_CONDITIONED_PREDICT2_VIDEO2WORLD_NET_2B = L(ActionConditionedMinimalV1LVGDiT)(
    max_img_h=240,
    max_img_w=240,
    max_frames=128,
    in_channels=16,
    out_channels=16,
    patch_spatial=2,
    patch_temporal=1,
    concat_padding_mask=True,
    # attention settings
    model_channels=2048,
    num_blocks=28,
    num_heads=16,
    atten_backend="minimal_a2a",
    # positional embedding settings
    pos_emb_cls="rope3d",
    pos_emb_learnable=True,
    pos_emb_interpolation="crop",
    use_adaln_lora=True,
    adaln_lora_dim=256,
    rope_h_extrapolation_ratio=3.0,
    rope_w_extrapolation_ratio=3.0,
    rope_t_extrapolation_ratio=1.0,
    extra_per_block_abs_pos_emb=False,
    rope_enable_fps_modulation=False,
    sac_config=L(SACConfig)(  # NOTE(bvh): this is for selective activation checkpointing
        every_n_blocks=1,
        mode="predict2_2b_720",
    ),
    # NOTE(bvh): add action dimension - this is the only difference between the two network configs
    # (i.e. PREDICT2_VIDEO2WORLD_NET_2B vs ACTION_CONDITIONED_PREDICT2_VIDEO2WORLD_NET_2B)
    action_dim=7*12,
    # NOTE(bvh): ^ seems hard-coded to single vector of size
    # 84 = 12 timesteps * 7 action values (= 6-DoF + gripper width)
    any4d_config=None,
)

# NOTE(bvh): we are NOT borrowing this one for Any4D (see cosmos_predict2/configs/base/defaults/model.py)
ACTION_CONDITIONED_PREDICT2_VIDEO2WORLD_PIPELINE_2B = Video2WorldPipelineConfig(
    adjust_video_noise=True,
    conditioner=L(ActionConditionedConditioner)(  # difference
        fps=L(ReMapkey)(
            dropout_rate=0.0,
            dtype=None,
            input_key="fps",
            output_key="fps",
        ),
        padding_mask=L(ReMapkey)(
            dropout_rate=0.0,
            dtype=None,
            input_key="padding_mask",
            output_key="padding_mask",
        ),
        text=L(TextAttr)(
            dropout_rate=0.2,
            input_key=["t5_text_embeddings"],
        ),
        use_video_condition=L(BooleanFlag)(
            dropout_rate=0.0,
            input_key="fps",
            output_key="use_video_condition",
        ),
        # NOTE: add additional action as condition
        action=L(ReMapkey)(  # difference
            input_key="action",
            output_key="action",
            dropout_rate=0.0,
            dtype=None,
        ),
    ),
    conditioning_strategy=str(ConditioningStrategy.FRAME_REPLACE),
    min_num_conditional_frames=1,
    max_num_conditional_frames=1,  # difference (2 vs 1)
    net=ACTION_CONDITIONED_PREDICT2_VIDEO2WORLD_NET_2B,
    precision="bfloat16",
    rectified_flow_t_scaling_factor=1.0,
    resize_online=True,
    resolution="720",
    ema=L(EMAConfig)(enabled=False),  # defaults to inference
    sigma_conditional=0.0001,
    sigma_data=1.0,
    state_ch=16,
    state_t=4,  # difference (24 vs 4)
    # NOTE(bvh): ^ this is number of latent frames (so only 13 pixel frames for this model)
    text_encoder_class="T5",
    tokenizer=L(TokenizerInterface)(
        chunk_duration=81,
        load_mean_std=False,
        name="tokenizer",
        # vae_pth="checkpoints/nvidia/Cosmos-Predict2-2B-Video2World/tokenizer/tokenizer.pth",
    ),
    # disable prompt refiner and guardrail for action conditional
    prompt_refiner_config=CosmosReason1Config(
        checkpoint_dir="checkpoints/nvidia/Cosmos-Reason1-7B",
        offload_model_to_cpu=True,
        enabled=False,  # difference (yes vs no)
    ),
    guardrail_config=CosmosGuardrailConfig(
        checkpoint_dir="checkpoints/",
        offload_model_to_cpu=True,
        enabled=False,  # difference (yes vs no)
    ),
)

PREDICT2_VIDEO2WORLD_NET_14B = None

PREDICT2_VIDEO2WORLD_PIPELINE_14B = None

# ^ last 2 = to force vscode to compare files correctly: