##########################################################################
#We implement a class in Sage for piecewise polynomial functions, using a
#newly developed canonical form.
#
#
#
#
#
#
#
#
#First of all, some tools
##########################################################################

from sage.rings.polynomial.real_roots import *


def IsGreaterRoot(P,i,intP,Q,j,intQ):
	"""
	True if the only root in intP is grater than the only root in intQ 
	and isolating intervals
	
	THIS IS A SUBROUTINE FOR THE MAIN METHODS OF THE PACKAGE, WHERE
	PRECONDITIONS ARE CONTROLLED. USE IT ALONE AT YOUR OWN RISK.
	"""
#	print [P,i,intP,Q,j,intQ]
	if intP[0]>=intQ[1]:
		return [True,intP,intQ]
	if intQ[0]>=intP[1]:
		return [False,intP,intQ]
	if intP[0]==intP[1]:
		if sign(Q.__call__(intP[0]))*sign(Q.__call__(intQ[1]))>0:
			return [True,intP,[intQ[0],intP[0]]]
		else:
			return [False, intP, [intP[0],intQ[1]]]
	intersInterv=[max(intP[0],intQ[0]),min(intP[1],intQ[1])]
	if sign(P.__call__(intersInterv[0]))*sign(P.__call__(intersInterv[1]))>0:
		if sign(P.__call__(intersInterv[1]))==sign(P.__call__(intP[1])):
		#This means that the root of P is to the left of the intersection interval, so is lesser
			return [False ,[intP[0],intersInterv[0]],intQ]
		else:
		#This means that the root of P is to the right of the intersection interval, so is greater
			return [True, [intersInterv[1],intP[1]],intQ]
	while not Q.__call__(intersInterv[0])*Q.__call__(intersInterv[1])>0:
	#Now we will always keep the root of P in the shrinking intersInterv
	#It is easy to see that, if the interval has length=0, we will not enter this loop
		midpt=(intersInterv[0]+intersInterv[1])/2
		if sign(P.__call__(intersInterv[1]))==sign(P.__call__(midpt)):
		#the root is in the first half of the interval
			intersInterv[1]=midpt
		else:
			intersInterv[0]=midpt
	if sign(Q.__call__(intersInterv[1]))==sign(Q.__call__(intQ[1])):
	#This means that the root of Q is to the left of the intersection interval, so is lesser
		return [True, intersInterv, [intQ[0],intersInterv[0]] ]
	else:
	#This means that the root of Q is to the right of the intersection interval, so is greater
		return [False, intersInterv, [intersInterv[1],intQ[1]] ]
	
		



##########################################################################
#Now, the classes
#
##########################################################################


class CanonicalPiecewisePolynomialFunction():
	"""
	Piecewise polynomial functions with rational coefficients
	in canonical form.
	"""
	def __init__(self,v):
		if not v.is_symbol():
			print 'Argument is supposed to be a variable'
			return
		self.var=v
		self.ring=PolynomialRing(QQ,v)
		self.Polys=[]
		self.indices=[]
		self.coeffs=[]
		self.interv=[]
		self.FullPol=self.ring(0)

		
	def __repr__(self):
		result=self.FullPol.__repr__()
		for i in range(len(self.Polys)):
			result+='\n+('+self.coeffs[i].__repr__()+')*C_'+str(self.indices[i])+'('+self.Polys[i].__repr__()+')'
		return result

	def __add__(self,other):
		if other in self.ring:
			result=CanonicalPiecewisePolynomialFunction(self.var)
			p=self.ring(other)
			result.FullPol=self.FullPol+p
			result.Polys=deepcopy(self.Polys)
			result.indices=deepcopy(self.indices)
			result.coeffs=deepcopy(self.coeffs)
			result.interv=deepcopy(self.interv)
			return result
		if not self.ring is other.ring:
			print 'Semipolynomials on different variables'
			return
