# Created by BVH, Jun 2026.
# AnyTask3 on Kubric5D (webbed): the 3-VIEW variant of anytask4_kubric.
# Differences vs anytask4_kubric: num_views=3 (3 views, no rgb3/cams3/depth3 entry) and
# use_views='all' (every sample uses all 3 views every step -- no random view subsampling). The
# all-views setup ensures no view-slot is starved (a suspected source of the 2-node artefacts).
# Each of the 3 views carries rgb + cams + depth; every batch randomly rolls cross-modal synthesis,
# dynamic view synthesis, forecasting, and pose estimation. Warm-started from MAnyView v2
# (s28c_dvs4_av40h): rgb + cams load (view 0), views 1-2 copy view-0 weights, depth is new.

import os
from hydra.core.config_store import ConfigStore
from custom.experiment.template import template_any4d_2b, S3_PRETRAINED_ROOT, S3_OUTPUT_ROOT

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

job = dict(  # this becomes JobConfig
    project='a4d2',
    group='example',
    name='at_kub3',
    prepend_datetime=True,  # becomes {now:%m-%d-%H-%M}_{config.job.name}
    local_root='',  # empty = _DEFAULT_LOCAL_ROOT (/any4d on DGX, /tmp/a4d2 on SM)
    s3_root=S3_OUTPUT_ROOT,  # default from template.py; set '' to disable S3 sync
)

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

any4d_config = dict(  # this is part of Any4DConfig, which inherits from Predict2Video2WorldModelConfig
    dataloader='unified_flex',  # the flexible dataloader supporting all tasks (incl cross_modal + pose_est)
    vae='a4d_vae',
    # NOTE(bvh): 3 views total (0, 1, 2) -- this is the 3-view variant (no rgb3/cams3/depth3).
    # V_max derives from these entries (1 + max view index = 3), so use_views='all' uses all 3.
    video_entries=[
        # first entry (view 0) = base / pretrained
        # NOTE(bvh): channel layout matches MAnyView v2 for rgb + cams (so they load), depth is appended.
        dict(
            # key: (channel_start, channel_end, in_proj_init, out_proj_init)
            # NOTE: there will be T*H*W tokens with this information
            rgb0=(0, 16, 'load', 'load'),  # do not change
            rgb0_input_mask=(16, 17, 'load', None),  # do not change
            rgb0_output_mask=(17, 18, 'load', None),  # do not change
            cams0=(18, 50, 'zero/load', 'zero/load'),
            cams0_input_mask=(50, 51, 'zero/load', None),
            cams0_output_mask=(51, 52, 'zero/load', None),
            depth0=(52, 68, 'zero/load', 'rand/load'),
            depth0_input_mask=(68, 69, 'zero/load', None),
            depth0_output_mask=(69, 70, 'zero/load', None),
        ),
        # later entries = new viewpoints
        dict(
            # key: (channel_start, channel_end, in_proj_init, out_proj_init)
            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),
            cams1=(18, 50, 'copy:cams0/load', 'copy:cams0/load'),
            cams1_input_mask=(50, 51, 'copy:cams0_input_mask/load', None),
            cams1_output_mask=(51, 52, 'copy:cams0_output_mask/load', None),
            # NOTE(bvh): depth is NOT in the MAnyView ckpt, so it cannot be copied from a source entry;
            # each view inits depth fresh (same strategy as depth0), like anytask3_ddad.
            depth1=(52, 68, 'zero/load', 'rand/load'),
            depth1_input_mask=(68, 69, 'zero/load', None),
            depth1_output_mask=(69, 70, 'zero/load', None),
        ),
        dict(
            rgb2=(0, 16, 'copy:rgb0/load', 'copy:rgb0/load'),
            rgb2_input_mask=(16, 17, 'copy:rgb0_input_mask/load', None),
            rgb2_output_mask=(17, 18, 'copy:rgb0_output_mask/load', None),
            cams2=(18, 50, 'copy:cams0/load', 'copy:cams0/load'),
            cams2_input_mask=(50, 51, 'copy:cams0_input_mask/load', None),
            cams2_output_mask=(51, 52, 'copy:cams0_output_mask/load', None),
            depth2=(52, 68, 'zero/load', 'rand/load'),
            depth2_input_mask=(68, 69, 'zero/load', None),
            depth2_output_mask=(69, 70, 'zero/load', None),
        ),
    ],
    num_views=3,
    video_concat_mode='view',
    video_proj_mode='per_view',
    view_timestep_mode='per_view',
    block_gate_fix=True,  # should be enabled in every new training run
    disable_risky_sharding=True,  # keep final_layer / t_embedder un-sharded to avoid multi-view hangs
    train_cond_aug_sigma_range=(0.001, 0.01),
    val_cond_aug_sigma=0.001,
    loss_weights=dict(
        rgb0=1.0, cams0=1.0, depth0=1.0,
        rgb1=1.0, cams1=1.0, depth1=1.0,
        rgb2=1.0, cams2=1.0, depth2=1.0,
    ),
    harmonize_streams=True,
    harmonize_frames=True,
    data_train_overrides=dict(
        frame_rate=24.0,  # NOTE(2026-06-17): beats baked metadata=10; drop after Kubric5D re-web at 24fps
        subsample=None,
        single_sample='random',
    ),
    data_val_overrides=dict(
        frame_rate=24.0,
        frame_stride=1,
        subsample=20,  # NOTE(bvh): for debugging non-aligned validation set sizes
        single_sample='first',
    ),
    use_views='all',  # always use all 3 views every step (no view-slot starvation)
    shuffle_cams2views=True,  # debias Any4D views vs the randomly subsampled dataset cameras
    reference_view=0,  # plucker rays re-anchored to view 0 (MAnyView convention)
    ref_view_rel_per_timestep=False,
    # NOTE(bvh): AnyTask = roll a flexible combination of tasks every batch. unified_flex keys are
    # pose_est / inv_dyn (NOT pose_estimation / inverse_dynamics); each is an independent Bernoulli.
    task_probs=dict(
        cross_modal=0.5,     # rgb <-> depth (needs >= 2 modalities)
        dyn_view_synth=0.5,  # predict held-out viewpoints (needs cams + V >= 2)
        forecast=0.5,        # predict future frames
        pose_est=0.5,        # predict held-out camera extrinsics (needs cams + V >= 2)
        # NOTE(bvh): action stuff (inv_dyn / policy / world_model) deprecated in unified_flex for now.
        world_model=0.0,
        inv_dyn=0.0,
        policy=0.0,
    ),
    train_directives=dict(
    ),
    val_directives=dict(
        # leave empty: randomly alternate tasks during validation to exercise the full mix
    ),
    track_metrics=dict(
        rgb0=['psnr', 'ssim', 'lpips'], cams0=['mse', 'mae'], depth0=['mse', 'mae'],
        rgb1=['psnr', 'ssim', 'lpips'], cams1=['mse', 'mae'], depth1=['mse', 'mae'],
        rgb2=['psnr', 'ssim', 'lpips'], cams2=['mse', 'mae'], depth2=['mse', 'mae'],
    ),
    train_visuals_interval=99,
    train_visuals_detail=1,
    val_visuals_detail=1,
    visuals_quality=8,
    viz_input_blacklist=[],
    viz_mask_border_width=2,  # NOTE(bvh): useful for temporally varying stuff such as forecasting
    val_num_steps=35,
)

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

