from collections import defaultdict

# The dict to be analyzed for all datasets and models. 
path_dict = defaultdict(dict)
# cora sage
store_key = ("cora", "SAGE")

name_model_pool = [("cora", "SAGE"), ("cora", "geomGCN"), ("citeseer", "SAGE"), ("citeseer", "geomGCN"), ("pubmed", "SAGE"), ("pubmed", "geomGCN"), ("a-computer", "SAGE"), ("a-computer", "exphormer"), ("a-photo", "SAGE"), ("a-photo", "exphormer"), ("ogbn-arxiv", "SAGE"), ("ogbn-arxiv", "DRGAT"), ("ogbn-products", "SAGE"), ("ogbn-products", "RevGNN-112")]

## the training configs of the SAGE models
train_conf_dict = defaultdict(dict)
# cora
train_conf_dict["cora"] = {
    "batch_size": 512,
    "epoch": 500,
    "lr": 0.01,
    "train_neighbor_size": 10,
    "hidden_layer": 3,
    "hidden_size": 64,
    "weight_decay": 0.0005,
    "dropout": 0.0
}

# citeseer
name = "citeseer"
train_conf_dict["citeseer"] = {# split_seed_pool = [0, 1, 2, 3, 4]
# split_seed_pool = [5, 6, 7, 8, 9]
    "batch_size": 512,
    "epoch": 500,
    "lr": 0.01,
    "train_neighbor_size": 5,
    "hidden_layer": 3,
    "hidden_size": 64,
    "weight_decay": 0.0005,
    "dropout": 0.0
}

# pubmed
train_conf_dict["pubmed"] = {
    "batch_size": 512,
    "epoch": 500,
    "lr": 0.01,
    "train_neighbor_size": 10,
    "hidden_layer": 3,
    "hidden_size": 64,
    "weight_decay": 0.0005,
    "dropout": 0.0
}

# a-computer
train_conf_dict["a-computer"] = {
    "batch_size": 512,
    "epoch": 500,
    "lr": 0.01,
    "train_neighbor_size": 5,
    "hidden_layer": 2,
    "hidden_size": 128,
    "weight_decay": 0.0005,
    "dropout": 0.0
}

# a-photo
train_conf_dict["a-photo"] = {
    "batch_size": 512,
    "epoch": 500,
    "lr": 0.01,
    "train_neighbor_size": 5,
    "hidden_layer": 2,
    "hidden_size": 128,
    "weight_decay": 0.0005,
    "dropout": 0.0
}

# reddit
train_conf_dict["reddit"] = {
    "batch_size": 1024,
    "epoch": 50,
    "lr": 0.001,
    "train_neighbor_size": 10,
    "hidden_layer": 2,
    "hidden_size": 256,
    "weight_decay": 0.0005,
    "dropout": 0.0
}

# flickr
train_conf_dict["flickr"] = {
    "batch_size": 1024,
    "epoch": 50,
    "lr": 0.001,
    "train_neighbor_size": 10,
    "hidden_layer": 2,
    "hidden_size": 256,
    "weight_decay": 0.0005,
    "dropout": 0.0
}

# OGBN-arxiv
train_conf_dict["ogbn-arxiv"] = {
    "batch_size": 4096,
    "epoch": 500,
    "lr": 0.01,
    "train_neighbor_size": 15,
    "hidden_layer": 3,
    "hidden_size": 256,
    "weight_decay": 0.0005,
    "dropout": 0.2
}

# OGBN-products GLNN
train_conf_dict["ogbn-products"] = {
    "batch_size": 512,
    "epoch": 200,
    "lr": 0.001,
    "train_neighbor_size": 10,
    "hidden_layer": 3,
    "hidden_size": 256,
    "weight_decay": 0.0005,
    "dropout": 0.3
}