#		result=deepcopy(self)
		result=CanonicalPiecewisePolynomialFunction(self.var)
		result.Polys=deepcopy(self.Polys)
		result.indices=deepcopy(self.indices)
		result.coeffs=deepcopy(self.coeffs)
		result.interv=deepcopy(self.interv)
		result.FullPol=self.FullPol+other.FullPol
		for i in range(len(other.Polys)):
			j=0
			contin=(j<len(self.indices))
			while contin:
				#print 'i,j==',i,j
				#print self.Polys,self.indices,other.Polys,other.indices
				if [self.Polys[j],self.indices[j]]==[other.Polys[i],other.indices[i]]:
					result.coeffs[j]+=other.coeffs[i]
					contin=False
				else:
					j+=1
					if j==len(self.indices):
						contin=False
			if j==len(self.indices):
				result.Polys.append(other.Polys[i])
				result.indices.append(other.indices[i])
				result.coeffs.append(other.coeffs[i])
				result.interv.append(other.interv[i])
		i=0
		while i<len(result.indices):
			if result.coeffs[i]==0:
				result.Polys.pop(i)
				result.indices.pop(i)
				result.coeffs.pop(i)
				result.interv.pop(i)
			else:
				i+=1
		return result.sort()

	def __radd__(self,other):
		if other in self.ring:
			return self+other



	def __mul__(self,other):
		if other in self.ring:
			if other==0:
				return CanonicalPiecewisePolynomialFunction(self.var) 
#			result=deepcopy(self)
			p=self.ring(other)
			result=CanonicalPiecewisePolynomialFunction(self.var)
			result.Polys=deepcopy(self.Polys)
			result.indices=deepcopy(self.indices)
			result.coeffs=deepcopy(self.coeffs)
			result.interv=deepcopy(self.interv)
			result.FullPol=self.FullPol*p
			for i in range(len(self.Polys)):
				result.coeffs[i]*=p
			return result
		if not self.ring is other.ring:
			print 'Semipolynomials on different variables'
			return
#		result=CanonicalPiecewisePolynomialFunction(self.var)
#		result.FullPol=self.FullPol*other.FullPol
		result=other*self.FullPol+self*other.FullPol+(-self.FullPol*other.FullPol)
#		if not other.FullPol==0:
#			for i in range(len(self.Polys)):
#				result.Polys.append(self.Polys[i])
#				result.indices.append(self.indices[i])
#				result.interv.append(self.interv[i])
#				result.coeffs.append(self.coeffs[i]*other.FullPol)
#		if not self.FullPol==0:
#			for j in range(len(other.Polys)):
#				result.Polys.append(other.Polys[j])
#				result.indices.append(other.indices[j])
#				result.interv.append(other.interv[j])
#				result.coeffs.append(other.coeffs[j]*self.FullPol)
#		print 'bucle'
		for i in range(len(self.Polys)):
			for j in range(len(other.Polys)):
				#print 'i,j=',i,j
				if self.Polys[i]==other.Polys[j]:
#					print 'caso0'
					aux=CiP(min(self.indices[i],other.indices[j]),self.Polys[i],self.var)
#					print '1'
					aux*=other.Polys[j]*other.coeffs[j]*self.coeffs[i]
#					print '2',aux,result
#					aux.coeffs[0]*=other.Polys[j]*other.coeffs[j]*self.coeffs[i]
					result+=aux
#					print '3'
				elif IsGreaterRoot(self.Polys[i],self.indices[i],self.interv[i],other.Polys[j],other.indices[j],other.interv[j])[0]:
#					print 'caso 1'
					aux=CiP(self.indices[i],self.Polys[i],self.var)
					aux*=other.Polys[j]*other.coeffs[j]*self.coeffs[i]
#					aux.coeffs[0]*=other.Polys[j]*other.coeffs[j]*self.coeffs[i]
					result+=aux
				else:
#					print 'caso 2'
					aux=CiP(other.indices[j],other.Polys[j],self.var)
					aux*=self.Polys[i]*self.coeffs[i]*other.coeffs[j]
#					aux.coeffs[0]*=self.Polys[i]*self.coeffs[i]*other.coeffs[j]
					result+=aux
		return result

	def __rmul__(self,other):
		if other in self.ring:
			return self*other

	def __neg__(self):
		return self*(-1)
		

	
	def sort(self):
		"""
		This method puts in order the structure of self. 
		It does not detect when two pairs (plynomial,number of root) are equal. 
		"""
