import logging
from queue import Queue
# import networkx as nx
from rdkit import Chem
import numpy as np

def mol_equal(smi_a, smi_b):
    return Chem.CanonSmiles(smi_a) == Chem.CanonSmiles(smi_b)

def mol_set_equal(smi_set_a, smi_set_b):
    a = set()
    for smi in smi_set_a:
        a.add(Chem.CanonSmiles(smi))
    b = set()
    for smi in smi_set_b:
        b.add(Chem.CanonSmiles(smi))

    return a == b

class SynRoute:
    def __init__(self, target_mol, succ_value, search_status):
        self.target_mol = target_mol
        self.mols = [target_mol]
        self.values = [None]
        self.templates = [None]
        self.parents = [-1]
        self.children = [None]
        self.optimal = False
        self.costs = {}

        self.succ_value = succ_value
        self.total_cost = 0
        self.length = 0
        self.search_status = search_status
        if self.succ_value <= self.search_status:
            self.optimal = True

    def _add_mol(self, mol, parent_id):
        self.mols.append(mol)
        self.values.append(None)
        self.templates.append(None)
        self.parents.append(parent_id)
        self.children.append(None)

        self.children[parent_id].append(len(self.mols)-1)

    def set_value(self, mol, value):
        assert mol in self.mols

        mol_id = self.mols.index(mol)
        self.values[mol_id] = value

    def add_reaction(self, mol, value, template, reactants, cost):
        assert mol in self.mols

        self.total_cost += cost
        self.length += 1

        parent_id = self.mols.index(mol)
        self.values[parent_id] = value
        self.templates[parent_id] = template
        self.children[parent_id] = []
        self.costs[parent_id] = cost

        for reactant in reactants:
            self._add_mol(reactant, parent_id)

    def serialize_reaction(self, idx):
        s = self.mols[idx]
        if self.children[idx] is None:
            return s
        s += '>%.4f>' % np.exp(-self.costs[idx])
        s += self.mols[self.children[idx][0]]
        for i in range(1, len(self.children[idx])):
            s += '.'
            s += self.mols[self.children[idx][i]]

        return s

    def serialize(self):
        s = self.serialize_reaction(0)
        for i in range(1, len(self.mols)):
            if self.children[i] is not None:
                s += '|'
                s += self.serialize_reaction(i)

        return s

    def contains_gt(self, gt_smi_list):
        if mol_set_equal([self.mols[0]], gt_smi_list):
            return True
        else:
            one_step_list = []
            if self.children[0] is not None:
                for c_ind in self.children[0]:
                    one_step_list.append(self.mols[c_ind])

            return mol_set_equal(one_step_list, gt_smi_list)

    def viz_route(self, viz_file):
        from graphviz import Digraph
        # pass
        G = Digraph('G', filename=viz_file)
        # G.attr(rankdir='LR')
        G.attr('node', shape='box')
        G.format = 'pdf'

        names = []
        for i in range(len(self.mols)):
            name = self.mols[i]
            # if self.templates[i] is not None:
            #     name += ' | %s' % self.templates[i]
            names.append(name)

        node_queue = Queue()
        node_queue.put((0,-1))   # target mol idx, and parent idx
        while not node_queue.empty():
            idx, parent_idx = node_queue.get()

            # G.add_node(names[idx])
            if parent_idx >= 0:
                G.edge(names[parent_idx], names[idx], label='cost')

            if self.children[idx] is not None:
                for c in self.children[idx]:
                    node_queue.put((c, idx))

        G.render()