"""Convert between math notation and internal notation.

Robert A McDougal, 2011
"""

from string import ascii_letters, digits
from errors import SyntaxError
import math
import copy

# use sympy to simplify expressions, if available
try:
    import sympy
    def _cas_simplify(s):
        try:
            result = str(sympy.sympify(s)).replace(' ', '')
            result = tokens_to_string(parse_string(result), use_cas=False)
            if is_enclosed_in_parens(s) and (not is_enclosed_in_parens(result)):
                return '(' + result + ')'
            else:
                return result
        except:
            return s
except ImportError:
    def _cas_simplify(s):
        return s


inf = float('inf')

class Token:
    """Generic Token."""
    
    def __init__(self, content, precedence = 0, num_args = 0):
        """Initialize the token.
        
        Parameters:
            content -- a string representing its content
            precedence -- priority in order of operations
            num_args -- number of parameters it takes
        """
        self.content = content
        self.precedence = precedence
        self.num_args = num_args
        self.last = None
        
    def value(self, random):
        """Return the current Token."""
        return self
    
    def __repr__(self):
        return self.content
    
    def __str__(self):
        return self.content


# TODO:  add more validation
class Number(Token):
    """Number Token."""
    
    def __init__(self, content):
        """Initialize the number.
        
        Parameters:
            content -- the value of the number
        
        Note:
        
        content could be a specific number or a probability distribution.
        
        Examples of valid content:
            1237.86
            12 [3]    <-- Normal r.v., mean = 12, standard deviation = 3
            
        """
        Token.__init__(self, str(content).strip(), 6, 0)
        

    def value(self, random):
        """Returns a new Number from current Number's probability distribution."""
        s = self.content
        try:
            if '[' not in s:
                v = float(s)
                return self
            
            i = s.find('[')
            j = s.find(']')
            
            if j + 1 != len(s):
                raise SyntaxError(s, 'Invalid number.')
            
            inside = s[i + 1 : j].strip()
            mean = s[ : i].strip()
            
            if len(mean) == 0:
                # uniform distribution
                parts = inside.split(':')
                if len(parts) != 2:
                    raise SyntaxError(s, 'Invalid number.')
                values = [float(v) for v in parts]
                return Number(random.uniform(min(values), max(values)))
            
            mean = float(mean)
            
            if inside[-1] == '%':
                std_dev = mean * float(inside[ : -1]) * .01
            else:
                std_dev = float(inside)
            
            return Number(random.gauss(mean, std_dev))
        except:
            raise SyntaxError(s, 'Invalid number.')

    def __float__(self):
        return float(self.content)

class Variable(Token):
    """Variable Token."""
    
    def __init__(self, content):
        """Initialize the variable.
        
        Parameters:
            content -- the name of the variable
        """
        Token.__init__(self, content, 6, 0)


class BinaryOperator(Token):
    """Binary Operator.
    
    When parsing, split at binary operators from right-to-left; otherwise
    left-to-right.
    """
    def __init__(self, content, precedence = 0):
        Token.__init__(self, content, precedence, 2)



class Function(Token):
    """Function Token."""
    
    def __init__(self, content, num_args = 1):
        """Initialize the function definition.
        
        Parameters:
            content -- the name of the function
            num_args -- the number of parameters it takes
        """
        Token.__init__(self, content, 4, num_args)


# common tokens
open_paren = Token('(', precedence = inf)
close_paren = Token(')', precedence = inf)
separator = Token(',', precedence = 0)
addition = BinaryOperator('+', precedence = 1)
subtraction = BinaryOperator('-', precedence = 1)
multiplication = BinaryOperator('*', precedence = 2)
division = BinaryOperator('/', precedence = 2)
exponentiation = BinaryOperator('^', precedence = 4)
negation = Token('-', precedence = 3, num_args = 1)
pi = Number(math.pi)

