from yacs.config import CfgNode as CN

from graphgym.register import register_config


def set_cfg_example(cfg):
    r'''
    This function sets the default config value for customized options
    :return: customized configuration use by the experiment.
    '''

    # ------------------------------------------------------------------------ #
    # Customized options
    # ------------------------------------------------------------------------ #

    # example argument
    cfg.example_arg = 'example'

    # example argument group
    cfg.example_group = CN()

    # then argument can be specified within the group
    cfg.example_group.example_arg = 'example'

    cfg.gnn.heads = 1
    cfg.gnn.concat = False
    cfg.gnn.add_attn_info = False
    cfg.gnn.convolve = False
    cfg.gnn.use_gcn_norm = False
    cfg.gnn.share_weights = False
    cfg.gnn.ln = True
    cfg.gnn.fixed_ansatz = False
    cfg.gnn.partial_ansatz = False
    cfg.gnn.lambda_policy = None
    cfg.gnn.gcn_mode = False
    cfg.gnn.freeze_pre_post = True
    cfg.gnn.eps = 0.0
    cfg.gnn.train_eps = False
    cfg.gnn.aggregators = ['mean', 'min', 'max', 'std']
    cfg.gnn.scalers = ['identity', 'amplification', 'attenuation']
    
    
    # Binary Contextual Stochastic Block Model (bcsbm)
    cfg.bcsbm = CN()
    cfg.bcsbm.n = 1000
    cfg.bcsbm.eps = 0.5
    cfg.bcsbm.p = 0.8
    cfg.bcsbm.q = 0.3
    cfg.bcsbm.mu_norm = 1.0
    cfg.bcsbm.d = 5
    cfg.bcsbm.std_dev = 1.0
    cfg.bcsbm.directed = False
    cfg.bcsbm.num_val = 100
    cfg.bcsbm.num_test = 100
    cfg.bcsbm.split = "full"
    cfg.bcsbm.version = None
    cfg.bcsbm.negative_slope = 0.2


register_config('example', set_cfg_example)
