calc.py

'''
	Copyright 2015 Teemu Piippo
	All rights reserved.

	Redistribution and use in source and binary forms, with or without
	modification, are permitted provided that the following conditions
	are met:

	1. Redistributions of source code must retain the above copyright
	   notice, this list of conditions and the following disclaimer.
	2. Redistributions in binary form must reproduce the above copyright
	   notice, this list of conditions and the following disclaimer in the
	   documentation and/or other materials provided with the distribution.
	3. The name of the author may not be used to endorse or promote products
	   derived from this software without specific prior written permission.

	THIS SOFTWARE IS PROVIDED BY THE AUTHOR ``AS IS'' AND ANY EXPRESS OR
	IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES
	OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED.
	IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR ANY DIRECT, INDIRECT,
	INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT
	NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
	DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
	THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
	(INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF
	THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
'''

import re
import cmath
import math
import random
import time
import operator
from copy import deepcopy

epsilon = 1e-10

# http://stackoverflow.com/a/2182437
class Enum (set):
	def __init__ (self, *args):
		super (Enum, self).__init__ (args)

	def __getattr__ (self, name):
		if name in self:
			return name
		raise AttributeError

	def __setattr__ (self, name, value):
		raise AttributeError

class Operator (object):
	def __init__ (self, name, symbol, operands, priority, function):
		self.name = name
		self.symbol = symbol
		self.operands = operands
		self.priority = priority
		self.function = function

	def __str__ (self):
		return '''<operator %s>''' % self.name

	def __repr__ (self):
		return self.__str__()

class FunctionCall (object):
	def __init__ (self, funcname, args):
		assert (type(args) is list)
		self.funcname = funcname
		self.args = args

	def __str__ (self):
		return '''<function %s (%s)>''' % (self.funcname, self.args)

	def __repr__ (self):
		return self.__str__()

def do_realf (func, *args):
	for x in args:
		if x.imag:
			raise ValueError ('%s called with a complex number' % func.__name__)

	return func (*[x.real for x in args])

def do_intf (func, *args):
	for x in args:
		if x.imag:
			raise ValueError ('%s called with a complex number' % func.__name__)

		if x.real - math.floor (x.real):
			raise ValueError ('%s called with a floating point number' % func.__name__)

	return func (*[int (x.real) for x in args])

def realf (func):
	return lambda *args: do_realf (func, *args)

def intf (func):
	return lambda *args: do_intf (func, *args)

Operators = []
OperatorData = {
	'lneg':		{ 'symbol': '!',  'operands': 1, 'priority': 5, 'function': lambda x: not x },
	'compl':	{ 'symbol': '~',  'operands': 1, 'priority': 5, 'function': intf (operator.inv) },
	'neg':		{ 'symbol': '-',  'operands': 1, 'priority': 5, 'function': lambda x: -x },
	'pow':		{ 'symbol': '**', 'operands': 2, 'priority': 10, 'function': lambda x, y: x ** y },
	'mul':		{ 'symbol': '*',  'operands': 2, 'priority': 50, 'function': lambda x, y: x * y },
	'div':		{ 'symbol': '/',  'operands': 2, 'priority': 50, 'function': lambda x, y: x / y },
	'mod':		{ 'symbol': '%',  'operands': 2, 'priority': 50, 'function': lambda x, y: math.fmod (x, y) },
	'add':		{ 'symbol': '+',  'operands': 2, 'priority': 100, 'function': lambda x, y: x + y },
	'sub':		{ 'symbol': '-',  'operands': 2, 'priority': 100, 'function': lambda x, y: x - y },
	'eq':		{ 'symbol': '==', 'operands': 2, 'priority': 500, 'function': lambda x, y: x == y },
	'neq':		{ 'symbol': '!=', 'operands': 2, 'priority': 500, 'function': lambda x, y: x != y },
	'lt':		{ 'symbol': '<',  'operands': 2, 'priority': 500, 'function': lambda x, y: x < y },
	'lteq':		{ 'symbol': '<=', 'operands': 2, 'priority': 500, 'function': lambda x, y: x <= y },
	'gt':		{ 'symbol': '>',  'operands': 2, 'priority': 500, 'function': lambda x, y: x > y },
	'gteq':		{ 'symbol': '>=', 'operands': 2, 'priority': 500, 'function': lambda x, y: x >= y },
	'btand':	{ 'symbol': '&',  'operands': 2, 'priority': 600, 'function': intf (operator.and_) },
	'btxor':	{ 'symbol': '^',  'operands': 2, 'priority': 601, 'function': intf (operator.xor) },
	'btor':		{ 'symbol': '|',  'operands': 2, 'priority': 602, 'function': intf (operator.or_) },
	'and':		{ 'symbol': '&&', 'operands': 2, 'priority': 603, 'function': lambda x, y: x and y },
	'or':		{ 'symbol': '||', 'operands': 2, 'priority': 604, 'function': lambda x, y: x or y },
}

