'''polymod.py  The class PolyMod manages polynomial expressions in one
variable, mod m.  
'''

# revisions:
#  3 param __pow__ allows 3 param pow with polynomial modulus  11/28/10
#  sparse constructor

from mod import Mod, ParamClass
from bases import intToBaseList, baseListToInt

VAR = 'x'

def PMod(m):
   '''Return a function that makes PolyMod objects for a particular modulus.
   Here m may be an integer, Mod or PolyMod.'''
   def PolyM(coef):
      return PolyMod(coef, m)
   return PolyM

class PolyMod(ParamClass):
    '''Class of polynomial expressions in one variable with
    coefficients in mod m. Binary operations +, - or *
    with another polynomial or number; / with numbers
    (if there is an inverse), ** with nonnegative integers. 

    Polynomials can be constructed from a string expression in normal
    math notation, with -, +, *, except with exponents preceded by ^ or **.

    >>> PM7 = PMod(7)
    >>> F = PM7([2, -1, 3])
    >>> F
    PolyMod('3x^2 + 6x + 2', 7)
    >>> print(F(5))
    2 mod 7
    >>> G = PM7('x')
    >>> print(G)
    x
    >>> G = PM7('-x')
    >>> print(G)
    6x
    >>> G = PM7('x + 3')
    >>> print(G)
    x + 3
    >>> G = PM7('-x^3 + 2(-x + 3)')
    >>> print(G)
    6x^3 + 5x + 6
    >>> H = PM7('2x - 1')
    >>> print(F + G)
    6x^3 + 3x^2 + 4x + 1
    >>> print(F*H)
    6x^3 + 2x^2 + 5x + 5
    >>> print(F + G/3)   
    2x^3 + 3x^2 + 3x + 4
    >>> print(H(F))
    6x^2 + 5x + 3
    >>> print(H(F(-1)))
    4 mod 7
    >>> print(H**3)
    x^3 + 2x^2 + 6x + 6

    Many polynomials with the same modulus can be created by generating
    a function with PMod(m).
    '''

    # invariant: fields coef, m
    #    coef is a list of coefficients mod m 
    #    The zero polynomial has just 0 as coefficient. 
    #    Otherwise include terms for degree 0 through the polynomial degree.
            
    def __init__(self, coef, m=None):
        ''' Allow coef to be:
          * A list of coefficients from degree zero through the highest degree.
          * A list of tuples (c, p) for term c*x**p, for a sparse polynomial.
          * A single constant
          * A string expression of integers, 'x' amd symbols in '+-*^()',
            allowing implict multiplication, like '3x^2 - 5x + 7'.
          * A polynomial to copy with the same modulus
        Allow m to be:
          integer modulus
          PolyMod, providing modulus
          None, in which case the modulus is inferred from coef
        Integer coefficients are converted to the modulus for the class.
        '''
        if isinstance(m, PolyMod):
            m = m.m  
        if isinstance(coef, str):
            assert isinstance(m, (int, long)) and m > 1, 'Bad modulus value'
            expr = ' '.join(explicitArithmeticTokens(coef))
            expr = expr.replace(VAR, 'PolyMod([0, 1], {0})'.format(m))
            coef = eval(expr) # result may be an int or PolyMod
        if isinstance(coef, PolyMod): # allow shallow copy
            if m is None:
               m = coef.m
            else: assert m == coef.m, 'Different moduli'
            self.coef = coef.coef
            self.m = m
            return
        if not isinstance(coef, list):
            coef = [coef]
        if not coef: # small point: allow empty list to mean 0
            coef = [0]
        if isinstance(coef[0], tuple):
           coef = sparse(coef)
        if isinstance(coef[0], Mod):
            if m is None:
                m = coef[0].m
            else: assert m == coef[0].m, 'Different moduli'
        assert isinstance(m, (int, long)) and m > 1, 'Bad modulus value'
        self.m = m
        self.coef = [Mod(c, self.m) for c in coef]
        while len(self.coef) > 1 and self.coef[-1] == 0: # remove high degree 0's
            self.coef.pop()

    def modulus(self):
        '''Return the modulus of coefficients that is set for the class.'''
        return self.m

    def __nonzero__(self):
        return self.coef[-1] != 0 # only 0 has has highest degree coef 0

    ### basic operations producing another PolyMod, PolyMod first
    def __add__(self, other):
        '''+ operation with another polynomial or number.'''
        other = self.tryLike(other)
        if other is None: return NotImplemented
        if len(self.coef) < len(other.coef):
            self, other = other, self
        coef = self.coef[:] # copy coef list for longer polynomial
        for (i, val) in enumerate(other.coef): # add corresp coef
            coef[i] += val
        return self.like(coef)

    def __sub__(self, other):
        '''- operation with another polynomial or number.'''
        other = self.tryLike(other)
        if other is None: return NotImplemented
        return self + other*-1

    def __mul__(self, other):
        '''* operation with another polynomial or number.'''
        other = self.tryLike(other)
        if other is None: return NotImplemented
        ds = self.degree()
        if ds == -1:# 0 poly
            return self
        do = other.degree()
        if do == -1: # 0 poly
            return other      
        coef = [Mod(0,self.m)] * (ds + do + 1) # space for new coef list
        for (i, sc) in enumerate(self.coef): # basic long multiplication
            for (j, oc) in enumerate(other.coef):
                coef[i+j] += sc*oc
        return PolyMod(coef)

    def __div__(self, invertible):
        '''/ operation with an invertible constant value.  '''
        invertible = self.coef[0].tryLike(invertible)
        if invertible is None: return NotImplemented
        return self*invertible.inverse()
        
    def __neg__(self):
        '''Negation operator -'''
        return self * -1

    def __pow__(self, n, modulus=None): 
        ''' power operator ** and pow function;
        n must evaluate to a non-negative integer.
        The pow function allows the modulus polynomial.
        '''
        if not isinstance(n, (int, long)) or n < 0:
            return NotImplemented
        result = PolyMod(1, self.m)
        s = self  # s holds the current square
        while n > 0:
           if n % 2 == 1: 
              result = s * result
              if modulus is not None:
                 result %= modulus
           s = s * s  # compute the next square
           if modulus is not None:
              s %= modulus
           n = n/2    # compute the next quotient
        return result        

    ### basic operations producing another Polynomial, with other operand first
    def __radd__(self, other):
        '''+ operation with Polynomial second.'''
        return self + other

    def __rsub__(self, other):
        '''- operator with the Polynomial second.'''
        return self*-1 + other
    
    def __rmul__(self, other):
        '''* operator with the Polynomial second.'''
        return self * other
    # no __rdiv__
    
    ### string conversions
    
    def __str__(self):
        '''Automatic conversion to an informal simplified string.
        '''
        if not self:
            return '0'  # only time a 0 term is printed
        p = self.degree()
        parts = [] 
        for i in range(p, -1, -1): # high degree to low
            val = int(self.coef[i])
            if val:    # avoid unneeded 0's and 1's, addition of negative
                s = str(val)    
                if i != p :
                    parts.append('+')
                if i > 0:
                    if s == '1':
                        s = ''
                    s += VAR
                if i > 1:
                    s += '^' + str(i)
                parts.append(s)
        return ' '.join(parts)

    def __repr__(self):
        """Returns a string representation that would evaluate to an
        equal Polynomial if typed into the shell.
        """