functions = {'exp': 1, 'sin': 1, 'cos': 1, 'tan': 1, 'atan': 1, 'acos': 1, 'asin': 1, 'log': 1,
             'log10': 1, 'abs': 1, 'heav': 1, 'sinh': 1, 'cosh': 1, 'tanh': 1, 'erf': 1,
             'erfc': 1, 'sign': 1, 'max': 2, 'min': 2, 'atan2': 2}



letters_and_numbers = ascii_letters + digits + '_'
numbers = digits + '.'
numbers_extended = numbers + ' [:%]'

def is_enclosed_in_parens(s):
    """Return True if s is an expression enclosed in parentheses.
    
    i.e. returns True  for (1+3-2)
             but False for (1+3)-(2)
    """
    if s[0] != '(' or s[-1] != ')':
        return False
    paren_count = 1
    for c in s[1 : -1]:
        if c == '(':
            paren_count += 1
        if c == ')':
            paren_count -= 1
            if paren_count == 0:
                return False
    if paren_count != 1:
        # should never get here
        raise SyntaxError(s, 'imbalanced parentheses')
    return True


def default_renderer(token, stack, user_functions, substitutions, renderer=None):
    """Express equation in human readable form.
    
    renderer points to current renderer.
    """
    if renderer is None:
        renderer = default_renderer
    if len(stack) < token.num_args:
        raise SyntaxError(token.content, 'Missing arguments')
    if token.num_args == 0:
        stack.append(token.content)
    elif isinstance(token, Function):
        args = []
        for i in range(token.num_args):
            arg = stack.pop()
            if is_enclosed_in_parens(arg):
                arg = arg[1 : -1]
            args.append(arg)
        args = args[: : -1]
        if token.content in functions:
            if len(args) != functions[token.content]:
                raise SyntaxError(token.content, '%s takes exactly %d arguments but %d were supplied' %
                                                 (token.content, functions[token.content], len(args)))
            stack.append(token.content + '(' + ','.join(args) + ')')
        elif token.content in user_functions:
            f_name = token.content
            f = user_functions[f_name]
            if len(args) != f['num_args']:
                raise SyntaxError(f_name, '%s takes exactly %d arguments but %d were supplied' % (f_name, f['num_args'], len(args)))
            var_names = f['var_names']
            f_val = []
            for t in f['expression']:
                if t.content in var_names:
                    t = '(%s)' % args[var_names.index(t.content)]
                f_val.append(t)
            stack.append(tokens_to_string(f_val, user_functions, substitutions, renderer=renderer))
        else:
            raise SyntaxError(stack, 'Undefined function: ' + token.content)
    elif token.num_args == 1:
        if token.content == '-':
            stack.append('(-' + stack.pop() + ')')
        else:
            # should never get here
            raise SyntaxError(token.content)
    elif isinstance(token, BinaryOperator):
        arg1 = stack.pop()
        arg2 = stack.pop()
        if token.content in ('/', '*', '^'):
            if not is_enclosed_in_parens(arg1):
                arg1 = '(' + arg1 + ')'
        if isinstance(token.last, Token) and token.content in ('+', '*', '-', '/') and token.precedence <= token.last.precedence:
            if isinstance(token.last, Number):
                # avoid negation
                if isinstance(token.last.last, Token) and token.precedence <= token.last.last.precedence:
                    if arg2[0] == '(' and arg2[-1] == ')' and arg2[1] != '-':
                        # strip parens
                        arg2 = arg2[1 : -1]
            elif arg1[0] == '(' and arg1[-1] == ')' and arg1[1] != '-' and token.content not in ('-', '/'):
                # strip parens from arg1                
                arg1 = arg1[1 : -1]
        stack.append('(' + arg2 + token.content + arg1 + ')')
    elif token.content in ('+', '*', '-', '/'):
        # common operations but with many arguments
        args = [stack.pop() for i in range(token.num_args)]
        # reverse order from stack
        args = args[::-1]
        stack.append('(' + token.content.join(args) + ')')
    else:
        raise SyntaxError(stack, 'unable to interpret expression: ' + str(token))


