# Trains Wan 2.1 Fun-V1.1-1.3B-Control on the rain_balanced subset, ELEVEN-CAMERA
# JOINT cross-view variant (Any4D-style): all 11 GAIA cameras are loaded as Any4D
# views and attend JOINTLY in the Wan DiT self-attention (concatenated token
# sequence). text (per-scene VLM caption) + canny edge map -> image, single-frame
# (T=1 / image mode).
#
# Companion to canny_1frame_wan_local.py (single-view M=1). Key diffs:
#   - 11 video_entries (view 0 = 'load'; views 1..10 = 'copy:rgb0/load' etc), so the
#     surgery clones view 0's patch-embedder / head init into the 10 new views. num_views=11.
#   - wan_gradient_checkpointing=True: the 11-view concatenated token sequence makes
#     per-block self-attention activations huge at 704x1280, so we trade recompute for
#     activation memory (non-reentrant torch.utils.checkpoint over the DiT block stack).
#   - loss_weights over all 11 rgb{v} streams; track_metrics on rgb0 (ego) only.
#   - data points at the 11-camera yamls (cameras_core_sample: 11).
# Everything else (structured noise, Wan-Control wiring, T=1, resolution) is copied
# verbatim from the single-view config.

import os

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

# Registers the 'adverse_weather_canny' data_library (AnyData Unified loader +
# precomputed UMT5 caption embeddings). Import for its side-effect BEFORE the
# template() call below resolves data_library via the registry.
import custom.dataset.adverse_weather_canny_dataset  # noqa: F401

NUM_VIEWS = 11


def _build_video_entries(num_views: int) -> list:
    """11 per-view 50-channel streams (Any4D JOINT cross-view layout).

    View 0 is the base/pretrained view ('load' = copy from the single-view ckpt).
    Views 1..N-1 clone view 0's init via 'copy:rgb0/load' (in-proj) and 'load'
    (out-proj). The 50-channel layout per view matches the single-view config:
      [0:16]  noisy RGB latent       [16:17] input_mask     [17:18] output_mask
      [18:34] reference latent       [34:50] edge_control latent (canny)
    """
    entries = [
        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),
        ),
    ]
    for v in range(1, num_views):
        entries.append({
            f'rgb{v}': (0, 16, 'copy:rgb0/load', 'copy:rgb0/load'),
            f'rgb{v}_input_mask': (16, 17, 'copy:rgb0_input_mask/load', None),
            f'rgb{v}_output_mask': (17, 18, 'copy:rgb0_output_mask/load', None),
            f'rgb{v}_ref': (18, 34, 'copy:rgb0_ref/load', None),
            f'rgb{v}_edge_control': (34, 50, 'copy:rgb0_edge_control/load', None),
        })
    return entries


job = dict(
    project='a4d2',
    group='canny_control',
    name='aw_canny_11view',
    prepend_datetime=True,
    # Local 8-GPU training; the in-container launcher (run_gaia_train.sh) overrides
    # local_root -> /workspace/Any4D/output_a4d and s3_root -> '' on the CLI.
    local_root='/workspace/Any4D/output_a4d',
    s3_root='',
)

wandb = dict(
    enabled=True,
    entity=os.environ.get('WANDB_ENTITY', 'yuzeng-tri'),
    project=os.environ.get('WANDB_PROJECT') or 'a4d2',
    num_validation_logs=8,
)

