import os
import shutil
import pytest
import numpy as np
import pandas as pd
import h5py
import tempfile

from bmtk.builder.network_adaptors.dm_network import DenseNetwork
# from bmtk.builder.network_adaptors.dm_network_orig import DenseNetworkOrig


@pytest.mark.parametrize('network_cls', [
    (DenseNetwork),
    # (DenseNetworkOrig)  # dropped support for original dense network
])
def test_save_nsyn_table(network_cls):
    np.random.seed(100)
    net = network_cls('NET1')
    net.add_nodes(N=10, position=[(0.0, 1.0, -1.0)]*10, cell_type='Scnna1', ei='e')
    net.add_nodes(N=10, position=[(0.0, 1.0, -1.0)]*10, cell_type='PV1', ei='i')
    net.add_nodes(N=10, position=[(0.0, 1.0, -1.0)]*10, tags=np.linspace(0, 100, 10), cell_type='PV2', ei='i')
    net.add_edges(source={'ei': 'i'}, target={'ei': 'e'}, connection_rule=lambda s, t: 1,
                  p1='e2i', p2='e2i')
    net.add_edges(source=net.nodes(cell_type='Scnna1'), target=net.nodes(cell_type='PV1'),
                  connection_rule=lambda s, t: 2, p1='s2p')
    net.build()
    nodes_h5 = tempfile.NamedTemporaryFile(suffix='.h5')
    nodes_csv = tempfile.NamedTemporaryFile(suffix='.csv')
    edges_h5 = tempfile.NamedTemporaryFile(suffix='.h5')
    edges_csv = tempfile.NamedTemporaryFile(suffix='.csv')

    net.save_nodes(nodes_h5.name, nodes_csv.name)
    net.save_edges(edges_h5.name, edges_csv.name)

    assert(os.path.exists(nodes_h5.name) and os.path.exists(nodes_csv.name))
    node_types_df = pd.read_csv(nodes_csv.name, sep=' ')
    assert(len(node_types_df) == 3)
    assert('cell_type' in node_types_df.columns)
    assert('ei' in node_types_df.columns)
    assert('positions' not in node_types_df.columns)

    nodes_h5 = h5py.File(nodes_h5.name, 'r')
    assert('node_id' in nodes_h5['/nodes/NET1'])
    assert(len(nodes_h5['/nodes/NET1/node_id']) == 30)
    assert(len(nodes_h5['/nodes/NET1/node_type_id']) == 30)
    assert(len(nodes_h5['/nodes/NET1/node_group_id']) == 30)
    assert(len(nodes_h5['/nodes/NET1/node_group_index']) == 30)

    node_groups = {nid: grp for nid, grp in nodes_h5['/nodes/NET1'].items() if isinstance(grp, h5py.Group)}
    for grp in node_groups.values():
        if len(grp) == 1:
            assert('position' in grp and len(grp['position']) == 20)

        elif len(grp) == 2:
            assert('position' in grp and len(grp['position']) == 10)
            assert('tags' in grp and len(grp['tags']) == 10)

        else:
            assert False

    assert(os.path.exists(edges_h5.name) and os.path.exists(edges_csv.name))
    edge_types_df = pd.read_csv(edges_csv.name, sep=' ')
    assert (len(edge_types_df) == 2)
    assert ('p1' in edge_types_df.columns)
    assert ('p2' in edge_types_df.columns)

    edges_h5 = h5py.File(edges_h5.name, 'r')
    assert('source_to_target' in edges_h5['/edges/NET1_to_NET1/indices'])
    assert('target_to_source' in edges_h5['/edges/NET1_to_NET1/indices'])
    assert(len(edges_h5['/edges/NET1_to_NET1/target_node_id']) == 300)
    assert(len(edges_h5['/edges/NET1_to_NET1/source_node_id']) == 300)

    # Check edges and node ids match up
    # warning, builder may not build edges in sequential order
    nid_idxs = np.sort(np.argwhere(edges_h5['/edges/NET1_to_NET1/target_node_id'][()] == 0).flatten())
    trg_ids = edges_h5['/edges/NET1_to_NET1/source_node_id'][nid_idxs]
    assert(np.all(trg_ids >= 10))
    assert(np.all(30 > trg_ids))
    edge_0 = nid_idxs[0]
    assert(edges_h5['/edges/NET1_to_NET1/edge_type_id'][edge_0] == 100)
    edge_id = edges_h5['/edges/NET1_to_NET1/edge_group_id'][edge_0]
    edge_idx = edges_h5['/edges/NET1_to_NET1/edge_group_index'][edge_0]
    assert(edges_h5['/edges/NET1_to_NET1'][str(edge_id)]['nsyns'][edge_idx] == 1)

    nid_idxs = np.sort(np.argwhere(edges_h5['/edges/NET1_to_NET1/target_node_id'][()] == 19).flatten())
    trg_ids = edges_h5['/edges/NET1_to_NET1/source_node_id'][nid_idxs]
    print(trg_ids)
    assert(np.all(trg_ids >= 0))
    assert(np.all(10 > trg_ids))
    edge_0 = nid_idxs[0]
    assert(edges_h5['/edges/NET1_to_NET1/edge_type_id'][edge_0] == 101)
    edge_id = edges_h5['/edges/NET1_to_NET1/edge_group_id'][edge_0]
    edge_idx = edges_h5['/edges/NET1_to_NET1/edge_group_index'][edge_0]
    assert(edges_h5['/edges/NET1_to_NET1'][str(edge_id)]['nsyns'][edge_idx] == 2)