#		result=deepcopy(self)
		i=0
		while i<len(self.Polys)-1:
			j=i+1
			while j<len(self.Polys):
				thing=IsGreaterRoot(self.Polys[i],self.indices[i],self.interv[i],self.Polys[j],self.indices[j],self.interv[j])
				if thing[0]:
					self.Polys[i],self.Polys[j]=self.Polys[j],self.Polys[i]
					self.indices[i],self.indices[j]=self.indices[j],self.indices[i]
					self.coeffs[i],self.coeffs[j]=self.coeffs[j],self.coeffs[i]
					self.interv[i],self.interv[j]=deepcopy(thing[2]),deepcopy(thing[1])
				else:
					self.interv[i],self.interv[j]=deepcopy(thing[1]),deepcopy(thing[2])
				j+=1
			i+=1
		return self	



	def of(self,other):
		"""
		Composition of piecewise polynomial functions
		"""
		result=CanonicalPiecewisePolynomialFunction(self.var)
		if self.FullPol in QQ:
			result=self.FullPol
		else:
			c=self.FullPol.coeffs()
			aux=deepcopy(other)
			result=c[0]+c[1]*aux
			for i in range(2,len(c)):
				aux*=other
#				print result, aux
				result+=c[i]*aux
		#We did what we needed with the full polynomial. Now we proceed with the C_i(P)
		#We will take advantage of the fact that all roots are sorted (thanks to the way we 
		#sum these functions)
		for i in range(len(self.indices)):
			aux=deepcopy(other)
			c=self.Polys[i].coeffs()
			CompWithPol=c[0]+c[1]*aux
			for j in range(2,len(c)):
				aux*=other
				CompWithPol+=c[j]*aux
			#Now we have composed with self.Poly[i]. Next step is computing the roots 
			#of such composition to compare their image by other with the
			#self.indices[i]th root of self.Polys[i]
			#result+=ComposeWithCiP(self.indices[i],self.interv[i],self.Polys[i],other)
#			print CompWithPol
			RootData=CompWithPol.real_roots()
#			print RootData
#			print CompWithPol
			if len(RootData[0])==0:
				a=0
			else:
				a=RootData[0][0][0]-1
			value=other.EvalCPWPF(a)
#			print a,'value=',value.n()
#			print self.var-value,1,[value,value],self.Polys[i],self.indices[i],self.interv[i],IsGreaterRoot(self.var-value,1,[value,value],self.Polys[i],self.indices[i],self.interv[i])
			if IsGreaterRoot(self.var-value,1,[value,value],self.Polys[i],self.indices[i],self.interv[i])[0]:
				is0=[False]
			else:
				is0=[True]
			for l in range(len(RootData[0])):
				if l==len(RootData[0])-1:
					a=RootData[0][l][1]+1
				else:
					a=(RootData[0][l][1]+RootData[0][l+1][0])/2
				value=other.EvalCPWPF(a)
#				print a,'value=',value.n()
				if IsGreaterRoot(self.var-value,1,[value,value],self.Polys[i],self.indices[i],self.interv[i])[0]:
					is0.append(False)
				else:
					is0.append(True)
#			print is0
			if is0[0]:
				resultadillo=0
			else:
				resultadillo=CompWithPol*self.coeffs[i]
#			print 'res',resultadillo
#			print 'com',CompWithPol
			for l in range(len(is0)-1):
				if not is0[l]==is0[l+1]:
					if is0[l]:#We are passing from 0 to nonzero
						auxPWF=deepcopy(CompWithPol)*CiP(RootData[2][l],RootData[1][l],self.var)*self.coeffs[i]
#						print 'aux',auxPWF
#						print 'RD',RootData[1][l]
						for k in range(len(auxPWF.coeffs)):
							auxPWF.coeffs[k]=auxPWF.coeffs[k].quo_rem(RootData[1][l])[0]
#						print 'aux',auxPWF
#						print 'res',resultadillo
						resultadillo+=auxPWF
					else: #We pass form nonzero to 0
						auxPWF=deepcopy(CompWithPol)*CiP(RootData[2][l],RootData[1][l],self.var)*self.coeffs[i]
#						print 'aux',auxPWF
#						print 'RD',RootData[1][l]
						for k in range(len(auxPWF.coeffs)):
							auxPWF.coeffs[k]=auxPWF.coeffs[k].quo_rem(-RootData[1][l])[0]
							#Observe the sign "-"
#						print 'aux',auxPWF
#						print 'res',resultadillo
						resultadillo+=auxPWF
