import os
import pytest

@pytest.mark.parametrize('data', ['cora', 'citeseer', 'pubmed', 'csv', 'reddit',
                                  'co-buy-computer', 'ogbn-arxiv', 'ogbn-products'])
def test_nodepred_data(data):
    os.system(f'dgl configure nodepred --data {data} --model gcn')
    assert os.path.exists(f'nodepred_{data}_gcn.yaml')

    custom_cfg = f'custom_{data}_gcn.yaml'
    os.system(f'dgl configure nodepred --data {data} --model gcn --cfg {custom_cfg}')
    assert os.path.exists(custom_cfg)

    custom_script = f'{data}_gcn.py'
    os.system(f'dgl export --cfg {custom_cfg} --output {custom_script}')
    assert os.path.exists(custom_script)

@pytest.mark.parametrize('model', ['gcn', 'gat', 'sage', 'sgc', 'gin'])
def test_nodepred_model(model):
    os.system(f'dgl configure nodepred --data cora --model {model}')
    assert os.path.exists(f'nodepred_cora_{model}.yaml')

    custom_cfg = f'custom_cora_{model}.yaml'
    os.system(f'dgl configure nodepred --data cora --model {model} --cfg {custom_cfg}')
    assert os.path.exists(custom_cfg)

    custom_script = f'cora_{model}.py'
    os.system(f'dgl export --cfg {custom_cfg} --output {custom_script}')
    assert os.path.exists(custom_script)

@pytest.mark.parametrize('data', ['cora', 'citeseer', 'pubmed', 'csv', 'reddit',
                                  'co-buy-computer', 'ogbn-arxiv', 'ogbn-products'])
def test_nodepred_ns_data(data):
    os.system(f'dgl configure nodepred-ns --data {data} --model gcn')
    assert os.path.exists(f'nodepred-ns_{data}_gcn.yaml')

    custom_cfg = f'ns-custom_{data}_gcn.yaml'
    os.system(f'dgl configure nodepred-ns --data {data} --model gcn --cfg {custom_cfg}')
    assert os.path.exists(custom_cfg)

    custom_script = f'ns-{data}_gcn.py'
    os.system(f'dgl export --cfg {custom_cfg} --output {custom_script}')
    assert os.path.exists(custom_script)


@pytest.mark.parametrize('model', ['gcn', 'gat', 'sage'])
def test_nodepred_ns_model(model):
    os.system(f'dgl configure nodepred-ns --data cora --model {model}')
    assert os.path.exists(f'nodepred-ns_cora_{model}.yaml')

    custom_cfg = f'ns-custom_cora_{model}.yaml'
    os.system(f'dgl configure nodepred-ns --data cora --model {model} --cfg {custom_cfg}')
    assert os.path.exists(custom_cfg)

    custom_script = f'ns-cora_{model}.py'
    os.system(f'dgl export --cfg {custom_cfg} --output {custom_script}')
    assert os.path.exists(custom_script)

@pytest.mark.parametrize('data', ['cora', 'citeseer', 'pubmed', 'csv', 'reddit',
                                  'co-buy-computer', 'ogbn-arxiv', 'ogbn-products', 'ogbl-collab',
                                  'ogbl-citation2'])
def test_linkpred_data(data):
    node_model = 'gcn'
    edge_model = 'ele'
    neg_sampler = 'global'
    custom_cfg = '_'.join([data, node_model, edge_model, neg_sampler]) + '.yaml'
    os.system('dgl configure linkpred --data {} --node-model {} --edge-model {} --neg-sampler {} --cfg {}'.format(
        data, node_model, edge_model, neg_sampler, custom_cfg))
    assert os.path.exists(custom_cfg)

    custom_script = '_'.join([data, node_model, edge_model, neg_sampler]) + '.py'
    os.system('dgl export --cfg {} --output {}'.format(custom_cfg, custom_script))
    assert os.path.exists(custom_script)

@pytest.mark.parametrize('node_model', ['gcn' ,'gat', 'sage', 'sgc', 'gin'])
def test_linkpred_node_model(node_model):
    data = 'cora'
    edge_model = 'ele'
    neg_sampler = 'global'
    custom_cfg = '_'.join([data, node_model, edge_model, neg_sampler]) + '.yaml'
    os.system('dgl configure linkpred --data {} --node-model {} --edge-model {} --neg-sampler {} --cfg {}'.format(
        data, node_model, edge_model, neg_sampler, custom_cfg))
    assert os.path.exists(custom_cfg)

    custom_script = '_'.join([data, node_model, edge_model, neg_sampler]) + '.py'
    os.system('dgl export --cfg {} --output {}'.format(custom_cfg, custom_script))
    assert os.path.exists(custom_script)


