from __future__ import annotations

import re

import networkx as nx
from rdkit import Chem as rdk
from rdkit.Chem.rdChemReactions import ChemicalReaction

from chemrecon.chem.gml.gml import remove_residues, GMLGraph


atom_regex = re.compile(r'([A-Z][a-z]*)([0-9])?([+\-])?')


# Molecule
def gml_to_rdkit_mol(gml: str) -> rdk.Mol:
    gmlgraph = GMLGraph(gml)
    rdk_mol = rdk.EditableMol(rdk.Mol())
    node_index: dict[int, int] = {}

    # Atoms
    for nx_index, nx_properties in gmlgraph.graph.nodes.items():
        # Add atom with element (charge?) given by the label, and atom map number given by the index
        atom = None
        label: str = nx_properties['label']
        if label == '*':
            atom = atom = rdk.AtomFromSmiles(f'[*]')
        else:
            match = atom_regex.match(label)
            charge = 0
            try:
                symbol = match.group(1)
                plusminus = match.group(3)
                if plusminus is None:
                    charge = 0
                else:
                    if plusminus == '+':
                        charge = 1
                    elif plusminus == '-':
                        charge = -1
                    else:
                        raise Exception()

                chargeabs = match.group(2)
                if chargeabs:
                    charge * int(chargeabs)

            except Exception:
                raise ValueError(f'Cannot convert S_GML: Atom label: {label}')
                # error.

            if symbol == 'R':
                # Radical
                # TODO
                atom = rdk.AtomFromSmiles(f'[*]')
            elif symbol.startswith('Alias'):
                # 'Alias in GML?'
                # TODO what does this mean exactly
                atom = rdk.AtomFromSmiles(f'[*]')
            else:
                atom = rdk.AtomFromSmiles(f'[{symbol}]')

            # Set charge based on match
            atom.SetFormalCharge(charge)

        # Atom is set
        if atom is None:
            raise ValueError(f'Cannot convert S_GML: Atom label: {label}')

        atom.SetAtomMapNum(nx_index)
        atom_index = rdk_mol.AddAtom(atom)
        node_index[nx_index] = atom_index

    # Edges
    for a, b in gmlgraph.graph.edges():
        index_a = node_index[a]
        index_b = node_index[b]
        label: str = gmlgraph.graph.get_edge_data(a, b)['label']
        match label:
            case '-':
                bondtype = rdk.BondType.SINGLE
            case '=':
                bondtype = rdk.BondType.DOUBLE
            case '#':
                bondtype = rdk.BondType.TRIPLE
            case ':':
                bondtype = rdk.BondType.AROMATIC
            case _:
                raise ValueError(f'Unknown bond type: {label}.')
        rdk_mol.AddBond(index_a, index_b, bondtype)

    # Mol is done, change to uneditable
    rdk_mol_final = rdk_mol.GetMol()
    return rdk_mol_final

# Reaction
def gml_to_rdkit_reaction(gml: str) -> ChemicalReaction:
    """ Relies on S_GML in the format outputted by the M-CSA script.
    """

    # TODO change to re-use above code?

    # Load two S_GML strings
    l_str = ''
    r_str = ''
    step: int = 0
    currentstr: str = ''
    for l in gml.splitlines():
        if l.startswith('graph') or l.startswith('left') or l.startswith('right'):
            if step == 1:
                l_str = currentstr
            elif step == 2:
                r_str = currentstr
            # Start new string
            step += 1
            currentstr = ''
        # Add to current string
        currentstr += l

    assert step == 2
    r_str = currentstr

    # Get the components and turn into MolInstances
        # TODO should not be necessary
    l_chemgraph = remove_residues(GMLGraph(l_str.replace('left', 'graph')))
    r_chemgraph = remove_residues(GMLGraph(r_str.replace('right', 'graph')))

    # For each component, add it
    rdk_r = ChemicalReaction()
    rdk_reaction = ChemicalReaction()

    for chemgraph, side in [[l_chemgraph, -1], [r_chemgraph, 1]]:
        for component in chemgraph.get_components_subgraphs():
            component: nx.Graph
            rdk_mol = rdk.EditableMol(rdk.Mol())
            node_index: dict[int, int] = {}

            # Atoms
            for nx_index, nx_properties in component.nodes.items():
                # Add atom with element (charge?) given by the label, and atom map number given by the index
                atom = None
                label: str = nx_properties['label']
                if label == '*':
                    atom = rdk.AtomFromSmiles(f'[*]')
                else:
                    match = atom_regex.match(label)
                    charge = 0
                    try:
                        symbol = match.group(1)
                        plusminus = match.group(3)
                        if plusminus is None:
                            charge = 0
                        else:
                            if plusminus == '+':
                                charge = 1
                            elif plusminus == '-':
                                charge = -1
                            else:
                                raise Exception()

                        chargeabs = match.group(2)
                        if chargeabs:
                            charge * int(chargeabs)

                    except Exception:
                        raise ValueError(f'Cannot convert S_GML: Atom label: {label}')
                        # error.
                    if symbol == 'R':
                        # Radical
                        # TODO
                        atom = rdk.AtomFromSmiles(f'[*]')
                    elif symbol.startswith('Alias'):
                        # 'Alias in GML?'
                        # TODO what does this mean exactly
                        atom = rdk.AtomFromSmiles(f'[*]')
                    else:
                        atom = rdk.AtomFromSmiles(f'[{symbol}]')


                if atom is None:
                    raise ValueError(f'Cannot convert S_GML: Atom label: {label}')
                atom.SetFormalCharge(charge)
                atom.SetAtomMapNum(nx_index)
                atom_index = rdk_mol.AddAtom(atom)
                node_index[nx_index] = atom_index

            # Edges
            for a, b in component.edges():
                index_a = node_index[a]
                index_b = node_index[b]
                label: str = component.get_edge_data(a, b)['label']
                match label:
                    case '-':
                        bondtype = rdk.BondType.SINGLE
                    case '=':
                        bondtype = rdk.BondType.DOUBLE
                    case '#':
                        bondtype = rdk.BondType.TRIPLE
                    case ':':
                        bondtype = rdk.BondType.AROMATIC
                    case _:
                        raise ValueError(f'Unknown bond type: {label}.')
                rdk_mol.AddBond(index_a, index_b, bondtype)

                pass

            # Mol is done, change to uneditable
            rdk_mol_final = rdk_mol.GetMol()
            # TODO ADD mol to reaction
            match side:
                case -1:
                    rdk_reaction.AddReactantTemplate(rdk_mol_final)
                case 1:
                    rdk_reaction.AddProductTemplate(rdk_mol_final)

    # Reaction now completed
    return rdk_reaction