# TODO: make expand user_functions
def rpn_renderer(token, stack, user_functions, substitutions):
    """Express equation in RPN form."""
    if len(stack) > 0:
        stack[0] += ' ' + token.content
    else:
        stack.append(token.content)


def c_renderer(token, stack, user_functions, substitutions):
    """Express equations in a form suitable for C."""
    if token.content == '^':
        arg1 = stack.pop()
        arg2 = stack.pop()
        stack.append('pow(' + arg2 + ',' + arg1 + ')')
    elif token.content == 'abs':
        stack.append('fabs(' + stack.pop() + ')')
    else:
        default_renderer(token, stack, user_functions, substitutions, c_renderer)

def matlab_renderer(token, stack, user_functions, substitutions):
    """Express equations in a form suitable for MATLAB."""
    if token.content == 'heav':
        stack.append('(' + stack.pop() + '> 0)')
    else:
        default_renderer(token, stack, user_functions, substitutions, matlab_renderer)
   

def simplify_expression(s):
    """Simplify a string, return result as a string."""
    return tokens_to_string(parse_string(s))

def simplify(tokens):
    """Preprocess a list of tokens by evaluating numerical expressions.
    
    Assumes numbers that are probabilities have been collapsed to specific
    values.
    """
    change = True
    while change:
        simplified = []
        change = False
        for i in range(len(tokens)):
            current = tokens[i]
            if isinstance(current, Token):
                symbol = current.content
                if symbol in ('cos', 'sin', 'tan', 'acos', 'asin', 'atan', 'sinh', 'cosh', 'tanh', 'log', 'log10', 'exp', 'heav', 'abs') and current.num_args == 1:
                    if len(simplified) < 1:
                        raise SyntaxError(token, 'not enough arguments')
                    last = simplified[-1]
                    if isinstance(last, Number):
                        arg = float(simplified.pop().content)
                        if symbol == 'cos':
                            current = Number(math.cos(arg))
                        elif symbol == 'sin':
                            current = Number(math.sin(arg))
                        elif symbol == 'tan':
                            current = Number(math.tan(arg))
                        elif symbol == 'acos':
                            current = Number(math.acos(arg))
                        elif symbol == 'asin':
                            current = Number(math.asin(arg))
                        elif symbol == 'atan':
                            current = Number(math.atan(arg))
                        elif symbol == 'cosh':
                            current = Number(math.cosh(arg))
                        elif symbol == 'sinh':
                            current = Number(math.sinh(arg))
                        elif symbol == 'tanh':
                            current = Number(math.tanh(arg))
                        elif symbol == 'log':
                            current = Number(math.log(arg))
                        elif symbol == 'log10':
                            current = Number(math.log10(arg))
                        elif symbol == 'exp':
                            current = Number(math.exp(arg))
                        elif symbol == 'abs':
                            current = Number(abs(arg))
                        elif symbol == 'heav':
                            if arg > 0:
                                current = Number(1)
                            else:
                                current = Number(0)
                        else:
                            raise Exception('Unsimplified token: ' + symbol)
                        change = True
                elif symbol in ('+', '*', '^', '/', '-', 'max', 'min', 'atan2') and current.num_args == 2:
                    if len(simplified) < 2:
                        raise SyntaxError(token, 'not enough arguments')
                    last = simplified[-1]
                    next_to_last = simplified[-2]
                    if isinstance(last, Number) and isinstance(next_to_last, Number):
                        arg2 = float(simplified.pop().content)
                        arg1 = float(simplified.pop().content)
                        if symbol == '+':
                            current = Number(arg1 + arg2)
                        elif symbol == '*':
                            current = Number(arg1 * arg2)
                        elif symbol == '^':
                            current = Number(arg1 ** arg2)
                        elif symbol == '/':
                            current = Number(arg1 / arg2)
                        elif symbol == '-':
                            current = Number(arg1 - arg2)
                        elif symbol == 'max':
                            current = Number(max(arg1, arg2))
                        elif symbol == 'min':
                            current = Number(min(arg1, arg2))
                        elif symbol == 'atan2':
                            current = Number(math.atan2(arg1, arg2))
                        else:
                            raise Exception('Should never get here.')
                        change = True
                        
                    # addition and subtraction of unary negation (note we know it's subtraction not negation here
                    # because current.num_args == 2)
                    elif isinstance(last, Token) and last.content == '-' and last.num_args == 1 and symbol in ('+', '-'):
                        simplified.pop()
                        if symbol == '+':
                            current = subtraction
                        else:
                            current = addition
                        change = True
                    elif symbol in ('+', '-') and isinstance(last, Number) and float(last) < 0:
                        old_value = float(last)
                        simplified[-1] = Number(-float(last))
                        if symbol == '+':
                            current = subtraction
                        else:
                            current = addition
                        change = True                    
                elif symbol == '-' and current.num_args == 1 and isinstance(simplified[-1], Number):
                    current = Number(-float(simplified.pop().content))
                    change = True
            simplified.append(current)
        tokens = simplified
    return simplified
                        

