import unittest
import doctest
import new

import line
from node import Node as N, Leaf as L
from line import generate_line, is_id, is_int, is_power


class TestLine(unittest.TestCase):

    def test_doctest(self):
        self.assertEqual(doctest.testmod(m=line)[0], 0)

    def test_empty(self):
        self.assertEquals(generate_line(None), '<empty expression>')

    def test_simple(self):
        l0, l1 = L(1), L(2)
        plus = N('+', l0, l1)
        self.assertEquals(generate_line(plus), '1 + 2')

    def test_parentheses(self):
        l0, l1 = L(1), L(2)
        plus = N('+', l0, l1)
        times = N('*', plus, plus)
        self.assertEquals(generate_line(times), '(1 + 2)(1 + 2)')

    def test_parentheses_no_equal_precedence(self):
        a, b, c = L('a'), L('b'), L('c')
        line = N('vv', a, N('^^', b, c))
        self.assertEquals(generate_line(line), 'a vv b ^^ c')
        line = N('^^', a, N('vv', b, c))
        self.assertEquals(generate_line(line), 'a ^^ (b vv c)')

    def test_parentheses_equal_precedence_right(self):
        l0, l1, l2, l3 = L(1), L(2), L(3), L(4)
        plus = N('+', l1, l2)
        minus = N('-', l0, plus)
        self.assertEquals(generate_line(minus), '1 - (2 + 3)')

        power = N('^', l0, N('_', l1, l2))
        self.assertEquals(generate_line(power), '1 ^ 2 _ 3')
        power = N('^', l0, N('^', l1, l2))
        self.assertEquals(generate_line(power), '1 ^ 2 ^ 3')
        power = N('^', N('^', l0, l1), l2)
        self.assertEquals(generate_line(power), '(1 ^ 2) ^ 3')
        power = N('^', l0, N('^', N('^', l1, l2), l3))
        self.assertEquals(generate_line(power), '1 ^ (2 ^ 3) ^ 4')

        power = N('^', l0, N('^', N('_', l1, l2), l3))
        self.assertEquals(generate_line(power), '1 ^ (2 _ 3) ^ 4')

    def test_parentheses_equal_precedence_left(self):
        a, b, c, d = L('a'), L('b'), L('c'), L('d')
        exp = N('*', N('/', N('*', a, b), c), d)
        self.assertEquals(generate_line(exp), '(ab) / c * d')

    def test_parentheses_unary(self):
        neg = N('-', N('+', L(1), L(2)))
        self.assertEquals(generate_line(neg), '-(1 + 2)')

    def test_parentheses_nary(self):
        l0, l1, l2 = L(1), L(2), L(3)
        plus = N('+', N('+', l0, l1), l2)
        self.assertEquals(generate_line(plus), '1 + 2 + 3')

    def test_function(self):
        sin = N('sin', N('*', L(2), L('x')))
        self.assertEquals(generate_line(sin), 'sin(2x)')

    def test_mod(self):
        l0, l1 = L(1), L(2)
        mod = N('mod', l1, l0)
        self.assertEquals(generate_line(mod), '2 mod 1')

    def test_multiplication_identifiers(self):
        a, b = L('a'), L('b')
        self.assertEquals(generate_line(N('*', a, b)), 'ab')
        self.assertEquals(generate_line(N('*', a, -b)), 'a * -b')

    def test_multiplication_constant_identifier(self):
        l0, a = L(2), L('a')
        mul = N('*', l0, a)
        self.assertEquals(generate_line(mul), '2a')

    def test_multiplication_identifier_constant(self):
        l0, a = L(2), L('a')
        mul = N('*', a, l0)
        self.assertEquals(generate_line(mul), 'a * 2')

    def test_multiplication_constants(self):
        l0, l1 = L(1), L(2)
        mul = N('*', l0, l1)
        self.assertEquals(generate_line(mul), '1 * 2')

    def test_nary(self):
        l0, l1, l2 = L(1), L(2), L(3)
        plus = N('+', l0, l1, l2)
        self.assertEquals(generate_line(plus), '1 + 2 + 3')

    def test_pow_basic(self):
        a, b, c = L('a'), L('b'), L('c')
        node_pow = N('^', a, N('+', b, c))
        self.assertEquals(generate_line(node_pow), 'a ^ (b + c)')

    def test_pow_intermediate1(self):
        # expression: (a(b+c))^(d+e)
        a, b, c, d, e = L('a'), L('b'), L('c'), L('d'), L('e')
        node_bc = N('+', b, c)
        node_de = N('+', d, e)
        node_mul = N('*', a, node_bc)
        node_pow = N('^', node_mul, node_de)
        self.assertEquals(generate_line(node_pow), '(a(b + c)) ^ (d + e)')

    def test_pow_intermediate2(self):
        # expression: a(b+c)^(d+e)
        a, b, c, d, e = L('a'), L('b'), L('c'), L('d'), L('e')
        node_bc = N('+', b, c)
        node_de = N('+', d, e)
        node_pow = N('^', node_bc, node_de)
        node_mul = N('*', a, node_pow)
        self.assertEquals(generate_line(node_mul), 'a(b + c) ^ (d + e)')

    def test_pow_negated_root(self):
        a, l2 = L('a'), L(2)
        power = -N('^', a, l2)
        self.assertEquals(generate_line(power), '-a ^ 2')
        power = N('^', -a, l2)
        self.assertEquals(generate_line(power), '(-a) ^ 2')

    def test_multiplication_sign(self):
        a, b, c, l2 = L('a'), L('b'), L('c'), L(2)
        mul = N('*', a, b)
        self.assertEquals(generate_line(mul), 'ab')
        mul = N('*', mul, c)
        self.assertEquals(generate_line(mul), 'abc')
        mul = N('*', N('*', a, l2), b)
        self.assertEquals(generate_line(mul), 'a * 2b')

        mul = -N('*', N('*', a, b), c)
        self.assertEquals(generate_line(mul), '-abc')
        mul = N('*', N('*', -a, b), c)
        self.assertEquals(generate_line(mul), '(-a)bc')

        mul = N('*', a, N('-', b, c))
        self.assertEquals(generate_line(mul), 'a(b - c)')
        mul = N('*', l2, N('-', b, c))
        self.assertEquals(generate_line(mul), '2(b - c)')

        mul = N('*', N('+', a, b), c)
        self.assertEquals(generate_line(mul), '(a + b)c')
        mul = N('*', N('+', a, b), l2)
        self.assertEquals(generate_line(mul), '(a + b)2')

        mul = N('*', N('+', a, b), N('+', c, l2))
        self.assertEquals(generate_line(mul), '(a + b)(c + 2)')

        mul = N('*', l2, a)
        self.assertEquals(generate_line(mul), '2a')
        mul = N('*', a, l2)
        self.assertEquals(generate_line(mul), 'a * 2')

        mul = N('*', l2, N('^', a, l2))
        self.assertEquals(generate_line(mul), '2a ^ 2')

        mul = N('*', l2, N('^', l2, a))
        self.assertEquals(generate_line(mul), '2 * 2 ^ a')

        mul = N('*', N('/', a, l2), N('+', a, b))
        self.assertEquals(generate_line(mul), 'a / 2 * (a + b)')

    def test_plus_to_minus(self):
        plus = N('+', L(1), -L(2))
        self.assertEquals(generate_line(plus), '1 - 2')

        l1, a, b, c = L(1), L('a'), L('b'), L('c')
        plus = N('+', l1, -N('*', N('*', a, b), c))
        self.assertEquals(generate_line(plus), '1 - abc')

    def test_helper_functions(self):
        l1, a = L(1), L('a')
        neg = -l1
        neg_a = -a
        plus = N('+', l1, a)
        power = N('^', a, l1)

        self.assertTrue(is_id(a))
        self.assertTrue(is_id(neg_a))
        self.assertFalse(is_id(neg))
        self.assertFalse(is_id(plus))

        self.assertTrue(is_int(l1))
        self.assertTrue(is_int(neg))
        self.assertFalse(is_int(neg_a))
        self.assertFalse(is_int(plus))

        self.assertTrue(is_power(power))
        self.assertFalse(is_power(l1))
        self.assertFalse(is_power(plus))

    def test_negated_operator(self):
        neg = -N('+', L(1), L(2))
        self.assertEquals(generate_line(neg), '-(1 + 2)')

        neg = -N('-', L(1), L(2))
        self.assertEquals(generate_line(neg), '-(1 - 2)')

        # FIXME: neg = N('+', L(1), N('+', L(1), L(2)))
        # FIXME: self.assertEquals(generate_line(neg), '1 + 1 + 2')

        neg = N('+', N('+', L(1), L(2)), L(3))
        self.assertEquals(generate_line(neg), '1 + 2 + 3')

        neg = N('+', L(1), N('+', L(1), L(2)))
        self.assertEquals(generate_line(neg), '1 + (1 + 2)')

        neg = N('+', L(1), -N('+', L(1), L(2)))
        self.assertEquals(generate_line(neg), '1 - (1 + 2)')

        neg = N('+', L(1), N('+', L(1), -L(2)))
        self.assertEquals(generate_line(neg), '1 + (1 - 2)')

        neg = -N('*', L(4), L('a'))
        self.assertEquals(generate_line(neg), '-4a')

        neg = N('*', L(4), -L('a'))
        self.assertEquals(generate_line(neg), '4 * -a')

        neg = -N('*', L(4), L(5))
        self.assertEquals(generate_line(neg), '-4 * 5')

        plus = N('+', L(1), -N('*', L(4), L(5)))
        self.assertEquals(generate_line(plus), '1 - 4 * 5')

        plus = N('+', L(1), -L(4))
        self.assertEquals(generate_line(plus), '1 - 4')

        plus = N('+', N('/', L('a'), L('b')), -N('/', L('c'), L('d')))
        self.assertEquals(generate_line(plus), 'a / b - c / d')

        mul = N('*', N('+', L('a'), L('b')), -N('+', L('c'), L('d')))
        self.assertEquals(generate_line(mul), '(a + b) * -(c + d)')

    def test_double_negation(self):
        neg = --L(1)
        self.assertEquals(generate_line(neg), '--1')
        neg = --N('*', L('x'), L(2))
        self.assertEquals(generate_line(neg), '--x * 2')
        neg = --N('^', L('x'), L(2))
        self.assertEquals(generate_line(neg), '--x ^ 2')

    def test_divide_fractions(self):
        a, b, c, d = L('a'), L('b'), L('c'), L('d')
        div = N('/', a, N('/', b, c))
        self.assertEquals(generate_line(div), 'a / (b / c)')
        div = N('/', N('/', a, b), N('/', c, d))
        self.assertEquals(generate_line(div), 'a / b / (c / d)')

    def test_prime(self):
        a, b, c, d = L('a'), L('b'), L('c'), L('d')
        root = N('*', a, N("'", b))
        self.assertEquals(generate_line(root), "a b'")
        root = N("'", -a)
        self.assertEquals(generate_line(root), "-a'")
        root = -N("'", a)
        self.assertEquals(generate_line(root), "-(a')")
        root = N("'", N('*', a, b))
        self.assertEquals(generate_line(root), "(ab)'")
        root = N("'", N('/', a, b))

    def test_function(self):
        root = N('sin', L('x'))
        self.assertEquals(generate_line(root), 'sin x')
        root = N('sin', N('+', L('x'), L(2)))
        self.assertEquals(generate_line(root), 'sin(x + 2)')
        root = N('dummyfunc', L('x'), L(2))
        self.assertEquals(generate_line(root), 'dummyfunc(x, 2)')

    def test_no_spacing(self):
        root = N('+', L('x'), L(2), no_spacing=True)
        self.assertEquals(generate_line(root), 'x+2')

    def test_explicit_parentheses(self):
        root = N('[]', L('x'))
        self.assertEquals(generate_line(root), '[x]')
        root = N('()', L('x'))
        self.assertEquals(generate_line(root), '(x)')
        root = N('{}', L('x'))
        self.assertEquals(generate_line(root), '{x}')

        root = N('^', N('[]', N('+', L('x'), L('y'))), L(2))
        self.assertEquals(generate_line(root), '[x + y] ^ 2')

    def test_abs(self):
        root = N('||', L('x'))
        self.assertEquals(generate_line(root), '|x|')
        root = N('||', N('+', L('x'), L(1)))
        self.assertEquals(generate_line(root), '|x + 1|')
        root = N('ln', N('||', L('x')))
        self.assertEquals(generate_line(root), 'ln|x|')

    def test_postprocess_str(self):
        root = N('int', N('^', L('x'), L(2)), L('x'))
        root.arity = lambda: 1
        root.postprocess_str = lambda s: s + ' dx'
        self.assertEquals(generate_line(root), 'int x ^ 2 dx')

    def test_concat_with_negation(self):
        root = N('*', -L(2), L('x'))
        self.assertEquals(generate_line(root), '(-2)x')
        root = N('*', N('*', L(3), -L(2)), L('x'))
        self.assertEquals(generate_line(root), '3 * -2x')
        root = N('*', L(3), -L(2), L('x'))
        self.assertEquals(generate_line(root), '3 * -2 * x')

    def test_first_child_negation(self):
        root = N('*', -L(1), L(2))
        self.assertEquals(generate_line(root), '(-1)2')
        root = -N('*', L(1), L(2))
        self.assertEquals(generate_line(root), '-1 * 2')
        root = N('/', -L(1), L(2))
        self.assertEquals(generate_line(root), '(-1) / 2')
        root = -N('/', L(1), L(2))
        self.assertEquals(generate_line(root), '-1 / 2')

    def test_postfix_brackets(self):
        root = N('*', L('x'), N("'", N('[]', N('^', L('x'), L(2)))))
        self.assertEquals(generate_line(root), "x[x ^ 2]'")

    def test_custom_line(self):
        root = N('*', L(1), L(2))
        root.custom_line = lambda: 'test'
        self.assertEquals(generate_line(root), 'test')

    def test_preprocess_str_exp(self):
        root = N('-', L(1))
        def addbrackets(self): self[0] = N('[]', self[0])
        root.preprocess_str_exp = new.instancemethod(addbrackets, root)
        self.assertEquals(generate_line(root), '-[1]')

    def test_division_by_multiplication(self):
        root = N('/', L(1), N('*', L('a'), L('b')))
        self.assertEquals(generate_line(root), '1 / (ab)')

        root = N('/', L(1), -N('*', L('a'), L('b')))
        self.assertEquals(generate_line(root), '1 / (-ab)')

        root = N('/', L(1), N('*', -L('a'), L('b')))
        self.assertEquals(generate_line(root), '1 / ((-a)b)')