"""Native Python defaults shared by generator training entries."""
import json
from model_resources import wan_base_root, wan_model_paths, wan_tokenizer_path


def set_training_defaults(parser, variant):
    model = wan_base_root()
    data = "data/vace_train_v2"
    caption_source = "teacher" if variant == "reward" else "grpo"
    defaults = dict(
        dataset_base_path=data,
        dataset_metadata_path=f"{data}/metadata.csv" if variant == "hoi" else f"{data}/metadata_caption_branch_{caption_source}.csv",
        data_file_keys="video,vace_video,vace_reference_image",
        extra_inputs="vace_video,vace_reference_image",
        height=480, width=480, num_frames=13, dataset_repeat=1, num_epochs=5,
        model_paths=json.dumps(wan_model_paths(model)),
        tokenizer_path=wan_tokenizer_path(model),
        use_gradient_checkpointing_offload=True,
        output_path=f"outputs/train_{variant}",
        learning_rate=1e-4 if variant == "hoi" else 3e-4,
    )
    if variant == "hoi":
        defaults.update(lora_base_model="vace", lora_target_modules="q,k,v,o,ffn.0,ffn.2", lora_rank=128)
    else:
        defaults.update(preset_lora_path="checkpoints/generator_hoi/ckpt.safetensors", preset_lora_model="vace", find_unused_parameters=True)
    parser.set_defaults(**defaults)
    # Upstream required options now have concrete package defaults.
    for action in parser._actions:
        if action.dest in defaults:
            action.required = False
    return parser