@pytest.mark.parametrize('edge_model', ['ele', 'bilinear'])
def test_linkpred_edge_model(edge_model):
    data = 'cora'
    node_model = 'gcn'
    neg_sampler = 'global'
    custom_cfg = '_'.join([data, node_model, edge_model, neg_sampler]) + '.yaml'
    os.system('dgl configure linkpred --data {} --node-model {} --edge-model {} --neg-sampler {} --cfg {}'.format(
        data, node_model, edge_model, neg_sampler, custom_cfg))
    assert os.path.exists(custom_cfg)

    custom_script = '_'.join([data, node_model, edge_model, neg_sampler]) + '.py'
    os.system('dgl export --cfg {} --output {}'.format(custom_cfg, custom_script))
    assert os.path.exists(custom_script)


@pytest.mark.parametrize('neg_sampler', ['global', 'persource', ''])
def test_linkpred_neg_sampler(neg_sampler):
    data = 'cora'
    node_model = 'gcn'
    edge_model = 'ele'
    custom_cfg = f'{data}_{node_model}_{edge_model}_{neg_sampler}.yaml'
    if neg_sampler == '':
        os.system('dgl configure linkpred --data {} --node-model {} --edge-model {} --cfg {}'.format(
            data, node_model, edge_model, custom_cfg))
    else:
        os.system('dgl configure linkpred --data {} --node-model {} --edge-model {} --neg-sampler {} --cfg {}'.format(
            data, node_model, edge_model, neg_sampler, custom_cfg))
    assert os.path.exists(custom_cfg)

    custom_script = f'{data}_{node_model}_{edge_model}_{neg_sampler}.py'
    os.system('dgl export --cfg {} --output {}'.format(custom_cfg, custom_script))
    assert os.path.exists(custom_script)

@pytest.mark.parametrize('data', ['csv', 'ogbg-molhiv', 'ogbg-molpcba'])
@pytest.mark.parametrize('model', ['gin', 'pna'])
def test_graphpred(data, model):
    os.system('dgl configure graphpred --data {} --model {}'.format(data, model))
    assert os.path.exists('graphpred_{}_{}.yaml'.format(data, model))

    custom_cfg = 'custom_{}_{}.yaml'.format(data, model)
    os.system('dgl configure graphpred --data {} --model {} --cfg {}'.format(data, model,
                                                                             custom_cfg))
    assert os.path.exists(custom_cfg)

    custom_script = '_'.join([data, model]) + '.py'
    os.system('dgl export --cfg {} --output {}'.format(custom_cfg, custom_script))
    assert os.path.exists(custom_script)

@pytest.mark.parametrize('recipe',
                         ['graphpred_hiv_gin.yaml',
                          'graphpred_hiv_pna.yaml',
                          'graphpred_pcba_gin.yaml',
                          'linkpred_cora_sage.yaml',
                          'linkpred_citation2_sage.yaml',
                          'linkpred_collab_sage.yaml',
                          'nodepred_citeseer_gat.yaml',
                          'nodepred_citeseer_gcn.yaml',
                          'nodepred_citeseer_sage.yaml',
                          'nodepred_cora_gat.yaml',
                          'nodepred_cora_gcn.yaml',
                          'nodepred_cora_sage.yaml',
                          'nodepred_pubmed_gat.yaml',
                          'nodepred_pubmed_gcn.yaml',
                          'nodepred_pubmed_sage.yaml',
                          'nodepred-ns_arxiv_gcn.yaml',
                          'nodepred-ns_product_sage.yaml'])
def test_recipe(recipe):
    # Remove all generated yaml files
    current_dir = os.listdir("./")
    for item in current_dir:
        if item.endswith(".yaml"):
            os.remove(item)

    os.system('dgl recipe get {}'.format(recipe))
    assert os.path.exists(recipe)

def test_node_cora():
    os.system('dgl configure nodepred --data cora --model gcn')
    os.system('dgl train --cfg nodepred_cora_gcn.yaml')
    assert os.path.exists('results')
    assert os.path.exists('results/run_0.pth')
    os.system('dgl configure-apply nodepred --cpt results/run_0.pth')
    assert os.path.exists('apply_nodepred_cora_gcn.yaml')
    os.system('dgl configure-apply nodepred --data cora --cpt results/run_0.pth --cfg apply.yaml')
    assert os.path.exists('apply.yaml')
    os.system('dgl apply --cfg apply.yaml')
    assert os.path.exists('apply_results/output.csv')
    os.system('dgl export --cfg apply.yaml --output apply.py')
    assert os.path.exists('apply.py')