#			print resultadillo
			result+=resultadillo
		return result



	def real_roots(self):
		intervals=[]
		PolyVanishing=[]
		NumberOfRealRoot=[]
		if self.FullPol==0:
			if len(self.Polys)==0:
				print 'All the line vanishes, it is the zero function.'
				return
			#if a whole segment is a root, we need a trick
			intervals.append([self.interv[0][0]-1,self.interv[0][0]])
			PolyVanishing.append(self.ring(0))
			NumberOfRealRoot.append(1)
		else:
			fact=self.FullPol.factor()[:]
			intervals.extend(real_roots(self.FullPol))
			for i in range(len(intervals)):
				intervals[i]=deepcopy(intervals[i][0])
			NumberOfCurrentRealRootForPoly=[0 for i in range(len(fact))]
			for i in range(len(intervals)):
				for j in range(len(fact)):
					if not sign(fact[j][0].subs(intervals[i][0]))==sign(fact[j][0].subs(intervals[i][1])):
						PolyVanishing.append(fact[j][0])
						NumberOfCurrentRealRootForPoly[j]+=1
						NumberOfRealRoot.append(NumberOfCurrentRealRootForPoly[j])
						break
			#Now we have factors for the full polynomial part
#		print 'empezamos'
#		print intervals,PolyVanishing
		j=0
		auxPol=self.FullPol
		for i in range(len(self.Polys)):
#			print i,'arriba',intervals			
			while j<len(intervals):
				if self.Polys[i].divides(PolyVanishing[j]) and NumberOfRealRoot[j]==self.indices[i]:
					condition=True
#					interval=[max(intervals[j][0][0],self.interv[i][0]),min(intervals[j][0][1],self.interv[i][1])]
				else:
					chain=IsGreaterRoot(PolyVanishing[j],NumberOfRealRoot[j],intervals[j],self.Polys[i],self.indices[i],self.interv[i])
					condition=chain[0]
#					interval=chain[2]
				if condition:
					for l in range(j,len(intervals)):
						intervals.pop(j)
						PolyVanishing.pop(j)
						NumberOfRealRoot.pop(j)
				else:
					intervals[j]=deepcopy(chain[1])
					j+=1
			#We now have removed all the roots that were too high
			auxPol+=self.coeffs[i]*self.Polys[i]
			#if i==len(self.Polys-1):
			#	sup=root_bounds(auxPol)
			#else:
			#	sup=self.interv[i+1][1]
			#This has been removed because we need to compute all real roots in 
			#order to assign its true order as roots of the irreducible components
			if auxPol==0:
				#if a whole segment is root, we need a trick
				if i==len(self.Polys)-1:
					sup=self.interv[i][1]+2
				else:
					sup=(self.interv[i+1][0]+2*self.interv[i][1])/3
				intervals.append([(self.interv[i][1]+sup)/2,sup])
				PolyVanishing.append(0)
				NumberOfRealRoot.append(0)
			else:
				intervals.extend(real_roots(auxPol))#,bounds=[self.interv[0],sup]))
				for l in range(j,len(intervals)):
					intervals[l]=deepcopy(intervals[l][0])
				fact=auxPol.factor()[:]
				NumberOfCurrentRealRootForPoly=[0 for k in range(len(fact))]
				for l in range(j,len(intervals)):
					for m in range(len(fact)):
						if not sign(fact[m][0].subs(intervals[l][0]))==sign(fact[m][0].subs(intervals[l][1])):
							PolyVanishing.append(fact[m][0])
							NumberOfCurrentRealRootForPoly[m]+=1
							NumberOfRealRoot.append(NumberOfCurrentRealRootForPoly[m])
							break
			#Now we must eliminate roots that are too low
			testnext=True
#			j=0
#			print i,'abajo',intervals 
			while testnext and (j<len(intervals)):
				if PolyVanishing[j].divides(self.Polys[i]) and NumberOfRealRoot[j]==self.indices[i]:
					condition=False
					interval=[max(intervals[j][0],self.interv[i][0]),min(intervals[j][1],self.interv[i][1])]
				else:
