# Created by BVH, Jul 2025.

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

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

job = dict(  # this becomes JobConfig
    project='a4d2',
    group='debug',
    name='wm1',
    prepend_datetime=True,  # becomes {now:%m-%d-%H-%M}_{config.job.name}
)

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

any4d_config = dict(  # this is part of Any4DConfig, which inherits from Predict2Video2WorldModelConfig
    transforms='default',  # with reference_camera=False
    dataloader='basile',
    vae='a4d_vae',
    video_entries=[
        # first entry = base / pretrained
        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, 'zero', None),  # do not change
        ),
        # later entries = new viewpoints
    ],
    lowdim_adaln_entries=dict(
        # key: (seq_start, seq_end, channel_start, channel_end, init)
        # NOTE: each token here exists only once, and they get flattened into one embedding
        action=(0, 25, 0, 20, 'zero'),
        action_input_mask=(0, 25, 20, 21, 'zero'),
    ),
    num_views=1,
    video_concat_mode='view',
    video_proj_mode='per_view',
    loss_weights=dict(
        rgb0=1.0,
    ),
    harmonize_streams=True,
    harmonize_frames=True,
    load_modals=['rgb', 'language', 'action'],  # cams is needed in vidar to plot actions
    task_probs=dict(
        cross_modal=0.0,
        dyn_view_synth=0.0,
        forecast=1.0,
        pose_estimation=0.0,
        inverse_dynamics=0.0,
        policy=0.0,
        world_model=0.8,
        # world_model=1.0,
    ),
    use_views='all',
    normalizer_stats='',
    relative_actions=False,
    action_multiplier=2.0,
    track_metrics=dict(
        rgb0=['psnr', 'ssim'],
    ),
    train_visuals_interval=99,
    train_visuals_detail=2,
    val_visuals_detail=2,
)

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

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

model = dict(  # this makes up parts of Predict2Video2WorldModelConfig and Predict2ModelManagerConfig
    # dit_path=f'{S3_PRETRAINED_PREFIX}/nvidia/Cosmos-Predict2-2B-Video2World/model-480p-10fps.pt',
    dit_path=f's3://tri-ml-sandbox-16011-us-west-2-datasets/sagemaker/cosmos-predict2/a4d2/debug/08-27-03-20_s9b_wm1_lbm/model/iter_000048000_003072000.pt',
    text_encoder_path=f'{S3_PRETRAINED_PREFIX}/google-t5/t5-11b',
    vae_path=f'{S3_PRETRAINED_PREFIX}/nvidia/Cosmos-Predict2-2B-Video2World/tokenizer/tokenizer.pth',
    fsdp_shard_size=8,
    ############
    run_validation=True,
    validation_iter=300,  # frequency / interval of validation runs
    max_val_iter=16,  # number of steps per validation run
    max_iter=150000,  # total number of training steps
    grad_accum_iter=2,
    context_parallel_size=1,
    device_monitor=0,
    manual_gc_iter=288,
    manual_gc_warm_up=-1,  # never disable automatic GC when actions are involved (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,
    s3_folder='s3://tri-ml-sandbox-16011-us-west-2-datasets/sagemaker/cosmos-predict2/',
    early_sanity_check=5,
)

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

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

dataset_train = dict(
    # config='custom/config/train/lbm24_all_1v.yaml',  # no pose
    config='custom/config/train/lbm40_most_1v.yaml',
    num_workers=6,
    batch_size=8,
)

dataset_val = dict(
    # config='custom/config/val/lbm24_some_1v.yaml',  # no pose
    config='custom/config/val/lbm40_some_1v.yaml',
    num_workers=0,  # can crash debugger otherwise
    batch_size=1,
)

metrics = dict()



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

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,
)