"""
********
UnitTest
OrientEdges
********

"""
__author__ = """Nicholas Cullen <ncullen.th@dartmouth.edu>"""

import os
import unittest
from os.path import dirname

import numpy as np
from external.pyBN.structure_learn.orient_edges import orient_edges_gs, orient_edges_pc


class OrientEdgesTestCase(unittest.TestCase):
    def setUp(self):
        pass

    def tearDown(self):
        pass

    def test_orient_edges_pc(self):
        e = {0: [1, 2], 1: [0], 2: [0]}
        b = {0: [], 1: {2: (0,)}, 2: {1: (0,)}}
        self.assertDictEqual(orient_edges_pc(e, b), {0: [1, 2], 1: [], 2: []})

    def test_orient_edges_gs(self):
        e = {0: [1, 2], 1: [0], 2: [0]}
        b = {0: [1, 2], 1: [0], 2: [0]}
        dpath = os.path.join(dirname(dirname(dirname(dirname(__file__)))), "data")
        path = os.path.join(dpath, "lizards.csv")
        data = np.loadtxt(path, dtype="int32", skiprows=1, delimiter=",")
        self.assertDictEqual(
            orient_edges_gs(e, b, data, 0.05), {0: [1, 2], 1: [], 2: []}
        )