##        if not hasattr(self, 'coef'):  # for debugging the creation process
##            return '<Polynomial init>'
        return "PolyMod('{0}', {1})".format(self, self.m)
    
    ### miscelaneous
    
    def __eq__(self, other):
        '''True only if the operands are identical as Polynomials.
        '''
        other = self.tryLike(other) # ! does make == operation NOT transitive!
        if other is None: return NotImplemented
        return self.coef == other.coef
        
    def __ne__(self, other):
        '''True only if the operands are not identical as Polynomials.
        '''
        return not self == other

    def __call__(self, value):
        '''Call the function with the variable replaced by value.
        If value is a number, the return value is a modular number.
        If value is a polynomial, the return value is a polynomial.
        '''
        if isinstance(value, (int, long, Mod)):
            value = Mod(value, self.m) 
        else:
            value = self.like(value)
        i = self.degree()
        tot = self.coef[i]
        for i in range(i-1, -1, -1):
            tot = tot * value + self.coef[i]
        return tot 

    def __divmod__(self, poly):
        '''Return (quotientPoly, remainderPoly)
        where self = quotientPoly*poly + remainderPoly
        and the degree of the remainder is less than for self.
        '''
        poly = self.tryLike(poly)
        if poly is None: 
            return NotImplemented
        degPoly = poly.degree()
        if degPoly == -1:
            raise ZeroDivisionError, 'PolyMod division by 0'
        zero = PolyMod([0], self)
        if degPoly == 0:  
            return (self*poly.coef[0].inverse(), zero)#may fail if not in field
        degQuot = self.degree() - degPoly       
        if degQuot < 0:
            return (zero, self)
        # basic long division follows
        qCoef = [None] * (degQuot+1) # locations for quotient coeficients 
        rCoef = self.coef[:] # will end up as remainder coefficients
        highCoefInv = poly.coef[-1].inverse() #may fail if not in field
        for i in range(degQuot, -1, -1):
            qc = rCoef[i+degPoly]*highCoefInv
            for j, c in enumerate(poly.coef):
                rCoef[i+j] -= c*qc
            qCoef[i] = qc
        return (self.like(qCoef), self.like(rCoef[:degPoly+1]))

    def __mod__(self, poly):
        ''' Return the remainder dividing self by poly, written self % poly.
        '''
        return divmod(self, poly)[1]

    def __hash__(self):
        ''' Hash value definition needed for use in dictionaries and sets.'''
        return hash(tuple([c.value for c in self.coef]))
        
    # public support functions
    
    def degree(self): 
        ''' Return the degree of the polynomial; -1 for the 0 polynomial.'''
        d = len(self.coef) - 1
        if d > 0 or self.coef[0] != 0:
            return d
        return -1

    def getCoef(self, degree):
        '''Return the coefficient for a specified power.'''
        if self.degree() >= degree >= 0:
            return self.coef[degree]
        return Mod(0, self.m)

    def __int__(self):
        '''canonical evaluation with ints for coef and modulus of self.'''
        return baseListToInt([c.value for c in self.coef], self.m)

    def likeFromInt(self, n):
        '''Convert an int code for coefficients to a PolyMod of same type.'''
        coef = intToBaseList(n, self.m)
        return PolyMod(coef, self)
            
    # support for like function
    def sameParam(self, val):
        '''True if val is also a PolyMod, and the moduli are the same'''
        return isinstance(val, PolyMod) and self.m == val.m
    
    ### End of class PolyMod ###########################################  