model = dict(  # this makes up parts of Predict2Video2WorldModelConfig and Predict2ModelManagerConfig
    # pretrained / vanilla cosmos:
    # dit_path=f'{S3_PRETRAINED_ROOT}/nvidia/Cosmos-Predict2-2B-Video2World/model-480p-10fps.pt',
    # MAnyView v2 (4 views + cams; rgb + cams load, depth added fresh):
    dit_path='s3://tri-ml-sandbox-16011-us-west-2-datasets/sagemaker/cosmos-predict2/a4d2/debug/01-03-14-18_s28c_dvs4_av40h_b2/model/iter_000059000_007552000.pt',
    text_encoder_path=f'{S3_PRETRAINED_ROOT}/google-t5/t5-11b',  # uncomment to enable
    # text_encoder_path='',  # uncomment to disable
    vae_path=f'{S3_PRETRAINED_ROOT}/nvidia/Cosmos-Predict2-2B-Video2World/tokenizer/tokenizer.pth',
    fsdp_shard_size=8,
    ############
    run_validation=True,
    skip_first_validation='delay3',  # to check VRAM
    validation_iter=500,  # frequency / interval of validation runs
    # max_val_iter=3,  # limit validation steps per dataset for faster iteration
    max_iter=60000,  # total number of training steps
    grad_accum_iter=1,
    context_parallel_size=1,
    device_monitor=0,
    manual_gc_iter=288,
    manual_gc_warm_up=-1,  # never disable automatic GC to be safe (weird VRAM issue)
    ############
    state_t=11,  # for noise level; = latent # frames for now
    tokenizer_chunk_duration=41,  # for VAE; = raw # frames for now
)

checkpoint = dict(
    save_iter=1000,
    early_sanity_check=5,
    # NOTE(bvh): full resume (model + optimizer + scheduler + iteration).
    # Set dit_path=None above when using this to avoid redundant loading.
    # resume='s3://...',
)

optimizer = dict(
    lr=2e-5,  # new active entry loss averaging
)

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=[
        # NOTE(bvh): if there are K uncommented lines here, run with at least K GPUs (RankDataset is
        # round-robin), else the surplus datasets are silently skipped.
        'custom/config/anydata/web/train/kubric_iked3_41.yaml',  # train split (scn00000-09599); 3-cam (matches num_views=3)
        # 'custom/config/anydata/web/debug/lbmall_ikela4_41.yaml',  # optional: add real+sim robot (rgb + cams, no depth)
    ],
    num_workers=4,
    batch_size=2,
)

dataset_val = dict(
    config=dict(
        Kubric5D='custom/config/anydata/web/val/kubric_iked3_41.yaml',  # val split (scn09600-09799)
        # LBM='custom/config/anydata/web/debug/lbmall_ikela4_41.yaml',  # optional: pairs with the commented train entry
    ),
    num_workers=0,
    batch_size=1,
)

metrics = None  # dict()  # disable because fragile

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

cs = ConfigStore.instance()

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

# Use the filename (without extension) as the experiment name
experiment_name = 'any4d_' + os.path.splitext(os.path.basename(__file__))[0]

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