# Mirror of rwm5_drive.py but with the Wan 2.1 Fun 1.3B InP image-to-video model
# as the base video diffusion backbone (instead of Cosmos-Predict2 2B).

import os
from hydra.core.config_store import ConfigStore
from custom.experiment.template import template_any4d_wan_1_3b

#################################################################

job = dict(
    project='a4d2',
    group='example',
    name='wm5_drive_wan',
    prepend_datetime=True,
)

wandb = dict(
    enabled=True,
    entity='tri',
    project='a4d2',
    num_validation_logs=5,
)

any4d_config = dict(
    dataloader='unified_anydrive',
    # Wan VAE wrapper (mirrors custom/vae/a4d_vae.py API but uses Wan2.1_VAE.pth).
    vae='a4d_vae_wan',
    use_wan_backbone=True,
    video_entries=[
        # Wan 2.1 Fun 1.3B InP DiT input layout (36 channels total, after _to_wan_36ch reshape):
        #   ch [0:16]  noisy rgb latent (actively denoised at all frames)
        #   ch [16:20] 4-channel mask (1 at given frames, 0 at target frames)
        #   ch [20:36] clean reference latent (VAE-encoded pixel video with zeros
        #              padding at non-given frames — NOT just the clean latent of frame 0)
        #
        # Any4D stream layout (34 channels), reshaped to Wan's 36 by Any4DWanDiT._to_wan_36ch:
        #   ch [0:16]  rgb latent (noisy target during training; fed to Wan's noisy_slot)
        #   ch [16:17] input_mask value — 1 at given frames, 0 elsewhere (broadcast to 4 ch inside _to_wan_36ch)
        #   ch [17:18] output_mask value — 0 at given frames, 1 elsewhere (kept for bookkeeping; not used by Wan)
        #   ch [18:34] reference latent (populated by a4d_vae_wan, fed to Wan's ref_slot)
        # first entry = base / pretrained
        dict(
            rgb0=(0, 16, 'load', 'load'),
            rgb0_input_mask=(16, 17, 'load', None),
            rgb0_output_mask=(17, 18, 'load', None),
            rgb0_ref=(18, 34, 'load', None),
        ),
        # later entries = new viewpoints
        dict(
            rgb1=(0, 16, 'copy:rgb0/load', 'copy:rgb0/load'),
            rgb1_input_mask=(16, 17, 'copy:rgb0_input_mask/load', None),
            rgb1_output_mask=(17, 18, 'copy:rgb0_output_mask/load', None),
            rgb1_ref=(18, 34, 'copy:rgb0_ref/load', None),
        ),
    ],
    lowdim_adaln_entries=dict(
        # NOTE(bvh): driving uses 2D trajectory instead of 20D action
        traj=(0, 21, 0, 2, 'zero/load'),
        traj_input_mask=(0, 21, 2, 3, 'zero/load'),
    ),
    num_views=2,
    video_concat_mode='view',
    video_proj_mode='per_view',
    view_timestep_mode='per_view',
    block_gate_fix=True,
    train_cond_aug_sigma_range=(0.001, 0.01),
    val_cond_aug_sigma=0.001,
    loss_weights=dict(
        rgb0=1.0,
        rgb1=1.0,
    ),
    harmonize_streams=True,
    harmonize_frames=True,
    load_modals=['rgb', 'traj'],
    data_train_overrides=dict(
        subsample=None,
        single_sample='random',
    ),
    data_val_overrides=dict(
        frame_stride=1,
        subsample=20,
        single_sample='first',
    ),
    use_views='all',  # TODO handle
    shuffle_cams2views=False,
    # NOTE(bvh): ^ False is important here to have views become consistent / semantically meaningful wrt cams!
    task_probs=dict(
        cross_modal=0.0,
        dyn_view_synth=0.0,
        forecast=1.0,
        pose_est=0.0,
        inv_dyn=0.0,
        policy=0.0,
        world_model=0.9,
        # text is 20% dropout, independently sampled (see config_video2world.py)
    ),
    train_directives=dict(
    ),
    val_directives=dict(
        tasks='forecast,world_model',
        perturb_traj='basile3',
        # NOTE(bvh): ^ key flag to generate counterfactual scenarios (at inference only, not training)
    ),
    action_multiplier=2.0,
    track_metrics=dict(
        rgb0=['psnr', 'ssim'],
        rgb1=['psnr', 'ssim'],
    ),
    train_visuals_interval=99,
    train_visuals_detail=1,
    val_visuals_detail=1,
    visuals_quality=8,
    viz_input_blacklist=[],
    viz_extra_modes=['anydrive1'],
    viz_mask_border_width=2,
    val_num_steps=35,
)

#################################################################

# Wan 2.1 Fun 1.3B InP checkpoint paths. Set these to your local / S3 mirrors of the
# PAI/Wan2.1-Fun-1.3B-InP release on ModelScope (the pipeline loader accepts both a
# directory with diffusion_pytorch_model*.safetensors shards and a single safetensors
# file).
WAN_CKPT_DIR = os.environ.get(
    'WAN_1_3B_INP_CKPT_DIR',
    's3://tri-ml-sandbox-16011-us-west-2-datasets/wan/Wan2.1-Fun-1.3B-InP',
)
WAN_VAE_PATH = os.environ.get(
    'WAN_VAE_PATH',
    f'{WAN_CKPT_DIR}/Wan2.1_VAE.pth',
)

model = dict(
    # Path (directory or file) to the Wan 2.1 Fun 1.3B InP DiT weights.
    dit_path=WAN_CKPT_DIR,
    # Disable text encoder (uses custom/vae/default_text.pkl). Set this to a UMT5-XXL dir
    # (e.g. f'{WAN_CKPT_DIR}/google/umt5-xxl') to use real text prompts instead.
    text_encoder_path='',
    vae_path=WAN_VAE_PATH,
    fsdp_shard_size=1,
    ############
    run_validation=True,
    skip_first_validation='delay3',
    validation_iter=500,
    max_val_iter=3,
    max_iter=80000,
    grad_accum_iter=1,
    context_parallel_size=1,
    device_monitor=0,
    manual_gc_iter=288,
    manual_gc_warm_up=-1,
    ############
    state_t=6,            # latent frames (Wan VAE: (21-1)/4+1 = 6)
    tokenizer_chunk_duration=21,  # raw frames fed to Wan VAE per chunk
)

checkpoint = dict(
    save_iter=1000,
    early_sanity_check=5,
)

optimizer = dict(
    lr=3.0e-5,
)

scheduler = dict(
    warm_up_steps=[500],
    cycle_lengths=[model['max_iter']],
    f_start=[0.01],
    f_max=[1.0],
    f_min=[0.05],
)

dataset_train = dict(
    config=[
        # 2-camera smoke-test variant (1 core + 1 aux). Created in this branch.
        'custom/config/anydata/web/debug/ddad_ike2_21_fixed.yaml',
    ],
    num_workers=4,
    batch_size=1,
)

dataset_val = dict(
    config=dict(
        DDAD  = 'custom/config/anydata/web/debug/ddad_ike2_21_fixed.yaml',
    ),
    num_workers=0,
    batch_size=1,
)

metrics = None

#################################################################

cs = ConfigStore.instance()

this_config = template_any4d_wan_1_3b(
    job, wandb, any4d_config, model,
    checkpoint, optimizer, scheduler,
    metrics, dataset_train, dataset_val)

experiment_name = 'any4d_' + os.path.splitext(os.path.basename(__file__))[0]

cs.store(
    group='experiment',
    package='_global_',
    name=experiment_name,
    node=this_config,
)