# This file is part of PSAMM.
#
# PSAMM is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# PSAMM 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 General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with PSAMM.  If not, see <http://www.gnu.org/licenses/>.
#
# Copyright 2014-2015  Jon Lund Steffensen <jon_steffensen@uri.edu>

"""Parser and representation of chemical formulas.

Chemical formulas (:class:`.Formula`) are represented as a number of
:class:`FormulaElements <.FormulaElement>` with associated counts. A
:class:`.Formula` is itself a :class:`.FormulaElement` so a formula can contain
subformulas. This allows some simple structure to be represented.
"""

from __future__ import unicode_literals

import re
from collections import Counter
import functools
import operator
import numbers

import six
from six import iteritems
from six.moves import reduce

from .expression.affine import Expression


class FormulaElement(object):
    """Base class representing elements of a formula"""

    def __add__(self, other):
        """Add formula elements creating subformulas"""
        if isinstance(other, FormulaElement):
            if self == other:
                return Formula({self: 2})
            return Formula({self: 1, other: 1})
        return NotImplemented

    def __radd__(self, other):
        return self + other

    def __or__(self, other):
        """Merge formula elements into one formula"""
        return Formula({self: 1}) | other

    def __ror__(self, other):
        return self | other

    def __mul__(self, other):
        """Multiply formula element by other"""
        return Formula({self: other})

    def __rmul__(self, other):
        return self * other

    def repeat(self, count):
        """Repeat formula element by creating a subformula"""
        return Formula({self: count})

    def variables(self):
        """Iterator over variables in formula element"""
        return iter([])

    def substitute(self, mapping):
        """Return formula element with substitutions performed"""
        return self


@six.python_2_unicode_compatible
@functools.total_ordering
class Atom(FormulaElement):
    """Represent an atom in a chemical formula

    >>> hydrogen = Atom('H')
    >>> oxygen = Atom('O')
    >>> str(oxygen | 2*hydrogen)
    'H2O'
    """

    def __init__(self, symbol):
        self._symbol = symbol

    @property
    def symbol(self):
        """Atom symbol

        >>> Atom('H').symbol
        'H'
        """

        return self._symbol

    def __hash__(self):
        return hash('Atom') ^ hash(self._symbol)

    def __eq__(self, other):
        return isinstance(other, Atom) and self._symbol == other._symbol

    def __ne__(self, other):
        return not self == other

    def __lt__(self, other):
        if isinstance(other, Atom):
            return self._symbol < other._symbol
        return NotImplemented

    def __str__(self):
        return self._symbol

    def __repr__(self):
        return str('Atom({})').format(repr(self._symbol))


@six.python_2_unicode_compatible
class Radical(FormulaElement):
    """Represents a radical or other unknown subformula"""

    def __init__(self, symbol):
        self._symbol = symbol

    @property
    def symbol(self):
        """Radical symbol

        >>> Radical('R1').symbol
        'R1'
        """

        return self._symbol

    def __hash__(self):
        return hash('Radical') ^ hash(self._symbol)

    def __eq__(self, other):
        return isinstance(other, Radical) and self._symbol == other._symbol

    def __ne__(self, other):
        return not self == other

    def __str__(self):
        return self._symbol

    def __repr__(self):
        return str('Radical({})').format(repr(self._symbol))


