from __future__ import division
from scipy import *

class Polynomial:
    """
    This class represents a polynomial in one variable.
    The attributes of the class are

      coeff -- array. The coefficients of the polynomial. The
               coefficient at index i is associated with x**i.
      var   -- string. The name of the polynomial variable.
    """

    # constructor
    def __init__(self, coefficients, variable='x'):
        self.coeff = array(coefficients)
        self.var = variable

    # call method: p(x)
    def __call__(self, x):
        return dot(self.coeff, x ** arange(len(self.coeff)))

    # convert to a string: str(p)
    #                      print p
    def __str__(self):
        result = ''
        for i in range(len(self.coeff)):
            if self.coeff[i] == 0:
                continue
            if i == 0:
                result += str(self.coeff[i]) + ' + '
            elif i == 1:
                result += str(self.coeff[i]) + self.var + ' + '
            else:
                result += str(self.coeff[i]) + self.var + '^' + str(i) + ' + '

        if result == '':
            return '0'
        else:
            return result[:-3]

    # add two polynomials: p + q
    def __add__(self, polynomial):
        if self.var != polynomial.var:
            raise ValueError, 'polynomial variables are different'
        result = zeros(max(len(self.coeff), len(polynomial.coeff)))
        result[:len(self.coeff)] += self.coeff
        result[:len(polynomial.coeff)] += polynomial.coeff
        return Polynomial(result, self.var)

    # negate a polynomial: -p
    def __neg__(self):
        return Polynomial(-self.coeff, self.var)

    # subtract two polynomials: p - q
    def __sub__(self, polynomial):
        return self + (-polynomial)

    # test of two polynomials are equal: p == q
    def __eq__(self, polynomial):
        return (self.coeff == polynomial.coeff) and (self.var == polynomial.var)

    # test of one polynomial is greater than another: p > q
    def __gt__(self, polynomial):
        raise TypeError, "polynomials cannot be compared"

    def degree(self):
        """
        Returns the degree of the polynomial.
        """
        return len(self.coeff)-1


p = Polynomial([1.5, 2.1, 3, 1, 0, 0, 5], 'x') # 1 + 2x + 3x^2 + x^3 + 5x^6
q = Polynomial([0, 0, 1], 'x') # x^2

print 'p(2.5) =', p(2.5)
print 'q(2.5) =', q(2.5)

print 'p =', p
print 'q =', q
print 'p-q =', p - q