from dgl import DGLGraph
import rdkit.Chem as Chem
from .chemutils import get_clique_mol, tree_decomp, get_mol, get_smiles, \
                       set_atommap, enum_assemble_nx, decode_stereo

class DGLMolTree(DGLGraph):
    def __init__(self, smiles):
        DGLGraph.__init__(self)
        self.smiles = smiles
        self.mol = get_mol(smiles)

        # Stereo Generation
        mol = Chem.MolFromSmiles(smiles)
        self.smiles3D = Chem.MolToSmiles(mol, isomericSmiles=True)
        self.smiles2D = Chem.MolToSmiles(mol)
        self.stereo_cands = decode_stereo(self.smiles2D)

        # cliques: a list of list of atom indices
        cliques, edges = tree_decomp(self.mol)
        root = 0
        for i, c in enumerate(cliques):
            cmol = get_clique_mol(self.mol, c)
            csmiles = get_smiles(cmol)
            self.add_node(
                    i,
                    smiles=csmiles,
                    mol=get_mol(csmiles),
                    clique=c,
                    )
            if min(c) == 0:
                root = i

        # The clique with atom ID 0 becomes root
        if root > 0:
            for attr in self.nodes[0]:
                self.nodes[0][attr], self.nodes[root][attr] = \
                        self.nodes[root][attr], self.nodes[0][attr]

        for _x, _y in edges:
            x = 0 if _x == root else root if _x == 0 else _x
            y = 0 if _y == root else root if _y == 0 else _y
            self.add_edge(x, y)
            self.add_edge(y, x)

        for i in self.nodes:
            self.nodes[i]['nid'] = i + 1
            if len(self[i]) > 1:    # Leaf node mol is not marked
                set_atommap(self.nodes[i]['mol'], self.nodes[i]['nid'])
            self.nodes[i]['is_leaf'] = (len(self[i]) == 1)

    # avoiding DiGraph.size()
    def treesize(self):
        return len(self.nodes)

    def _recover_node(self, i, original_mol):
        node = self.nodes[i]

        clique = []
        clique.extend(node['clique'])
        if not node['is_leaf']:
            for cidx in node['clique']:
                original_mol.GetAtomWithIdx(cidx).SetAtomMapNum(node['nid'])

        for j in self[i]:
            nei_node = self.nodes[j]
            clique.extend(nei_node['clique'])
            if nei_node['is_leaf']: # Leaf node, no need to mark
                continue
            for cidx in nei_node['clique']:
                # allow singleton node override the atom mapping
                if cidx not in node['clique'] or len(nei_node['clique']) == 1:
                    atom = original_mol.GetAtomWithIdx(cidx)
                    atom.SetAtomMapNum(nei_node['nid'])

        clique = list(set(clique))
        label_mol = get_clique_mol(original_mol, clique)
        node['label'] = Chem.MolToSmiles(Chem.MolFromSmiles(get_smiles(label_mol)))
        node['label_mol'] = get_mol(node['label'])

        for cidx in clique:
            original_mol.GetAtomWithIdx(cidx).SetAtomMapNum(0)

        return node['label']

    def _assemble_node(self, i):
        neighbors = [self.nodes[j] for j in self[i] if self.nodes[j]['mol'].GetNumAtoms() > 1]
        neighbors = sorted(neighbors, key=lambda x: x['mol'].GetNumAtoms(), reverse=True)
        singletons = [self.nodes[j] for j in self[i] if self.nodes[j]['mol'].GetNumAtoms() == 1]
        neighbors = singletons + neighbors

        cands = enum_assemble_nx(self, i, neighbors)

        if len(cands) > 0:
            self.nodes[i]['cands'], self.nodes[i]['cand_mols'], _ = list(zip(*cands))
            self.nodes[i]['cands'] = list(self.nodes[i]['cands'])
            self.nodes[i]['cand_mols'] = list(self.nodes[i]['cand_mols'])
        else:
            self.nodes[i]['cands'] = []
            self.nodes[i]['cand_mols'] = []

    def recover(self):
        for i in self.nodes:
            self._recover_node(i, self.mol)

    def assemble(self):
        for i in self.nodes:
            self._assemble_node(i)
