import numpy

def integrate(t,y):
#integrate function y (N,n) or (N,) or a scalar tabulated at t (N,) values

#compose time and value shift arrays
   t0=[x for x in t]
   t1=[x for x in t]
#t0.insert(0,0)
   t0=t0[:-1]
   t1=t1[1:]
   t0=numpy.array(t0)
   t1=numpy.array(t1)
   try:
      n=y.shape[1]
      y0=numpy.matrix.copy(y)
      y1=numpy.matrix.copy(y)
   except IndexError:
      try:
         n=len(y)
         y0=numpy.outer(numpy.ones(t.shape),y)
         y1=numpy.outer(numpy.ones(t.shape),y)
      except TypeError:
         n=1
         y0=y*numpy.outer(numpy.ones(t.shape),numpy.ones(n))
         y1=y*numpy.outer(numpy.ones(t.shape),numpy.ones(n))

   y0=numpy.delete(y0,-1,axis=0)
   y1=numpy.delete(y1,0,axis=0)
   

#integral (trapezoid) updates
#(N,n)
   dt=numpy.outer(t1-t0,numpy.ones(n))
   #print('{} {}'.format(dt.shape,(y0+y1).shape))
   d=0.5*dt*(y1+y0)

#use cumsum to compute integrals
#(N-1,n)
   return numpy.cumsum(d,axis=0)


def integrateFull(t,Y,axis=1):
    #calculate integrals of Y along axis, assume value i of each vector along axis to be calculated at ti

#the formula for trapezoid is 
#2*I=y[n]*t[n]-y[0]*t[0]
#+[y[0]*t[1]+y[1]*t[2]+...+y[n-2]*t[n-1]]+
#-[y[1]*t[0]+y[2]*t[1]+...+y[n-1]*t[n-2]]=
#y*(t1-t0)
#where
#t0=[t[0],t[0],...,t[n-3],t[n-2]]
#t1=[t[1],t[2],...,t[n-1],t[n-1]]
    t0=numpy.concatenate((t[0:1],t[:-1]))
    t1=numpy.concatenate((t[1:],t[-1:]))
    #this works for arbitrary dimension of Y, so some Einstein magic has to be applied
    v='abcdef'
    n=len(Y.shape)
    einsumStr=v[:n]+','+v[axis]+'->'+v[:n]
    print('einsum [{}] {} {}'.format(einsumStr,Y.shape,t0.shape))
    T=numpy.einsum(einsumStr,numpy.ones(Y.shape),t1-t0)
    return 0.5*numpy.sum(Y*T,axis)


def skipRedundant(model):
   redundant=[r for r in model.scaled]
   redundant.append('total')
   lut=model.lut
   v_sorted=numpy.sort([lut[x] for x in redundant])[::-1]
#n=len(Q['lut'])
#sort by value
   sMap={x:lut[x] for x in lut if x not in redundant}
#sMap=dict(sorted(redundantMap.items(),key=lambda item:item[1]))
   updatedMap={list(sMap.keys())[i]:i for i in range(len(sMap))}
#print(sMap)
#print(updatedMap)
#return v_sorted
   return updatedMap



def removeScaled(model,M,mode='matrix'):
   originalMap=model.lut
   updatedMap=skipRedundant(model)
   n=len(updatedMap)
#print(skip)
   if mode=='matrix':
      A=numpy.zeros((n,n))
      for c in updatedMap:
         i=updatedMap[c]
         i0=originalMap[c]
         for c1 in updatedMap:
            j=updatedMap[c1]
            j0=originalMap[c1]
            A[i,j]=M[i0,j0]
      return A
   if mode=='columnsOnly':
      A=numpy.zeros((M.shape[0],n))
      for c in updatedMap:
         i=updatedMap[c]
         i0=originalMap[c]
         A[:,i]=M[:,i0]
      return A
   if mode=='vector':
      v=numpy.zeros(n)
      for c in updatedMap:
         i=updatedMap[c]
         i0=originalMap[c]
         v[i]=M[i0]
      return v
   return numpy.zeros(0)

def getSE(s1):
   s2=numpy.multiply(s1,s1)
   return numpy.sqrt(numpy.dot(s2,numpy.ones(s1.shape[2])))

def removeImVector(v,eps):
   for i in range(len(v)):
      if numpy.abs(numpy.imag(v[i]))<eps*numpy.abs(numpy.real(v[i])):
         v[i]=numpy.real(v[i])
   return v

