# *****************************************************************************
#
# This program 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 2 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 General Public License for more
# details.
#
# You should have received a copy of the GNU General Public License along with
# this program; if not, write to the Free Software Foundation, Inc.,
# 59 Temple Place, Suite 330, Boston, MA  02111-1307  USA
#
# Module authors:
#   Markus Zolliker <markus.zolliker@psi.ch>
#   Anik Stark <anik.stark@psi.ch>
#   https://stackoverflow.com/questions/43836866/safely-evaluate-simple-string-equation
#
# *****************************************************************************

import math
import ast
import operator as op

class MathParser:
    _operators2method = {
        ast.Add: op.add,
        ast.Sub: op.sub,
        ast.BitXor: op.xor,
        ast.Or:  op.or_,
        ast.And: op.and_,
        ast.Mod:  op.mod,
        ast.Mult: op.mul,
        ast.Div:  op.truediv,
        ast.Pow:  op.pow,
        ast.FloorDiv: op.floordiv,
        ast.Lt: op.lt,
        ast.Gt: op.gt,
        ast.LtE: op.le,
        ast.GtE: op.ge,
        ast.Eq: op.eq,
        ast.NotEq: op.ne,
        ast.Not: op.not_,
        ast.USub: op.neg,
        ast.UAdd: lambda a:a}

    def __init__(self, math=True, **kwargs):
        self._vars = kwargs
        if not math:
            self._alt_name = self._no_alt_name

    def _Name(self, name):
        try:
            return  self._vars[name]    # look up in user-provided dict
        except KeyError:
            return self._alt_name(name)     # fall back to math functions

    @staticmethod
    def _alt_name(name):
        if name.startswith("_"):    # prevent access to hidden names
            raise NameError(f"{name!r}") 
        try:
            return  getattr(math, name)
        except AttributeError:
            raise NameError(f"{name!r}")

    @staticmethod
    def _no_alt_name(name):
        raise NameError(f"{name!r}") 

    def eval_(self, node):
        if isinstance(node, ast.Expression):
            return self.eval_(node.body)
        if isinstance(node, ast.Constant):    # return the number
            return node.value
        if isinstance(node, ast.Name):    # return variable or math function
            return self._Name(node.id) 
        if isinstance(node, ast.BinOp):    # evaluate binary operations
            method = self._operators2method[type(node.op)]
            return method( self.eval_(node.left), self.eval_(node.right))
        if isinstance(node, ast.Compare):    # evaluate binary operations
            left = self.eval_(node.left)
            for oper, value in zip(node.ops, node.comparators):
                method = self._operators2method[type(oper)]
                right = self.eval_(value)
                if not method(left, right):
                    return False
                left = right
            return True
        if isinstance(node, ast.UnaryOp):    # handle operators
            method = self._operators2method[type(node.op)]
            return method( self.eval_(node.operand) )
        if isinstance(node, ast.Attribute):    # handle attributes (e.g. math.cos)
            return getattr(self.eval_(node.value), node.attr)        
        if isinstance(node, ast.Call):    # evaluate the function and its arguments, calls function
            return self.eval_(node.func)( 
                      *(self.eval_(a) for a in node.args),
                      **{k.arg:self.eval_(k.value) for k in node.keywords})
        raise TypeError(node)

    def calculate(self, expr, **kwargs):
        self._vars.update(kwargs)
        return  self.eval_(ast.parse(expr, mode='eval'))
