import unittest

import networkx as nx

from Flowsheet_Class.flowsheet import Flowsheet


class TestSFILESctrl(unittest.TestCase):
    """Class performs an unittest to test the SFILES 2.0 control extension on specific test cases."""

    def test_case_1a(self):
        # 1) Measuring point in/at unit operation
        test_case = "1a"
        edges = [("IO-1", "tank-1"), ("tank-1", "IO-2"), ("tank-1", "C-1/TIR"), ("tank-1", "C-2/LIR")]
        self.SFILESctrl(test_case, edges)

    def test_case_1b(self):
        test_case = "1b"
        edges = [("IO-1", "tank-1"), ("C-1/LIR", "v-1", {"tags": {"signal": ["not_next_unitop"]}}),
                 ("tank-1", "C-1/LIR"), ("v-1", "IO-2"), ("tank-1", "v-1")]
        self.SFILESctrl(test_case, edges)

    def test_case_2(self):
        # 2) Measuring point at material stream
        test_case = "2"
        edges = [("IO-1", "C-1/FC"), ("C-1/FC", "v-1", {"tags": {"signal": ["next_unitop"]}}), ("v-1", "IO-2")]
        self.SFILESctrl(test_case, edges)

    def test_case_3(self):
        # 3) Cascade control
        test_case = "3"
        edges = [("IO-1", "tank-1"), ("tank-1", "C-1/LC"),
                 ("C-1/LC", "C-2/FC", {"tags": {"signal": ["not_next_unitop"]}}),
                 ("tank-1", "C-2/FC"), ("C-2/FC", "v-1", {"tags": {"signal": ["next_unitop"]}}), ("v-1", "IO-2")]
        self.SFILESctrl(test_case, edges)

    def test_case_A(self):
        # A)
        test_case = "A"
        edges = [("IO-1", "C-1/FC"), ("C-1/FC", "v-1", {"tags": {"signal": ["next_unitop"]}}), ("v-1", "IO-2")]
        self.SFILESctrl(test_case, edges)

    def test_case_B(self):
        # B)
        test_case = "B"
        edges = [("IO-1", "C-1/F"), ("C-1/F", "C-2/FFC", {"tags": {"signal": ["not_next_unitop"]}}),
                 ("C-2/FFC", "v-1", {"tags": {"signal": ["not_next_unitop"]}}), ("C-1/F", "v-1"), ("v-1", "IO-2"),
                 ("IO-3", "C-3/F"), ("C-3/F", "C-2/FFC", {"tags": {"signal": ["not_next_unitop"]}}),
                 ("C-3/F", "IO-4")]
        self.SFILESctrl(test_case, edges)

    def test_case_C(self):
        # C)
        test_case = "C"
        edges = [("IO-1", "C-1/T"), ("C-1/T", "IO-2"), ("IO-3", "C-2/FQC"),
                 ("C-1/T", "C-2/FQC", {"tags": {"signal": ["not_next_unitop"]}}),
                 ("C-2/FQC", "v-1", {"tags": {"signal": ["next_unitop"]}}), ("v-1", "IO-4")]
        self.SFILESctrl(test_case, edges)

    def test_case_D(self):
        # D)
        test_case = "D"
        edges = [("IO-1", "hex-1/1"), ("hex-1/1", "C-1/TC"), ("C-1/TC", "IO-2"), ("IO-3", "C-2/FC"),
                 ("C-2/FC", "v-1", {"tags": {"signal": ["next_unitop"]}}), ("v-1", "hex-1/2"),
                 ("hex-1/2", "IO-4"), ("C-1/TC", "C-2/FC", {"tags": {"signal": ["not_next_unitop"]}})]
        self.SFILESctrl(test_case, edges)

    def test_case_E(self):
        # E)
        test_case = "E"
        edges = [("IO-1", "splt-1"), ("splt-1", "v-1"), ("v-1", "IO-2"), ("splt-1", "C-1/FC"),
                 ("C-1/FC", "v-2", {"tags": {"signal": ["next_unitop"]}}), ("v-2", "IO-3")]
        self.SFILESctrl(test_case, edges)

    def test_case_F(self):
        # F)
        test_case = "F"
        edges = [("IO-1", "tank-1"), ("tank-1", "IO-2"), ("tank-1", "C-1/FC"),
                 ("C-1/FC", "v-1", {"tags": {"signal": ["next_unitop"]}}), ("v-1", "IO-3")]
        self.SFILESctrl(test_case, edges)

    def test_case_G(self):
        # G)
        test_case = "G"
        edges = [("IO-1", "tank-1"), ("tank-1", "C-1/LC"), ("tank-1", "splt-1"), ("splt-1", "C-2/FC"),
                 ("C-2/FC", "v-1", {"tags": {"signal": ["next_unitop"]}}), ("v-1", "IO-2"), ("splt-1", "v-2"),
                 ("v-2", "IO-3"), ("C-1/LC", "v-2", {"tags": {"signal": ["not_next_unitop"]}})]
        self.SFILESctrl(test_case, edges)

    def test_case_HX_4(self):
        # HX_4)
        test_case = "HX_4"
        edges = [("IO-1", "C-1/TC"), ("C-1/TC", "hex-1/1"), ("hex-1/1", "IO-2"), ("IO-3", "C-2/FC"),
                 ("C-2/FC", "v-1", {"tags": {"signal": ["next_unitop"]}}), ("v-1", "hex-1/2"), ("hex-1/2", "IO-4"),
                 ("C-1/TC", "C-2/FC", {"tags": {"signal": ["not_next_unitop"]}})]
        self.SFILESctrl(test_case, edges)

    def test_case_Umpumpanalge(self):
        # Umpumpanlage
        test_case = "Umpumpanlage"
        edges = [("IO-1", "v-1"), ("v-1", "tank-1"), ("tank-1", "v-2"), ("v-2", "pp-1"), ("pp-1", "v-3"),
                 ("v-3", "C-4/FRC"), ("C-4/FRC", "splt-1"), ("splt-1", "v-4"), ("v-4", "v-5"), ("v-5", "v-6"),
                 ("v-6", "mix-1"), ("mix-1", "v-7"), ("splt-1", "v-8"), ("v-8", "mix-1"), ("v-7", "tank-2"),
                 ("tank-2", "v-9"), ("v-9", "IO-2"), ("tank-1", "C-1/TIR"), ("tank-1", "C-2/LIR"), ("pp-1", "C-3/M"),
                 ("v-8", "C-5/H"), ("tank-2", "C-6/TIR"), ("tank-2", "C-7/PICA"),
                 ("C-4/FRC", "v-5", {"tags": {"signal": ["not_next_unitop"]}})]
        self.SFILESctrl(test_case, edges)

    def test_case_Rectification(self):
        # Rectification)
        test_case = "Rectification"
        edges = [("IO-1", "C-1/FC"), ("C-1/FC", "v-1"), ("v-1", "hex-1/1"), ("hex-1/1", "C-2/TC"), ("C-2/TC", "dist-1"),
                 ("dist-1", "C-3/PC"), ("dist-1", "C-4/LC"), ("C-1/FC", "v-1", {"tags": {"signal": ["next_unitop"]}}),
                 ("C-2/TC", "v-2", {"tags": {"signal": ["not_next_unitop"]}}), ("IO-2", "v-2"),
                 ("v-2", "hex-1/2"), ("hex-1/2", "IO-3"), ("dist-1", "hex-2", {"tags": {"col": ["tout"]}}),
                 ("hex-2", "tank-1"), ("tank-1", "C-5/LC"), ("tank-1", "splt-1"), ("splt-1", "v-3"), ("v-3", "dist-1"),
                 ("C-5/LC", "v-3", {"tags": {"signal": ["not_next_unitop"]}}), ("splt-1", "C-6/FC"),
                 ("C-6/FC", "v-4", {"tags": {"signal": ["next_unitop"]}}), ("v-4", "IO-4"), ("tank-1", "v-5"),
                 ("v-5", "IO-5"), ("C-3/PC", "v-5", {"tags": {"signal": ["not_next_unitop"]}}),
                 ("dist-1", "splt-2", {"tags": {"col": ["bout"]}}), ("splt-2", "v-6"), ("v-6", "IO-6"),
                 ("C-4/LC", "v-6", {"tags": {"signal": ["not_next_unitop"]}}), ("splt-2", "hex-3/1"),
                 ("hex-3/1", "dist-1"), ("IO-7", "C-7/FC"), ("C-7/FC", "v-7", {"tags": {"signal": ["next_unitop"]}}),
                 ("v-7", "hex-3/2"), ("hex-3/2", "IO-8")]
        self.SFILESctrl(test_case, edges)

    def test_case_H(self):
        # Test case H
        test_case = "H"
        edges = [("IO-1", "pp-1"), ("pp-1", "C-1/FC"), ("C-1/FC", "v-1", {"tags": {"signal": ["next_unitop"]}}),
                 ("v-1", "hex-1"), ("hex-1", "mix-1"), ("mix-1", "r-1"),
                 ("r-1", "mix-2", {"tags": {"col": ["tout"]}}), ("mix-2", "mix-1"), ("IO-2", "mix-2"),
                 ("r-1", "C-2/TC", {"tags": {"col": ["bout"]}}), ("C-2/TC", "dist-1"),
                 ("dist-1", "pp-2", {"tags": {"col": ["bout"]}}), ("pp-2", "hex-2"), ("hex-2", "splt-1"),
                 ("splt-1", "dist-1"), ("splt-1", "C-3/FC"), ("C-3/FC", "v-2", {"tags": {"signal": ["next_unitop"]}}),
                 ("v-2", "IO-3"), ("dist-1", "C-4/LC"), ("C-4/LC", "C-3/FC", {"tags": {"signal": ["not_next_unitop"]}}),
                 ("C-2/TC", "C-1/FC", {"tags": {"signal": ["not_next_unitop"]}}),
                 ("dist-1", "IO-4", {"tags": {"col": ["tout"]}}), ("IO-5", "r-1"), ("IO-6", "tank-1"),
                 ("tank-1", "pp-3"), ("pp-3", "C-5/FC"), ("C-5/FC", "v-3", {"tags": {"signal": ["next_unitop"]}}),
                 ("v-3", "hex-3"), ("hex-3", "mix-1")]
        self.SFILESctrl(test_case, edges)

    def SFILESctrl(self, test_case, edges):
        graph = nx.DiGraph()
        graph.add_edges_from(edges)

        flowsheet = Flowsheet()
        flowsheet.state = graph
        flowsheet.convert_to_sfiles()
        sfilesctrl1 = flowsheet.sfiles
        flowsheet.create_from_sfiles(sfilesctrl1, overwrite_nx=True)
        flowsheet.convert_to_sfiles()
        sfilesctrl2 = flowsheet.sfiles

        sfiles = flowsheet.convert_sfilesctrl_to_sfiles()
        flowsheet.sfiles = sfiles
        flowsheet.create_from_sfiles(sfiles, overwrite_nx=True)
        flowsheet.convert_to_sfiles()
        sfiles2 = flowsheet.sfiles

        if sfilesctrl1 == sfilesctrl2 and sfiles == sfiles2:
            print(
                "-----------------------------------------------------------------------------------------------------")
            print("Test case ", test_case)
            print("Conversion back successful")
            print("SFILESctrl: ", sfilesctrl1)
            print("SFILES: ", sfiles)
        else:
            print(
                "-----------------------------------------------------------------------------------------------------")
            print("Test case ", test_case)
            print("Conversion back produced a different SFILES string. Input:", sfilesctrl1, "Output:", sfilesctrl2)
            print("Conversion back produced a different SFILES string. Input:", sfiles, "Output:", sfiles2)

        self.assertEqual(sfilesctrl1, sfilesctrl2, "Not correct!")
        self.assertEqual(sfiles, sfiles2, "Not correct!")


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