@pytest.mark.parametrize('network_cls', [
    (DenseNetwork),
    # (DenseNetworkOrig)
])
def test_save_weights(network_cls):
    net = network_cls('NET1')
    net.add_nodes(N=100, position=[(0.0, 1.0, -1.0)]*100, cell_type='Scnna1', ei='e')
    net.add_nodes(N=100, position=[(0.0, 1.0, -1.0)]*100, cell_type='PV1', ei='i')
    net.add_nodes(N=100, position=[(0.0, 1.0, -1.0)]*100, tags=np.linspace(0, 100, 100), cell_type='PV2', ei='i')
    cm = net.add_edges(source={'ei': 'i'}, target={'ei': 'e'}, connection_rule=lambda s, t: 3,
                       p1='e2i', p2='e2i')  # 200*100 = 60000 edges
    cm.add_properties(names=['segment', 'distance'], rule=lambda s, t: [1, 0.5], dtypes=[int, float])

    net.add_edges(source=net.nodes(cell_type='Scnna1'), target=net.nodes(cell_type='PV1'),
                  connection_rule=lambda s, t: 2, p1='s2p')  # 100*100 = 20000'

    net.build()
    net_dir = tempfile.mkdtemp()
    net.save_nodes('tmp_nodes.h5', 'tmp_node_types.csv', output_dir=net_dir)
    net.save_edges('tmp_edges.h5', 'tmp_edge_types.csv', output_dir=net_dir)

    edges_h5 = h5py.File('{}/tmp_edges.h5'.format(net_dir), 'r')
    assert(net.nedges == 80000)
    assert(len(edges_h5['/edges/NET1_to_NET1/0/distance']) == 60000)
    assert(len(edges_h5['/edges/NET1_to_NET1/0/segment']) == 60000)
    assert(len(edges_h5['/edges/NET1_to_NET1/1/nsyns']) == 10000)
    assert(edges_h5['/edges/NET1_to_NET1/0/distance'][0] == 0.5)
    assert(edges_h5['/edges/NET1_to_NET1/0/segment'][0] == 1)
    assert(edges_h5['/edges/NET1_to_NET1/1/nsyns'][0] == 2)


