"""Save the initialized Any4D+Wan DiT as a single-file .pt checkpoint.

Builds Any4DWanPipeline via `build_wan_pipeline` (which loads Wan pretrained weights
into Any4DWanDiT) and saves pipe.dit.state_dict() with a `net.` prefix. Any4D-added
parameters (view_embs, action_embedder, sattn/xattn projections) remain at their
zero-init values so the checkpoint is a ControlNet-style identity relative to Wan.

Usage (inside cosmos-predict2-local docker):
    python -m custom.eval.save_init_checkpoint_wan \\
        --dit_path  /path/to/Wan2.1-Fun-1.3B-InP \\
        --vae_path  /path/to/Wan2.1_VAE.pth \\
        --out       /path/to/init_any4d_wan.pt
"""

import argparse
import copy
import os
import sys

import torch
import torch.nn as nn

ANY4D_ROOT = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
if ANY4D_ROOT not in sys.path:
    sys.path.insert(0, ANY4D_ROOT)
os.chdir(ANY4D_ROOT)


def _init_dist_single_gpu():
    if torch.distributed.is_initialized():
        return
    import random
    torch.distributed.init_process_group(
        backend='gloo',
        init_method=f'tcp://127.0.0.1:{random.randint(29000, 30000)}',
        world_size=1, rank=0,
    )


def _build_wan_pipe_cfg(vae_path: str):
    """Return a fresh copy of ANY4D_WAN_PIPELINE_1_3B with the tokenizer pointed at vae_path."""
    import copy as _copy
    from cosmos_predict2.configs.base.defaults.model import ANY4D_WAN_PIPELINE_1_3B

    pipe_cfg = _copy.deepcopy(ANY4D_WAN_PIPELINE_1_3B)
    pipe_cfg.tokenizer.vae_pth = vae_path
    if pipe_cfg.extra_nets is None:
        pipe_cfg.extra_nets = nn.ModuleDict()
    return pipe_cfg


def build_pipeline(dit_path: str, vae_path: str):
    """Build Any4DWanPipeline with Wan weights loaded. Returns (pipe, cfg)."""
    _init_dist_single_gpu()

    from custom.any4d.a4d_config import Any4DConfig
    from custom.any4d import a4d_surgery
    from custom.wan.a4d_pipe_wan import build_wan_pipeline

    pipe_cfg = _build_wan_pipe_cfg(vae_path)

    # Mirror the Any4D config used in rwm3_drive40h_wan.py — 3 views, traj, Wan ref entries.
    cfg = Any4DConfig(
        vidar_active=True, any4d_active=True, use_wan_backbone=True,
        num_views=3,
        video_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)),
            dict(rgb1=(0, 16, 'copy:rgb0/load', 'copy:rgb0/load'),
                 rgb1_input_mask=(16, 17, 'copy:rgb0_input_mask/load', None),
                 rgb1_output_mask=(17, 18, 'copy:rgb0_output_mask/load', None),
                 rgb1_ref=(18, 34, 'copy:rgb0_ref/load', None)),
            dict(rgb2=(0, 16, 'copy:rgb0/load', 'copy:rgb0/load'),
                 rgb2_input_mask=(16, 17, 'copy:rgb0_input_mask/load', None),
                 rgb2_output_mask=(17, 18, 'copy:rgb0_output_mask/load', None),
                 rgb2_ref=(18, 34, 'copy:rgb0_ref/load', None)),
        ],
        lowdim_adaln_entries=dict(
            traj=(0, 41, 0, 2, 'zero/load'),
            traj_input_mask=(0, 41, 2, 3, 'zero/load'),
        ),
        video_concat_mode='view', video_proj_mode='per_view',
        view_timestep_mode='per_view', block_gate_fix=True,
        harmonize_streams=True, harmonize_frames=True,
        load_modals=['rgb', 'traj'],
        pipe_config=pipe_cfg,
    )
    cfg = a4d_surgery.validate_config(cfg)

    print(f'[build] loading Wan weights from {dit_path} ...')
    pipe = build_wan_pipeline(
        pipe_cfg,
        dit_path=dit_path,
        text_encoder_path='',
        vae_path=vae_path,
        any4d_config=cfg,
    )
    print(f'[build] done. DiT has {sum(p.numel() for p in pipe.dit.parameters()):,} params')
    return pipe, cfg


def save_init_checkpoint(pipe, out_path: str):
    sd = {f'net.{k}': v.detach().cpu() for k, v in pipe.dit.state_dict().items()}
    os.makedirs(os.path.dirname(os.path.abspath(out_path)) or '.', exist_ok=True)
    torch.save(sd, out_path)
    print(f'[save] wrote {out_path}')
    print(f'[save] {len(sd)} tensors, {sum(v.numel() for v in sd.values()):,} params')


def main():
    p = argparse.ArgumentParser()
    p.add_argument('--dit_path', type=str, required=True,
                   help='Wan 2.1 Fun 1.3B InP dir or .safetensors file')
    p.add_argument('--vae_path', type=str, required=True)
    p.add_argument('--out', type=str, required=True)
    args = p.parse_args()

    pipe, cfg = build_pipeline(args.dit_path, args.vae_path)
    save_init_checkpoint(pipe, args.out)


if __name__ == '__main__':
    main()