<commit_before>from __future__ import annotations

from typing import Dict, Iterable, List, Set, Tuple, Union

from selfies.grammar_rules import find_element, get_num_from_bond, \
    parse_atom_symbol

ATOM_TYPE = 1
BRANCH_TYPE = 2
RING_TYPE = 3


def kekulize_parser(smiles_gen: Iterable[Tuple[str, str, int]]) \
        -> Iterable[Tuple[str, str, int]]:
    """Kekulizes a SMILES in the form of an iterable.

    This method intercepts the output of ``encoder._parse_smiles``, and
    acts as filter that kekulizes the SMILES. The motivation for having
    this setup is that string parsing and concatenation is minimized,
    as the parsing is already done by ``_parse_smiles``.

    Reference: https://depth-first.com/articles/2020/02/10/a-comprehensive
               -treatment-of-aromaticity-in-the-smiles-language/

    :param smiles_gen: an iterator returned by ``encoder._parse_smiles``.
    :return: an iterator representing the kekulized SMILES, in the same
        format as that returned by ``encoder._parse_smiles``.
    """

    # save to list, so the iterator can be used across multiple functions
    # change elements from tuple -> list to allow in-place modifications
    smiles_symbols = list(map(list, smiles_gen))

    mol_graph = MolecularGraph(smiles_symbols)

    rings = {}
    _build_molecular_graph(mol_graph, smiles_symbols, rings)

    if mol_graph.aro_indices:
        _kekulize(mol_graph)

    for x in mol_graph.smiles_symbols:  # return as iterator
        yield tuple(x)


def _build_molecular_graph(graph: MolecularGraph,
                           smiles_symbols: List[List[Union[str, int]]],
                           rings: Dict[int, Tuple[int, int]],
                           prev_idx: int = -1,
                           curr_idx: int = -1) -> int:
    """From the iterator returned by ``encoder._parse_smiles``, builds
    a graph representation of the molecule.

    This is done by iterating through ``smiles_symbols``, and then adding bonds
    to the molecular graph. Note that ``smiles_symbols`` is mutated in this
    method, for convenience.

    :param graph: the MolecularGraph to be added to.
    :param smiles_symbols: a list created from the iterator returned
        by ``encoder._parse_smiles``.
    :param rings: an, initially, empty dictionary used to keep track of
        rings to be made.
    :param prev_idx:
    :param curr_idx:
    :return: the last index of ``smiles_symbols`` that was processed.
    """

    while curr_idx + 1 < len(smiles_symbols):

        curr_idx += 1
        _, symbol, symbol_type = smiles_symbols[curr_idx]

        if symbol_type == ATOM_TYPE:
            if prev_idx >= 0:
                graph.add_bond(prev_idx, curr_idx, curr_idx)
            prev_idx = curr_idx

        elif symbol_type == BRANCH_TYPE:
            if symbol == '(':
                curr_idx = _build_molecular_graph(graph, smiles_symbols, rings,
                                                  prev_idx, curr_idx)
            else:
                break

        else:
            if symbol in rings:
                left_idx, left_bond_idx = rings.pop(symbol)
                right_idx, right_bond_idx = prev_idx, curr_idx

                # we mutate one bond index to be '', so that we
                # can faithfully represent the bond to be localized at
                # one index. For example, C=1CCCC=1 --> C1CCCC=1.

                if smiles_symbols[left_bond_idx][0] != '':
                    bond_idx = left_bond_idx
                    smiles_symbols[right_bond_idx][0] = ''
                else:
                    bond_idx = right_bond_idx
                    smiles_symbols[left_bond_idx][0] = ''

                graph.add_bond(left_idx, right_idx, bond_idx)
            else:
                rings[symbol] = (prev_idx, curr_idx)

    return curr_idx


def _kekulize(mol_graph: MolecularGraph) -> None:
    """Kekulizes the molecular graph.

    :param mol_graph: a molecular graph to be kekulized.
    :return: None.
    """

    mol_graph.prune_to_pi_subgraph()

    visited = set()
    for i in mol_graph.get_nodes_by_num_edges():
        success = mol_graph.dfs_assign_bonds(i, visited, set(), set())
        if not success:
            raise ValueError("Kekulization Failed.")

    mol_graph.write_to_smiles_symbols()


# Aromatic Helper Methods and Classes

# key = aromatic SMILES element, value = number of valence electrons
# Note: wild card '*' not supported currently
_aromatic_valences = {
    'b': 3, 'c': 4, 'n': 5, 'p': 5, 'as': 5,
    'o': 6, 's': 6, 'se': 6
}


