import argparse

def parse_args():
    parser = argparse.ArgumentParser()

    # Dataset
    parser.add_argument('--dataset', type=str, default='Cora',
                        help='Dataset Name')
    parser.add_argument('--diffusion_type', type=str, default='GraphMaker',
                        help='[DiGress, GraphMaker]')
    parser.add_argument('--imb_rate', type=float, default=0.05,
                        help='The Imbalance Rate')
    parser.add_argument('--lambda_', type=float, default=10,
                        help='The Importance of Synthetic Graphs')
    parser.add_argument('--imb_ratio', type=float, default=20,
                        help='Imbalance Ratio')
    parser.add_argument('--num_graphs', type=float, default=5,
                        help='# of Synthetic Graphs')
    # Architecture
    parser.add_argument('--net', type=str, default='GAT',
                        help='Architecture name')
    parser.add_argument('--n_layer', type=int, default=2,
                        help='the number of layers')
    parser.add_argument('--feat_dim', type=int, default=64,
                        help='Feature dimension')
    # GATpython 4.retrain_classifier.py

    parser.add_argument('--n_head', type=int, default=4,
                        help='the number of heads in GAT') 
    # Imbalance Loss
    parser.add_argument('--loss_type', type=str, default='ce',
                        help='Loss type')
    # Method
    parser.add_argument('--ens', action='store_false',
                        help='Mixing node')
    # Hyperparameter for our approach
    parser.add_argument('--keep_prob', type=float, default=0.01,
                        help='Keeping Probability')
    parser.add_argument('--pred_temp', type=float, default=10,
                        help='Prediction temperature')
    parser.add_argument('--warmup', type=int, default=10,
                        help='warmup')
    args = parser.parse_args()

    return args