import argparse


def get_args() -> argparse.Namespace:
    """ Get the arguments.

    Returns:
        argparse.Namespace: The arguments.
    """

    parser = argparse.ArgumentParser()

    parser.add_argument(
        '-c',
        '--config',
        type=str,
        required=True,
        help='The configuration file path.',
    )
    parser.add_argument(
        '-mc',
        '--model_config',
        type=str,
        required=False,
        help='The model configuration file path (only for MoLoS model).',
    )
    parser.add_argument(
        '-dsc',
        '--deepspeed_config',
        type=str,
        required=True,
        help='The DeepSpeed configuration file path.',
    )
    parser.add_argument(
        '-lc',
        '--lora_config',
        type=str,
        required=False,
        help='The LoRA configuration file path.',
    )
    parser.add_argument(
        '-m',
        '--mode',
        type=str,
        choices=[
            'chat',
            'infer',
            'test',
            'train',
        ],
        required=True,
        help='The mode.',
    )
    parser.add_argument(
        '-mt',
        '--model_type',
        default='molos',
        type=str,
        choices=[
            'molos',
            'mosle',
            'vanilla',
        ],
        help='The model type.',
    )
    parser.add_argument(
        '-is_oriup',
        '--is_original_upcycling',
        action='store_true',
        help='Initialize the model with original upcycling (only for MoLoS).',
    )
    parser.add_argument(
        '-cei',
        '--chosen_expert_index',
        default=-2,
        type=int,
        help='The chosen expert index (only for MoLoS model testing).',
    )
    parser.add_argument(
        '-v',
        '--version',
        type=str,
        help='The version of this execution.',
    )
    parser.add_argument(
        '--local_rank',
        default=0,
        type=int,
        help='The local rank. (Auto-filled by DeepSpeed)',
    )

    return parser.parse_args()