@pytest.mark.parametrize('network_cls', [
    (DenseNetwork),
    # (DenseNetworkOrig)
])
def test_save_multinetwork(network_cls):
    net1 = DenseNetwork('NET1')
    net1.add_nodes(N=100, position=[(0.0, 1.0, -1.0)] * 100, cell_type='Scnna1', ei='e')
    net1.add_edges(source={'ei': 'e'}, target={'ei': 'e'}, connection_rule=5, ctype_1='n1_rec')
    net1.build()

    net2 = DenseNetwork('NET2')
    net2.add_nodes(N=10, position=[(0.0, 1.0, -1.0)] * 10, cell_type='PV1', ei='i')
    net2.add_edges(connection_rule=10, ctype_1='n2_rec')
    net2.add_edges(source=net1.nodes(), target={'ei': 'i'}, connection_rule=1, ctype_2='n1_n2')
    net2.add_edges(target=net1.nodes(cell_type='Scnna1'), source={'cell_type': 'PV1'}, connection_rule=2,
                   ctype_2='n2_n1')
    net2.build()

    net_dir = tempfile.mkdtemp()
    net1.save_edges(output_dir=net_dir)
    net2.save_edges(output_dir=net_dir)

    n1_n1_fname = '{}/{}_{}'.format(net_dir, 'NET1', 'NET1')
    edges_h5 = h5py.File(n1_n1_fname + '_edges.h5', 'r')
    assert(len(edges_h5['/edges/NET1_to_NET1/target_node_id']) == 100*100)
    assert(len(edges_h5['/edges/NET1_to_NET1/0/nsyns']) == 100*100)
    assert(edges_h5['/edges/NET1_to_NET1/0/nsyns'][0] == 5)
    edge_types_csv = pd.read_csv(n1_n1_fname + '_edge_types.csv', sep=' ')
    assert(len(edge_types_csv) == 1)
    assert('ctype_2' not in edge_types_csv.columns.values)
    assert(edge_types_csv['ctype_1'].iloc[0] == 'n1_rec')

    n1_n2_fname = '{}/{}_{}'.format(net_dir, 'NET1', 'NET2')
    edges_h5 = h5py.File(n1_n2_fname + '_edges.h5', 'r')
    assert(len(edges_h5['/edges/NET1_to_NET2/target_node_id']) == 100*10)
    assert(len(edges_h5['/edges/NET1_to_NET2/0/nsyns']) == 100*10)
    assert(edges_h5['/edges/NET1_to_NET2/0/nsyns'][0] == 1)
    edge_types_csv = pd.read_csv(n1_n2_fname + '_edge_types.csv', sep=' ')
    assert(len(edge_types_csv) == 1)
    assert('ctype_1' not in edge_types_csv.columns.values)
    assert(edge_types_csv['ctype_2'].iloc[0] == 'n1_n2')

    n2_n1_fname = '{}/{}_{}'.format(net_dir, 'NET2', 'NET1')
    edges_h5 = h5py.File(n2_n1_fname + '_edges.h5', 'r')
    assert(len(edges_h5['/edges/NET2_to_NET1/target_node_id']) == 100*10)
    assert(len(edges_h5['/edges/NET2_to_NET1/0/nsyns']) == 100*10)
    assert(edges_h5['/edges/NET2_to_NET1/0/nsyns'][0] == 2)
    edge_types_csv = pd.read_csv(n2_n1_fname + '_edge_types.csv', sep=' ')
    assert(len(edge_types_csv) == 1)
    assert('ctype_1' not in edge_types_csv.columns.values)
    assert(edge_types_csv['ctype_2'].iloc[0] == 'n2_n1')

    n2_n2_fname = '{}/{}_{}'.format(net_dir, 'NET2', 'NET2')
    edges_h5 = h5py.File(n2_n2_fname + '_edges.h5', 'r')
    assert(len(edges_h5['/edges/NET2_to_NET2/target_node_id']) == 10*10)
    assert(len(edges_h5['/edges/NET2_to_NET2/0/nsyns']) == 10*10)
    assert(edges_h5['/edges/NET2_to_NET2/0/nsyns'][0] == 10)
    edge_types_csv = pd.read_csv(n2_n2_fname + '_edge_types.csv', sep=' ')
    assert(len(edge_types_csv) == 1)
    assert('ctype_2' not in edge_types_csv.columns.values)
    assert(edge_types_csv['ctype_1'].iloc[0] == 'n2_rec')


if __name__ == '__main__':
    test_save_nsyn_table(DenseNetwork)
    # test_save_weights(DenseNetwork)
    # test_save_multinetwork(DenseNetwork)
