#! /usr/bin/env python

"""
File: Polynomial_vec.py

Copyright (c) 2016 Taylor Patti

License: MIT

This class provides a vectorized version of the example class
Polynomial from the Springer textbook webpage.

"""

import numpy

class Polynomial(object):
    """Initializes a polynomial object with a list of coefficients."""
    def __init__(self, coefficients):
        self.coeff = numpy.zeros(len(coefficients))
        for i in xrange(len(coefficients)):
            self.coeff[i] = coefficients[i]

    def __call__(self, x):
        """Evaluate the polynomial."""
        s = 0
        x_array = numpy.zeros(len(self.coeff))
        for i in xrange(len(x_array)):
            x_array[i] = x**i
        return numpy.dot(self.coeff, x_array)

    def __add__(self, other):
        """Return self + other as Polynomial object."""
        if len(self.coeff) > len(other.coeff):
            val = len(other.coeff)
            marker = 'self'
        else:
            val = len(self.coeff)
            marker = 'other'
        result_coeff = numpy.add(self.coeff[0:val], other.coeff[0:val])
        if marker == 'self':
            result_coeff = numpy.append(result_coeff, self.coeff[val : ])
        else:
            result_coeff = numpy.append(result_coeff, other.coeff[val : ])
        return Polynomial(result_coeff)

    def __mul__(self, other):
        """Multiplies the polynomial by another polynomial."""
        c = self.coeff
        d = other.coeff
        M = len(c) - 1
        N = len(d) - 1
        result_coeff = numpy.zeros(M+N+1)
        for i in range(0, M+1):
            for j in range(0, N+1):
                result_coeff[i+j] += c[i]*d[j]
        return Polynomial(result_coeff)

    def differentiate(self):
        """Differentiate this polynomial in-place."""
        n = len(self.coeff)
        self.coeff[:-1] = numpy.linspace(1, n-1, n-1)*self.coeff[1:]
        self.coeff = self.coeff[:-1]

    def derivative(self):
        """Copy this polynomial and return its derivative."""
        dpdx = Polynomial(self.coeff[:])  # make a copy
        dpdx.differentiate()
        return dpdx

    def __str__(self):
        """Expresses the polynomial as a string."""
        s = ''
        for i in range(0, len(self.coeff)):
            if self.coeff[i] != 0:
                s += ' + %g*x^%d' % (self.coeff[i], i)
        # Fix layout
        s = s.replace('+ -', '- ')
        s = s.replace('x^0', '1')
        s = s.replace(' 1*', ' ')
        s = s.replace('x^1 ', 'x ')
        #s = s.replace('x^1', 'x') # will replace x^100 by x^00
        if s[0:3] == ' + ':  # remove initial +
            s = s[3:]
        if s[0:3] == ' - ':  # fix spaces for initial -
            s = '-' + s[3:]
        return s

    def simplestr(self):
        """Expresses the polynomial as an even simpler string."""
        s = ''
        for i in range(0, len(self.coeff)):
            s += ' + %g*x^%d' % (self.coeff[i], i)
        return s


def test_Polynomial():
    """Ensures that all of the polynomial functions work properly by comparing the function
    result to the result which you would get by preforming these calculations by hand or in a separate
    computation of this kind."""
    p1 = Polynomial([1, -1])
    p2 = Polynomial([0, 1, 0, 0, -6, -1])
    p3 = p1.__add__(p2)
    p3_exact = Polynomial([1, 0, 0, 0, -6, -1])
    msg = 'p1 = %s, p2 = %s\np3=p1+p2 = %s\nbut wrong p3 = %s'%\
          (p1.coeff, p2.coeff, p3_exact.coeff, p3.coeff)
    assert numpy.allclose(p3.coeff, p3_exact.coeff, rtol=1E-14), msg
    # Note __add__ applies lists only, here with integers, so
    # == for comparing lists is not subject to round-off errors

    p4 = p1.__mul__(p2)
    # p4.coeff becomes a numpy array, see __mul__
    p4_exact = Polynomial([0, 1, -1,  0, -6,  5,  1])
    msg = 'p1 = %s, p2 = %s\np4=p1*p2 = %s\ngot wrong p4 = %s'%\
          (p1, p2, p4_exact, p4)
    assert numpy.allclose(p4.coeff, p4_exact.coeff), msg

    p5 = p2.derivative()
    p5_exact = Polynomial([1, 0, 0, -24, -5])
    msg = 'p2 = %s\np5 = p2.derivative() = %s\ngot wrong p5 = %s'%\
          (p2, p5_exact, p5)
    assert numpy.allclose(p5.coeff, p5_exact.coeff), msg

    p6 = Polynomial([0, 1, 0, 0, -6, -1])  # p2
    p6.differentiate()
    p6_exact = p5_exact
    msg = 'p6 = %s\p6.differentiate() = %s\ngot wrong p6 = %s'%\
          (p2, p6_exact, p6)
    assert numpy.allclose(p6.coeff, p6_exact.coeff), msg

if __name__ == '__main__':
    import sys
    if len(sys.argv) >= 2 and sys.argv[1] == 'verify':
        test_Polynomial()