def _join_with_commas(stack):
    return ','.join(stack)


def tokens_to_string(tokens, user_functions={}, substitutions={}, renderer=default_renderer, use_cas=True):
    """Convert the list of tokens to a string.
    
    Parameters:
        tokens -- a list of objects of type Token in RPN order
        substitutions -- dict of variable names and values
        renderer -- rule for converting the RPN Token objects to a string
    
    Note:
    
    Use a different renderer for different output formats. For example, the
    default_renderer writes exponentiation as a^b, but in C one must use
    pow(a, b).
    
    Assumes numbers that are probabilities have been collapsed to specific
    values.
    """
    
    # apply substitutions
    # TODO: add way to catch circular loops
    change = True
    while change:
        final_tokens = []
        change = False
        for t in tokens:
            if isinstance(t, Token) and t.content in substitutions:
                final_tokens += substitutions[t.content]
                change = True
            else:
                final_tokens.append(t)
        tokens = final_tokens
        
    # bubble sort +, * together 
    change = True
    sort_ops = (addition, multiplication)
    while change:
        change = False
        for i in range(4, len(tokens)):
            if tokens[i] == tokens[i - 2] and tokens[i] in sort_ops and isinstance(tokens[i - 1], Token) and tokens[i - 1].num_args == 0:
                change = True
                tokens[i - 1], tokens[i - 2] = tokens[i - 2], tokens[i - 1]
    

    # simplify the resulting expression
    tokens = simplify(tokens)
    
    # render it in the desired format
    stack = []
    last = None
    for t in tokens:
        t = copy.copy(t)
        if isinstance(t, Token):
            t.last = last
            renderer(t, stack, user_functions, substitutions)
        elif isinstance(t, str):
            stack.append(t)
        else:
            raise SyntaxError(tokens, 'parsing error')
        last = t
    if len(stack) != 1:
        raise SyntaxError(tokens, 'parsing error')
        
    if use_cas:
        return _cas_simplify(stack[0])
    else:
        return stack[0]


def parse_string(s):
    """Parse the string s, return a list of Token in RPN order."""
    return parse(tokenize(str(s)))


    
def parse(tokens):
    """Take a list of Token in math order, return them sorted in RPN order."""
    parsed_tokens = _parse(tokens)
    return [t for t in parsed_tokens if t is not separator]

def _parse(tokens):
    """Helper function for parse.
    
    Does all the work, except for removing separators."""
    if len(tokens) < 2:
        return tokens

    rpn = []

    paren_count = 0
    min_precedence = inf
    min_at = -1

    for i, t in enumerate(tokens):
        if t is open_paren:
            paren_count += 1
        elif t is close_paren:
            paren_count -= 1
            if paren_count < 0:
                raise SyntaxError(tokens, 'Unbalanced parentheses')
        elif t.precedence <= min_precedence and paren_count == 0:
            if isinstance(t, BinaryOperator) or t.precedence < min_precedence:
                min_precedence = t.precedence
                min_at = i

    if min_precedence < inf:
        current = tokens[min_at]
        if current.num_args == 1:
            if min_at == 0:
                return parse(tokens[1 : ]) + [current]
            else:
                return parse(tokens[0 : min_at - 1]) + parse(tokens[min_at + 1 : ]) + [current] + [tokens[min_at - 1]]
            return parse(tokens[min_at + 1 : ]) + [current] + parse(tokens[0 : min_at])
        return parse(tokens[0 : min_at]) + parse(tokens[min_at + 1 : ]) + [tokens[min_at]]
    else:
        if tokens[0] == open_paren and tokens[-1] == close_paren:
            return parse(tokens[1 : -1])
        else:
            # should never be here?
            raise SyntaxError(tokens)


