import os
import unittest

from pytlex_core.Graph import Graph
from pytlex_core.Partitioner import partition_graph
from pytlex_tests import Testing_utilities as test_ut, Testing_utilities
from pytlex_tests.jTLEX_Data import jTLEX_Counts

"""
Tester for the Partitioner class. 
Each method partitions a different TimeML Graph.
Partitions are then checked for correctness.
"""
class PartitionerTest(unittest.TestCase):
    error_message = 'Erroneous or unexpected output'

    def test_partition_graph_wsj_0006(self):
        graph = test_ut.wsj_0006()

        results = partition_graph(graph)
        print("Filename: wsj_0006")
        for main_partition in results["main_graphs"]:
            print('Main Partition nodes:')
            for n in main_partition.nodes:
                print(n.get_id_str())

        print("\n")
        for sub_partition in results["subordination_graphs"]:
            print('Subordination Partition nodes:')
            for n in sub_partition.nodes:
                print(n.get_id_str())
        print("\n")

        print("S_links:")
        for l in results["s_links"]:
            print(f'{l.start_node.get_id_str()} -> {l.related_to_node.get_id_str()}')

    def test_partition_graph_wsj_1073(self):
        graph = test_ut.wsj_1073()

        results = partition_graph(graph)
        print("\nFilename: wsj_1073")
        for main_partition in results["main_graphs"]:
            print('Main Partition nodes:')
            for n in main_partition.nodes:
                print(n.get_id_str())
        print("\n")

        for sub_partition in results["subordination_graphs"]:
            print('Subordination Partition nodes:')
            for n in sub_partition.nodes:
                print(n.get_id_str())
        print("\n")

        print("S_links:")
        for l in results["s_links"]:
            print(f'{l.start_node.get_id_str()} -> {l.related_to_node.get_id_str()}')

    def test_partition_graph_wsj_0555(self):
        graph = test_ut.wsj_0555()

        results = partition_graph(graph)
        print("\nFilename: wsj_0555")
        for main_partition in results["main_graphs"]:
            print('Main Partition nodes:')
            for n in main_partition.nodes:
                print(n.get_id_str())
        print("\n")
        for sub_partition in results["subordination_graphs"]:
            print('Subordination Partition nodes:')
            for n in sub_partition.nodes:
                print(n.get_id_str())
        print("\n")
        print("S_links:")
        for l in results["s_links"]:
            print(f'{l.start_node.get_id_str()} -> {l.related_to_node.get_id_str()}')

    def test_partition_lengths_vs_jtlex(self):
        filepath = r"../pytlex_data/TimeBankCorpus"
        for subdir, dirs, files in os.walk(filepath):
            for filename in files:
                if filename not in Testing_utilities.failed_parse:
                    graph = Graph(filepath=filepath + "/" + filename)
                    self.assertEqual(
                        len(graph.main_graphs + graph.subordination_graphs),
                        jTLEX_Counts.jTLEX_partition_counts[filename]
                    )


if __name__ == '__main__':
    unittest.main()
