from argparse import ArgumentParser


def continual_args(parser: ArgumentParser):
    """Adds continual learning arguments to a parser.

    Args:
        parser (ArgumentParser): parser to add dataset args to.
    """
    # base continual learning args
    parser.add_argument("--num_tasks", type=int, default=2)
    parser.add_argument("--task_idx", type=int, required=True)

    SPLIT_STRATEGIES = ["class", "data", "domain"]
    parser.add_argument("--split_strategy", choices=SPLIT_STRATEGIES, type=str, required=True)

    # distillation args
    parser.add_argument("--distiller", type=str, default=None)

    # my methods args
    parser.add_argument("--prune_perc", type=float, default=0.5)
    parser.add_argument("--retrain_start_epoch", type=int, default=500)
    parser.add_argument("--my_method", type=str, default=None)
    parser.add_argument("--ea", action="store_true")
    parser.add_argument("--addition_images", type=int, default=12500)
    parser.add_argument("--addition_classes", type=int, default=1000)
    parser.add_argument("--dataselect", type=str, default=None, nargs="+")
    parser.add_argument("--sample_method", type=str, default='random')
    parser.add_argument("--distill_lam", type=float, default=0.5)