for name, data in OperatorData.iteritems():
	Operators.append (Operator (name=name, symbol=data['symbol'], operands=data['operands'],
		priority=data['priority'], function=data['function']))

OperatorSymbols={}
for op in Operators:
	if op.symbol in OperatorSymbols:
		OperatorSymbols[op.symbol].append (op)
	else:
		OperatorSymbols[op.symbol] = [op]

Constants = {
	'pi': cmath.pi,
	'e': cmath.e,
	'phi': 1.6180339887498948482,
	'epsilon': epsilon,
}

random.seed (time.time())

Functions = {
	'round':	{ 'function': realf (round) },
	'floor':	{ 'function': realf (math.floor) },
	'ceil':		{ 'function': realf (math.ceil) },
	'fact':		{ 'function': intf (math.factorial) },
	'abs':		{ 'function': realf (math.fabs) },
	'degrees':	{ 'function': realf (math.degrees) },
	'radians':	{ 'function': realf (math.radians) },
	'erf':		{ 'function': realf (math.erf) },
	'erfc':		{ 'function': realf (math.erfc) },
	'gamma':	{ 'function': realf (math.gamma) },
	'lgamma':	{ 'function': realf (math.lgamma) },
	'sqrt':		{ 'function': cmath.sqrt },
	'ln':		{ 'function': cmath.log },
	'log':		{ 'function': cmath.log10 },
	'sin':		{ 'function': cmath.sin },
	'sinh':		{ 'function': cmath.sinh },
	'asin':		{ 'function': cmath.asin },
	'asinh':	{ 'function': cmath.asinh },
	'cos':		{ 'function': cmath.cos },
	'cosh':		{ 'function': cmath.cosh },
	'acos':		{ 'function': cmath.acos },
	'acosh':	{ 'function': cmath.acosh },
	'tan':		{ 'function': cmath.tan },
	'tanh':		{ 'function': cmath.tanh },
	'atan':		{ 'function': cmath.atan },
	'atanh':	{ 'function': cmath.atanh },
	'exp':		{ 'function': cmath.exp },
	'phase':	{ 'function': cmath.phase },
	'lg':		{ 'function': lambda x: cmath.log (x, 2) },
	're':		{ 'function': lambda x: x.real },
	'im':		{ 'function': lambda x: x.imag },
	'conjugate':{ 'function': lambda x: x.conjugate() },
	'rand':		{ 'function': random.random },
}

Tokens = ['(', ')']

# Symbol table
SymbolType = Enum ('CONSTANT', 'FUNCTION', 'OPERATOR', 'TOKEN')
SymbolTypes = {}
Symbols = []

for name, value in Constants.iteritems():
	Symbols.append (name)
	SymbolTypes[name] = SymbolType.CONSTANT

for name, data in Functions.iteritems():
	Symbols.append (name)
	SymbolTypes[name] = SymbolType.FUNCTION

for op in Operators:
	if op.symbol not in Symbols:
		Symbols.append (op.symbol)
	SymbolTypes[op.symbol] = SymbolType.OPERATOR

for name in Tokens:
	SymbolTypes[name] = SymbolType.TOKEN

Symbols += Tokens
Symbols = sorted (Symbols, key=lambda x: len (x), reverse=True)

