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='full'):
   if method=='full':
      obj=_solveMatrix(model,tmax,nt,t0,y0,sIn)
   if method=='separately':
      obj=_solveMatrixSeparately(model,tmax,nt,t0,y0,sIn)

   return obj['t'],obj['sol'],obj['se'],obj['s1']

def _solveMatrix(model,tmax,nt=201,t0=0,y0=numpy.array([]), sIn=numpy.array([])):
   t=numpy.linspace(0,tmax-t0,nt)

#solve system Q at time points t for initial values y0
#SS is gradient of system nxn matrix M with respect to m parameters p 
   SS=model.fSS(0)
#large (M+1)*N by (M+1)*N matrix
   m=len(model.lutSE)
   A=removeScaled(model,model.M(0))
#adjusted n where only non-terminal(scaled) compartments are counted
   n=A.shape[0]
   SM=numpy.zeros(((m+1)*n,(m+1)*n))
   for k in range(m):
#A is on diagonal
      SM[k*n:(k+1)*n,k*n:(k+1)*n]=A
#on the far right, gradient matrix SS sits
      SM[k*n:(k+1)*n,-n:]=removeScaled(model,SS[k])
#for original solutions, matrix A is used
      SM[-n:,-n:]=A
#matrix SM is now of full rank and can be diagonalized
#print('{}/{}'.format(numpy.linalg.matrix_rank(SM),(m+1)*n))

#RHS
#SU has shape mxn
#we have to remove scaled columns
#Add plain rhs
   SU=removeScaled(model,model.Su(0),mode='columnsOnly').flatten()
   u=removeScaled(model,model.u(0),mode='vector')
   print('{} {}'.format(SU.shape,u.shape))
   SU=numpy.concatenate((SU,u))
   print(SU.shape)
#diagonalize matrix
   lam,P=numpy.linalg.eig(SM)
#no effect down to eps=1e-2
   #eps=1e-4
   #print(f'Using eps={eps}')
   #lam=removeImVector(lam,eps)
   #P=removeImMatrix(P,eps)
   P1=numpy.linalg.inv(P)
   #P1=removeImMatrix(P1,eps)
   D=numpy.diag(lam)

#verify
#print(numpy.max(P @ D @ P1-M))

#shift input to diagonal system
   v=P1 @ SU
   #also no effect
   #maxV=numpy.max(numpy.abs(v))
   #v=v*(numpy.abs(v)>eps*maxV)
   #v=removeImVector(v,eps)

   #print(f'v={v}')
   #print(f'lam={lam}')

#set initial values if missing
   if not y0.size:
      y0p=numpy.zeros(n)
   else:
#should move this to diagonal space
      y0p=removeScaled(model,y0,mode='vector')

   if not sIn.size:
      s0=numpy.zeros(n*m)
   else:
#sIn is n by m matrix, so transpose before reshaping
      s0=numpy.transpose(sIn)
      s0=removeScaled(model,s0,mode='columnsOnly')
      s0=s0.flatten()
   print('s0 {} y0p {}'.format(s0.shape,y0p.shape))
   S0=numpy.concatenate((s0,y0p))
   S0= P1 @ S0   

#present results as a (n*(m+1) x N array)
#n number of variables/compartments
#N number of time points
   W=numpy.outer(lam,t)
   tE=numpy.ones(t.shape)
   V0=numpy.outer(S0,tE)
   V=numpy.outer(v/lam,tE)
   y1=(V0+V)*numpy.exp(W)-V

#convert back to true system
   y=P @ y1

#reinsert scaled into array
#start with y of shape (m+1)*n by N where N is number of time-points and n are only non-scaled compartments
#should end up with y of shape (m+1)*n by N, where n is now number of ALL compartments
   originalMap=model.lut
   updatedMap=skipRedundant(model)
   n0=len(originalMap)
   n=len(updatedMap)
   N=y.shape[1]
   sOut=numpy.zeros((N,n0,m+1))
   for k in range(m+1):
      for c in updatedMap:
         i=updatedMap[c]
         i0=originalMap[c]
         sOut[:,i0,k]=y[k*n+i,:]

#print('Shape: {}, n={} N={} m={}'.format(yout.shape,n0,N,m))
#equivalent of Sout (N,n,m)
#return numpy.transpose(yout).reshape(N,n0,m+1)
   sOut=calculateScaled(model,t,sOut,sIn,y0)
   sol=sOut[:,:,-1]
   s1=sOut[:,:,:-1]
   se=getSE(s1)
   return {'t':t+t0,'sol':sol,'se':se,'s1':s1,'P':P,'D':D,'SM':SM}




def calculateScaled(model,t,sOut,sIn=numpy.zeros(0),y0=numpy.zeros(0)):

