# import mod
import molfoundry as mf
import networkx as nx
import time

import re

# from equilibrator_assets.generate_compound import create_compounds
from prettify import red, warn

# each vertex has attribute "label" which stores string label
# each edge has attribute "bond" which stores bond order (-, =, etc.)

class Graph:
    def __init__(self, graph, left=False, verbose=False):
        self._verbose = verbose
        if isinstance(graph, nx.Graph):
            self._nx_graph = graph
            self._gml_string  = self.nx_graph_to_GML_string(graph)
        elif isinstance(graph, str):
            self._gml_string = graph
            self._nx_graph = self.GML_to_nx_graph(self._gml_string)
        elif isinstance(graph, mf.graph.Graph):
            self._nx_graph = self.mf_to_nx_graph(graph)
            self._gml_string = self.nx_graph_to_GML_string(self._nx_graph)
        elif isinstance(graph, mf.rule.Rule):
            if left:
                self._nx_graph = self.rule_to_nx_graph(graph, left=True)
                self._gml_string = self.nx_graph_to_GML_string(self._nx_graph)
            else:
                self._nx_graph = self.rule_to_nx_graph(graph, left=False)
                self._gml_string = self.nx_graph_to_GML_string(self._nx_graph)
        else:
            print(red(f"Error: type not found for {type(graph)}"))

        # up to the caller to check the number of components created
        ccs = [self._nx_graph.subgraph(c).copy() for c in nx.connected_components(self._nx_graph)]
        ccs = sorted(ccs, key=len, reverse=True)
        self._num_components = len(ccs)
        if self._num_components > 1:
            self._components = [Graph(c) for c in ccs]
        else:
            self._components = [self]

        if self._num_components == 1:
            self._mod_graph = self.nx_graph_to_mf(self._nx_graph)

    def rule_to_nx_graph(self, rule, left):
        g = nx.Graph()

        # vertices that don't change
        for v in rule.vertices:
            if v.isContext:
                g.add_node(int(v.id), label=v.leftLabel, modID=int(v.id))

        # bonds that don't change
        for e in rule.edges:
            if e.isContext:
                g.add_edge(e.source.id, e.target.id, bond=e.leftLabel)

        if left:
            # vertices/edges which are destroyed
            verts = [v for v in rule.vertices if v.isDeleted]
            edges = [e for e in rule.edges if e.isDeleted]
            
        else:
            # vertices/edges which are created
            verts = [v for v in rule.vertices if v.isCreated]
            edges = [e for e in rule.edges if e.isCreated]

        # vertices which change label
        verts.extend([v for v in rule.vertices if v.isChanged])
        # bonds which change label
        edges.extend([e for e in rule.edges if e.isChanged])

        for v in verts:
            vLabel = v.leftLabel if left else v.rightLabel
            g.add_node(v.id, label=vLabel, modID=v.id)

        for e in edges:
            eLabel = e.leftLabel if left else e.rightLabel
            g.add_edge(e.source.id, e.target.id, bond=eLabel)

        return g
    
    def mf_to_nx_graph(self, mfGraph):
        g = nx.Graph()

        for v in mfGraph.vertices:
            g.add_node(int(v.id), label=str(v.stringLabel), modID=int(v.id))

        for e in mfGraph.edges:
            g.add_edge(int(e.source.id), int(e.target.id), bond=e.stringLabel)
        
        return g

    def nx_graph_to_mf(self, nxGraph):
        try:
            return mf.Graph.fromGMLString(self.nx_graph_to_GML_string(nxGraph))
        except: # graph is not connected probably
            if self._verbose:
                print(warn("Error converting nxGraph to mod graph. Likely graph is not connected. This will not affect rule generation."))
            return None

    def GML_to_nx_graph(self, gml_string):
        g = nx.Graph()
        lines = gml_string.split("\n")
        for line in lines:
            tokens = line.split()
            if len(tokens) < 2:
                continue
            if tokens[0] == "node":
                m = re.search(r"label\s+\"(.*)\"", line)
                l = m.group(1)
                m = re.search(r"id\s+(\d+)", line)
                mid = m.group(1)
                #(_,_,_, mid, _, l, _) = tokens
                g.add_node(int(mid), label=l, modID=int(mid))
            elif tokens[0] == "edge":
                m = re.search(r"source\s+(\d+)\s+target\s+(\d+)\s+label\s+\"(.*)\"", line)
                u = m.group(1)
                v = m.group(2)
                bondOrder = m.group(3)
                #(_, _, _, u, _, v, _, bondOrder, _ ) = tokens
                g.add_edge(int(u), int(v), bond=bondOrder)
        return g

    def nx_graph_to_GML_string(self, nxGraph):
        out = []
        out.append("graph [")

        out.extend([f"\t\tnode [ id {nxGraph.nodes[node]['modID']} label \"{nxGraph.nodes[node]['label']}\" ]" for node in nxGraph.nodes])
        out.extend([f"\t\tedge [ source {u} target {v} label \"{nxGraph[u][v]['bond']}\" ]"
            for (u,v) in nxGraph.edges])

        out.append("]")

        return "\n".join(out)

    @property
    def nx_graph(self):
        return self._nx_graph

    @property
    def mod_graph(self):
        if self.num_components == 1:
            return self._mod_graph
        else:
            return mf.Graph.fromGMLString(self.gml_string)

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

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

    @property
    def edges(self):
        return self._nx_graph.edges

    @property
    def nodes(self):
        return self._nx_graph.nodes

    @property
    def num_components(self):
        return self._num_components

    @property
    def connected_components(self):
        return self._components

    def mod_print(self):
        self._mod_graph.print()

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

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