# Single-camera Wan-Fun-InP experiment on the existing GAIA webbed reference data
# Uses Wan2.1-Fun-1.3B-InP first-and-last frame conditioning recipe

import os

from cosmos_predict2.callbacks.grad_clip import GradClipCallback
from custom.experiment.template import template_any4d_wan_inp_canny_1_3b
from hydra.core.config_store import ConfigStore
from imaginaire.lazy_config import LazyCall as L

job = dict(
    project='a4d2',
    group='inp',
    name='wm3_gaia_wan_iked1_41',
    prepend_datetime=True,
    # SageMaker (AD2 gaia-wfm account): leave local_root empty so it resolves to the
    # SM scratch default (/tmp/a4d2), and sync checkpoints/visuals to the AD2 bucket the
    # SM execution role can write. (Override via the env/job dict for other clusters.)
    local_root='',
    s3_root='s3://gaia-e2e-wfm-datasets/any4d',
)

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

any4d_config = dict(
    transforms='default',
    dataloader='unified_anydrive_vid2vid',
    vae='a4d_vae_wan',
    use_wan_backbone=True,
    wan_variant='inp_canny',
    wan_inp_single_end_frame=True,  # end conditioning = a single real frame (orig Wan-Fun-InP)
    video_entries=[
        # Single view with a Canny edge-control channel. 50-ch per-view stream
        # (reshaped to the 52-ch canny DiT input by Any4DWanInPCannyDiT._to_wan_36ch):
        #   [0:16]  noisy rgb latent      [16:17] input_mask   [17:18] output_mask
        #   [18:34] gray-padded ref latent          [34:50] canny edge-control latent
        # The edge_control video is derived from RGB via Canny in the dataloader
        # (compute_canny_edge_control), then VAE-encoded like RGB.
        dict(
            rgb0=(0, 16, 'load', 'load'),
            rgb0_input_mask=(16, 17, 'load', None),
            rgb0_output_mask=(17, 18, 'load', None),
            rgb0_ref=(18, 34, 'load', None),
            rgb0_edge_control=(34, 50, 'load', None),
        ),
    ],
    num_views=1,
    video_concat_mode='view',
    video_proj_mode='per_view',
    view_timestep_mode='per_view',
    view_emb_std=0.06,
    block_gate_fix=True,
    loss_weights=dict(rgb0=1.0),
    harmonize_streams=True,
    harmonize_frames=True,
    load_modals=['rgb'],
    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.0,
    ),
    use_views='all',
    shuffle_cams2views=False,
    normalizer_stats='',
    relative_actions=False,
    action_multiplier=2.0,
    # InP first+last frame conditioning.
    train_directives=dict(cond_frames_raw=1, cond_last_frame=True),
    val_directives=dict(
        tasks='forecast',
        cond_frames_raw=1,
        cond_last_frame=True,
        remove_cond=True,
    ),
    data_train_overrides=dict(frame_stride=1, resize=[-16, 1280], length=9),  # 720p, 9 frames
    data_val_overrides=dict(frame_stride=1, subsample=8, resize=[-16, 1280], length=9),  # 720p, 9 frames
    track_metrics=dict(
        rgb0=['psnr', 'ssim'],
    ),
    train_visuals_interval=99,
    train_visuals_detail=1,
    val_visuals_detail=1,
    visuals_quality=8,
    viz_input_blacklist=[],
    viz_mask_border_width=2,
    val_num_steps=35,
)

WAN_INP_CKPT_DIR = os.environ.get(
    'WAN_1_3B_INP_CKPT_DIR',
    # Default to the AD2-bucket S3 mirror so the SageMaker execution role can read it
    # (the model_manager syncs s3:// dirs to local at runtime). Override the env var to
    # point at a local copy on other clusters.
    's3://gaia-e2e-wfm-datasets/any4d/pretrained/wan/Wan2.1-Fun-1.3B-InP',
)
WAN_VAE_PATH = os.environ.get(
    'WAN_VAE_PATH',
    f'{WAN_INP_CKPT_DIR}/Wan2.1_VAE.pth',
)

model = dict(
    dit_path=WAN_INP_CKPT_DIR,
    text_encoder_path='',
    vae_path=WAN_VAE_PATH,
    fsdp_shard_size=8,
    ema_enabled=True,  # Karras power EMA (rate 0.1); validation runs on EMA weights
    run_validation=True,
    skip_first_validation='delay3',
    validation_iter=500,
    max_val_iter=3,
    # T=17 raw -> Tl=5 latent (1 + 4*4). Matches the YAML's length=17, which is
    # constrained by the iked11 tarball window size (32) and multi_tarfiles=1.
    max_iter=800000,
    grad_accum_iter=1,
    context_parallel_size=1,
    device_monitor=0,
    manual_gc_iter=288,
    manual_gc_warm_up=-1,
    state_t=3,            # 9 raw frames -> (9-1)/4+1 = 3 latent frames
    tokenizer_chunk_duration=9,
)

checkpoint = dict(
    save_iter=1000,
    early_sanity_check=5,
)

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

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

dataset_train = dict(
    # NOTE: the iked1 (single random camera) YAML lives at web/, not web/debug/.
    config='custom/config/anydata/web/gaia_iked1_41.yaml',
    num_workers=4,
    prefetch_factor=4,
    batch_size=1,
)

dataset_val = dict(
    config=dict(
        GaiaWebbed1Cam='custom/config/anydata/web/gaia_iked1_41.yaml',
    ),
    num_workers=0,
    batch_size=1,
)

metrics = None

cs = ConfigStore.instance()

from cosmos_predict2.configs.base.defaults import optimizer as _opt_mod  # noqa: E402
from cosmos_predict2.utils.optim_instantiate_dtensor import get_base_optimizer  # noqa: E402
from imaginaire.lazy_config import PLACEHOLDER  # noqa: E402

_opt_mod.FusedAdamWConfig = L(get_base_optimizer)(
    model=PLACEHOLDER,
    lr=1e-4,
    weight_decay=0.1,
    betas=[0.9, 0.99],
    optim_type='adamw',
    eps=1e-8,
)

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

this_config['trainer']['callbacks']['grad_clip'] = L(GradClipCallback)(
    max_norm=1.0, log_every_n=10,
)

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

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