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
import string
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]

def sgn (x):
	return cmp (x, 0)

Constants = {
	'pi': cmath.pi,
	'e': cmath.e,
	'phi': 1.6180339887498948482,
	'epsilon': epsilon,
	'potato': cmath.sqrt(8) / 3,
}

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 },
	'sgn':		{ 'function': realf (sgn) },
}

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)
PreferredBase = 10

def set_preferred_base(x):
	global PreferredBase
	PreferredBase = x

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

def realPrint (x):
	print x

Attributes = {
	'hex': lambda self: self.set_preferred_base (0x10),
	'binary': lambda self: self.set_preferred_base (0b10),
	'decimal': lambda self: self.set_preferred_base (10),
}

Attributes['bin'] = Attributes['binary']
Attributes['dec'] = Attributes['decimal']
AttributeNames = sorted ([key for key in Attributes], key=lambda x:len(x), reverse=True)

class Calculator (object):
	def __init__ (self):
		self.preferred_base = None

	def set_preferred_base (self, x):
		self.preferred_base = x

	def trim_spaces (self, expr):
		return re.sub ('\s+', '', expr.strip())

	def parse_attributes (self, expr):
		if expr[0] != '<':
			return expr

		i = 1
		while expr[i] != '>':
			if expr[i] == '|':
				i += 1

			for name in AttributeNames:
				if expr[i:i + len(name)] == name:
					Attributes[name] (self)
					i += len(name)

					if expr[i] == '>':
						break

					if expr[i] != '|':
						raise ValueError ('malformed attributes')
					break
			else:
				raise ValueError ('bad attributes: %s' % expr[i:])

		return expr[i + 1:]

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

		# Possibly it's hexadecimal
		if expr[0:2].lower() == '0x':
			base = 0x10
			digits = string.hexdigits
			digitname = 'hexit'
			i = 2
		elif expr[0:2].lower() == '0b':
			base = 0b10
			digits = ['0', '1']
			digitname = 'bit'
			i = 2

		if base != 10:
			if not self.preferred_base:
				self.preferred_base = base
		
			# Skip leading zeroes
			while i < len (expr) and expr[i] == '0':
				i += 1

			startingIndex = i
			while i < len (expr) and expr[i] in digits:
				i += 1

			if i < len(expr) and expr[i] == '.':
				raise ValueError ('non-decimal floating point numbers are not supported')

			if i == startingIndex:
				if i < len (expr):
					raise ValueError ('not a valid %s "%s" in %s' % (digitname, expr[i], expr[0:i+1]))
				else:
					raise ValueError ('end of expression encountered while parsing '
						'base-%d literal' % base)

			return (complex (int (expr[startingIndex:i], base), 0), i)

		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 (self, expr):
		for sym in Symbols:
			if expr[:len (sym)] == sym:
				return sym

		return None

	def tokenize (self, expr):
		i=0
		tokens = []

		while i < len(expr):
			sym = self.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 = self.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 process_parens (self, 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
		self.process_parens (expr)

	def process_functions (self, expr):
		"""Processes functions in-place"""
		i = 0
		while i < len (expr):
			if type (expr[i]) is list:
				self.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 (self, x):
		# Operands can be either lists (which also mean parens, thus can be single-expressions)
		# or complex numbers
		return type(x) in [complex, list, FunctionCall]

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

			if op.operands == numOperands:
				return op

		# Pass 2: by symbol
		for op in Operators:
			if op.symbol == sym:
				return op

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

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

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

			if type (expr[i]) is FunctionCall:
				self.process_functions (expr[i].args)
				self.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 self.is_operand (expr[i - 1]):
				args.append (expr[i - 1])
				argIndices.append (i - 1)

			if i - 1 < len(expr) and self.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 self.is_operand (arg):
					numOperands += 1

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

	def find_priority_operator (self, 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 evaluate (self, expr, verbose=False):
		printFunc = realPrint if verbose else lambda x:None
		printFunc (self.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 (self.tabs + 'Evaluating sub-expression: %s' % (sym))
				self.tabs += '\t'
				expr[i] = self.evaluate (list (sym), verbose)
				self.tabs = self.tabs[:-1]
				printFunc (self.tabs + '-> %s' % expr[i])

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

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

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

			op, i = self.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 (self.tabs + 'Processing: (%s, %d) with args %s' % (op, i, args))
			expr[i] = op.function (*args)
			printFunc (self.tabs + '-> %s' % expr[i])

			for i2 in argIndices:
				del expr[i2]

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

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

		return expr[0]

	def repr_number (self, x):
		"""Returns a string representation for a real number"""
		base = self.preferred_base if self.preferred_base else 10
		if math.fabs (x - math.floor(x)) < epsilon and base != 10:
			digits='0123456789abcdef'
			assert base <= len (digits), '''base %d is too large''' % base

			divisor = base
			rep = ''
			x = int (x)
			runaway = 0

			while x > 0:
				runaway += 1
				if runaway > 1000:
					raise ValueError('runaway triggered')

				i = (x % divisor) / (divisor / base)
				x -= i * (divisor / base)
				rep += digits[i]
				divisor *= base

			rep += 'x' if base == 16 else 'b'
			rep += '0'
			return rep[::-1]

		rep = '%.10f' % x

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

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

		return rep

	def repr_imaginary (self, x):
		rep = self.repr_number (x)

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

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

		return rep + 'i'

	def represent (self, 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' % (self.repr_number (x.real),
					'+' if x.imag >= 0 else '-',
					self.repr_imaginary (math.fabs (x.imag)))
			else:
				# Pure imaginary number
				return self.repr_imaginary (x.imag)
		else:
			# Real number
			return self.repr_number (x.real)

	def calc (self, expr, verbose=False):
		self.state = {}
		self.tabs = ''
		expr = self.trim_spaces (expr)
		expr = self.parse_attributes (expr)
		expr = self.tokenize (expr)
		self.process_parens (expr)
		self.process_functions (expr)
		self.process_operators (expr)
		result = self.evaluate (expr, verbose)
		return self.represent (result)