#					print 'tamoaqui',PolyVanishing[j],NumberOfRealRoot[j],intervals[j],self.Polys[i],self.indices[i],self.interv[i]
					chain=IsGreaterRoot(PolyVanishing[j],NumberOfRealRoot[j],intervals[j],self.Polys[i],self.indices[i],self.interv[i])
					condition=not chain[0]
					interval=chain[1]
				if condition:
					intervals.pop(j)
					PolyVanishing.pop(j)
					NumberOfRealRoot.pop(j)
				else:
					intervals[j]=deepcopy(interval)
					testnext=False
		return [intervals, PolyVanishing, NumberOfRealRoot]




#	def EvalFunction(self,number):
		

	def EvalCPWPF(self,number):
		if self.FullPol in QQ:
			result=self.Fullpol
		else:
			result=self.FullPol.__call__(number)
		for i in range(len(self.Polys)):
			if self.coeffs[i] in QQ:
				cffcnt=self.coeffs[i]
			else:
				cffcnt=self.coeffs[i].__call__(number)
			if number>self.interv[i][1]:
				result+=self.Polys[i].__call__(number)*cffcnt
			elif number>self.interv[i][0]:
				if sign(self.Polys[i].__call__(number))==(-1)^(self.indices[i]+self.Polys[i].degree()):
					result+=self.Polys[i].__call__(number)*cffcnt
		return result




class CiP(CanonicalPiecewisePolynomialFunction):
	"""
	This class is for C_i(P)
	"""
	def __init__(self,i,P,v):
		if not v.is_symbol():
			print 'Third argument is supposed to be a variable'
			return
		if not P in PolynomialRing(QQ,v):
			print 'Second argument is expected to be an univariate polynomial in the third argument with rational coefficients'
			return
		if not i in NN:
			print 'First argument is expected to be a nonnegative integer' 
			return
		self.var=v
		self.ring=PolynomialRing(QQ,self.var)
		p=self.ring(P)
		#Since we want this to be done soon, we do not factor the polynomial P yet
		#Maybe for the next version...
		if not p.is_irreducible():
#			print 'second argument is supposed to be an irreducible polynomial'
			interv=real_roots(p)
			if i>len(interv):
			#We have zero
				self.Polys=[]
				self.indices=[]
				self.coeffs=[]
				self.interv=[]
				self.FullPol=0
				return
			elif i==0:
				self.Polys=[]
				self.indices=[]
				self.coeffs=[]
				self.interv=[]
				self.FullPol=p		
				return
			fac=p.factor()[:]
			for l in range(len(interv)):
				interv[l]=deepcopy(interv[l][0])
#			print p,fac,i,interv
			for k in range(len(fac)):
#				print k, interv[i-1]
				if fac[k][0].__call__(interv[i-1][0])*fac[k][0].__call__(interv[i-1][1])<0:
					l=k
			inew=1
			for k in range(i-1):
				if fac[l][0].__call__(interv[k][0])*fac[l][0].__call__(interv[k][1])<0:
					inew+=1
#			print p.quo_rem(fac[l][0])[0]*CiP(inew,fac[l][0],self.var)
			aux=p.quo_rem(fac[l][0])[0]*CiP(inew,fac[l][0],self.var)
			self.Polys=aux.Polys
			self.indices=aux.indices
			self.coeffs=aux.coeffs
			self.interv=aux.interv
			self.FullPol=aux.FullPol	
			return #p.quo_rem(fac[l][0])[0]*CiP(inew,fac[l][0],self.var)
		#der=P.diff(self.var)
		#g=gcd(P,der)
		#q=P.quo_rem(g)[0]
		lc=p.leading_coefficient()
		interv=real_roots(p)
		for l in range(len(interv)):
			interv[l]=deepcopy(interv[l][0])
		if i>len(interv):
			#We have zero
			self.Polys=[]
			self.indices=[]
			self.coeffs=[]
			self.interv=[]
			self.FullPol=0
		elif i==0:
			self.Polys=[]
			self.indices=[]
			self.coeffs=[]
			self.interv=[]
			self.FullPol=p
		else:
			self.Polys=[p/lc]
			self.indices=[i]
			self.coeffs=[self.ring(lc)]
			self.interv=[[interv[i-1][0],interv[i-1][1]]]
			self.FullPol=0

print 'Introduce CiP(i=integer,p=polynomial,x=variable) to create a piecewise polynomial funtion'
print 'valued 0 before the i-th root of p(x) and as p after such root'
print 'Before, you need to have the variable defined, and p must be irreducible'
print 'One can sum, multiply and compose this kind of objects'