def parse_number (expr):
	"""Tries to parse a number from the expression. Returns (value, length) on success."""
	i = 0
	value = None

	if expr[0] == 'i':
		# Special case, we just have 'i' -- need special handling here because otherwise this would
		# call float('') in the complex number routine, which throws an error.
		# TODO: maybe 'i' can be a symbol instead?
		return (1j, 1)

	# Try parse increasingly long substrings until we are unable to create a float or complex number
	# from it.
	try:
		while i < len (expr):
			if expr[i] == 'i':
				value = complex (0.0, float (expr[:i]))
			else:
				value = complex (float (expr [:i + 1]), 0.0)
			i += 1

		return (value, i)
	except ValueError:
		if i != 0:
			# Got a number (the error happened when we tried to parse past the number)
			return (value, i)
		else:
			# The error happened on the first character. So this is not a number.
			return None

def parse_symbol (expr):
	for sym in Symbols:
		if expr[:len (sym)] == sym:
			return sym

	return None

def tokenize (expr):
	expr = re.sub ('\s+', '', expr.strip())
	i=0
	tokens = []

	while i < len(expr):
		sym = parse_symbol (expr[i:])

		if sym:
			symtype = SymbolTypes[sym]
			if symtype == SymbolType.CONSTANT:
				tokens.append (Constants[sym])
			else:
				tokens.append (sym)

			i += len(sym)
			continue

		result = parse_number (expr[i:])
		if result:
			num, length = result
			tokens.append (num)
			i += length
			continue

		raise ValueError ("""bad expression, couldn't parse: %s""" % expr[i:])

	return tokens

def rindex (li, value):
	return (len(li) - 1) - li[::-1].index(value)

def process_parens (expr):
	"""Processes parentheses of expr into sublists in-place.
	[1.0, '*', '(', 3.5, '+', 1j, ')']
	-> [1.0, '*', [3.5, '+', 1j]]"""
	if '(' not in expr and ')' not in expr:
		return

	try:
		parenStart = rindex (expr, '(')
		parenEnd = expr.index (')', parenStart)
	except ValueError:
		raise ValueError ("""mismatched parentheses in expression: %s""" % expr)

	subexpr = expr[parenStart + 1:parenEnd]
	del expr[parenStart + 1:parenEnd + 1]
	expr[parenStart] = subexpr
	process_parens (expr)

def process_functions (expr):
	"""Processes functions in-place"""
	i = 0
	while i < len (expr):
		if type (expr[i]) is list:
			process_functions (expr[i])

		if (type(expr[i]) is not str) or (expr[i] not in Functions):
			i += 1
			continue

		# Ensure that arguments follow
		if (i + 1 >= len (expr)) or (type (expr[i + 1]) is not list):
			raise ValueError ("""function %s used without arguments""" % expr[i])

		args = expr[i + 1]
		del expr[i + 1]
		expr[i] = FunctionCall (expr[i], args)
		i += 1

def is_operand (x):
	# Operands can be either lists (which also mean parens, thus can be single-expressions)
	# or complex numbers
	return type(x) in [complex, list]

def find_fitting_operator (sym, numOperands):
	# Pass 1: exact numOperands match
	for op in Operators:
		if op.symbol != sym:
			continue

		if op.operands == numOperands:
			# print '%s, %d -> %s (1st pass)' % (sym, numOperands, op)
			return op

	# Pass 2: by symbol
	for op in Operators:
		if op.symbol == sym:
			# print '%s, %d -> %s (2nd pass)' % (sym, numOperands, op)
			return op

	raise ValueError ('''unknown operator %s!''' % sym)

