""" Code to handle S_GML from Juri's M-CSA tool
"""
from __future__ import annotations

import re
from typing import Optional

import matplotlib.pyplot as plt
import networkx as nx
from rdkit import Chem as rdk

from chemrecon.chem.mol import Mol

# Regex patterns
rule_pattern = re.compile(
    r'^rule \[(?:ruleID \".*\")?(?:left \[(.*?)\])?context \[(.*?)\](?:right \[(.*?)\])?\]$'
)
label_sub = re.compile(
    r'id (\d+) label \"([A-Z][a-z]?\d?\+?\-?|Amino\(..?, ..., \d+, [\w|*]\)|Alias\(.,.\)|\*)\"'
)
label_rep = r'id \g<1> label "\g<1>: \g<2>"'


def sanitise_gml(gml: str) -> str:
    r = f'{gml}'
    r.replace('\n', '').replace('\t', '')
    return r


class GMLException(Exception):
    pass


class GMLEmptyException(Exception):
    pass


class GMLEntry:
    entry_index: int
    proposals: list[list[GMLStep]]
    map: dict[int, int]  # For proposal 0

    def __init__(self, index: int):
        self.entry_index = index
        self.proposals = []

    def add_step(self, proposal_index: int, step: GMLStep):
        # Extend proposals list
        while len(self.proposals) < proposal_index:
            self.proposals.append([])

        # Add step
        self.proposals[proposal_index - 1].append(step)

    def set_map(self, aam: dict[int, int]):
        self.map = aam
        sanity_check(self.left_graph(), self.right_graph(), self.map)

    # Getters always get proposal 0 (best star rating in M-CSA)
    def left_graph(self) -> GMLGraph:
        return self.proposals[0][0].left_graph()

    def right_graph(self) -> GMLGraph:
        return self.proposals[0][-1].right_graph()

    def atom_atom_map(self) -> dict[int, int]:
        return self.map

class GMLStep:
    step_index: int
    gml: str
    lhs: GMLGraph
    rhs: GMLGraph

    # Pure gml
    left_gml: str
    context_gml: str
    right_gml: str

    def __init__(self, step_index: int, gml: str):
        self.step_index = step_index
        self.gml = str(gml).replace('\n', '').replace('\t', '')

        # Make S_GML for educt and product sides
        match = re.match(rule_pattern, self.gml)
        self.left_gml = match.group(1)
        self.context_gml = match.group(2)
        self.right_gml = match.group(3)

        lhs = (f'graph [\n {self.left_gml} {self.context_gml} \n]'
               if self.left_gml else f'graph [\n{self.context_gml}\n]')
        rhs = (f'graph [\n {self.context_gml} {self.right_gml} \n]'
               if self.right_gml else f'graph [\n{self.context_gml}\n]')

        self.lhs = GMLGraph(
            gml = lhs
        )
        self.rhs = GMLGraph(
            gml = rhs
        )

    def left_graph(self) -> GMLGraph:
        return self.lhs

    def right_graph(self) -> GMLGraph:
        return self.rhs


class GMLGraph:
    graph: nx.Graph
    gml: str
    component_subsets: list[frozenset[int]]
    component_subgraphs: list[nx.Graph]

    def __init__(self, gml: str, graph: nx.Graph = None):
        self.gml = sanitise_gml(gml)

        # TODO raise exception of called with empty graph

        if graph:
            self.graph = graph
        else:
            # Make graph from gml
            self.graph = nx.parse_gml(
                gml,
                label = 'id'
            )
        self.component_subgraphs = []
        self.component_subsets = []

    def get_component_subsets(self) -> list[frozenset[int]]:
        if not self.component_subsets:
            self.component_subsets = list(
                frozenset(s)
                for s in nx.connected_components(self.graph)
            )
        return self.component_subsets

    def get_components_subgraphs(self) -> list[nx.Graph]:
        if self.component_subgraphs:
            return self.component_subgraphs
        else:
            if not self.component_subsets:
                # Make subsets
                self.get_component_subsets()
            # Set the subgraphs
            self.component_subgraphs = [
                nx.subgraph(self.graph, compset) for compset in self.component_subsets
            ]
            return self.component_subgraphs

    def remove_hydrogens(self):
        g_ = self.graph.copy()
        to_remove: list = list()
        for n in g_.nodes():
            if g_.nodes[n]['label'] == 'H' and g_.degree[n] > 0:
                to_remove.append(n)
        g_.remove_nodes_from(to_remove)
        self.graph = g_

    def get_subgraph(self, indices: set[int]) -> nx.Graph:
        return self.graph.subgraph(indices)

    def draw(self):
        # Draw with networkx and matplotlib - need to display atom numbers on label!
        plt.figure(1, figsize = (8, 8), dpi = 240)
        pos = nx.spring_layout(self.graph)

        # Graph structure and node labels
        nx.draw(
            self.graph, pos,
            with_labels = True,
            node_size = 50,
            font_size = 6,
            labels = {
                k: f"{k}: {n['label']}" for k, n in self.graph.nodes.items()
            }
        )

        # Edge labels
        edict = dict()
        for (u, v, d) in self.graph.edges.data('label'):
            edict[(u, v)] = d

        nx.draw_networkx_edge_labels(
            self.graph, pos, edict, font_color = 'blue', font_size = 6
        )

        # Display graph
        plt.show()
        # TODO