def removeImMatrix(A,eps):
   f=A.ravel()  
   f=removeImVector(f,eps)
   return f.reshape(A.shape)
 
def solveMatrix(model,tmax,nt=201,t0=0,y0=numpy.array([]), sIn=numpy.array([]),method='fraction'):
#method=fraction sets output of excrement compartments to fractions of total ingested chemical
#method=derivate sets output to amount excreted per unit time (in native units, ng/min)
   obj=_solveMatrix(model,tmax,nt,t0,y0,sIn,method)
   return obj['t'],obj['sol'],obj['se'],obj['s1']

def _solveMatrix(model,tmax,nt=201,t0=0,y0=numpy.array([]), sIn=numpy.array([]),method='fraction'):
#sIn (n,m)
#y0 (n,)
   t=numpy.linspace(0,tmax-t0,nt)
   n=len(model.lut)
   m=len(model.lutSE)
   N=len(t)
   
   if not y0.size:
      y0=numpy.zeros(n)

   iT=model.lut['total']
   u0=y0[iT]
   solver=solveUnitObj(model,t,u0)


   if not sIn.size:
      sIn=numpy.zeros([n,m])
   
  
   sOut=numpy.zeros([N,n,m+1])
#y is (N,n)
   y=solver.solve(y0,model.u(0),method=method)   
   sOut[:,:,-1]=y
#(N,n,1)
   y1=numpy.expand_dims(y,axis=2)


#SS (m,n,n)
   SS=model.fSS(0)
#SU (m,n)
   SU=model.Su(0)  #m by n
#u (n0,)
   u=model.u(0) 
   for p in model.lutSE:
      k=model.lutSE[p]
#SS[k] is (n,n)
#y is (N,n)
#rhs is (N,n)=(n,n) @ (N,n,1)
      rhs=(SS[k] @ y1).squeeze()
#yIn is (n,)
      yIn=sIn[:,k]
#return in (N,n)
      sOut[:,:,k]=solver.solveS(SS[k],y,y0,u=u,z0=sIn[:,k],r=SU[k],method=method)


      
   sol=sOut[:,:,-1]
   s1=sOut[:,:,:-1]
   se=getSE(s1)
   return {'t':t+t0,'sol':sol,'se':se,'s1':s1}


class solveUnitObj:

   def __init__(self,model,t,u0):

      self.model=model
      self.t=t
      self.u0=u0
#some dimensions
      self.originalMap=model.lut
      self.updatedMap=skipRedundant(model)
      self.n0=len(self.originalMap)
      self.n=len(self.updatedMap)
 
      self.N=len(t)

      self.M=self.model.M(0)
      
      self.A=removeScaled(model,self.M)

#diagonalize matrix
      self.lam,self.P=numpy.linalg.eig(self.A)
      self.P1=numpy.linalg.inv(self.P)
      D=numpy.diag(self.lam)
      self.D1=numpy.diag(1/self.lam)


#next are all (N,n,n)
      self.tE=numpy.ones(t.shape)
      E=numpy.einsum('a,bc',self.tE,numpy.eye(self.n))
      QT=numpy.einsum('a,bc',t,numpy.eye(self.n))
      QTL=numpy.einsum('a,bc',t,D)
      self.W=E*numpy.exp(QTL)
      invL=numpy.einsum('a,bc',self.tE,self.D1)
      self.Q=(self.W-E)*invL

#we have exp(lam*(t-t')) as Green's function
#construct matrix W1 (N,N,n,n) where the last part is diag(lam)
#W[i,j,k,l]=delta(k,l)*exp(lam_k*(t_i-t'_j))
#delta(i,j) is Kronecker's delta
      T=numpy.outer(t,self.tE)-numpy.outer(self.tE,t)
      T1=numpy.einsum('ab,cd',T,numpy.eye(self.n))
      F=numpy.einsum('ab,cd',T>=0,numpy.eye(self.n))
      L1=numpy.einsum('ab,cd',numpy.ones(T.shape),D)
      self.W1=numpy.zeros(F.shape,dtype='complex')
      self.W1[F>0]=numpy.exp(L1[F>0]*T1[F>0])

