# import mod
import molfoundry as mf
import networkx as nx
import random as rd

# from mod import BondType
from graph import Graph
from prettify import header, blue, green, red, warn, bold, underline
from typing import Dict, List

class Reaction:
    '''
    Can either call constructor with Reaction(leftGraph, rightGraph) or 
                                     Reaction(rule) (where rule is type mod.Rule)
    '''
    def __init__(self, rule=None, leftGraph=None, rightGraph=None, isMinimal=False, name="no name", context=0):
        self._name: str = rule.name if rule else name

        if leftGraph and rightGraph:
            self._leftGraph = leftGraph
            self._rightGraph = rightGraph
        else:
            self._leftGraph = Graph(rule, left=True)
            self._rightGraph = Graph(rule, left=False)

        self._educts = self._leftGraph.connected_components
        self._products = self._rightGraph.connected_components

        self._mod_rule = rule if rule else self.get_mod_rule()

        self._is_minimal: bool = isMinimal
        self._maximum_subrule: Reaction = self
        self._minimum_subrule: Reaction = self if isMinimal else self.set_minimum_subrule(add_context=int(context))
        self._shells: Reaction = {}

        self._gml_string = self.to_ruleGML_string()
        self._mod_rule_with_core = mf.Rule.fromGMLString(self._gml_string)

    # check that all connected components are used in the reaction core
    def check_if_proper(self):
        core_left = self._minimum_subrule.left_graph.nx_graph

    def get_mod_rule(self):
        return mf.Rule.fromGMLString(self.to_ruleGML_string(get_from_mod=False))

    def to_ruleGML_string(self, get_from_mod=False):

            if get_from_mod:
                return self._mod_rule.getGMLString()

            leftGraph = self._leftGraph.nx_graph
            rightGraph = self._rightGraph.nx_graph

            # NODES
            label_changing_nodes = [node for node in leftGraph.nodes if node in rightGraph.nodes and \
                                    leftGraph.nodes[node]['label']!=rightGraph.nodes[node]['label']]
            shared_nodes = [node for node in leftGraph.nodes if node in rightGraph.nodes and node not in label_changing_nodes]
            only_right_nodes = [node for node in rightGraph.nodes if node not in shared_nodes and node not in label_changing_nodes]
            only_left_nodes = [node for node in leftGraph.nodes if node not in shared_nodes and node not in label_changing_nodes]

            # EDGES
            shared_edges = []
            for (u, v, attrs1) in leftGraph.edges(data=True):
                for (p,q,attrs2) in rightGraph.edges(data=True):
                    if {u,v}=={p,q} and attrs1['bond']==attrs2['bond']:
                        shared_edges.append({u,v})
            shared_edges_list = [list(s) for s in shared_edges]

            bond_changing_edges = []
            for (u,v,attrs) in leftGraph.edges(data=True):
                for (p,q,attrs2) in rightGraph.edges(data=True):
                    if {u,v}=={p,q} and attrs['bond']!=attrs2['bond']:
                        bond_changing_edges.append({u,v})
            bond_changing_edges_list = [list(s) for s in bond_changing_edges]

            only_left_edges = []
            for (u,v) in leftGraph.edges:
                if {u,v} not in shared_edges and {u,v} not in bond_changing_edges:
                    only_left_edges.append({u,v})
            only_left_edges_list = [list(s) for s in only_left_edges]

            only_right_edges = []
            for (u,v) in rightGraph.edges:
                if {u,v} not in shared_edges and {u,v} not in bond_changing_edges:
                    only_right_edges.append({u,v})
            only_right_edges_list = [list(s) for s in only_right_edges]


            out = []
            out.append("rule [")
            out.append(f'\truleID "{self._name}"')

            out.append("\tleft [")
            out.extend([f"\t\tnode [ id {leftGraph.nodes[node]['modID']} label \"{leftGraph.nodes[node]['label']}\" ]" for node in only_left_nodes])
            out.extend([f"\t\tnode [ id {leftGraph.nodes[node]['modID']} label \"{leftGraph.nodes[node]['label']}\" ]" for node in label_changing_nodes])
            out.extend([f"\t\tedge [ source {leftGraph.nodes[u]['modID']} target {leftGraph.nodes[v]['modID']} label \"{leftGraph.edges[u,v]['bond']}\" ]" for [u,v] in only_left_edges_list])
            out.extend([f"\t\tedge [ source {leftGraph.nodes[u]['modID']} target {leftGraph.nodes[v]['modID']} label \"{leftGraph.edges[u,v]['bond']}\" ]" for [u,v] in bond_changing_edges_list])
            out.append("\t]")


            out.append("\tcontext [")
            out.extend([f"\t\tnode [ id {leftGraph.nodes[node]['modID']} label \"{leftGraph.nodes[node]['label']}\" ]" for node in shared_nodes])
                                                                # could take from left or right, arbitrarily chose leftGraph
            out.extend([f"\t\tedge [ source {leftGraph.nodes[u]['modID']} target {leftGraph.nodes[v]['modID']} label \"{leftGraph.edges[u,v]['bond']}\" ]" for [u,v] in shared_edges_list])
            out.append("\t]")

            out.append("\tright [")
            out.extend([f"\t\tnode [ id {rightGraph.nodes[node]['modID']} label \"{rightGraph.nodes[node]['label']}\" ]" for node in only_right_nodes])
            out.extend([f"\t\tnode [ id {rightGraph.nodes[node]['modID']} label \"{rightGraph.nodes[node]['label']}\" ]" for node in label_changing_nodes])
            out.extend([f"\t\tedge [ source {rightGraph.nodes[u]['modID']} target {rightGraph.nodes[v]['modID']} label \"{rightGraph.edges[u,v]['bond']}\" ]" for [u,v] in only_right_edges_list])
            out.extend([f"\t\tedge [ source {rightGraph.nodes[u]['modID']} target {rightGraph.nodes[v]['modID']} label \"{rightGraph.edges[u,v]['bond']}\" ]" for [u,v] in bond_changing_edges_list])
            out.append("\t]")

            out.append("]")

            return "\n".join(out)

    def set_minimum_subrule(self, add_context=0):
        core_left_nodes = []
        core_right_nodes = []
        left_graph = self._leftGraph.nx_graph
        left_nodes = list(left_graph.nodes)
        right_graph = self._rightGraph.nx_graph
        right_nodes = list(right_graph.nodes)

        # label changing nodes
        left_labels = { v : left_graph.nodes[v]['label'] for v in left_nodes }
        right_labels = { v : right_graph.nodes[v]['label'] for v in right_nodes }
        for v in set(left_nodes+right_nodes):
            if v in left_labels and v in right_labels and right_labels[v]!=left_labels[v]:
                core_left_nodes.append((v, left_graph.nodes[v]))
                core_right_nodes.append((v, right_graph.nodes[v]))

        # node in leftGraph, not in rightGraph
        for v in left_nodes:
            if v not in right_nodes:
                core_left_nodes.append((v, left_graph.nodes[v]))
        # node in rightGraph, not in leftGraph
        for v in right_nodes:
            if v not in left_nodes:
                core_right_nodes.append((v, right_graph.nodes[v]))
                    
        # edge in leftGraph, no edge in rightGraph
        core_left_edges = []
        for u,v,attrs in left_graph.edges(data=True):
            try:
                right_graph[u][v]
            except KeyError:
                core_left_edges.append((u,v,attrs))
                u_tuple = (u, left_graph.nodes[u])
                v_tuple = (v, left_graph.nodes[v])
                if u_tuple not in core_left_nodes: core_left_nodes.append(u_tuple)
                if v_tuple not in core_left_nodes: core_left_nodes.append(v_tuple)

        # edge in rightGraph, no edge in left graph
        core_right_edges = []
        for u,v,attrs in right_graph.edges(data=True):
            try:
                left_graph[u][v]
                # bond type changes; both left and right must have these nodes
                if left_graph[u][v]['bond'] != right_graph[u][v]['bond']:
                    core_left_edges.append((u,v,left_graph[u][v]))
                    core_right_edges.append((u,v,attrs))
                    left_u = (u, left_graph.nodes[u])
                    left_v = (v, left_graph.nodes[v])
                    right_u = (u, right_graph.nodes[u])
                    right_v = (v, right_graph.nodes[v])

                    if left_u not in core_left_nodes: core_left_nodes.append(left_u)
                    if left_v not in core_left_nodes: core_left_nodes.append(left_v)
                    if right_u not in core_right_nodes: core_right_nodes.append(right_u)
                    if right_v not in core_right_nodes: core_right_nodes.append(right_v)

            except KeyError:
                core_right_edges.append((u,v,attrs))
                u_tuple = (u, right_graph.nodes[u])
                v_tuple = (v, right_graph.nodes[v])
                if u_tuple not in core_right_nodes: core_right_nodes.append(u_tuple)
                if v_tuple not in core_right_nodes: core_right_nodes.append(v_tuple)
        
        # nothing changes
        if core_left_nodes==[] and core_right_nodes==[] and not core_left_edges and not core_right_edges:
            print(red(f"Cannot set minimum subrule for reaction {self._name} as no bonds change."))
            return False


        # ADDING EXTRA NODES FOR CONTEXT
        while add_context>0:
            added=False
            int_nodes = [v for v,_ in core_left_nodes]
            for v in int_nodes:
                if left_graph.degree[v] > 1:
                    nbs = list(left_graph.neighbors(v))
                    for neighbour in nbs:
                        if neighbour not in int_nodes:
                            # print(blue(f"added context: {left_graph.nodes[neighbour]} | {right_graph.nodes[neighbour]}"))
                            core_left_nodes.append((neighbour, {'label':left_graph.nodes[neighbour]['label'], 'modID':left_graph.nodes[neighbour]['modID']}))
                            core_right_nodes.append((neighbour, {'label':right_graph.nodes[neighbour]['label'], 'modID':right_graph.nodes[neighbour]['modID']}))
                            add_context-=1
                            added=True
            if added==False:
                break

        core_left_graph = nx.Graph()
        core_left_graph.add_nodes_from(core_left_nodes)
        core_left_graph.add_edges_from(core_left_edges)
        
        core_right_graph = nx.Graph()
        core_right_graph.add_nodes_from(core_right_nodes)
        core_right_graph.add_edges_from(core_right_edges)

        new_left_graph = Graph(core_left_graph)
        new_right_graph = Graph(core_right_graph)

        # if parent reaction does not synthesize nor fragment, minimum subrule should not either
        # aka. minimum subrule should stay connected if parent reaction is connected
        # solution: take the shortest path between connected components
        '''if len(self._educts)==1 and len(self._products)==1 and (new_left_graph.num_components>1 or new_right_graph.num_components>1):
            comp_nodes = []
            if new_left_graph.num_components > 1:
                ccps = new_left_graph.connected_components
            else:
                ccps = new_right_graph.connected_components

            for component in ccps:
                acceptable_nodes = [ n for n in list(component.nx_graph.nodes()) if (n in left_graph and n in right_graph) ]
                comp_nodes.append(acceptable_nodes[0])
            
            tsp = nx.approximation.traveling_salesman_problem
            nodes_to_add = tsp(left_graph, nodes=comp_nodes, cycle=False)

            for n in nodes_to_add:
                core_left_graph.add_node(n, label=left_graph.nodes[n]['label'], modID=left_graph.nodes[n]['modID'])
                core_right_graph.add_node(n, label=right_graph.nodes[n]['label'], modID=right_graph.nodes[n]['modID'])

                for u in left_graph.neighbors(n):
                    if u in core_left_graph.nodes():
                        core_left_graph.add_edge(u,n, bond=left_graph[u][n]['bond'])

                for u in right_graph.neighbors(n):
                    if u in core_right_graph.nodes():
                        core_right_graph.add_edge(u,n, bond=right_graph[u][n]['bond'])
                
            new_left_graph = Graph(core_left_graph)
            new_right_graph = Graph(core_right_graph)'''

        # relabel core nodes in original graph
        for n in core_left_graph.nodes:
            for g in [self._leftGraph, core_left_graph]:
                label = g.nodes[n]['label']
                if 'core' not in label:
                    g.nodes[n]['label'] = f"{label}_core"

        for n in core_right_graph.nodes:
            for g in [self._rightGraph, core_right_graph]:
                label = g.nodes[n]['label']
                if 'core' not in label:
                    g.nodes[n]['label'] = f"{label}_core"

        return Reaction(leftGraph=new_left_graph, rightGraph=new_right_graph, isMinimal=True, name=(str(self._name)+"_min"))
    
    def copy(self, exact=False):
        #if exact:
        #    new_left_graph = Graph(self._leftGraph.gml_string)
        #    new_right_graph = Graph(self._rightGraph.gml_string)
        #else:
        new_left_graph = Graph(self._leftGraph.gml_string)
        new_right_graph = Graph(self._rightGraph.gml_string)
        return Reaction(leftGraph=new_left_graph, rightGraph=new_right_graph, name=self._name)

    def get_shell(self, k:int):
        if k<=0:
            return self._minimum_subrule
        try:
            reac_k = self._shells[k]
            return reac_k
        except KeyError:
            pass 

        prev_shell = self.get_shell(k-1).copy()
        lgraph = prev_shell.left_graph.nx_graph
        rgraph = prev_shell.right_graph.nx_graph
        max_lgraph = self._leftGraph.nx_graph
        max_rgraph = self._rightGraph.nx_graph

        lnodes = list(lgraph.nodes)
        rnodes = list(rgraph.nodes)
        for n in lnodes:
            for neigh in list(max_lgraph.neighbors(n)):
                lgraph.add_node(neigh, label=max_lgraph.nodes[neigh]['label'], modID=max_lgraph.nodes[neigh]['modID'])
                lgraph.add_edge(n, neigh, bond=max_lgraph[n][neigh]['bond'])
        
        for n in rnodes:
            for neigh in list(max_rgraph.neighbors(n)):
                rgraph.add_node(neigh, label=max_rgraph.nodes[neigh]['label'], modID=max_rgraph.nodes[neigh]['modID'])
                rgraph.add_edge(n, neigh, bond=max_rgraph[n][neigh]['bond'])
        new_shell = Reaction(leftGraph=Graph(lgraph), rightGraph=Graph(rgraph), isMinimal=True, name=f"{self._name}_k{k}")
        self._shells[k] = new_shell
        return new_shell

    @property
    def size(self):
        return self.left_graph.nx_graph.number_of_nodes()

    @property
    def minimum_subrule(self):
        return self._minimum_subrule

    @property
    def maximum_subrule(self):
        return self._maximum_subrule

    @property
    def mod_rule(self):
        return self._mod_rule

    @property
    def educts(self):
        return self._educts

    @property
    def products(self):
        return self._products

    @property
    def left_graph(self):
        return self._leftGraph

    @property
    def right_graph(self):
        return self._rightGraph

    @property
    def name(self):
        return self._name

    @name.setter
    def name(self, val):
        self._name = val

    @property
    def gml(self):
        return self._gml_string

    @property
    def mod_rule_with_core(self):
        return self._mod_rule_with_core

    def is_minimal(self):
        return self._is_minimal

    def __eq__(self, other: 'Reaction') -> bool:
        return hash(self) == hash(other)

    def __lt__(self, other):
        return hash(self) < hash(other)

    def __le__(self, other):
        return hash(self) <= hash(other)

    def __gt__(self, other):
        return hash(self) > hash(other)

    def __ge__(self, other):
        return hash(self) >= hash(other)

    def __ne__(self, other):
        return not self==other

    def __str__(self) -> str:
        return self._name

    def __repr__(self) -> str:
        return self._name

    def __hash__(self) -> int:
        return (hash(self.left_graph)-hash(self.right_graph))


    def _self_check_bond_counts(self):
        for nid in self._leftGraph.nx_graph.nodes:
            left_count = count_bonds(self._leftGraph.nx_graph, nid)
            right_count = count_bonds(self._rightGraph.nx_graph, nid)
            if left_count != right_count:
                print(red(f"BOND COUNT HAS CHANGED! from {left_count} to {right_count}"))

def count_bonds(graph: nx.Graph, v_id):
    count = 0
    for (u, v) in graph.edges([v_id]):
        bt = graph[u][v]['bond']
        if bt == BondType.Single:
            count+=1
        elif bt == BondType.Double:
            count+=2
        elif bt == BondType.Triple:
            count+=3
        elif bt == BondType.Aromatic:
            count+=1.5
        else:
            print(red("BOND DETECTED OTHER THAN EXPECTED"))
    return count