def tokenize(s):
    """Return a list of Token in the order of the math equation string s."""
    tokens = []
    length = len(s)
    next_unary = True
    no_numbers = False
    token = None

    i = j = 0
    while i < length:
        c = s[i]
        if c in (' ', '\t'):
            # whitespace; ignore
            j = i + 1
        else:
            if i + 1 >= length and c in ('(', '+', '*', '^', '/', '-', '['):
                raise SyntaxError(s, 'Cannot end expression with ' + c)
            if c in ascii_letters:
                j = i
                while j < length and s[j] in letters_and_numbers:
                    j += 1
                token = s[i : j]

                if token == 'pi':
                    token = pi
                else:
                    token = Variable(token)
            elif c in numbers_extended:
                j = i
                inside_bracket = False
                while j < length and (s[j] in numbers_extended or s[j] in ('e', 'E') or ((s[j] == '+' or s[j] == '-') and s[j - 1]
 in ('e', 'E')) or inside_bracket):
                    if s[j] == ']' and not inside_bracket:
                        raise SyntaxError(s, '] without matching [')
                    if inside_bracket:
                        if s[j] == '[':
                            raise SyntaxError(s, 'Cannot have nested [')
                        elif s[j] == ']':
                            inside_bracket = False
                    elif s[j] == '[':
                        inside_bracket = True
                    j += 1
                token = Number(s[i : j])
            elif c == '(':
                token = open_paren
                j = i + 1
            elif c == ')':
                token = close_paren
                j = i + 1
            elif c == ',':
                token = separator
                j = i + 1
            elif c == '+':
                token = addition
                j = i + 1
            elif c == '-':
                if next_unary:
                    token = negation
                else:
                    token = subtraction
                j = i + 1
            elif c == '*':
                if s[i + 1] == '*':
                    token = exponentiation
                    j = i + 2
                else:
                    token = multiplication
                    j = i + 1
            elif c == '/':
                token = division
                j = i + 1
            elif c == '^':
                token = exponentiation
                j = i + 1
            else:
                raise SyntaxError(s, 'Symbol does not belong here: ' + c)
            next_unary = (token.num_args > 0) or (token is open_paren) or (token is separator)
            no_numbers = (token.num_args == 0) and (token is not open_paren)
            if token is None:
                raise SyntaxError(s, 'Unknown token: ' + c)
            tokens.append(token)
            token = None
        i = j
    
    # distinguish between functions and variables
    for i in range(len(tokens) - 1):
        t = tokens[i]
        if isinstance(t, Variable) and tokens[i + 1] == open_paren:
            paren_count = 1
            num_args = 1
            for t2 in tokens[i + 2 :]:
                if t2 is open_paren:
                    paren_count += 1
                elif t2 is close_paren:

                    paren_count -= 1
                    if paren_count == 0:
                        break
                elif t2 is separator and paren_count == 1:
                    num_args += 1
            if paren_count > 0:
                raise SyntaxError(s, '( without matching )')
                
            tokens[i] = Function(t.content, num_args)

    # make sure no repeated separators or empty parentheses
    last_token = None
    for t in tokens:
        if t is separator and last_token is separator:
            raise SyntaxError(s, 'repeated commas')
        elif t is close_paren and last_token is open_paren:
            raise SyntaxError(s, 'empty parentheses')
        last_token = t
        
    return tokens
