# -*- coding: utf-8 -*-
#
#  Copyright 2018-2020 Ramil Nugmanov <nougmanoff@protonmail.com>
#  Copyright 2019 Artem Mukanov <nostro32@mail.ru>
#  This file is part of CGRtools.
#
#  CGRtools is free software; you can redistribute it and/or modify
#  it under the terms of the GNU Lesser General Public License as published by
#  the Free Software Foundation; either version 3 of the License, or
#  (at your option) any later version.
#
#  This program is distributed in the hope that it will be useful,
#  but WITHOUT ANY WARRANTY; without even the implied warranty of
#  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
#  GNU Lesser General Public License for more details.
#
#  You should have received a copy of the GNU Lesser General Public License
#  along with this program; if not, see <https://www.gnu.org/licenses/>.
#
from collections import defaultdict
from itertools import permutations
from io import StringIO, TextIOWrapper
from logging import warning
from pathlib import Path
from re import split, compile, fullmatch
from traceback import format_exc
from typing import Union, List
from warnings import warn
from ._CGRrw import CGRRead, parse_error
from ..containers import MoleculeContainer, CGRContainer, ReactionContainer
from ..exceptions import IncorrectSmiles, IsChiral, NotChiral, ValenceError


# tokens structure:
# (type: int, value)
# types:
# 0: atom
# 1: bond
# 2: open chain (
# 3: close chain )
# 4: dot bond .
# 5: in bracket raw data []
# 6: closure number
# 7: raw closure number
# 8: aromatic atom
# 9: up down bond
# 10: dynamic bond
# 11: dynamic atom

replace_dict = {'-': 1, '=': 2, '#': 3, ':': 4, '~': 8, '.': None, '(': 2, ')': 3}
charge_dict = {'+': 1, '+1': 1, '++': 2, '+2': 2, '+3': 3, '+++': 3, '+4': 4, '++++': 4,
               '-': -1, '-1': -1, '--': -2, '-2': -2, '-3': -3, '---': -3, '-4': -4, '----': -4}
dynamic_bonds = {'.>-': (None, 1), '.>=': (None, 2), '.>#': (None, 3), '.>:': (None, 4), '.>~': (None, 8),
                 '->.': (1, None), '->=': (1, 2), '->#': (1, 3), '->:': (1, 4), '->~': (1, 8),
                 '=>.': (2, None), '=>-': (2, 1), '=>#': (2, 3), '=>:': (2, 4), '=>~': (2, 8),
                 '#>.': (3, None), '#>-': (3, 1), '#>=': (3, 2), '#>:': (3, 4), '#>~': (3, 8),
                 ':>.': (4, None), ':>-': (4, 1), ':>=': (4, 2), ':>#': (4, 3), ':>~': (4, 8),
                 '~>.': (8, None), '~>-': (8, 1), '~>=': (8, 2), '~>#': (8, 3), '~>:': (8, 4)}
dyn_charge_dict = {'-4': -4, '-3': -3, '-2': -2, '--': -2, '-': -1, '0': 0, '+': 1, '++': 2, '+2': 2, '+3': 3, '+4': 4}
tmp = {f'{i}>{j}': (x, ('charge', y - x)) for (i, x), (j, y) in permutations(dyn_charge_dict.items(), 2) if x != y}
dyn_charge_dict = {k: (v,) for k, v in dyn_charge_dict.items()}
dyn_charge_dict.update(tmp)
dyn_radical_dict = {'*': (True,), '*>^': (True, ('radical', None)), '^>*': (False, ('radical', None))}

atom_re = compile(r'([1-9][0-9]{0,2})?([A-IK-PR-Zacnops][a-ik-pr-vy]?)(@@|@)?(H[1-4]?)?([+-][1-4+-]?)?(:[0-9]{1,4})?')
dyn_atom_re = compile(r'([1-9][0-9]{0,2})?([A-IK-PR-Z][a-ik-pr-vy]?)([+-0][1-4+-]?(>[+-0][1-4+-]?)?)?([*^](>[*^])?)?')
delimiter = compile(r'[=:]')


