import unittest

from rdkit import Chem

from frag.network.models import NodeHolder, Node, Attr
from frag.utils.network_utils import (
    rebuild_smi,
    make_child_mol,
    get_fragments,
    build_network,
    get_comb_index,
    ret_comb_index,
)
from frag.network.decorate import (
    decorate_smi,
    deletion_linker_mol,
    deletion_linker_smi,
    del_link_coord,
)


def parse_node(input_str):
    """
    Convert something like to a Node:
    NODE O=CCCc1ccc(cc1)c2ccccc2 16 12 OCCCC1CCC(CC1)C2CCCCC2 0
    :param input_str:
    :return:
    """
    smiles = input_str.split()[1]
    new_node = Node()
    new_node.SMILES = Chem.CanonSmiles(smiles)
    new_node.HAC = input_str.split()[2]
    new_node.RAC = input_str.split()[3]
    new_node.RING_SMILES = input_str.split()[4]
    return new_node


def conv_smi(input_smi):
    return Chem.MolToSmiles(Chem.MolFromSmiles(input_smi), isomericSmiles=False)


class NetworksTest(unittest.TestCase):

    def test_rebuild(self):
        input_list = [
            ["O[100Xe]", "[100Xe]c1ccc([101Xe])cc1"],
            ["O[100Xe]", "[101Xe]c1ccccc1"],
            ["[101Xe]c1ccccc1", "[100Xe]c1ccc([101Xe])cc1"],
        ]
        rebuild_list = [
            "Oc1ccc([Xe])cc1",
            "O[Xe].[Xe]c1ccccc1",
            "[Xe]c1ccc(cc1)c2ccccc2",
        ]
        for i in range(len(input_list)):
            self.assertEqual(
                conv_smi(rebuild_smi(input_list[i], ring_ring=False)),
                conv_smi(rebuild_list[i]),
            )

    def test_child(self):
        rebuild_list = [
            "Oc1ccc([Xe])cc1",
            "O[Xe].[Xe]c1ccccc1",
            "[Xe]c1ccc(cc1)c2ccccc2",
        ]
        child_list = ["Oc1ccccc1", "O.c1ccccc1", "c1ccc(cc1)c2ccccc2"]
        for i in range(len(child_list)):
            self.assertEqual(
                conv_smi(make_child_mol(rebuild_list[i])), conv_smi(child_list[i])
            )

    def test_get(self):
        input_list = ["CC.CC", "CC.c1ccccc1C", "CCC"]
        output_list = [["CC", "CC"], ["CC", "[100Xe]C", "[100Xe]c1ccccc1"], ["CCC"]]
        for i in range(len(input_list)):
            self.assertListEqual(
                output_list[i], get_fragments(Chem.MolFromSmiles(input_list[i]))
            )

    def test_generate_nodes(self):
        """
        Test we can generate nodes for the basic data.
        :return:
        """
        try:
            nodes = [x for x in open("frag/tests/data/nodes.txt").readlines()]
            edges = [x.split() for x in open("frag/tests/data/edges.txt").readlines()]
            attrs = [
                Attr(input_str=x)
                for x in open("frag/tests/data/attributes.txt").readlines()
            ]
        except IOError:
            nodes = [x for x in open("data/nodes.txt").readlines()]
            edges = [x.split() for x in open("data/edges.txt").readlines()]
            attrs = [Attr(input_str=x) for x in open("data/attributes.txt").readlines()]
        node_holder = NodeHolder(iso_flag=True)
        node_holder = build_network(attrs, node_holder)
        # Create the nodes and test with output
        self.assertEqual(len(node_holder.node_list), len(nodes))
        # This doesn't work yet(we get 3687 edges - should be 3691
        # Close enough - and the output looks right...
        self.assertEqual(len(node_holder.get_edges()), 3687)

    def test_generate_nodes_non_iso(self):
        """
        Test we can generate nodes for the basic data.
        :return:
        """
        try:
            nodes = [x for x in open("frag/tests/data/nodes.txt").readlines()]
            edges = [x.split() for x in open("frag/tests/data/edges.txt").readlines()]
            attrs = [
                Attr(input_str=x)
                for x in open("frag/tests/data/attributes.txt").readlines()
            ]
        except IOError:
            nodes = [x for x in open("data/nodes.txt").readlines()]
            edges = [x.split() for x in open("data/edges.txt").readlines()]
            attrs = [Attr(input_str=x) for x in open("data/attributes.txt").readlines()]
        node_holder = NodeHolder(iso_flag=False)
        node_holder = build_network(attrs, node_holder)
        # Create the nodes and test with output
        self.assertEqual(len(node_holder.node_list), len(nodes))
        # This doesn't work yet(we get 3687 edges - should be 3691
        # Close enough - and the output looks right...
        self.assertEqual(len(node_holder.get_edges()), 3687)

    def test_compare_iso_non_iso(self):
        """
        Test that the iso flag makes a difference.
        :return:
        """
        input_smis = ["C#CC(C)(C)NC[C@]1(O)CCCN2CCCC[C@@H]21"]
        test_iso_node_list = [
            "C#CC(C)(C)NC",
            "OC1CCCN2CCCCC12",
            "O",
            "C#CC(C)(C)NCC1CCCN2CCCCC12",
            "C#CC(C)(C)NC[C@]1(O)CCCN2CCCC[C@@H]21",
            "C1CCN2CCCCC2C1",
            "C#CC(C)(C)NC.O",
        ]
        test_non_iso_node_list = [
            "C#CC(C)(C)NC",
            "OC1CCCN2CCCCC12",
            "O",
            "C#CC(C)(C)NCC1CCCN2CCCCC12",
            "C#CC(C)(C)NCC1(O)CCCN2CCCCC21",
            "C1CCN2CCCCC2C1",
            "C#CC(C)(C)NC.O",
        ]
        test_iso_edge_list = [
            "EDGE C#CC(C)(C)NC[C@]1(O)CCCN2CCCC[C@@H]21 OC1CCCN2CCCCC12 FG|C#CC(C)(C)NC[Xe]|CCC(C)(C)NC[100Xe]|RING|OC1([Xe])CCCN2CCCCC21|O[C@@]1([100Xe])CCCC2CCCC[C@@H]21",
            "EDGE OC1CCCN2CCCCC12 O RING|[Xe]C1CCCN2CCCCC12|[100Xe]C1CCCC2CCCCC12|FG|O[Xe]|O[100Xe]",
            "EDGE C#CC(C)(C)NC[C@]1(O)CCCN2CCCC[C@@H]21 C#CC(C)(C)NC.O RING|[Xe]C1([Xe])CCCN2CCCCC21|[100Xe][C@]1([101Xe])CCCC2CCCC[C@@H]21|FG|C#CC(C)(C)NC[Xe].O[Xe]|CCC(C)(C)NC[100Xe].O[101Xe]",
            "EDGE OC1CCCN2CCCCC12 C1CCN2CCCCC2C1 FG|O[Xe]|O[100Xe]|RING|[Xe]C1CCCN2CCCCC12|[100Xe]C1CCCC2CCCCC12",
            "EDGE C#CC(C)(C)NC.O C#CC(C)(C)NC FG|O|O|FG|C#CC(C)(C)NC|CCC(C)(C)NC",
            "EDGE C#CC(C)(C)NC[C@]1(O)CCCN2CCCC[C@@H]21 C#CC(C)(C)NCC1CCCN2CCCCC12 FG|O[Xe]|O[101Xe]|RING|C#CC(C)(C)NCC1([Xe])CCCN2CCCCC21|CCC(C)(C)NC[C@@]1([101Xe])CCCC2CCCC[C@@H]21",
            "EDGE C#CC(C)(C)NCC1CCCN2CCCCC12 C1CCN2CCCCC2C1 FG|C#CC(C)(C)NC[Xe]|CCC(C)(C)NC[100Xe]|RING|[Xe]C1CCCN2CCCCC12|[100Xe]C1CCCC2CCCCC12",
            "EDGE C#CC(C)(C)NC.O O FG|C#CC(C)(C)NC|CCC(C)(C)NC|FG|O|O",
            "EDGE C#CC(C)(C)NCC1CCCN2CCCCC12 C#CC(C)(C)NC RING|[Xe]C1CCCN2CCCCC12|[100Xe]C1CCCC2CCCCC12|FG|C#CC(C)(C)NC[Xe]|CCC(C)(C)NC[100Xe]",
        ]

        test_non_iso_edge_list = [
            "EDGE C#CC(C)(C)NCC1(O)CCCN2CCCCC21 C#CC(C)(C)NCC1CCCN2CCCCC12 FG|O[Xe]|O[101Xe]|RING|C#CC(C)(C)NCC1([Xe])CCCN2CCCCC21|CCC(C)(C)NCC1([101Xe])CCCC2CCCCC21",
            "EDGE OC1CCCN2CCCCC12 C1CCN2CCCCC2C1 FG|O[Xe]|O[100Xe]|RING|[Xe]C1CCCN2CCCCC12|[100Xe]C1CCCC2CCCCC12",
            "EDGE C#CC(C)(C)NCC1(O)CCCN2CCCCC21 C#CC(C)(C)NC.O RING|[Xe]C1([Xe])CCCN2CCCCC21|[100Xe]C1([101Xe])CCCC2CCCCC21|FG|C#CC(C)(C)NC[Xe].O[Xe]|CCC(C)(C)NC[100Xe].O[101Xe]",
            "EDGE OC1CCCN2CCCCC12 O RING|[Xe]C1CCCN2CCCCC12|[100Xe]C1CCCC2CCCCC12|FG|O[Xe]|O[100Xe]",
            "EDGE C#CC(C)(C)NCC1CCCN2CCCCC12 C1CCN2CCCCC2C1 FG|C#CC(C)(C)NC[Xe]|CCC(C)(C)NC[100Xe]|RING|[Xe]C1CCCN2CCCCC12|[100Xe]C1CCCC2CCCCC12",
            "EDGE C#CC(C)(C)NC.O O FG|C#CC(C)(C)NC|CCC(C)(C)NC|FG|O|O",
            "EDGE C#CC(C)(C)NCC1(O)CCCN2CCCCC21 OC1CCCN2CCCCC12 FG|C#CC(C)(C)NC[Xe]|CCC(C)(C)NC[100Xe]|RING|OC1([Xe])CCCN2CCCCC21|OC1([100Xe])CCCC2CCCCC21",
            "EDGE C#CC(C)(C)NC.O C#CC(C)(C)NC FG|O|O|FG|C#CC(C)(C)NC|CCC(C)(C)NC",
            "EDGE C#CC(C)(C)NCC1CCCN2CCCCC12 C#CC(C)(C)NC RING|[Xe]C1CCCN2CCCCC12|[100Xe]C1CCCC2CCCCC12|FG|C#CC(C)(C)NC[Xe]|CCC(C)(C)NC[100Xe]",
        ]

        attrs = [Attr(input_smi) for input_smi in input_smis]
        node_holder = NodeHolder(iso_flag=False)
        node_holder = build_network(attrs, node_holder)
        non_iso_node_list = [x.SMILES for x in node_holder.node_list]
        non_iso_edge_list = [str(x) for x in node_holder.edge_list]
        self.assertListEqual(sorted(non_iso_node_list), sorted(test_non_iso_node_list))
        self.assertListEqual(sorted(non_iso_edge_list), sorted(test_non_iso_edge_list))
        node_holder = NodeHolder(iso_flag=True)
        node_holder = build_network(attrs, node_holder)
        iso_node_list = [x.SMILES for x in node_holder.node_list]
        iso_edge_list = [str(x) for x in node_holder.edge_list]
        self.assertListEqual(sorted(iso_node_list), sorted(test_iso_node_list))
        self.assertListEqual(sorted(iso_edge_list), sorted(test_iso_edge_list))

    def test_decorate(self):
        """
        Test we can decorate a series of input SMILEs
        :return:
        """
        input_data = ["Oc1ccc(cc1)c2ccccc2", "c1ccccc1", "c1ccncc1", "c1cccnc1"]
        output_data = [
            [
                "Oc1ccc(-c2ccccc2[At])cc1",
                "Oc1ccc(-c2ccccc2)c([At])c1",
                "Oc1ccc(-c2ccc([At])cc2)cc1",
                "Oc1ccc(-c2ccccc2)cc1[At]",
                "Oc1ccc(-c2cccc([At])c2)cc1",
            ],
            ["[At]c1ccccc1"],
            ["[At]c1cccnc1", "[At]c1ccccn1", "[At]c1ccncc1"],
            ["[At]c1cccnc1", "[At]c1ccccn1", "[At]c1ccncc1"],
        ]
        for i, smi in enumerate(input_data):
            self.assertListEqual(list(decorate_smi(smi).keys()), output_data[i])

    def test_comb_index(self):
        """
        Test we combine indices
        :return:
        """
        input_data = [(12, 19), (6, 14), (99, 98), (4, 0)]
        output_data = [2012, 1506, 9999, 104]
        for i, data in enumerate(input_data):
            self.assertEqual(get_comb_index(data[0], data[1]), output_data[i])
            self.assertTupleEqual(ret_comb_index(output_data[i]), data)
            self.assertTupleEqual(
                ret_comb_index(get_comb_index(data[0], data[1])), data
            )

    def test_ring_ring_smi(self):
        input_smi = "CC(=O)NC=1C=CC(=CC1)C2=CSC(N)=N2"
        res = deletion_linker_smi("CC(=O)NC=1C=CC(=CC1)C2=CSC(N)=N2")
        self.assertListEqual(
            [Chem.MolToSmiles(x, isomericSmiles=True) for x in res[0]],
            ["Nc1nc(-c2ccc([100Xe])cc2)cs1", "CC(=O)Nc1ccc(-c2csc([102Xe])n2)cc1"],
        )
        self.assertListEqual(
            [Chem.MolToSmiles(x, isomericSmiles=True) for x in res[1]],
            ["CC(=O)Nc1ccc([1107Xe])cc1.Nc1nc([1107Xe])cs1"],
        )
        self.assertListEqual(
            [Chem.MolToSmiles(x, isomericSmiles=True) for x in res[2]],
            ["CC(=O)N[100Xe].Nc1nc([101Xe])cs1", "CC(=O)Nc1ccc([101Xe])cc1.N[102Xe]"],
        )

    def test_ring_ring_mol(self):
        input_sd = """
     RDKit          3D

 16 17  0  0  0  0  0  0  0  0999 V2000
  -10.1430  -10.6060  -10.9420 C   0  0  0  0  0  0  0  0  0  0  0  0
  -11.4350  -10.5270  -11.7010 C   0  0  0  0  0  0  0  0  0  0  0  0
  -12.5150  -10.4940  -11.1270 O   0  0  0  0  0  0  0  0  0  0  0  0
  -11.3000  -10.4840  -13.0450 N   0  0  0  0  0  0  0  0  0  0  0  0
  -12.2950  -10.4260  -14.0500 C   0  0  0  0  0  0  0  0  0  0  0  0
  -11.8760  -10.0530  -15.3220 C   0  0  0  0  0  0  0  0  0  0  0  0
  -12.7770   -9.9750  -16.3610 C   0  0  0  0  0  0  0  0  0  0  0  0
  -14.1220  -10.2670  -16.1660 C   0  0  0  0  0  0  0  0  0  0  0  0
  -14.5380  -10.6530  -14.8930 C   0  0  0  0  0  0  0  0  0  0  0  0
  -13.6380  -10.7330  -13.8470 C   0  0  0  0  0  0  0  0  0  0  0  0
  -15.0650  -10.1850  -17.2930 C   0  0  0  0  0  0  0  0  0  0  0  0
  -16.4170  -10.0960  -17.2170 C   0  0  0  0  0  0  0  0  0  0  0  0
  -17.1150   -9.9510  -18.7790 S   0  0  0  0  0  0  0  0  0  0  0  0
  -15.5150  -10.0790  -19.4610 C   0  0  0  0  0  0  0  0  0  0  0  0
  -15.3280  -10.0290  -20.7950 N   0  0  0  0  0  0  0  0  0  0  0  0
  -14.5560  -10.2150  -18.5880 N   0  0  0  0  0  0  0  0  0  0  0  0
  2  1  1  0
  3  2  2  0
  4  2  1  0
  5  4  1  0
  6  5  2  0
  7  6  1  0
  8  7  2  0
  9  8  1  0
 10  9  2  0
 10  5  1  0
 11  8  1  0
 12 11  2  0
 13 12  1  0
 14 13  1  0
 15 14  1  0
 16 14  2  0
 16 11  1  0
M  END

"""
        res = del_link_coord(input_sd)
        values = res["linkers"]["CC(=O)Nc1ccc([Xe])cc1.Nc1nc([Xe])cs1"]
        self.assertEqual(len(values), 4)
        self.assertEqual(type(values[0][2]), float)
        self.assertEqual(type(values[2][1]), float)