#a matrix with lam_i in rows LR(i,j)=lam(i)
#(N,n,n)
      LR=numpy.einsum('a,b,c',self.tE,self.lam,numpy.ones(self.n))
      LRT=numpy.einsum('a,b,c',t,self.lam,numpy.ones(self.n))
#a matrix with lam_i in columns LC(i,j)=lam(j)
#(N,n,n)
      LC=numpy.einsum('a,b,c',self.tE,numpy.ones(self.n),self.lam)
      LCT=numpy.einsum('a,b,c',t,numpy.ones(self.n),self.lam)

#a matrix with exp(lam_i) in rows DR(i,j)=exp(lam(i))
#(N,n,n)
      DR=numpy.exp(LRT)
#a matrix with exp(lam_j) in columns DC(i,j)=exp(lam(j))
#(N,n,n)
      DC=numpy.exp(LCT)

#a diagonal matrix with t*exp(D*t) on diagonal
#(N,n,n)
      self.H=numpy.zeros(self.W.shape,dtype='complex')
      self.H[E>0]=QT[E>0]*numpy.exp(QTL[E>0])
#off diagonal is exp(lam(j))-exp(lam(i))/(lam(j)-lam(i))
      self.H[E==0]=(DC[E==0]-DR[E==0])/(LC[E==0]-LR[E==0])
      
      
#


   def solve(self,y0=numpy.zeros(0),u=numpy.zeros(0),method='fraction'):

#solve system
#y0 (n,) initial values
#u (n,) the time independent input


#use Green function theorem, since impulse response is simple
#yGREEN=exp(lam*(t-t0)
   
#matrix solution works for reduced space where all
#non-trivial components are reduced

#typical variables are denoted with trailing p (for prime)
#and have a smaller size (fewer compartments) than
#full variables


#get the matrix and remove trivial components
#remove trivial components from non-homogeneous parts
      if not u.size:
         u=numpy.zeros(self.n0)

      up=removeScaled(self.model,u,mode='vector')

#remove trivial components from initial values
      if not y0.size:
         y0=numpy.zeros(u.shape)
      y0p=removeScaled(self.model,y0,mode='vector')

#(N,n)=(n,n) @ (N,n,n) @ (n,n) @ (n,) + (n,n) @ (N,n,n) @ (n,n) @ (n,)
      y= self.P @ self.W @ self.P1 @ y0p + self.P @ self.Q @ self.P1 @ up

#back to original system, add trivial compartments
#(N,n)
      yOut=numpy.zeros((self.N,self.n0))

#shift compartments from modified index i (from y) to original index i0 (to yOut)
      for c in self.updatedMap:
         i=self.updatedMap[c]
         i0=self.originalMap[c]
         yOut[:,i0]=y[:,i]
  
#here we need original, full sized y0,u,rhs, so make sure they are not overwritten by reduced versions
#content determined by method
      if method=='fraction':
         return getScaled(self.model,self.t,yOut,self.u0,y0,u)
      if method=='derivative':
         return setScaledDerivative(self.model,yOut,u)

   def solveS(self,S,y,y0=numpy.zeros(0),u=numpy.zeros(0),\
      z0=numpy.zeros(0),r=numpy.zeros(0),method='fraction'):

#solve gradient system (S) at time points t 
#y (N,n) is solution of the system 
#y0 (n,) initial values of solution  
#u (n,) is the non-homogenous (but constant) part of equation
#z0 (n,) initial values of gradients
#r (n,) is the time independent gradient of nonhomogenous part / input (grad u)

#use Green function theorem, since impulse response is simple
#yGREEN=exp(lam*(t-t0) and integrate it for exponentail solutions 
   
#matrix solution works for reduced space where all
#non-trivial components are reduced

#typical variables are denoted with trailing p (for prime)
#and have a smaller size (fewer compartments) than
#full variables


#get the matrix and remove trivial components

#remove trivial components from initial values
      if not y0.size:
         y0=numpy.zeros(self.n0)
      y0p=removeScaled(self.model,y0,mode='vector')

#remove trivial components from non-homogenuous part of base equation
      if not u.size:
         u=numpy.zeros(y0.shape)
      up=removeScaled(self.model,u,mode='vector')

#remove trivial components from gradient of non-homogenous part
      if not r.size:
         r=numpy.zeros(self.n0)
      rp=removeScaled(self.model,r,mode='vector')

