# parameter sets of model thp

realworld_training_hyperparameter_list = {
    "file_name": "start.py",
    "single": [
        "--no_seed",
        "--dataloader_name", "generic",
        "--model_name", "thp",
        "--lr", "0.002",
        "--save_mode", "best",
        "--lr_sched",
        "--op_name", "AdamW",
        "--optim_config", "optimizer.yml",
        "--n_cycles", "0.5",
        "--wandb"
    ],
    'index':
    {
        'multiple':[
            "--dataset_name", ["stackoverflow", "retweet", "mooc", "bookorder"],
            "--dataloader_config", ["stackoverflow/thp_dl.yml", "retweet/thp_dl.yml", "mooc/thp_dl.yml", "bookorder/thp_dl.yml"],
            "--model_config", ["stackoverflow/thp.yml", "retweet/thp.yml", "mooc/thp.yml", "bookorder/thp.yml"],
            "--n_training_steps", ["200000", "400000", "400000", "20000"],
            "--n_evaluation_steps", ["2000", "4000", "4000", "500"],
            "--n_report_steps", ["2000", "4000", "4000", "4000"],
            "-tb", ["32", "32", "32", "8"],
            "-eb", ["32", "32", "32", "8"],
            "--n_warmup_steps", ["40000", "80000", "80000", "4000"],
        ]
    }
}

syn_training_hyperparameter_list = {
    "file_name": "start.py",
    "single": [
        "--no_seed",
        "--dataloader_name", "generic",
        "--n_training_steps", "100000",
        "--n_evaluation_steps", "2000",
        "--n_report_steps", "2000",
        "-tb", "32",
        "-eb", "128",
        "--n_warmup_steps", "20000",
        "--model_name", "thp",
        "--lr", "0.002",
        "--save_mode", "best",
        "--lr_sched",
        "--op_name", "AdamW",
        "--optim_config", "optimizer.yml",
        "--model_config", "syn/thp.yml",
        "--n_cycles", "0.5",
        "--wandb"
    ],
    "index":{
        'multiple':[
            "--dataset_name", ["hawkes_1_v2", "hawkes_2_v2", "poisson_v2", "self_correct_v2", "stationary_renewal_v2"],
        ]
    }
}

realworld_plot_hyperparameter_list = {
    "file_name": "start.py",
    "single": [
        "--seed", "32",
        "--model_name", "thp",
        "--lr", "0.002",
        "--dataloader_name", "generic",
        "--figure_count", "1",
        # "--train",
        # "--evaluation",
        "--test",
        "--used_dataloader_config", "thp_dl.yml",
        "--resolution", "200",
    ],
    'index': {
        'multiple':[
            "--dataset_name", ["retweet", "mooc", "bookorder", "stackoverflow"],
            "--model_config", ["retweet/thp.yml", "mooc/thp.yml", "bookorder/thp.yml", "stackoverflow/thp.yml"],
            "--n_training_steps", ["400000", "400000","20000", "200000"],
            "--dataloader_config", ["retweet/plot.yml", "mooc/plot.yml", "bookorder/plot.yml", "stackoverflow/plot.yml"],
            "--used_batch_size", ["32", "32", "8", "32"],
        ]
    },
    "counting": {
        'index':{
            'multiple':[
                "--plot_type", ["intensity", "probability", "debug", "debug", "debug"],
                "--task_name", ['graph', 'graph', 'graph', 'mae_and_f1', 'mae_e_and_f1']
            ]
        }
    }
}

syn_plot_hyperparameter_list = {
    "file_name": "start.py",
    "single": [
        "--seed", "32",
        "--model_name", "thp",
        "--lr", "0.002",
        "--dataloader_name", "generic",
        "--figure_count", "1",
        "--n_training_steps", "100000",
        # "--train",
        # "--evaluation",
        "--test",
        "--resolution", "200",
        "--used_batch_size", "32",
        "--dataloader_config", "syn/plot.yml",
        "--model_config", "syn/thp.yml",
    ],
    'index': {
        'multiple':[
            "--dataset_name", ["hawkes_1_v2", "hawkes_2_v2", "poisson_v2", "self_correct_v2", "stationary_renewal_v2"],
        ]
    },
    "counting": {
        'index':{
            "multiple":[
                "--plot_type", ["intensity", "probability", "debug", "debug"],
                "--task_name", ['graph', 'graph', 'graph', 'spearman_and_l1']
            ]
        }
    }
}

training_hyperparameter = {
    'real': realworld_training_hyperparameter_list,
    'syn': syn_training_hyperparameter_list
}

plot_hyperparameter = {
    'real': realworld_plot_hyperparameter_list,
    'syn': syn_plot_hyperparameter_list
}