class SMILESRead(CGRRead):
    """SMILES separated per lines files reader. Works similar to opened file object. Support `with` context manager.
    On initialization accept opened in text mode file, string path to file,
    pathlib.Path object or another buffered reader object.

    Line should be start with SMILES string and optionally continues with space/tab separated list of
    `key:value` [or `key=value`] data if `header=None`. For example::

        C=C>>CC id:123 key=value

    if `header=True` then first line of file should be space/tab separated list of keys including smiles column key.
    For example::

        ignored_smi_key key1 key2
        CCN 1 2

    Also possible to pass list of keys (without smiles_pseudo_key) for mapping space/tab separated list
    of SMILES and values: `header=['key1', 'key2'] # order depended`.

    For reactions . [dot] in bonds should be used only for molecules separation.
    """
    def __init__(self, file, header=None, ignore_stereo=False, **kwargs):
        """
        :param ignore: Skip some checks of data or try to fix some errors.
        :param remap: Remap atom numbers started from one.
        :param store_log: Store parser log if exists messages to `.meta` by key `CGRtoolsParserLog`.
        :param ignore_stereo: Ignore stereo data.
        """
        if isinstance(file, str):
            self._file = open(file)
            self.__is_buffer = False
        elif isinstance(file, Path):
            self._file = file.open()
            self.__is_buffer = False
        elif isinstance(file, (TextIOWrapper, StringIO)):
            self._file = file
            self.__is_buffer = True
        else:
            raise TypeError('invalid file. TextIOWrapper, StringIO subclasses possible')
        super().__init__(**kwargs)
        self.__file = iter(self._file.readline, '')

        if header is True:
            self.__header = next(self.__file).split()[1:]
        elif header:
            if not isinstance(header, (list, tuple)) or not all(isinstance(x, str) for x in header):
                raise TypeError('expected list (tuple) of strings')
            self.__header = header
        else:
            self.__header = None

        self.__ignore_stereo = ignore_stereo
        self._data = self.__data()

    def __data(self):
        file = self._file
        parse = self.parse
        seekable = file.seekable()
        pos = file.tell() if seekable else None
        for n, line in enumerate(self.__file):
            x = parse(line)
            if x is None:
                yield parse_error(n, pos, self._format_log())
                if seekable:
                    pos = file.tell()
            else:
                yield x

    @classmethod
    def create_parser(cls, *args, **kwargs):
        """
        Create SMILES parser function configured same as SMILESRead object.
        """
        obj = object.__new__(cls)
        obj._SMILESRead__header = None
        obj._SMILESRead__ignore_stereo = False
        super(SMILESRead, obj).__init__(*args, **kwargs)
        return obj.parse

    def close(self, force=False):
        """
        Close opened file.

        :param force: Force closing of externally opened file or buffer.
        """
        if not self.__is_buffer or force:
            self._file.close()

    def __enter__(self):
        return self

    def __exit__(self, _type, value, traceback):
        self.close()

    def read(self) -> List[Union[MoleculeContainer, CGRContainer, ReactionContainer]]:
        """
        Parse whole file.

        :return: List of parsed molecules or reactions.
        """
        return list(iter(self))

    def __iter__(self):
        return (x for x in self._data if not isinstance(x, parse_error))

    def __next__(self):
        return next(iter(self))

    def parse(self, smiles: str) -> Union[MoleculeContainer, CGRContainer, ReactionContainer, None]:
        """SMILES string parser."""
        self._flush_log()
        smi, *data = smiles.split()
        if not smi:
            self._info('empty smiles')
            return
        elif self.__header is None:
            meta = {}
            for x in data:
                try:
                    k, v = split(delimiter, x, 1)
                    meta[k] = v
                except ValueError:
                    self._info(f'invalid metadata entry: {x}')
        else:
            meta = dict(zip(self.__header, data))

        if '>' in smi and (smi[smi.index('>') + 1] in '>([' or smi[smi.index('>') + 1].isalpha()):
            record = {'reactants': [], 'reagents': [], 'products': [], 'meta': meta, 'title': ''}
            try:
                reactants, reagents, products = smi.split('>')
            except ValueError:
                self._info('invalid SMIRKS')
                return

            try:
                if reactants:
                    for x in reactants.split('.'):
                        if not x:
                            if self._ignore:
                                self._info('two dots in line ignored')
                            else:
                                self._info('two dots in line')
                                return
                        else:
                            record['reactants'].append(self.__parse_tokens(x))
                if products:
                    for x in products.split('.'):
                        if not x:
                            if self._ignore:
                                self._info('two dots in line ignored')
                            else:
                                self._info('two dots in line')
                                return
                        else:
                            record['products'].append(self.__parse_tokens(x))
                if reagents:
                    for x in reagents.split('.'):
                        if not x:
                            if self._ignore:
                                self._info('two dots in line ignored')
                            else:
                                self._info('two dots in line')
                                return
                        else:
                            record['reagents'].append(self.__parse_tokens(x))
            except ValueError:
                self._info(f'record consist errors:\n{format_exc()}')
                return

            try:
                container = self._convert_reaction(record)
            except ValueError:
                self._info(f'record consist errors:\n{format_exc()}')
            else:
                if self._store_log:
                    log = self._format_log()
                    if log:
                        container.meta['CGRtoolsParserLog'] = log
                return container
        else:
            try:
                record = self.__parse_tokens(smi)
            except ValueError:
                self._info(f'line: {smi}\nconsist errors:\n{format_exc()}')
                return

            record['meta'] = meta
            try:
                container = self._convert_structure(record)
            except ValueError:
                self._info(f'record consist errors:\n{format_exc()}')
            else:
                if self._store_log:
                    log = self._format_log()
                    if log:
                        container.meta['CGRtoolsParserLog'] = log
                return container

    def _convert_molecule(self, molecule, mapping):
        mol = super()._convert_molecule(molecule, mapping)
        if self.__ignore_stereo or not molecule['stereo_atoms'] and not molecule['stereo_bonds']:
            return mol

        st = mol._stereo_tetrahedrons
        sa = mol._stereo_allenes
        sat = mol._stereo_allenes_terminals
        ctt = mol._stereo_cis_trans_terminals

        order = {mapping[n]: [mapping[m] for m in ms] for n, ms in molecule['order'].items()}

        stereo = []
        for n, s in molecule['stereo_atoms'].items():
            n = mapping[n]
            if n in st:
                stereo.append((mol.add_atom_stereo, n, order[n], s))
            elif n in sa:
                t1, t2 = sat[n]
                env = sa[n]
                n1 = next(x for x in order[t1] if x in env)
                n2 = next(x for x in order[t2] if x in env)
                stereo.append((mol.add_atom_stereo, n, (n1, n2), s))

        stereo_bonds = {mapping[n]: {mapping[m]: s for m, s in ms.items()}
                        for n, ms in molecule['stereo_bonds'].items()}
        seen = set()
        for n, ns in stereo_bonds.items():
            if n in seen:
                continue
            if n in ctt:
                nm = ctt[n]
                m = nm[1] if nm[0] == n else nm[0]
                if m in stereo_bonds:
                    seen.add(m)
                    n2, s2 = stereo_bonds[m].popitem()
                    n1, s1 = ns.popitem()
                    stereo.append((mol.add_cis_trans_stereo, n, m, n1, n2, s1 == s2))

        while stereo:
            fail_stereo = []
            old_stereo = len(stereo)
            for f, *args in stereo:
                try:
                    f(*args, clean_cache=False)
                except NotChiral:
                    fail_stereo.append((f, *args))
                except IsChiral:
                    pass
                except ValenceError:
                    self._info('structure has errors, stereo data skipped')
                    mol.flush_cache()
                    break
            else:
                stereo = fail_stereo
                if len(stereo) == old_stereo:
                    break
                del mol.__dict__['_MoleculeStereo__chiral_centers']
                continue
            break
        return mol

    @staticmethod
    def _raw_tokenize(smiles):
        token_type = token = None
        tokens = []
        for s in smiles:
            if s == '[':  # open complex token
                if token_type == 5:  # two opened [
                    raise IncorrectSmiles('[...[')
                elif token:
                    tokens.append((token_type, token))
                elif token_type == 7:  # empty closure
                    raise IncorrectSmiles('invalid closure')
                token = []
                token_type = 5
            elif s == ']':  # close complex token
                if token_type != 5:
                    raise IncorrectSmiles(']..]')
                elif not token:
                    raise IncorrectSmiles('empty [] brackets')
                tokens.append((5, ''.join(token)))
                token_type = token = None
            elif s in '()':
                if token_type == 5:
                    raise IncorrectSmiles('brackets in atom/bond token')
                elif token:
                    tokens.append((token_type, token))
                    token = None
                elif token_type == 7:  # empty closure
                    raise IncorrectSmiles('invalid closure')
                elif token_type == 2:  # barely opened
                    raise IncorrectSmiles('(( or ()')
                token_type = replace_dict[s]
                tokens.append((token_type, None))
            elif token_type == 5:  # grow token with brackets. skip validation
                token.append(s)
            elif s.isnumeric():  # closures
                if token_type == 7:  # % already found. collect number
                    if not token and s == '0':
                        raise IncorrectSmiles('number starts with 0')
                    token.append(s)
                else:
                    if s == '0':
                        raise IncorrectSmiles('number starts with 0')
                    elif token:
                        tokens.append((token_type, token))
                        token = None
                    token_type = 6
                    tokens.append((6, int(s)))
            elif s == '%':
                if token:
                    tokens.append((token_type, token))
                elif token_type == 7:
                    raise IncorrectSmiles('%%')
                token_type = 7
                token = []
            elif s in '=#:-~':  # bonds found
                if token:
                    tokens.append((token_type, token))
                    token = None
                token_type = 1
                tokens.append((1, replace_dict[s]))
            elif s in r'\/':
                if token:
                    tokens.append((token_type, token))
                    token = None
                token_type = 9
                tokens.append((9, s == '/'))  # Up is true
            elif s == '.':
                if token:
                    tokens.append((token_type, token))
                    token = None
                token_type = 4
                tokens.append((4, None))
            elif s in 'NOPSFI':  # organic atoms
                if token:
                    tokens.append((token_type, token))
                    token = None
                token_type = 0
                tokens.append((0, s))
            elif s in 'cnops':  # aromatic ring atom
                if token:
                    tokens.append((token_type, token))
                    token = None
                token_type = 8
                tokens.append((8, s.upper()))
            elif s in 'CB':  # flag possible Cl or Br
                if token:
                    tokens.append((token_type, token))
                token_type = 0
                token = s
            elif token_type == 0:
                if s == 'l':
                    if token == 'C':
                        tokens.append((0, 'Cl'))
                        token = None
                    else:
                        raise IncorrectSmiles('invalid element Bl')
                elif s == 'r':
                    if token == 'B':
                        tokens.append((0, 'Br'))
                        token = None
                    else:
                        raise IncorrectSmiles('invalid smiles for Cr')
                else:
                    raise IncorrectSmiles('invalid smiles')
            else:
                raise IncorrectSmiles('invalid smiles')

        if token_type == 5:
            raise IncorrectSmiles('atom description has not finished')
        elif token_type == 2:
            raise IncorrectSmiles('not closed')
        elif token:
            tokens.append((token_type, token))  # %closure or C or B
        return [(6, int(''.join(x[1]))) if x[0] == 7 else x for x in tokens]  # composite closures folding

    @classmethod
    def _fix_tokens(cls, tokens):
        out = []
        for token_type, token in tokens:
            if token_type in (0, 8):  # simple atom
                out.append((token_type, {'element': token, 'charge': 0, 'isotope': None, 'is_radical': False,
                                         'mapping': 0, 'x': 0., 'y': 0., 'z': 0., 'hydrogen': 0, 'stereo': None}))
            elif token_type == 5:
                if '>' in token:  # dynamic bond or atom
                    if len(token) == 3:  # bond only possible
                        try:
                            out.append((10, dynamic_bonds[token]))
                        except KeyError:
                            raise IncorrectSmiles('invalid dynamic bond token')
                    else:  # dynamic atom token
                        out.append((11, cls.__dynatom_parse(token)))
                elif '*' in token:  # CGR atom radical mark
                    out.append((11, cls.__dynatom_parse(token)))
                else:  # atom token
                    out.append(cls.__atom_parse(token))
            else:  # as is types: 1, 2, 3, 4, 6, 9
                out.append((token_type, token))
        return out

    @staticmethod
    def __atom_parse(token):
        # [isotope]Element[element][@[@]][H[n]][+-charge][:mapping]
        match = fullmatch(atom_re, token)
        if match is None:
            raise IncorrectSmiles('atom token invalid')
        isotope, element, stereo, hydrogen, charge, mapping = match.groups()

        if isotope:
            isotope = int(isotope)

        if stereo:
            stereo = stereo == '@'

        if hydrogen:
            if len(hydrogen) > 1:
                hydrogen = int(hydrogen[1:])
            else:
                hydrogen = 1
        else:
            hydrogen = 0

        if charge:
            try:
                charge = charge_dict[charge]
            except KeyError:
                raise IncorrectSmiles('charge token invalid')
        else:
            charge = 0

        if mapping:
            try:
                mapping = int(mapping[1:])
            except ValueError:
                raise IncorrectSmiles('invalid mapping token')
        else:
            mapping = 0

        if element in ('c', 'n', 'o', 'p', 's', 'as', 'se'):
            _type = 8
            element = element.capitalize()
        else:
            _type = 0
        return _type, {'element': element, 'charge': charge, 'isotope': isotope, 'is_radical': False,
                       'mapping': mapping, 'x': 0., 'y': 0., 'z': 0., 'hydrogen': hydrogen, 'stereo': stereo}

    @staticmethod
    def __dynatom_parse(token):
        # [isotope]Element[element][+-charge[>+-charge]][*^[>*^]]
        match = fullmatch(dyn_atom_re, token)
        if match is None:
            raise IncorrectSmiles('atom token invalid')
        isotope, element, charge, _, is_radical, _ = match.groups()

        if isotope:
            isotope = int(isotope)

        if charge:
            try:
                charge, *cgr = dyn_charge_dict[charge]
            except KeyError:
                raise IncorrectSmiles('charge token invalid')
        else:
            charge = 0
            cgr = []

        if is_radical:
            try:
                is_radical, *dyn = dyn_radical_dict[is_radical]
            except KeyError:
                raise IncorrectSmiles('invalid dynamic radical token')
            else:
                cgr.extend(dyn)
        else:
            is_radical = False

        return {'element': element, 'charge': charge, 'isotope': isotope, 'is_radical': is_radical,
                'mapping': 0, 'x': 0., 'y': 0., 'z': 0., 'cgr': cgr}

    def __parse_tokens(self, smiles):
        tokens = self._raw_tokenize(smiles)
        tokens = self._fix_tokens(tokens)
        return self._parse_tokens(tokens)

    def _parse_tokens(self, tokens):
        strong_cycle = not self._ignore
        t1 = tokens[0][0]
        if t1 == 2:
            if tokens[1][0] not in (0, 8, 11):
                raise IncorrectSmiles('not atom started')
        elif t1 not in (0, 8, 11):
            raise IncorrectSmiles('not atom started')

        atoms = []
        bonds = []
        order = defaultdict(list)
        atoms_types = []
        atom_num = 0
        last_num = 0
        stack = []
        cycles = {}
        used_cycles = set()
        cgr = []
        stereo_bonds = defaultdict(dict)
        stereo_atoms = {}
        hydrogens = {}
        previous = None

        for token_type, token in tokens:
            if token_type == 2:  # ((((((
                if previous:
                    if previous[0] != 4:
                        raise IncorrectSmiles('bond before side chain')
                    previous = None
                stack.append(last_num)
            elif token_type == 3:  # ))))))
                if previous:
                    raise IncorrectSmiles('bond before closure')
                try:
                    last_num = stack.pop()
                except IndexError:
                    raise IncorrectSmiles('close chain more than open')
            elif token_type in (1, 4, 9, 10):  # bonds. only keeping for atoms connecting
                if previous:
                    raise IncorrectSmiles('2 bonds in a row')
                elif not atoms:
                    raise IncorrectSmiles('started from bond')
                previous = (token_type, token)
            elif token_type == 6:  # cycle
                if previous and previous[0] == 4:
                    raise IncorrectSmiles('dot-cycle pattern invalid')
                elif token not in cycles:
                    if token in used_cycles:
                        if strong_cycle:
                            raise IncorrectSmiles('reused closure number')
                        else:
                            self._info(f'reused closure number: {token}')
                    else:
                        used_cycles.add(token)
                    cycles[token] = (last_num, previous, len(order[last_num]))
                    order[last_num].append(None)  # Reserve a table
                else:
                    a, ob, ind = cycles[token]
                    if ob:
                        if not previous:
                            bt, b = ob
                            if bt == 9:  # closure open is \/ bonded
                                stereo_bonds[a][last_num] = b
                                bt = b = 1
                            elif strong_cycle:
                                raise IncorrectSmiles('not equal cycle bonds')
                        else:
                            bt, b = previous
                            obt, ob = ob
                            if bt == 9:  # \/ bonds can be unequal
                                if obt == 9:
                                    stereo_bonds[a][last_num] = ob
                                elif ob != 1:
                                    raise IncorrectSmiles('not equal cycle bonds')
                                stereo_bonds[last_num][a] = b
                                bt = b = 1
                            elif obt == 9:
                                if b != 1:
                                    raise IncorrectSmiles('not equal cycle bonds')
                                stereo_bonds[a][last_num] = ob
                            elif b != ob:
                                raise IncorrectSmiles('not equal cycle bonds')
                    elif previous:
                        bt, b = previous
                        if bt == 9:  # stereo \/
                            stereo_bonds[last_num][a] = b
                            bt = b = 1
                        elif strong_cycle:
                            raise IncorrectSmiles('not equal cycle bonds')
                    else:
                        bt = 1
                        b = 4 if atoms_types[last_num] == atoms_types[a] == 8 else 1

                    if bt == 1:
                        bonds.append((last_num, a, b))
                    else:  # bt == 10
                        bonds.append((last_num, a, 8))
                        cgr.append(((last_num, a), 'bond', b))
                    order[a][ind] = last_num
                    order[last_num].append(a)
                    del cycles[token]
                previous = None
            else:  # atom
                if atoms:
                    if not previous:
                        bt = 1
                        b = 4 if atoms_types[last_num] == token_type == 8 else 1
                    else:
                        bt, b = previous

                    if bt == 1:
                        bonds.append((atom_num, last_num, b))
                    elif bt == 9:
                        bonds.append((atom_num, last_num, 1))
                        stereo_bonds[last_num][atom_num] = b
                        stereo_bonds[atom_num][last_num] = not b
                    elif bt == 10:
                        bonds.append((atom_num, last_num, 8))
                        cgr.append(((atom_num, last_num), 'bond', b))
                    order[last_num].append(atom_num)
                    order[atom_num].append(last_num)

                if token_type == 11:
                    cgr.extend((atom_num, *x) for x in token.pop('cgr'))
                else:
                    stereo = token.pop('stereo')
                    if stereo is not None:
                        stereo_atoms[atom_num] = stereo
                    hydrogen = token.pop('hydrogen')
                    if hydrogen:
                        hydrogens[atom_num] = hydrogen

                atoms.append(token)
                atoms_types.append(token_type)

                last_num = atom_num
                atom_num += 1
                previous = None

        if stack:
            raise IncorrectSmiles('number of ( does not equal to number of )')
        elif cycles:
            raise IncorrectSmiles('cycle is not finished')
        elif previous:
            raise IncorrectSmiles('bond on the end')

        stereo_bonds = {n: ms for n, ms in stereo_bonds.items() if len(ms) == 1 or len(ms) == set(ms.values())}
        mol = {'atoms': atoms, 'bonds': bonds, 'order': order,
               'stereo_bonds': stereo_bonds, 'stereo_atoms': stereo_atoms, 'hydrogens': hydrogens}
        if cgr or any(x == 11 for x in atoms_types):
            mol['cgr'] = cgr
        return mol


class SMILESread:
    def __init__(self, *args, **kwargs):
        warn('SMILESread deprecated. Use SMILESRead instead', DeprecationWarning)
        warning('SMILESread deprecated. Use SMILESRead instead')
        self.__obj = SMILESRead(*args, **kwargs)

    def __getattr__(self, item):
        return getattr(self.__obj, item)

    def __iter__(self):
        return iter(self.__obj)

    def __next__(self):
        return next(self.__obj)

    def __enter__(self):
        return self.__obj.__enter__()

    def __exit__(self, _type, value, traceback):
        return self.__obj.__exit__(_type, value, traceback)


__all__ = ['SMILESRead', 'SMILESread']
