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

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

job = dict(
    project="finetune",
    group="sandbox",
    name='multi',
)

modules = dict(
    arch="default",
    transforms="fit_batch",
    dataloader="default",
    vae="identity",
)

wandb = dict(
    # enabled=False,
    enabled=True,
    entity='tri',
    project='lightning_logs',
)

model = dict(
    remove_dit=True,
    model='predict2_video2world_fsdp_2b_regression',
    # run_validation=False,
    run_validation=True,
    validation_iter=500,
    max_iter=200000,
    fsdp_shard_size=8,
    extra_nets=dict(
        vggt=dict(),
    ),
)

checkpoint = dict(
    save_iter=500,
    # pretrained='s3://tri-ml-sandbox-16011-us-west-2-datasets/sagemaker/cosmos-predict2/checkpoints/finetune/sandbox/2025-8-16-1h56m29s--nograd/extra/iter_000006000_003072000.pt',
    # pretrained='s3://tri-ml-sandbox-16011-us-west-2-datasets/sagemaker/cosmos-predict2/checkpoints/finetune/sandbox/2025-8-16-18h53m24s--multi/extra/iter_000000005_000000320.pt',
    # pretrained='s3://tri-ml-sandbox-16011-us-west-2-datasets/sagemaker/cosmos-predict2/checkpoints/finetune/sandbox/2025-8-16-19h29m54s--multi/extra/iter_000000000_000000000.pt',
    # pretrained='s3://tri-ml-sandbox-16011-us-west-2-datasets/sagemaker/cosmos-predict2/checkpoints/finetune/sandbox/2025-8-16-19h38m29s--multi/extra/iter_000000005_000000040.pt',
    # pretrained='s3://tri-ml-sandbox-16011-us-west-2-datasets/sagemaker/cosmos-predict2/checkpoints/finetune/sandbox/2025-8-16-20h17m37s--multi/extra/iter_000000010_000000080.pt',
    # pretrained='s3://tri-ml-sandbox-16011-us-west-2-datasets/sagemaker/cosmos-predict2/checkpoints/finetune/sandbox/2025-8-16-21h18m15s--multi/extra/iter_000000005_000000040.pt'
    # pretrained='s3://tri-ml-sandbox-16011-us-west-2-datasets/sagemaker/cosmos-predict2/checkpoints/finetune/sandbox/2025-8-16-23h26m13s--multi/extra/iter_000000000_000000000.pt',
    # pretrained='s3://tri-ml-sandbox-16011-us-west-2-datasets/sagemaker/cosmos-predict2/checkpoints/finetune/sandbox/2025-8-17-1h4m32s--multi/extra/iter_000000000_000000000.pt',
)

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

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

dataset_train = dict(
    config='custom/config/train/multi/multi.yaml',
    num_workers=16, # 24,
    batch_size=8, # 8,
)

dataset_val = dict(
    config={
        # 'OmniMMT': 'custom/config/val/multi/omni.yaml',
        # 'TartanAir': 'custom/config/val/multi/tartan.yaml',
        # 'DDAD': 'custom/config/val/multi/ddad.yaml',
        'PDv2': 'custom/config/val/multi/pdv2.yaml',
        # 'LBM': 'custom/config/val/multi/lbm.yaml',
    },
    num_workers=8,
    batch_size=1,
)

metrics = dict(
    depth=dict(
        min_depth=0.0,
        max_depth=80.0,
        crop='garg',
    )
)

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

cs = ConfigStore.instance()
multi = template_vidar_2b(
    job, wandb, modules, model, 
    checkpoint, optimizer, scheduler, metrics, 
    dataset_train, dataset_val,
)
items = [multi] # Add recipes defined below in this list

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

for _item in items:
    experiment_name = [name.lower() for name, value in globals().items() if value is _item][0]
    cs.store(
        group="experiment",
        package="_global_",
        name=experiment_name,
        node=_item,
    )