def _capitalize(atom_symbol: str) -> str:
    """Capitalizes the element portion of an aromatic SMILES atom symbol,
    converting it into a standard SMILES atom symbol.

    :param atom_symbol: an aromatic SMILES atom symbol.
    :return: the capitalized ``atom_symbol``.
    """

    s, _ = find_element(atom_symbol)
    return atom_symbol[:s] + atom_symbol[s].upper() + atom_symbol[s + 1:]


def _is_aromatic(atom_symbol: str) -> bool:
    """Checks whether a SMILES atom symbol is an aromatic SMILES atom symbol.

    An aromatic SMILES atom symbol is indicated by an element substring
    that is not capitalized.

    :param atom_symbol: a SMILES atom symbol.
    :return: True, if ``atom_symbol`` is an aromatic atom symbol,
        and False otherwise.
    """

    s, e = find_element(atom_symbol)

    if e == len(atom_symbol):  # optimization to prevent string copying
        element = atom_symbol
    else:
        element = atom_symbol[s: e]

    if element[0].isupper():  # check if element is capitalized
        return False

    if element not in _aromatic_valences:
        raise ValueError(f"Kekulization Failed: aromatic symbol {atom_symbol} "
                         f"not recognized.")
    return True


def _in_pi_subgraph(atom_symbol: str, bonds: Tuple[str]) -> bool:
    """Checks whether a SMILES atom symbol should be a node in the pi
    subgraph, based on its bonds.

    More specifically, an atom should be a node in the pi subgraph if it has
    an unpaired valence electron, and thus, is able to make a double bond.

    Reference: https://depth-first.com/articles/2020/02/10/a-comprehensive
               -treatment-of-aromaticity-in-the-smiles-language/

    :param atom_symbol: a SMILES atom symbol representing an atom.
    :param bonds: the bonds connected to ``atom_symbol``.
    :return: True if ``atom_symbol`` should be included in the pi subgraph,
        and False otherwise.
    """

    atom, h_count, charge = parse_atom_symbol(atom_symbol)

    used_electrons = 0
    for b in bonds:
        used_electrons += get_num_from_bond(b)

    if (atom == 'c') and (h_count == 0) and (len(bonds) == 2):  # e.g. c1ccccc1<commit_msg>kekulization bug fix<commit_after>from __future__ import annotations

from typing import Dict, Iterable, List, Set, Tuple, Union

from selfies.grammar_rules import find_element, get_num_from_bond, \
    parse_atom_symbol

ATOM_TYPE = 1
BRANCH_TYPE = 2
RING_TYPE = 3


def kekulize_parser(smiles_gen: Iterable[Tuple[str, str, int]]) \
        -> Iterable[Tuple[str, str, int]]:
    """Kekulizes a SMILES in the form of an iterable.

    This method intercepts the output of ``encoder._parse_smiles``, and
    acts as filter that kekulizes the SMILES. The motivation for having
    this setup is that string parsing and concatenation is minimized,
    as the parsing is already done by ``_parse_smiles``.

    Reference: https://depth-first.com/articles/2020/02/10/a-comprehensive
               -treatment-of-aromaticity-in-the-smiles-language/

    :param smiles_gen: an iterator returned by ``encoder._parse_smiles``.
    :return: an iterator representing the kekulized SMILES, in the same
        format as that returned by ``encoder._parse_smiles``.
    """

    # save to list, so the iterator can be used across multiple functions
    # change elements from tuple -> list to allow in-place modifications
    smiles_symbols = list(map(list, smiles_gen))

    mol_graph = MolecularGraph(smiles_symbols)

    rings = {}
    _build_molecular_graph(mol_graph, smiles_symbols, rings)

    if mol_graph.aro_indices:
        _kekulize(mol_graph)

    for x in mol_graph.smiles_symbols:  # return as iterator
        yield tuple(x)


