"""Best hyperparameters found."""
import torch

MWE_GCN_proteins = {
    'num_ew_channels': 8,
    'num_epochs': 2000,
    'in_feats': 1,
    'hidden_feats': 10,
    'out_feats': 112,      
    'n_layers': 3,
    'lr': 2e-2,
    'weight_decay': 0,
    'patience': 1000,
    'dropout': 0.2,           
    'aggr_mode': 'sum',     ## 'sum' or 'concat' for the aggregation across channels
    'ewnorm': 'both'
    }

MWE_DGCN_proteins = {
    'num_ew_channels': 8,
    'num_epochs': 2000,
    'in_feats': 1,
    'hidden_feats': 10,
    'out_feats': 112,     
    'n_layers': 2,
    'lr': 1e-2,
    'weight_decay': 0,
    'patience': 300,
    'dropout': 0.5,           
    'aggr_mode': 'sum',   
    'residual': True,
    'ewnorm': 'none'
    }


def get_exp_configure(args):
    if (args['model'] == 'MWE-GCN'):
        return MWE_GCN_proteins
    elif (args['model'] == 'MWE-DGCN'):
        return MWE_DGCN_proteins