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

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

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

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

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

model = dict(
    model='predict2_video2world_fsdp_2b_regression',
    # run_validation=False,
    run_validation=True,
    validation_iter=5,
    max_iter=200000,
    fsdp_shard_size=8,
    net=dict(
        # in_channels=4*768,
        # out_channels=4*768,
        in_channels=1024,
        out_channels=1024,
    ),
    extra_nets=dict(
        # dav2=dict(),
        vggt=dict(),
        ),
    remove_dit=True
)

checkpoint = dict(
    save_iter=1000,
    # pretrained='s3://tri-ml-sandbox-16011-us-west-2-datasets/sagemaker/cosmos-predict2/checkpoints/finetune/sandbox/2025-7-31-20h31m38s--ratio/model/iter_000036000_055296000.pt',
    # resume='s3://tri-ml-sandbox-16011-us-west-2-datasets/sagemaker/cosmos-predict2/checkpoints/finetune/sandbox/2025-7-31-20h31m38s--ratio/model/iter_000036000_055296000.pt',
    # resume='s3://tri-ml-sandbox-16011-us-west-2-datasets/sagemaker/cosmos-predict2/checkpoints/finetune/sandbox/2025-8-11-16h29m40s--omni/extra/iter_000000003_000000024.pt',
    # resume="s3://tri-ml-sandbox-16011-us-west-2-datasets/sagemaker/cosmos-predict2/checkpoints/finetune/sandbox/2025-8-11-18h45m13s--pointcloud/extra/iter_000000100_000000800.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/tartan.yaml',
    config='custom/config/train/omni.yaml',
    num_workers=8,
    batch_size=1,
)

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

metrics = dict(
    depth=dict(
        min_depth=0.0,
        max_depth=200.0,
    )
)

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

cs = ConfigStore.instance()
sandbox = template_vidar_2b(
    job, wandb, modules, model, 
    checkpoint, optimizer, scheduler, metrics, 
    dataset_train, dataset_val,
)
items = [sandbox] # 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,
    )