#!/usr/bin/python
"""Symbolic differentiation.

Just what I managed to livecode in half an hour as an example of
doctest-driven development.

(Then another few minutes later.)

>>> x = Var('x')
>>> x
x
>>> x.diff(x)
1
>>> y = Var('y')
>>> x.diff(y)
0
>>> (x*x).diff(x)
(2 * x)
>>> (x+x).diff(x)
2
>>> Constant(0) == Constant(0)
True
>>> (x*y).diff(x)
y
>>> x.subst({x: Constant(3)})
3
>>> (x * Constant(2)).subst({x: Constant(3)})
6
>>> (x * x + Constant(2)).subst({x: Constant(5)})
27
>>> (x * x).subst({x: Var('a') + Var('b')})
((a + b) * (a + b))

"""
import doctest

class Formula:
    def __mul__(self, other):
        return Product(self, other).simplify()

    def __add__(self, other):
        return Sum(self, other).simplify()

    def simplify(self):
        return self

class Constant(Formula):
    def __init__(self, value):
        self._value = value
        
    def __repr__(self):
        return repr(self._value)

    def __eq__(self, other):
        return isinstance(other, Constant) and self._value == other._value

    def diff(self, independent_variable):
        return Constant(0)

    def subst(self, replacements):
	return self

class Product(Formula):
    def __init__(self, a, b):
        self._a = a
        self._b = b

    def diff(self, independent_variable):
        return (self._a.diff(independent_variable) * self._b
                + self._a * self._b.diff(independent_variable))

    def __repr__(self):
        return "(%s * %s)" % (repr(self._a), repr(self._b))

    def simplify(self):
        if isinstance(self._b, Constant) and not isinstance(self._a, Constant):
            return self._b * self._a
        elif (isinstance(self._a, Constant)
              and isinstance(self._b, Constant)):
            return Constant(self._a._value * self._b._value)
        elif self._a == Constant(0):
            return Constant(0)
        elif self._a == Constant(1):
            return self._b
        return self

    def subst(self, replacements):
	return self._a.subst(replacements) * self._b.subst(replacements)

class Sum(Formula):
    def __init__(self, a, b):
        self._a = a
        self._b = b

    def __repr__(self):
        return "(%s + %s)" % (repr(self._a), repr(self._b))

    def diff(self, independent_variable):
        return self._a.diff(independent_variable) + self._b.diff(independent_variable)

    def simplify(self):
        if (isinstance(self._a, Product)
            and isinstance(self._b, Product)
            and self._a._b is self._b._b):
            return (self._a._a + self._b._a) * self._a._b
        elif (isinstance(self._a, Constant)
              and isinstance(self._b, Constant)):
            return Constant(self._a._value + self._b._value)
        elif self._a == Constant(0):
            return self._b
        elif self._b == Constant(0):
            return self._a
        elif self._a == self._b:
            return Constant(2) * self._a
        return self

    def subst(self, replacements):
	return self._a.subst(replacements) + self._b.subst(replacements)

class Var(Formula):
    def __init__(self, name):
        """
        
        Arguments:
        - `name`: name of the variable
        """
        self._name = name

    def __repr__(self):
        return self._name

    def diff(self, independent_variable):
        return Constant(1) if self is independent_variable else Constant(0)

    def subst(self, replacements):
	return replacements.get(self, self)
    

if __name__ == '__main__':
    doctest.testmod()