from __future__ import division

from .transformation import Transformation
from pyparsing import (Literal, CaselessLiteral, Word, Combine, Group, Optional,
                       ZeroOrMore, Forward, nums, alphas, oneOf)
import math
import re
import operator

__author__ = 'Paul McGuire'
__version__ = '$Revision: 0.0 $'
__date__ = '$Date: 2009-03-20 $'
__source__ = '''http://pyparsing.wikispaces.com/file/view/fourFn.py
http://pyparsing.wikispaces.com/message/view/home/15549426
'''
__note__ = '''
All I've done is rewrap Paul McGuire's fourFn.py as a class, so I can use it
more easily in other places.
'''


class Custom(Transformation):
    """
    Most of this code comes from the fourFn.py pyparsing example
    """

    title = "Custom equation"
    key = "Math equation"
    fields = {
        "equation": {"name": "Equation", "type": "string", "help": "The equation to evaluate. Column values should be entered as {COLUMN NAME}",
                     "required": True, "input": "text", "default": ""},
        "output": {"name": "Output column", "type": "string", "input": "text", "required": True,
                   "help": "The name of the (newly created) column that contains the results", "default": ""},
    }

    def __init__(self, arguments: dict, sample_size: int, example: dict = None):
        """
        Initialize the transformation with the given parameters.

        expop   :: '^'
        multop  :: '*' | '/'
        addop   :: '+' | '-'
        integer :: ['+' | '-'] '0'..'9'+
        atom    :: PI | E | real | fn '(' expr ')' | '(' expr ')'
        factor  :: atom [ expop factor ]*
        term    :: factor [ multop factor ]*
        expr    :: term [ addop term ]*

        Arguments:
            arguments {dict} -- The arguments
        """
        super().__init__(arguments, sample_size, example)
        self.equation = arguments["equation"]
        self.output = arguments["output"]

        point = Literal(".")
        e = CaselessLiteral("E")
        fnumber = Combine(Word("+-" + nums, nums) +
                          Optional(point + Optional(Word(nums))) +
                          Optional(e + Word("+-" + nums, nums)))
        ident = Word(alphas, alphas + nums + "_$")
        plus = Literal("+")
        minus = Literal("-")
        mult = Literal("*")
        div = Literal("/")
        mod = Literal("%")
        lpar = Literal("(").suppress()
        rpar = Literal(")").suppress()
        addop = plus | minus
        multop = mult | div | mod
        expop = Literal("^")
        pi = CaselessLiteral("PI")
        expr = Forward()
        atom = ((Optional(oneOf("- +")) +
                 (ident + lpar + expr + rpar | pi | e | fnumber).setParseAction(self.push_first))
                | Optional(oneOf("- +")) + Group(lpar + expr + rpar)
                ).setParseAction(self.push_u_minus)
        # by defining exponentiation as "atom [ ^ factor ]..." instead of
        # "atom [ ^ atom ]...", we get right-to-left exponents, instead of left-to-right
        # that is, 2^3^2 = 2^(3^2), not (2^3)^2.
        factor = Forward()
        factor << atom + \
        ZeroOrMore((expop + factor).setParseAction(self.push_first))
        term = factor + \
               ZeroOrMore((multop + factor).setParseAction(self.push_first))
        expr << term + \
        ZeroOrMore((addop + term).setParseAction(self.push_first))
        # addop_term = ( addop + term ).setParseAction( self.push_first )
        # general_term = term + ZeroOrMore( addop_term ) | OneOrMore( addop_term)
        # expr <<  general_term
        self.bnf = expr
        # map operator symbols to corresponding arithmetic operations
        epsilon = 1e-12
        self.opn = {"+": operator.add,
                    "-": operator.sub,
                    "*": operator.mul,
                    "%": operator.mod,
                    "/": operator.truediv,
                    "^": operator.pow}
        self.expr_stack = None
        self.fn = {"sin": math.sin,
                   "sinh": math.sinh,
                   "cos": math.cos,
                   "cosh": math.cosh,
                   "tan": math.tan,
                   "tanh": math.tanh,
                   "exp": math.exp,
                   "sqrt": math.sqrt,
                   "radians": math.radians,
                   "degrees": math.degrees,
                   "sign": lambda x: 0 if x == 0 else x / abs(x),
                   "log": math.log10,
                   "ln": math.log,
                   "abs": abs,
                   "trunc": lambda a: int(a),
                   "round": round,
                   "floor": math.floor,
                   "ceil": math.ceil,
                   "sgn": lambda a: abs(a) > epsilon and cmp(a, 0) or 0}

    def push_first(self, strg, loc, toks):
        self.expr_stack.append(toks[0])

    def push_u_minus(self, strg, loc, toks):
        if toks and toks[0] == '-':
            self.expr_stack.append('unary -')

    def evaluate_stack(self, s):
        op = s.pop()
        if op == 'unary -':
            return -self.evaluate_stack(s)
        if op in "+-*/^%":
            op2 = self.evaluate_stack(s)
            op1 = self.evaluate_stack(s)
            return self.opn[op](op1, op2)
        elif op == "PI":
            return math.pi  # 3.1415926535
        elif op == "E":
            return math.e  # 2.718281828
        elif op in self.fn:
            return self.fn[op](self.evaluate_stack(s))
        elif op[0].isalpha():
            return 0
        else:
            return float(op)

    def eval(self, num_string, parse_all=True):
        self.expr_stack = []
        results = self.bnf.parseString(num_string, parse_all)
        val = self.evaluate_stack(self.expr_stack[:])
        return val

    def __call__(self, row, index: int):
        """This class is called on each row.

        Arguments:
            row {dict} -- The complete row

        Returns:
            dict -- The row, including the extra output column
        """
        row[self.output] = self.eval(re.sub(r'{(\w+)}', lambda x: str(row.get(x.group(1), 0)), self.equation))

        return row, index

