# Created by BVH & VG, Jul 2025.

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

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

job = dict(
    project='a4d2',
    group='cvpr',
    name='debug',
    prepend_datetime=True,
)

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

any4d_config = dict(
    transforms='default_fitorigin',
    dataloader='identity',
    vae='',
    train_directives=dict(
    ),
    val_directives=dict(
        tasks='dyn_view_synth',
        num_pred_views=1,
        remove_cond=True,
    ),
)

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

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

model = dict(
    dit_path=None,
    text_encoder_path=None, 
    vae_path=None, 
    fsdp_shard_size=8,
    ############
    # run_validation=False,
    run_validation=True,
    validation_iter=500,
    max_iter=50000,
    grad_accum_iter=1,
    context_parallel_size=1,
    device_monitor=0,
    manual_gc=488,
    ############
    remove_dit=True,
    model='predict2_video2world_fsdp_2b_regression',
    extra_nets=dict(
        vggt=dict(aaa=4,bbb=8),
    ),
)

checkpoint = dict(
    save_iter=1000,
    s3_folder='s3://tri-ml-sandbox-16011-us-west-2-datasets/sagemaker/cosmos-predict2/',
)

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

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

dataset_train = dict(
    config='custom/experiment/regression/config/ddad_train.yaml',
    num_workers=8,
    batch_size=1,
)

dataset_val = dict(
    config='custom/experiment/regression/config/ddad_val.yaml',
    num_workers=1,
    batch_size=1,
)

metrics = dict(
)

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

cs = ConfigStore.instance()

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

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

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