#update sOut w/ calculated scaled values
#sIn,yIn are initial values
#sIn is of shape (n,m)
#yIn is (n,)
#sOut is (N,n,m+1)
   lutSE=model.lutSE
   lut=model.lut

#if reading a model, keep the read lut
#lutSE=Q['lutSE']
#lut=Q['lut']

   m=len(lutSE)
   n=len(lut)
   N=len(t)
   
   if not sIn.size:
      sIn=numpy.zeros((n,m))

   if not y0.size:
      y0=numpy.zeros(n)

   #add column for initial values
   sIn=numpy.c_[sIn,y0]


#reshape to n*(m+1) by N
   y=numpy.zeros((n*(m+1),N))

   for k in range(m+1):
      for i in range(n):
         y[n*k+i,:]=sOut[:,i,k]

#full version of SM

   SS=model.fSS(0)
   A=model.M(0)
   SM=numpy.zeros(((m+1)*n,(m+1)*n))
   for k in range(m):
#A is on diagonal
      SM[k*n:(k+1)*n,k*n:(k+1)*n]=A
#on the far right, gradient matrix SS sits
      SM[k*n:(k+1)*n,-n:]=model.SS[k]

   SM[-n:,-n:]=A

#full version of RHS
   SU=model.Su(0).flatten()
   u=model.u(0)
#n*(m+1) by N
   SU=numpy.outer(numpy.concatenate((SU,u)),numpy.ones(t.shape))

#integral, n*(m+1) by N
   fI=integrate(t,y)
   fU=integrate(t,SU)

#apply couplings
   fI= SM @ fI

#update values; scale compartments to total exposure
   iT=lut['total']
   fTotal=y0[iT]+fU[m*n+iT,:]
   for k in range(m+1):
      for c in model.scaled:
         i=lut[c]
#sOut[1:,i,k]=fI[k*n+i,:]+fU[k*n+i,:]
         sOut[1:,i,k]=(sIn[i,k]*y0[iT]+fI[k*n+i,:]+fU[k*n+i,:])/fTotal
         sOut[0,i,k]=sIn[i,k]
#set total for solution only
   sOut[1:,iT,-1]=fTotal
   sOut[0,iT,-1]=y0[iT]
   return sOut

def _solveMatrixSeparately(model,tmax,nt=201,t0=0,y0=numpy.array([]), sIn=numpy.array([])):
#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=solveUnit(model,t,y0,u0,model.u(0))   
   y=solver.solve(y0,model.u(0))   
   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]=solveUnit(model,t,yIn,u0,u,rhs)
      #sOut[:,:,k]=solver.solve(sIn[:,k],SU[k],rhs)
      sOut[:,:,k]=solver.solveS(SS[k],y1,y0,u=u,z0=sIn[:,k],r=SU[k])


#      break
      
   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.A=removeScaled(model,model.M(0))

#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),rhs=numpy.zeros(0)):

#solve system Q at time points t for initial values y0
#u is the time independent nohomogeneous part
#rhs is tabulated non-homogeneous part (N,n)
#u0 is initial accumulated xenobiote


#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')

#overload w/ rhs
#RHS has a time component
#typically of shape n by N
      if rhs.size:
#to N by n shape
#      rhs=rhs.transpose()
#remove trivial components
         rhsp=removeScaled(self.model,rhs,mode='columnsOnly')
#to (N,n,1)
         rhsp=numpy.expand_dims(rhsp,axis=2) 
      else:
         rhsp=numpy.zeros(0)
   #print(numpy.linalg.matrix_rank(A))



#(n,n) @ (N,n,n) @ (n,n) @ (n,)
#y1 is (N,n)
#sum works for (N,n)+(n,)
   #print(up)
   #print(P1 @ up)
   #print( (W-E)*invL @ P1 @ up)
      y= self.P @ self.W @ self.P1 @ y0p + self.P @ self.Q @ self.P1 @ up

      if rhsp.size:

#time dependent component
#(a,b,c,e)=(a,b,c,d) @ (b,d,e)
#(N,N,n,1)=(n,n) @ (N,N,n,n) @ (n,n) @ (N,n,1)
#Y1 is (N,N,n,1)
         Y1=numpy.real(self.P @ self.W1 @ self.P1 @ rhsp)
#to apply Green's theorem, integral along axis=1 should be performed
#(N,n,1)
      
         Y2=integrateFull(self.t,Y1,axis=1)
#to remove last dimension, squeeze at the end
#(N,n), same as y1
         y+=Y2.squeeze()
     
#back to original system, add trivial compartments
      yOut=numpy.zeros((self.N,self.n0))
   
      for c in self.updatedMap:
         i=self.updatedMap[c]
         i0=self.originalMap[c]
         yOut[:,i0]=y[:,i]
  
      #return yOut
