# Created by VG & BVH, Jul 2025.

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

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

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

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

modules = dict(  # this is part of Predict2Video2WorldModelConfig
    # arch='default',
    transforms='default',
    dataloader='default',
    vae='a4d_vae',
    # load_modals=['rgb', 'language', 'action'],
    # load_modals=['rgb', 'language', 'action', 'depth', 'points', 'cams'],
    train_visuals_interval=39,
    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',
    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=2,
    ############
    validation_iter=500,  # frequency / interval of validation runs
    max_val_iter=4,  # number of steps per validation run
    max_iter=25000,  # total number of training steps
    context_parallel_size=1,
    device_monitor=0,
    manual_gc=288,
    ############
    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/humanoid.yaml',
    num_workers=4,
    batch_size=2,
)

dataset_val = dict(
    config='custom/config/val/humanoid.yaml',
    num_workers=2,
    batch_size=1,
)

metrics = dict(
)

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

cs = ConfigStore.instance()

this_config = template_vidar_2b(
    job, wandb, modules, model, 
    checkpoint, optimizer, scheduler, metrics, 
    dataset_train, dataset_val,
)

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

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