import argparse
from .config import resolve_args


def parse_args():
    parser = argparse.ArgumentParser(description="Train Multics")
    parser.add_argument(
        "--task",
        type=str,
        default="train",
        required=False,
        choices=["train"],
        help="Train the hierarchical model.",
    )
    parser.add_argument(
        "--dataset_path",
        type=str,
        default=None,
        required=True,
        help="The path of the Dataset.",
    )
    parser.add_argument(
        "--output_path", type=str, default="./", help="Path to save the model."
    )
    parser.add_argument(
        "--ckpt_local_dir",
        type=str,
        default=None,
        help="If set, torch.save writes checkpoints here instead of the Lightning checkpoint directory.",
    )
    parser.add_argument(
        "--run_version",
        type=str,
        default=None,
        help="Lightning log directory name, for example version_76. If omitted, the next local version_N under the output path is used.",
    )
    parser.add_argument(
        "--text_encoder_path", type=str, default=None, help="Path of text encoder."
    )
    parser.add_argument(
        "--image_encoder_path", type=str, default=None, help="Path of image encoder."
    )
    parser.add_argument("--vae_path", type=str, default=None, help="Path of VAE.")
    parser.add_argument("--dit_path", type=str, default=None, help="Path of DiT.")
    parser.add_argument(
        "--steps_per_epoch", type=int, default=500, help="Number of steps per epoch."
    )
    parser.add_argument(
        "--dataloader_num_workers",
        type=int,
        default=4,
        help="Number of subprocesses to use for data loading. 0 means that the data will be loaded in the main process.",
    )
    parser.add_argument(
        "--learning_rate", type=float, default=1e-05, help="Learning rate."
    )
    parser.add_argument(
        "--accumulate_grad_batches",
        type=int,
        default=1,
        help="The number of batches in gradient accumulation.",
    )
    parser.add_argument(
        "--gradient_clip_val",
        type=float,
        default=1.0,
        help="Global gradient-norm clip. 0 disables clipping.",
    )
    parser.add_argument("--max_epochs", type=int, default=1, help="Number of epochs.")
    parser.add_argument(
        "--max_steps",
        type=int,
        default=-1,
        help="Maximum optimizer steps for this run; -1 disables the step limit.",
    )
    parser.add_argument(
        "--training_strategy",
        type=str,
        default="deepspeed_stage_1",
        choices=["auto", "deepspeed_stage_1", "deepspeed_stage_2", "deepspeed_stage_3"],
        help="Training strategy",
    )
    parser.add_argument(
        "--use_gradient_checkpointing",
        default=False,
        action="store_true",
        help="Whether to use gradient checkpointing.",
    )
    parser.add_argument(
        "--use_gradient_checkpointing_offload",
        default=False,
        action="store_true",
        help="Whether to use gradient checkpointing offload.",
    )
    parser.add_argument(
        "--use_swanlab",
        default=False,
        action="store_true",
        help="Whether to use SwanLab logger.",
    )
    parser.add_argument(
        "--swanlab_mode", default=None, help="SwanLab mode (cloud or local)."
    )
    parser.add_argument(
        "--val_infer",
        action="store_true",
        help="Run causal inference on the first --val_num_samples training rows and log the generated and reference videos.",
    )
    parser.add_argument(
        "--val_num_samples", type=int, default=4, help="Number of leading training samples used for validation."
    )
    parser.add_argument(
        "--val_every_n_steps",
        type=int,
        default=500,
        help="Run validation every this many training batches.",
    )
    parser.add_argument(
        "--val_num_inference_steps",
        type=int,
        default=50,
        help="Denoising steps used during validation inference.",
    )
    parser.add_argument(
        "--val_causal_num_frames",
        type=int,
        default=None,
        help="Frames to generate during validation. Defaults to target_rgb_frames.",
    )
    parser.add_argument(
        "--val_output_dir",
        type=str,
        default="outputs/validation",
        help="Directory for validation MP4 files.",
    )
    parser.add_argument(
        "--val_grid_fps", type=int, default=15, help="Frames per second for saved validation videos."
    )
    parser.add_argument("--metadata_file_name", type=str, default="metadata.csv")
    parser.add_argument("--resume_ckpt_path", type=str, default=None)
    parser.add_argument(
        "--train_step_offset",
        type=int,
        default=None,
        help="Completed optimizer updates before this run. Default: infer from the resume checkpoint's parent hparams.yaml and stepN filename.",
    )
    parser.add_argument(
        "--num_players",
        type=int,
        default=4,
        help="Number of player streams in the precomputed dataset.",
    )
    parser.add_argument(
        "--use_simplex_rope",
        action="store_true",
        default=True,
        help="Use simplex 4D RoPE instead of learned player embeddings.",
    )
    parser.add_argument(
        "--simplex_pool_size",
        type=int,
        default=4,
        help="Simplex vertex-pool size. Training draws num_players vertices from this pool.",
    )
    parser.add_argument(
        "--simplex_gamma", type=float, default=1.0, help="Rotation strength of the simplex player axis."
    )
    parser.add_argument(
        "--simplex_train_permute",
        action="store_true",
        help="Randomly permute player slots at each step, together with the sampled vertices.",
    )
    parser.add_argument(
        "--no_simplex_train_permute",
        action="store_false",
        dest="simplex_train_permute",
        help="Do not permute player streams. Vertex sampling stays on its own switch.",
    )
    parser.add_argument(
        "--simplex_train_sample_vertex_ids",
        action="store_true",
        help="Sample P vertex IDs from the simplex pool without replacement, without reordering player streams.",
    )
    parser.add_argument(
        "--training_stage",
        type=str,
        default="stage2",
        choices=["stage2"],
        help=argparse.SUPPRESS,
    )
    parser.add_argument(
        "--loss_type",
        type=str,
        default="diffusion_forcing",
        choices=["diffusion_forcing"],
        help="Training objective. Fixed to diffusion forcing.",
    )
    parser.add_argument(
        "--attention_mode",
        type=str,
        default="hier_hub",
        choices=["hier_hub"],
        help="Attention layout. Fixed to hierarchical hubs.",
    )
    parser.add_argument("--local_attn_size", type=int, default=7)
    parser.add_argument("--local_attn_size_after", type=int, default=21)
    parser.add_argument("--local_attn_switch_step", type=int, default=10000)
    parser.add_argument("--num_frame_per_block", type=int, default=7)
    parser.add_argument(
        "--num_local_hub_tokens",
        type=int,
        default=5,
        help="hier_hub: local hub tokens K_L.",
    )
    parser.add_argument(
        "--num_global_hub_tokens",
        type=int,
        default=6,
        help="hier_hub: global hub tokens K_G.",
    )
    parser.add_argument(
        "--hier_no_local_hub",
        action="store_true",
        help="Ablation: remove local hubs and connect every player directly to the global hub.",
    )
    parser.add_argument(
        "--hier_local_hub_mode",
        type=str,
        default="shared",
        choices=["shared", "per_group"],
        help="hier_hub local seed parameterization: shared uses [K_L,D] for all dynamic group slots; per_group learns independent [N,K_L,D] seeds.",
    )
    parser.add_argument(
        "--hier_direct_schedule",
        type=str,
        default="none",
        choices=["none", "uniform", "explicit"],
        help="hier_hub: selected layers with same-block same-group direct player attention.",
    )
    parser.add_argument(
        "--hier_direct_frac",
        type=float,
        default=0.0,
        help="hier_hub uniform schedule only: fraction in [0,1] of DiT layers that enable direct same-group, same-latent-block player-player attention. Ignored when --hier_direct_schedule is none or explicit; same-player self-history is always controlled by the base mask.",
    )
    parser.add_argument(
        "--hier_direct_layers",
        type=str,
        default=None,
        help="hier_hub explicit zero-based layer ids, e.g. 2,5,8,11.",
    )
    parser.add_argument("--group_path_column", type=str, default="group_path")
    parser.add_argument(
        "--hier_group_format",
        type=str,
        default="labels",
        choices=["labels", "membership"],
        help="hier_hub grouping tensor format. labels: partition ids in [0,N), no player↔global (default, bit-identical to previous runs). membership: cover with membership[T,N,P] + global_only[T,P].",
    )
    parser.add_argument(
        "--hier_group_mode",
        type=str,
        default="csv",
        choices=["csv", "random_static", "random_block"],
        help="hier_hub grouping: CSV labels, one fixed random permutation, or one permutation per latent block.",
    )
    parser.add_argument("--hier_group_seed", type=int, default=0)
    parser.add_argument(
        "--hier_random_preserve_sizes",
        type=int,
        default=1,
        choices=[0, 1],
        help="Random grouping: preserve each source block's group-size histogram when set to 1.",
    )
    parser.add_argument(
        "--group_label_player_ids",
        type=str,
        default=None,
        help="1-based original player ids for group label columns; default 1..P.",
    )
    parser.add_argument(
        "--attn_impl",
        type=str,
        default="sdpa",
        choices=["sdpa", "flex", "hier_tiled_ste"],
        help="Self-attention kernel: sdpa, flex, or hier_tiled_ste. hier_tiled_ste uses tiled forward and Flex backward, and requires hier_hub.",
    )
    parser.add_argument(
        "--use_camera_pose",
        action="store_true",
        help="Add relative camera conditioning to self-attention Q and K. The CSV must contain precomputed camera_pose_path* or camera12_path*.",
    )
    parser.add_argument(
        "--camera_dim",
        type=int,
        default=12,
        help="Camera-pose embedding size. The default is the 12-D relative camera.",
    )
    parser.add_argument(
        "--camera_dropout_p",
        type=float,
        default=0.1,
        help="Shared dropout on the camera Q/K adapter. Used only while training. Default 0.1.",
    )
    parser.add_argument(
        "--unfreeze_self_attn",
        action="store_true",
        help="Train the self-attention weights. On by default.",
    )
    parser.add_argument(
        "--causal_forcing_training",
        action="store_true",
        default=False,
        help="Causal-forcing block training. A supervised block is chosen at random; earlier blocks are teacher-forced from clean latents. Inference context noise should stay at 0 when this flag is active.",
    )
    parser.add_argument(
        "--clean_prefix_prob",
        type=float,
        default=0.0,
        help="Probability of training on a clean prefix and supervising only the current block. 0 disables it. A typical range is 0.5 to 0.8.",
    )
    parser.add_argument(
        "--clean_prefix_context_noise",
        type=int,
        default=0,
        help="Context timestep used for clean-prefix blocks. Keep it equal to the inference context noise. 0 means fully clean.",
    )
    parser.add_argument(
        "--generated_history_prob",
        type=float,
        default=0.0,
        help="Fraction of batches that train on generated history: a ground-truth prefix, one to three generated blocks, and loss only on the next ground-truth block. 0 disables it.",
    )
    parser.add_argument(
        "--df_per_frame_noise",
        action="store_true",
        help="Sample an independent noise level for each latent frame instead of sharing one level across a block.",
    )
    parser.add_argument(
        "--target_rgb_frames",
        type=int,
        default=81,
        help="RGB frames in the supervised clip. Must be 4n+1. Default 81. The diffusion loss and action slices use this window.",
    )
    parser.add_argument(
        "--blank_prompt_emb_path",
        type=str,
        default=None,
        help="Cached embedding of the empty text prompt. Defaults to asset/blank_string_prompt_emb.pth. If the file is missing, training encodes the empty string and writes the cache.",
    )
    parser.add_argument(
        "--text_prompt_dropout_p",
        type=float,
        default=0.0,
        help="Probability of replacing a sample's text embedding with the empty-prompt embedding. 0 disables it. The bundled training configs use 0.1.",
    )
    parser.add_argument(
        "--vae_latent_channels",
        type=int,
        default=48,
        help="VAE latent channels. 48 for the Wan 5B VAE.",
    )
    parser.add_argument(
        "--first_frame_dropout_p",
        type=float,
        default=0.0,
        help="Probability of dropping the clean first frame during training and applying loss to every frame. Inference still uses a clean first frame. 0 disables it.",
    )
    parser.add_argument(
        "--first_frame_noise_p",
        type=float,
        default=0.0,
        help="Probability of noising the first latent frame of the anchor. 0 disables it.",
    )
    parser.add_argument(
        "--first_frame_noise_corrupt_ratio",
        type=float,
        default=1 / 3,
        help="Upper bound of the uniform noise scale applied to the anchor's first frame. Default 1/3.",
    )
    parser.add_argument(
        "--first_frame_noise_min_seg",
        type=int,
        default=1,
        help="Apply anchor corruption only at or after this segment index.",
    )
    parser.add_argument(
        "--perf_log_every_n_steps",
        type=int,
        default=100,
        help="Print averaged timing and the bottleneck every N steps.",
    )
    parser.add_argument(
        "--no_perf_timing",
        action="store_true",
        help="Disable step timing. Skips the extra CUDA sync and the perf logs.",
    )
    args = resolve_args(parser, "train")
    if args.hier_group_format == "membership" and args.attn_impl == "hier_tiled_ste":
        raise ValueError(
            "hier_group_format=membership is not supported with --attn_impl hier_tiled_ste; use flex or sdpa"
        )
    return args