def sparse(pairs):
    '''Convert list of (c, i) to list with element i = c, rest 0.'''
    deg = max([pr[1] for pr in pairs])
    coef = [0]*(deg+1)
    for c, i in pairs:
        coef[i] = c
    return coef

def explicitArithmeticTokens(mathString):
    '''Return a list of tokens, converting a mathematical notation
    string. Allowed symbols:  single letter lowercase variable
    symbols, integers, and symbols in '()+-*/^,'.  Multication becomes
    explicit and exponentiation is changed from ^ to **. Example '3x +
    3(2x - 1/7)5^2' would produce ['3', '*', 'x', '+', '3', '*', '(',
    '2', '*', 'x', '-', '1', '/', '7', ')', '*', '5', '**', '2'].    
    '''
    ops = '^*-+/,'
    s = mathString.replace('**', '^')
    for sym in '()abcdefghijklmnopqrstuvwxyz' + ops:
        s = s.replace(sym, ' ' + sym + ' ')
    tokens = s.split()
    lastToken = ''
    newTokens = []
    for token in tokens:
        if lastToken not in ops + '(' and token not in ops + ')': 
            newTokens.append('*')
        if token == '^':
              newTokens.append('**')
        else:
              newTokens.append(token)
        lastToken = token
    return newTokens