[((56, 16, 56, 28), 'pyparsing.Literal', 'Literal', ({(56, 24, 56, 27): '"""."""'}, {}), "('.')", False, 'from pyparsing import Literal, CaselessLiteral, Word, Combine, Group, Optional, ZeroOrMore, Forward, nums, alphas, oneOf\n'), ((57, 12, 57, 32), 'pyparsing.CaselessLiteral', 'CaselessLiteral', ({(57, 28, 57, 31): '"""E"""'}, {}), "('E')", False, 'from pyparsing import Literal, CaselessLiteral, Word, Combine, Group, Optional, ZeroOrMore, Forward, nums, alphas, oneOf\n'), ((61, 16, 61, 50), 'pyparsing.Word', 'Word', ({(61, 21, 61, 27): 'alphas', (61, 29, 61, 49): "alphas + nums + '_$'"}, {}), "(alphas, alphas + nums + '_$')", False, 'from pyparsing import Literal, CaselessLiteral, Word, Combine, Group, Optional, ZeroOrMore, Forward, nums, alphas, oneOf\n'), ((62, 15, 62, 27), 'pyparsing.Literal', 'Literal', ({(62, 23, 62, 26): '"""+"""'}, {}), "('+')", False, 'from pyparsing import Literal, CaselessLiteral, Word, Combine, Group, Optional, ZeroOrMore, Forward, nums, alphas, oneOf\n'), ((63, 16, 63, 28), 'pyparsing.Literal', 'Literal', ({(63, 24, 63, 27): '"""-"""'}, {}), "('-')", False, 'from pyparsing import Literal, CaselessLiteral, Word, Combine, Group, Optional, ZeroOrMore, Forward, nums, alphas, oneOf\n'), ((64, 15, 64, 27), 'pyparsing.Literal', 'Literal', ({(64, 23, 64, 26): '"""*"""'}, {}), "('*')", False, 'from pyparsing import Literal, CaselessLiteral, Word, Combine, Group, Optional, ZeroOrMore, Forward, nums, alphas, oneOf\n'), ((65, 14, 65, 26), 'pyparsing.Literal', 'Literal', ({(65, 22, 65, 25): '"""/"""'}, {}), "('/')", False, 'from pyparsing import Literal, CaselessLiteral, Word, Combine, Group, Optional, ZeroOrMore, Forward, nums, alphas, oneOf\n'), ((66, 14, 66, 26), 'pyparsing.Literal', 'Literal', ({(66, 22, 66, 25): '"""%"""'}, {}), "('%')", False, 'from pyparsing import Literal, CaselessLiteral, Word, Combine, Group, Optional, ZeroOrMore, Forward, nums, alphas, oneOf\n'), ((71, 16, 71, 28), 'pyparsing.Literal', 'Literal', ({(71, 24, 71, 27): '"""^"""'}, {}), "('^')", False, 'from pyparsing import Literal, CaselessLiteral, Word, Combine, Group, Optional, ZeroOrMore, Forward, nums, alphas, oneOf\n'), ((72, 13, 72, 34), 'pyparsing.CaselessLiteral', 'CaselessLiteral', ({(72, 29, 72, 33): '"""PI"""'}, {}), "('PI')", False, 'from pyparsing import Literal, CaselessLiteral, Word, Combine, Group, Optional, ZeroOrMore, Forward, nums, alphas, oneOf\n'), ((73, 15, 73, 24), 'pyparsing.Forward', 'Forward', ({}, {}), '()', False, 'from pyparsing import Literal, CaselessLiteral, Word, Combine, Group, Optional, ZeroOrMore, Forward, nums, alphas, oneOf\n'), ((81, 17, 81, 26), 'pyparsing.Forward', 'Forward', ({}, {}), '()', False, 'from pyparsing import Literal, CaselessLiteral, Word, Combine, Group, Optional, ZeroOrMore, Forward, nums, alphas, oneOf\n'), ((67, 15, 67, 27), 'pyparsing.Literal', 'Literal', ({(67, 23, 67, 26): '"""("""'}, {}), "('(')", False, 'from pyparsing import Literal, CaselessLiteral, Word, Combine, Group, Optional, ZeroOrMore, Forward, nums, alphas, oneOf\n'), ((68, 15, 68, 27), 'pyparsing.Literal', 'Literal', ({(68, 23, 68, 26): '""")"""'}, {}), "(')')", False, 'from pyparsing import Literal, CaselessLiteral, Word, Combine, Group, Optional, ZeroOrMore, Forward, nums, alphas, oneOf\n'), ((58, 26, 58, 49), 'pyparsing.Word', 'Word', ({(58, 31, 58, 42): "'+-' + nums", (58, 44, 58, 48): 'nums'}, {}), "('+-' + nums, nums)", False, 'from pyparsing import Literal, CaselessLiteral, Word, Combine, Group, Optional, ZeroOrMore, Forward, nums, alphas, oneOf\n'), ((60, 39, 60, 62), 'pyparsing.Word', 'Word', ({(60, 44, 60, 55): "'+-' + nums", (60, 57, 60, 61): 'nums'}, {}), "('+-' + nums, nums)", False, 'from pyparsing import Literal, CaselessLiteral, Word, Combine, Group, Optional, ZeroOrMore, Forward, nums, alphas, oneOf\n'), ((76, 43, 76, 68), 'pyparsing.Group', 'Group', ({(76, 49, 76, 67): 'lpar + expr + rpar'}, {}), '(lpar + expr + rpar)', False, 'from pyparsing import Literal, CaselessLiteral, Word, Combine, Group, Optional, ZeroOrMore, Forward, nums, alphas, oneOf\n'), ((74, 26, 74, 38), 'pyparsing.oneOf', 'oneOf', ({(74, 32, 74, 37): '"""- +"""'}, {}), "('- +')", False, 'from pyparsing import Literal, CaselessLiteral, Word, Combine, Group, Optional, ZeroOrMore, Forward, nums, alphas, oneOf\n'), ((76, 27, 76, 39), 'pyparsing.oneOf', 'oneOf', ({(76, 33, 76, 38): '"""- +"""'}, {}), "('- +')", False, 'from pyparsing import Literal, CaselessLiteral, Word, Combine, Group, Optional, ZeroOrMore, Forward, nums, alphas, oneOf\n'), ((59, 52, 59, 62), 'pyparsing.Word', 'Word', ({(59, 57, 59, 61): 'nums'}, {}), '(nums)', False, 'from pyparsing import Literal, CaselessLiteral, Word, Combine, Group, Optional, ZeroOrMore, Forward, nums, alphas, oneOf\n')]