#remove trivial components from initial values of gradient
      if not z0.size:
         z0=numpy.zeros(r.shape)
      z0p=removeScaled(self.model,z0,mode='vector')


      Sp=removeScaled(self.model,S)
#Spp is nearly diagonal in (n,n)
      Spp=self.P1 @ Sp @ self.P
#converted to time space for H multiplication
#(N,n,n)
      HS= numpy.einsum('a,bc',self.tE,Spp) * self.H

#(N,n) = (n,n) @ (N,n,n) @ (n,n) @ (n,)

      z1 = self.P @ self.W @ self.P1 @ z0p 
      z2 = self.P @ self.Q @ self.P1 @ rp

#intermediate variable for speed-up
#(n,) = (n,n) @ (n,n) @(n,)
      _f=self.D1 @ self.P1 @ up

      #combine above for speed-up
#(N,n) = (n,n) @ (N,n,n) @ ( (n,n) @ (n,) + (n,) )
      z3 = self.P @ HS @ (self.P1 @ y0p + _f)
#(N,n) = (n,n) @ (n,n) @ (N,n,n) @ (n,)
      z4 = self.P @ self.Q @ Spp @ _f
     
#(N,n)
      z=z1+z2+(z3-z4)

#back to original system, add trivial compartments
      sOut=numpy.zeros((self.N,self.n0))
   
      for c in self.updatedMap:
         i=self.updatedMap[c]
         i0=self.originalMap[c]
         sOut[:,i0]=z[:,i]
  
#here we need full sized y0,u,rhs, so make sure they are not overwritten by reduced versions
#(N,n) = (n,n) @ (N,n,1)
      rhs = (S @ numpy.expand_dims(y,axis=2)).squeeze()
      if method=='fraction':
         return setScaled(self.model,self.t,sOut,self.u0,z0,r,rhs)
      if method=='derivative':
         return setScaledDerivative(self.model,sOut,r,rhs)


 

def setScaled(model,t,y,u0=0,y0=numpy.zeros(0),u=numpy.zeros(0),rhs=numpy.zeros(0)):
#update y with scaled containers, which were excluded from direct calculation to make M of full rank
#u has shape (n,)
#rhs has shape (N,n)
#y is (N,n)
   M=model.M(0)

#integrate input
#(N-1,n)
   fI=integrate(t,y)

#assemble non-homogenous part of equation
   if not rhs.size:
      rhs=numpy.zeros(y.shape[1])
   if u.size:
#rhs is either (n,) or (N,n) so adding (n,) should work
      rhs+=u

#(N-1,n) even if rhs is (n,)
   fU=integrate(t,rhs)

#apply coupling
#(N-1,n)
   fI=(M @ numpy.expand_dims(fI,axis=2)).squeeze()

#find initial values
#(n,)
   if not y0.size:
      y0=numpy.zeros(u.shape)

#get starting value for fTotal
#(N-1,)
   fTotal=getTotal(model,t,u0)

#update solutions with scaled compartments
   for c in model.scaled:
      i=model.lut[c]
      y[1:,i]=(u0*y0[i]+fI[:,i]+fU[:,i])/fTotal[1:]
      y[0,i]=y0[i]

   iT=model.lut['total']
   y[:,iT]=fTotal
 
#(N,n)
   
   return y


def setScaledDerivative(model,y,u=numpy.zeros(0),rhs=numpy.zeros(0)):
#update y with scaled containers, which were excluded from direct calculation to make M of full rank. 
#Rhs is added to the argument list so that the same routine can be used for Jacobi calculation
#M is (n,n), the system matrix, M=model.M(0)
#y is (N,n), the solution to the coupled compartments (ie. excluded excrement)
#u is (n,), the input part
#rhs is (N,n), for Jacobi, set it to (grad M) @ y, for k-th component of gradient

#assemble non-homogenous part of equation
   if not rhs.size:
#(N,n)
      rhs=numpy.zeros(y.shape)

   if u.size:
#rhs is (n,n) so adding (n,) should work, rhs remains (N,n)
      rhs+=u

#apply couplings and inputs
#(n,n)
   M=model.M(0)
#(N,n)=[(n,n) @ (N,n,1)] + (N,n)
   fI=( M @ numpy.expand_dims(y,axis=2) ).squeeze() + rhs

#only set values back to y for excrement compartments (scaled)
   for c in model.scaled:
      i=model.lut[c]
      y[:,i]=fI[:,i]

#(N,n)
   return y