any4d_config = dict(
    transforms='default',
    # AnyData "Unified" loader + precomputed caption embeddings (captions_t5_all),
    # registered in custom/dataset/adverse_weather_canny_dataset.py.
    data_library='adverse_weather_canny',
    dataloader='unified_anydrive_vid2vid',
    vae='a4d_vae_wan',
    use_wan_backbone=True,
    wan_variant='control',
    # PPD structured noise (https://arxiv.org/abs/2512.05106): replace ε with
    # phase-preserved noise so the clean latent's geometry is baked into the
    # diffusion forward process. cutoff_scale = mean of exponential dist;
    # PPD reference is 10.0. (Copied verbatim from the single-view config.)
    use_structured_noise=True,
    structured_noise_cutoff_scale=10.0,  # Exp(scale) for per-iter cutoff_radius (PPD ref)
    structured_noise_channels=4,
    structured_noise_seed=0,
    val_structured_noise=True,
    val_structured_noise_cutoff=10.0,
    # 11-view JOINT cross-view. Stream layout per view is 50 channels:
    #   [0:16]  noisy RGB latent           (xt / target)
    #   [16:17] input_mask                 (Any4D bookkeeping)
    #   [17:18] output_mask                (Any4D bookkeeping)
    #   [18:34] reference latent           (zeros for T2V mode; forced zero in DiT)
    #   [34:50] edge_control latent        (canny VAE-encoded by a4d_vae_wan)
    video_entries=_build_video_entries(NUM_VIEWS),
    num_views=NUM_VIEWS,
    # Gradient checkpointing over the Wan DiT block stack: the 11-view concatenated
    # token sequence makes self-attention activations dominate VRAM at 704x1280.
    wan_gradient_checkpointing=True,
    # SAC 'mm_only' (cosmos scheme): save matmuls + attention, recompute only cheap ops ->
    # NO attention recompute -> much faster than full-block checkpointing, at higher VRAM
    # (we had ~70GB free). Applied to every block (period=1); recompute is cheap.
    wan_sac_mode='mm_only',
    wan_grad_ckpt_period=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={f'rgb{v}': 1.0 for v in range(NUM_VIEWS)},
    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,
    # T=1 (single-frame): unified_anydrive's forecast task tries to randomly pick
    # cond_frames from range(1, (T*2)//3, 4) which is empty for T=1. Force
    # cond_frames_raw=0 so the whole frame is a generation target (text+canny only).
    train_directives=dict(cond_frames_raw=0),
    val_directives=dict(
        tasks='forecast',
        cond_frames_raw=0,
        remove_cond=True,
    ),
    data_train_overrides=dict(frame_stride=1, resize=[704, 1280]),  # 704x1280 (16:9, /16)
    # subsample=64 so the val set has enough samples to fill max_val_iter=8 across
    # dp_world=8 (8 samples shown, not 1). single_sample='first' keeps it deterministic.
    data_val_overrides=dict(frame_stride=1, subsample=64, resize=[704, 1280]),
    # Ego view only, to limit clutter across 11 views.
    track_metrics=dict(
        rgb0=['psnr', 'ssim'],
    ),
    # Train-visual saving runs on all ranks and is memory/IO-heavy; firing it every 99
    # steps stalled the box for ~2.3h once (host OOM-killed a dataloader worker -> crash).
    # 500 matches the validation cadence (validation_iter=500) -- 5x fewer visual passes.
    train_visuals_interval=500,
    train_visuals_detail=1,
    val_visuals_detail=1,
    visuals_quality=8,
    viz_input_blacklist=[],
    viz_mask_border_width=2,
    val_num_steps=35,
)

WAN_CTRL_CKPT_DIR = os.environ.get(
    'WAN_1_3B_CONTROL_CKPT_DIR',
    '/workspace/Any4D/checkpoints/wan/Wan2.1-Fun-V1.1-1.3B-Control',
)
WAN_VAE_PATH = os.environ.get(
    'WAN_VAE_PATH',
    f'{WAN_CTRL_CKPT_DIR}/Wan2.1_VAE.pth',
)

model = dict(
    dit_path=WAN_CTRL_CKPT_DIR,
    # Empty -> no UMT5 text encoder is loaded; the dataset (adverse_weather_canny)
    # injects precomputed UMT5 caption embeddings per sample, which
    # a4d_vae_wan.encode_text_as_needed consumes via its fast-path. Saves ~11GB
    # VRAM/GPU (enabling a larger batch). default_text_wan.pkl covers any
    # caption missing from the t5 store.
    text_encoder_path='',
    vae_path=WAN_VAE_PATH,
    fsdp_shard_size=8,
    # Karras EDM2 power-function EMA (rate=σ_rel=0.1 default). Builds a frozen dit_ema that
    # the training loop averages into; validation runs under ema_scope() so it evaluates the
    # EMA weights. EMA is seeded from the loaded init checkpoint by the checkpointer.
    ema_enabled=True,
    run_validation=True,
    # 'delay3' => first validation runs at step 3 (sanity check that EMA + validation work),
    # instead of the usual pre-training pass. Subsequent validations follow validation_iter.
    skip_first_validation='delay3',
    validation_iter=500,
    max_val_iter=8,  # weather-transfer: log ~8 clear-scene x rain/snow-prompt examples
    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=1,                   # T=1: single-frame image-mode training (Tl=1)
    tokenizer_chunk_duration=1,  # raw frames per VAE chunk = 1
)

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],
)

_AW_CANNY_YAML = 'custom/config/anydata/local/adverse_weather_canny_11view.yaml'
# Validation uses a DIFFERENT split: clear (non-rain/snow) scenes + rain/snow prompt
# override, for weather-transfer eval. 11-camera mirror of the single-view val yaml.
_AW_CANNY_VAL_YAML = 'custom/config/anydata/local/adverse_weather_canny_11view_val_transfer.yaml'

dataset_train = dict(
    config=_AW_CANNY_YAML,
    # 8 ranks x 8 workers x prefetch held a lot of host RAM with no swap to absorb a
    # spike; T=1 samples are tiny so 4 workers keep the GPUs fed at far lower peak RAM.
    num_workers=4,
    prefetch_factor=2,
    batch_size=1,  # overridden by the launcher's BATCH env (probed for largest fit)
)

dataset_val = dict(
    config=dict(
        AdverseWeatherCanny=_AW_CANNY_VAL_YAML,
    ),
    num_workers=0,
    batch_size=1,
)

metrics = None

cs = ConfigStore.instance()

this_config = template_any4d_wan_control_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,
)