def _build_molecular_graph(graph: MolecularGraph,
                           smiles_symbols: List[List[Union[str, int]]],
                           rings: Dict[int, Tuple[int, int]],
                           prev_idx: int = -1,
                           curr_idx: int = -1) -> int:
    """From the iterator returned by ``encoder._parse_smiles``, builds
    a graph representation of the molecule.

    This is done by iterating through ``smiles_symbols``, and then adding bonds
    to the molecular graph. Note that ``smiles_symbols`` is mutated in this
    method, for convenience.

    :param graph: the MolecularGraph to be added to.
    :param smiles_symbols: a list created from the iterator returned
        by ``encoder._parse_smiles``.
    :param rings: an, initially, empty dictionary used to keep track of
        rings to be made.
    :param prev_idx:
    :param curr_idx:
    :return: the last index of ``smiles_symbols`` that was processed.
    """

    while curr_idx + 1 < len(smiles_symbols):

        curr_idx += 1
        _, symbol, symbol_type = smiles_symbols[curr_idx]

        if symbol_type == ATOM_TYPE:
            if prev_idx >= 0:
                graph.add_bond(prev_idx, curr_idx, curr_idx)
            prev_idx = curr_idx

        elif symbol_type == BRANCH_TYPE:
            if symbol == '(':
                curr_idx = _build_molecular_graph(graph, smiles_symbols, rings,
                                                  prev_idx, curr_idx)
            else:
                break

        else:
            if symbol in rings:
                left_idx, left_bond_idx = rings.pop(symbol)
                right_idx, right_bond_idx = prev_idx, curr_idx

                # we mutate one bond index to be '', so that we
                # can faithfully represent the bond to be localized at
                # one index. For example, C=1CCCC=1 --> C1CCCC=1.

                if smiles_symbols[left_bond_idx][0] != '':
                    bond_idx = left_bond_idx
                    smiles_symbols[right_bond_idx][0] = ''
                else:
                    bond_idx = right_bond_idx
                    smiles_symbols[left_bond_idx][0] = ''

                graph.add_bond(left_idx, right_idx, bond_idx)
            else:
                rings[symbol] = (prev_idx, curr_idx)

    return curr_idx


def _kekulize(mol_graph: MolecularGraph) -> None:
    """Kekulizes the molecular graph.

    :param mol_graph: a molecular graph to be kekulized.
    :return: None.
    """

    mol_graph.prune_to_pi_subgraph()

    visited = set()
    for i in mol_graph.get_nodes_by_num_edges():
        success = mol_graph.dfs_assign_bonds(i, visited, set(), set())
        if not success:
            raise ValueError("Kekulization Failed.")

    mol_graph.write_to_smiles_symbols()


# Aromatic Helper Methods and Classes

# key = aromatic SMILES element, value = number of valence electrons
# Note: wild card '*' not supported currently
_aromatic_valences = {
    'b': 3, 'c': 4, 'n': 5, 'p': 5, 'as': 5,
    'o': 6, 's': 6, 'se': 6
}


def _capitalize(atom_symbol: str) -> str:
    """Capitalizes the element portion of an aromatic SMILES atom symbol,
    converting it into a standard SMILES atom symbol.

    :param atom_symbol: an aromatic SMILES atom symbol.
    :return: the capitalized ``atom_symbol``.
    """

    s, _ = find_element(atom_symbol)
    return atom_symbol[:s] + atom_symbol[s].upper() + atom_symbol[s + 1:]


def _is_aromatic(atom_symbol: str) -> bool:
    """Checks whether a SMILES atom symbol is an aromatic SMILES atom symbol.

    An aromatic SMILES atom symbol is indicated by an element substring
    that is not capitalized.

    :param atom_symbol: a SMILES atom symbol.
    :return: True, if ``atom_symbol`` is an aromatic atom symbol,
        and False otherwise.
    """

    s, e = find_element(atom_symbol)

    if e == len(atom_symbol):  # optimization to prevent string copying
        element = atom_symbol
    else:
        element = atom_symbol[s: e]

    if element[0].isupper():  # check if element is capitalized
        return False

    if element not in _aromatic_valences:
        raise ValueError(f"Kekulization Failed: aromatic symbol {atom_symbol} "
                         f"not recognized.")
    return True


def _in_pi_subgraph(atom_symbol: str, bonds: Tuple[str]) -> bool:
    """Checks whether a SMILES atom symbol should be a node in the pi
    subgraph, based on its bonds.

    More specifically, an atom should be a node in the pi subgraph if it has
    an unpaired valence electron, and thus, is able to make a double bond.

    Reference: https://depth-first.com/articles/2020/02/10/a-comprehensive
               -treatment-of-aromaticity-in-the-smiles-language/

    :param atom_symbol: a SMILES atom symbol representing an atom.
    :param bonds: the bonds connected to ``atom_symbol``.
    :return: True if ``atom_symbol`` should be included in the pi subgraph,
        and False otherwise.
    """

    atom, h_count, charge = parse_atom_symbol(atom_symbol)

    used_electrons = 0
    for b in bonds:
        used_electrons += get_num_from_bond(b)

    # e.g. c1ccccc1
    if (atom == 'c') and (h_count == charge == 0) and (len(bonds) == 2):
{'fraction_non_alphanumeric': 0.08820482462631345, 'fraction_numerical': 0.00680775492082285, 'mean_word_length': 2.517699115044248, 'pattern_counts': {'":': 0, '<': 5, '<?xml version=': 0, '>': 21, 'https://': 4, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 4, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '5068696', 'n_tokens_mistral': 4297, 'n_tokens_neox': 4006, 'n_words': 1300}