# ChemGraph generation
def chemgraph_from_mol(mol: Mol) -> GMLGraph:
    # Based on an RDK mol, create a ChemGraph object.
    struct_graph = nx.Graph()
    for atom in mol.mol.GetAtoms():
        # Make the label based on element and charge
        fcharge = atom.GetFormalCharge()
        symbol = atom.GetSymbol()

        # Do not include hydrogen
        if symbol == 'H' and len(atom.GetNeighbors()) > 0:
            continue

        # Make label
        label: str = 'ERR'
        match fcharge:
            case 0:
                label = symbol
            case -1:
                label = f'{symbol}-'
            case 1:
                label = f'{symbol}+'
            case _ if fcharge < 0:
                label = f'{symbol}{fcharge}-'
            case _ if fcharge >= 0:
                label = f'{symbol}{fcharge}+'

        struct_graph.add_node(
            atom.GetIdx(),
            label = label
        )

    for bond in mol.mol.GetBonds():
        label: str = 'ERR'
        match bond.GetBondType():
            case rdk.BondType.SINGLE:
                label = '-'
            case rdk.BondType.DOUBLE:
                label = '='
            case rdk.BondType.TRIPLE:
                label = '#'
            case rdk.BondType.AROMATIC:
                label = ':'
            case _:
                raise RuntimeError(f'Unimplemented bond-type {bond.GetBondType()}')

        struct_graph.add_edge(
            bond.GetBeginAtomIdx(),
            bond.GetEndAtomIdx(),
            label = label
        )

    # Finalize and add the graph to the list of graphs to match
    cgraph = GMLGraph(
        gml = '',
        graph = struct_graph
    )
    return cgraph


# Entry serialisation
# ----------------------------------------------------------------------------------------------------------------------
def graph_to_gml(graph: nx.Graph, relabel_map: Optional[dict[int, int]] = None) -> str:
    """ Only returns 'inner' gml.
    """
    if relabel_map is None:
        relabel_map = {x: x for x in graph.nodes}

    s = '[\n'
    # s = 'graph [\n'
    # Nodes
    for n_index in graph.nodes:
        label = graph.nodes[n_index]['label']
        s += f'\tnode [ id {relabel_map[n_index]} label "{label}"]\n'

    # Edges
    for x, y in graph.edges:
        label = graph.edges[(x, y)]['label']
        s += f'\tedge [ source {relabel_map[x]} target {relabel_map[y]} label "{label}" ]\n'
    s += '\n]'
    return s

# Reaction graph sanitation and handling -------------------------------------------------------------------------------
def remove_residues(
    chemgraph: GMLGraph
) -> GMLGraph:
    """ Given a chemgraph representing either side of a reaction, return the same graph with identified
        catalytic enzyme residues removed.
    """
    g: nx.Graph = chemgraph.graph
    remove_nodes: set = set()  #

    components = list(nx.connected_components(g))
    for component in components:
        for v in component:
            l: str = g.nodes[v]['label']
            if l.startswith('Amino'):
                # Remove the entire component if it contains an amino acid
                remove_nodes = remove_nodes.union(component)
                break

    # Remove the nodes
    g_ = g.copy(as_view = False)
    g_.remove_nodes_from(remove_nodes)
    if g_.number_of_nodes() == 0:
        raise GMLEmptyException
    return GMLGraph(gml = '', graph = g_)


# Misc
# --------------------------------------------------------------------------------------------------------------
# Matching functions for graph similarity
def nmatch(n1, n2):
    try:
        return n1['label'] == n2['label']
    except KeyError:
        return 1


def ematch(e1, e2):
    try:
        return e1['label'] == e2['label']
    except KeyError:
        return 1

def sanity_check(l_graph: GMLGraph, r_graph: GMLGraph, map: dict[int, int]):
    for l_index, r_index in map.items():
        l_label: str = l_graph.graph.nodes[l_index]['label']
        r_label: str = r_graph.graph.nodes[r_index]['label']
        if l_label.removesuffix('-').removesuffix('+') != r_label.removesuffix('-').removesuffix('+'):
            assert False, 'Wrong match'