### testing only below
    
def doctests():
    '''
    >>> PM7 = PMod(7)
    >>> F = PM7([2, -1, 3])
    >>> G = PM7('-x^3 + 2(-x + 3)')
    >>> H = PM7('2x - 1')
    >>> print(H(H))
    4x + 4
    >>> print(F - 2*G)
    2x^3 + 3x^2 + 3x + 4
    >>> X = PM7('x')
    >>> print(X)
    x
    >>> print(0*X)
    0
    >>> ZERO = PM7(0)
    >>> print(ZERO*G)
    0
    >>> print(-3*X)
    4x
    >>> print X**3
    x^3
    >>> print (2*X+3)**2
    4x^2 + 5x + 2
    >>> print X/2
    4x
    >>> G == 3*X + 5
    False
    >>> H == 1*H
    True
    >>> G != 3*X + 5
    True
    >>> H != 1*H
    False
    >>> G == 3
    False
    >>> TWO = PM7(2)
    >>> TWO == 2
    True
    >>> G != 3
    True
    >>> TWO != 2
    False
    >>> 3 == G
    False
    >>> 2 == TWO
    True
    >>> 3 != G
    True
    >>> 2 != TWO
    False
    >>> (X + 2)*H == H*2 - -H*X
    True
    >>> print(X(3))
    3 mod 7
    >>> print(H.getCoef(-1))
    0 mod 7
    >>> print(H.getCoef(0))
    6 mod 7
    >>> print(H.getCoef(1))
    2 mod 7
    >>> print(H.getCoef(2))
    0 mod 7
    >>> print(divmod(X**5, X**2))
    (PolyMod('x^3', 7), PolyMod('0', 7))
    >>> print(divmod(2*X**2 + 9*X - 5, X+5))
    (PolyMod('2x + 6', 7), PolyMod('0', 7))
    >>> print(divmod(X**4 - 1, X-1))
    (PolyMod('x^3 + x^2 + x + 1', 7), PolyMod('0', 7))
    >>> print(divmod(X**4 - 1, X**3 + X**2 + X + 1))
    (PolyMod('x + 6', 7), PolyMod('0', 7))
    >>> print(divmod(X**4 + X**2 - 2*X - 5, X**3 + X**2 + X + 1))
    (PolyMod('x + 6', 7), PolyMod('x^2 + 5x + 3', 7))
    >>> print((X**4 + X**2 - 2*X - 5) % (X**3 + X**2 + X + 1))
    x^2 + 5x + 3
    >>> print (G*H) % G
    0
    >>> Poly8 = PMod(8)
    >>> S = Poly8(X)
    Traceback (most recent call last):
    ...
    AssertionError: Different moduli
    >>> S = Poly8('x')
    >>> S
    PolyMod('x', 8)
    >>> print((3*S - 2)/5)
    7x + 6
    >>> print(6*S/6)
    Traceback (most recent call last):
    ...
    ValueError: Value not invertible
    >>> G.likeFromInt(int(H)) == H
    True
    >>> print PolyMod([(5, 0), (2, 1), (5, 7), (3, 12)], 7)
    3x^12 + 5x^7 + 2x + 5
    >>> H**14 % F == pow(H, 14, F)
    True
    '''
  
if __name__ == '__main__':
    import doctest
    doctest.testmod() #verbose=True)