#here we need full sized y0,u,rhs, so make sure they are not overwritten by reduced versions
      return getScaled(self.model,self.t,yOut,self.u0,y0,u,rhs)

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

#solve gradient system (S) at time points t 
#y (N,n,1) 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

#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

#z=P @ W @ P1 @ z0 + P @ Q @ P1 r + 
# + P @ (S * H) @ P1 @ y0 + P @ (S * H) @ D1 @ P1 @ u +P @ S @ Q @ D1 @ P1 @ u
#(n,n) @ (N,n,n) @ (n,n) @ (n,)
#y1 is (N,n)
#sum works for (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,)
      _f=self.D1 @ self.P1 @ up

      #combine above for speed-up
#(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,)
      z4 = self.P @ self.Q @ Spp @ _f
     
      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]
  
      #return yOut
#here we need full sized y0,u,rhs, so make sure they are not overwritten by reduced versions
      #y = self.P @ self. W @ self.P1 @ y + self.P @ self.Q @ self.P1 @ u
      #(N,n0) = (n0,n0) @ (N,n0,1)
      rhs = (S @ y).squeeze()
      return getScaled(self.model,self.t,sOut,self.u0,z0,r,rhs)


 

def solveUnit(model,t,y0=numpy.zeros(0),u0=0,u=numpy.zeros(0),rhs=numpy.zeros(0)):
#solve system Q at time points t for initial values y0
#u is the time independent nohomogeneous part
#rhs is tabulated non-homogeneous part (N,n)
#u0 is initial accumulated xenobiote

#use Green function theorem, since impulse response is simple
#yGREEN=exp(lam*(t-t0)
   
#some dimensions
   originalMap=model.lut
   updatedMap=skipRedundant(model)
   n0=len(originalMap)
   n=len(updatedMap)
   N=len(t)

#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
   A=removeScaled(model,model.M(0))
#remove trivial components from non-homogeneous parts
   if not u.size:
      u=numpy.zeros(n0)

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

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

#overload w/ rhs
#RHS has a time component
#typically of shape n by N
   if rhs.size:
#to N by n shape
#      rhs=rhs.transpose()
#remove trivial components
      rhsp=removeScaled(model,rhs,mode='columnsOnly')
#to (N,n,1)
      rhsp=numpy.expand_dims(rhsp,axis=2) 
   else:
      rhsp=numpy.zeros(0)
   #print(numpy.linalg.matrix_rank(A))

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

#next are all (N,n,n)
   W=numpy.exp(numpy.einsum('a,bc',t,numpy.diag(lam)))
   tE=numpy.ones(t.shape)
   invL=numpy.einsum('a,bc',tE,numpy.diag(1/lam))
   E=numpy.ones(W.shape)

#(n,n) @ (N,n,n) @ (n,n) @ (n,)
#y1 is (N,n)
#sum works for (N,n)+(n,)
   #print(up)
   #print(P1 @ up)
   #print( (W-E)*invL @ P1 @ up)
   y= y0p + P @ ((W-E)*invL) @ P1 @ up

   if rhsp.size:

#time dependent component
#t-t'
#we have exp(lam*(t-t')) as Green's function
#construct matrix (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,tE)-numpy.outer(tE,t)
      T1=numpy.einsum('ab,cd',T,numpy.eye(n))
      F=numpy.einsum('ab,cd',T>=0,numpy.eye(n))
      L1=numpy.einsum('ab,cd',numpy.ones(T.shape),numpy.diag(lam))
      W1=numpy.zeros(F.shape,dtype='complex')
      W1[F>0]=numpy.exp(L1[F>0]*T1[F>0])
#(a,b,c,e)=(a,b,c,d) @ (b,d,e)
#(N,N,n,1)=(n,n) @ (N,N,n,n) @ (n,n) @ (N,n,1)
#Y1 is (N,N,n,1)
      Y1=P @ W1 @ P1 @ rhsp
#to apply Green's theorem, integral along axis=1 should be performed
#(N,n,1)
      
      Y2=integrateFull(t,Y1,axis=1)
#to remove last dimension, squeeze at the end
#(N,n), same as y1
      y+=Y2.squeeze()
     
#back to original system, add trivial compartments
   yOut=numpy.zeros((N,n0))
   
   for c in updatedMap:
      i=updatedMap[c]
      i0=originalMap[c]
      yOut[:,i0]=y[:,i]
   
#here we need full sized y0,u,rhs, so make sure they are not overwritten by reduced versions
   return getScaled(model,t,yOut,u0,y0,u,rhs)

def getTotal(model,t,u0=0):
   u=model.u(0)
   iT=model.lut['total']
   U=u[iT]
#*numpy.ones(t.shape)
   fU=numpy.zeros(t.shape)
   fU[1:]=u0+integrate(t,U).squeeze()
   fU[0]=u0
   return fU


def getScaled(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