def process_operators (expr):
	"""Processes operators"""
	i = 0

	# Find all operators in this expression
	while i < len (expr):
		if type (expr[i]) is list:
			process_functions (expr[i])
			process_operators (expr[i])

		if type (expr[i]) is FunctionCall:
			process_functions (expr[i].args)
			process_operators (expr[i].args)

		if (type(expr[i]) is not str) or (expr[i] not in OperatorSymbols):
			i += 1
			continue

		args = []
		argIndices = []
		if i > 0 and is_operand (expr[i - 1]):
			args.append (expr[i - 1])
			argIndices.append (i - 1)

		if i - 1 < len(expr) and is_operand (expr[i + 1]):
			args.append (expr[i + 1])
			argIndices.append (i + 1)

		# Resolve operators with the same symbol based on operand count
		numOperands = 0
		for arg in args:
			if is_operand (arg):
				numOperands += 1

		expr[i] = find_fitting_operator (expr[i], numOperands)
		i += 1

def find_priority_operator (expr):
	"""Finds the operator with most priority in the expression"""
	bestOp = None
	bestOpIndex = -1

	for i in range (0, len (expr)):
		sym = expr[i]

		if type (sym) is not Operator:
			continue

		if not bestOp or sym.priority < bestOp.priority:
			bestOp = sym
			bestOpIndex = i

	return (bestOp, bestOpIndex)

def realPrint (x):
	print x

tabs=''
def evaluate (expr, verbose=False):
	global tabs
	printFunc = realPrint if verbose else lambda x:None
	printFunc (tabs + 'Preprocess: %s' % expr)

	# If there are sub-expressions in here, those need to be evaluated first
	i = 0
	while i < len (expr):
		sym = expr[i]

		if type (sym) is list and sym:
			printFunc (tabs + 'Evaluating sub-expression: %s' % (sym))
			tabs += '\t'
			expr[i] = evaluate (list (sym), verbose)
			tabs = tabs[:-1]
			printFunc (tabs + '-> %s' % expr[i])

		# If there are function calls, evaluate those
		if type (sym) is FunctionCall:
			tabs += '\t'
			if sym.args:
				sym.args = [evaluate (sym.args, verbose)]
			tabs = tabs[:-1]

			printFunc (tabs + 'Evaluating function call: %s' % (sym))
			expr[i] = Functions[sym.funcname]['function'] (*sym.args)
			printFunc (tabs + '-> %s' % expr[i])
			
		i += 1

	printFunc (tabs + 'Evaluate: %s' % expr)
	runaway = 0
	while True:
		runaway += 1
		if runaway > 1000:
			raise ValueError ('infinite loop detected')

		op, i = find_priority_operator (expr)
		if not op:
			break

		if op.operands == 2:
			argIndices = [i - 1, i + 1]
		else:
			argIndices = [i + 1]

		args = [expr[x] for x in argIndices]
		argIndices = sorted (argIndices, reverse=True)
		printFunc (tabs + 'Processing: (%s, %d) with args %s' % (op, i, args))
		expr[i] = op.function (*args)
		printFunc (tabs + '-> %s' % expr[i])

		for i2 in argIndices:
			del expr[i2]

	printFunc (tabs + 'Result: %s' % expr[0])

	if len(expr) != 1:
		printFunc (tabs + 'Bad expression detected, tokens: %s' % expr)
		raise ValueError ('malformed expression')

	return expr[0]

def repr_number (x):
	"""Returns a string representation for a real number"""
	rep = '%.10f' % x

	if '.' in rep:
		while rep[-1] == '0':
			rep = rep[:-1]

		if rep[-1] == '.':
			rep = rep[:-1]

	return rep

def repr_imaginary (x):
	rep = repr_number (x)

	if rep == '1':
		return 'i'

	if rep == '-1':
		return '-i'

	return rep + 'i'

def represent (x):
	"""Returns a string representation of a float or complex number"""
	if math.fabs (x.imag) > epsilon:
		if math.fabs (x.real) > epsilon:
			# Complex number
			return '%s %s %s' % (repr_number (x.real),
				'+' if x.imag >= 0 else '-',
				repr_imaginary (math.fabs (x.imag)))
		else:
			# Pure imaginary number
			return repr_imaginary (x.imag)
	else:
		# Real number
		return repr_number (x.real)

def calc (expr, verbose=False):
	expr = tokenize (expr)
	process_parens (expr)
	process_functions (expr)
	process_operators (expr)
	return represent (evaluate (expr, verbose))