@six.python_2_unicode_compatible
class Formula(FormulaElement):
    """Representation of a chemial formula

    This is represented as a number of
    :class:`FormulaElements <.FormulaElement>` with associated counts.

    >>> f = Formula({Atom('C'): 6, Atom('H'): 12, Atom('O'): 6})
    >>> str(f)
    'C6H12O6'
    """

    def __init__(self, values={}):
        self._values = {}
        self._variables = set()

        for element, value in iteritems(values):
            if not isinstance(element, FormulaElement):
                raise ValueError('Not a formula element: {}'.format(
                    repr(element)))
            if element != Formula() and value != 0:
                self._values[element] = value

            if callable(getattr(value, 'variables', None)):
                for var in value.variables():
                    self._variables.add(var)
            for var in element.variables():
                self._variables.add(var)

    def substitute(self, mapping):
        result = self.__class__()
        for element, value in iteritems(self._values):
            if callable(getattr(value, 'substitute', None)):
                value = value.substitute(mapping)
                if isinstance(value, int) and value <= 0:
                    raise ValueError(
                        'Expression evaluated to non-positive number')
            # TODO does not merge correctly with subformulas
            result += value * element.substitute(mapping)
        return result

    def flattened(self):
        """Return formula where subformulas have been flattened

        >>> str(Formula.parse('(CH2)(CH2)2').flattened())
        'C3H6'
        """

        stack = [(self, 1)]
        result = {}
        while len(stack) > 0:
            var, value = stack.pop()
            if isinstance(var, Formula):
                for sub_var, sub_value in iteritems(var._values):
                    stack.append((sub_var, value*sub_value))
            else:
                if var in result:
                    result[var] += value
                else:
                    result[var] = value
        return Formula(result)

    def variables(self):
        return iter(self._variables)

    def items(self):
        """Iterate over (:class:`.FormulaElement`, value)-pairs"""
        return iteritems(self._values)

    def is_variable(self):
        return len(self._variables) > 0

    def __contains__(self, element):
        return element in self._values

    def __str__(self):
        """Return formula represented using Hill notation system

        >>> str(Formula({Atom('C'): 6, Atom('H'): 12, Atom('O'): 6}))
        'C6H12O6'
        """

        def hill_sorted_elements(values):
            def element_sort_key(pair):
                element, value = pair
                if isinstance(element, Atom):
                    return 0, element.symbol
                elif isinstance(element, Radical):
                    return 1, element.symbol
                else:
                    return 2, None

            if Atom('C') in values:
                yield Atom('C'), values[Atom('C')]
                if Atom('H') in values:
                    yield Atom('H'), values[Atom('H')]
                for element, value in sorted(
                        iteritems(values), key=element_sort_key):
                    if element not in (Atom('C'), Atom('H')):
                        yield element, value
            else:
                for element, value in sorted(
                        iteritems(values), key=element_sort_key):
                    yield element, value

        s = ''
        for element, value in hill_sorted_elements(self._values):
            def grouped(element, value):
                return '({}){}'.format(element, value if value != 1 else '')

            def nongrouped(element, value):
                return '{}{}'.format(element, value if value != 1 else '')

            if isinstance(element, Radical):
                if len(element.symbol) == 1:
                    s += nongrouped(element, value)
                else:
                    s += grouped(element, value)
            elif isinstance(element, Atom):
                s += nongrouped(element, value)
            else:
                s += grouped(element, value)
        return s

    def __repr__(self):
        return str('Formula({})').format(repr(self._values))

    def __or__(self, other):
        """Merge formulas into one formula"""
        if isinstance(other, Formula):
            values = Counter(self._values)
            values.update(other._values)
            return Formula(values)
        elif isinstance(other, FormulaElement):
            return self | Formula({other: 1})
        return NotImplemented

    def __mul__(self, other):
        """Multiply formula element by other"""
        values = {key: value*other for key, value in iteritems(self._values)}
        return Formula(values)

    def __hash__(self):
        h = hash('Formula')
        for element, value in iteritems(self._values):
            h ^= hash(element) ^ hash(value)
        return h

    def __eq__(self, other):
        return isinstance(other, Formula) and self._values == other._values

    def __ne__(self, other):
        return not self == other

    @classmethod
    def parse(cls, s):
        """Parse a formula string (e.g. C6H10O2)"""

        scanner = re.compile(r'''
            (\s+) |         # whitespace
            (\(|\)) |       # group
            ([A-Z][a-z]*) | # element
            (\d+) |         # number
            ([a-z]) |       # variable
            (\Z) |          # end
            (.)             # error
        ''', re.DOTALL | re.VERBOSE)

        def transform_subformula(form):
            '''Extract radical if subformula is a singleton with a radical'''
            if isinstance(form, dict) and len(form) == 1:
                # A radical in a singleton subformula is interpreted as a
                # numbered radical.
                element, value = next(iteritems(form))
                if isinstance(element, Radical):
                    return Radical('{}{}'.format(element.symbol, value))
            return form

        stack = []
        formula = {}
        expect_count = False

        def close(formula, count=1):
            if len(stack) == 0:
                raise ValueError('Unbalanced parenthesis group in formula')
            subformula = transform_subformula(formula)
            if isinstance(subformula, dict):
                subformula = Formula(subformula)

            formula = stack.pop()
            if subformula not in formula:
                formula[subformula] = 0
            formula[subformula] += count
            return formula

        for match in re.finditer(scanner, s):
            (whitespace, group, element, number, variable, end,
                error) = match.groups()

            if error is not None:
                raise ValueError('Invalid token in formula string: {}'.format(
                    match.group(0)))
            elif whitespace is not None:
                continue
            elif group is not None and group == '(':
                if expect_count:
                    formula = close(formula)
                stack.append(formula)
                formula = {}
                expect_count = False
            elif group is not None and group == ')':
                if expect_count:
                    formula = close(formula)
                expect_count = True
            elif element is not None:
                if expect_count:
                    formula = close(formula)
                stack.append(formula)
                if element in 'RX':
                    formula = Radical(element)
                else:
                    formula = Atom(element)
                expect_count = True
            elif number is not None and expect_count:
                formula = close(formula, int(number))
                expect_count = False
            elif variable is not None and expect_count:
                formula = close(formula, Expression(variable))
                expect_count = False
            elif end is not None:
                if expect_count:
                    formula = close(formula)
            else:
                raise ValueError('Invalid token in formula string:'
                                 ' {}'.format(match.group(0)))

        if len(stack) > 0:
            raise ValueError('Unbalanced parenthesis group in formula')

        return Formula(formula)

    @classmethod
    def balance(cls, lhs, rhs):
        """Return formulas that need to be added to balance given formulas

        Given complete formulas for right side and left side of a reaction,
        calculate formulas for the missing compounds on both sides. Return
        as a left, right tuple. Formulas can be flattened before balancing
        to disregard grouping structure.
        """

        def missing(formula, other):
            for element, value in iteritems(formula._values):
                if element not in other._values:
                    yield value*element
                else:
                    delta = value - other._values[element]
                    if isinstance(delta, numbers.Number) and delta > 0:
                        yield delta*element

        return (reduce(operator.or_, missing(rhs, lhs), Formula()),
                reduce(operator.or_, missing(lhs, rhs), Formula()))


if __name__ == '__main__':
    import doctest
    doctest.testmod()
