import numpy
import scipy.integrate
#scipy solve_ivp interface using system of n equations
#infrastructure for sensitivity analysis of y with respect
#to model parameters p also integrated
#
#General form of equation:
#  dy/dt=M*y+u
#where y is an n-vector and M is nxn Jacobi matrix 
#
#implementation of system has to provide attributes:
# - n, number of variables
# - m, number of parameters
# - M(t), nxn, a time dependent Jacobi matrix,
# - u(t), 1xn, a time dependent source term
# - fS(t) mxn, a time dependent array of derivatives 
#           fS(t)[p,:] = d(M*y)/dp (t)
#   where y is internal interpolated y-only ivp solution
# - fSY(t,y), mxn, a time dependent array of derivatives:
#           fS(t)[p,:] = d(M*y)/dp (t)
#   where current y is supplied at each step
# - Su(t), mxn, a (time dependent) array of partial derivatives:
#           Su(t)[p,i] = du_i/dp (t)


def dfdy(t,y,system):
    dfdy=system.M(t).dot(y)+system.u(t)
    return dfdy

def jacobi(t,y,system):
    return system.M(t)

#SE post calculation
def dfdyS(t,S,system):
    #unwrap S to NxM where M is number of parameters
    mS=numpy.reshape(S,(system.n,system.m))
    mOut=system.M(t).dot(mS)+system.fS(t)
    return numpy.ravel(mOut)

def jacobiSE(t,S,system):
    N=system.n*(system.m)
    fJ=numpy.zeros((N,N))
    #print('fJ shape {}'.format(fJ.shape))
    for i in range(system.m):
        fJ[i*system.n:(i+1)*system.n,i*system.n:(i+1)*system.n]=system.M(t)
    return fJ

#SE simultaeneous calculation
def dfdySFull(t,S,system):
    #unwrap S to NxM where M is number of parameters
    mS=numpy.reshape(S,(system.n,system.m+1))
    #system.fS(y,t) is NxM matrix where M are parameters
    y=mS[:,0]
    mOut=system.M(t,y).dot(mS)+system.fSY(y,t)
    #add u derivatives
    mOut[:,1:]+=numpy.transpose(system.Su(t))
    try:
      system.iPrint+=1
    except AttributeError:
      system.iPrint=0
    if system.iPrint>1000:
       print('At t={:.2f}'.format(t),flush=True)
       system.iPrint=0
    return numpy.ravel(mOut)

def jacobiSEFull(t,S,system):
   return system.jacobiFull(t) 

def solveSimultaneous(model,tmax,atol,rtol,method='LSODA',nt=201,\
      t0=0,y0=numpy.array([]), sIn=numpy.array([])):

   t_eval=numpy.linspace(t0,tmax,nt)


   s0=numpy.zeros((model.n,model.m+1))

   if sIn.size!=0:
      s0[:,1:]=sIn
   if y0.size==0:
      y0=numpy.zeros(model.n)

   #set initial condition
   s0[:,0]=y0
   s0=s0.ravel()

   sol=scipy.integrate.solve_ivp(dfdySFull,[t0, tmax],s0, args=(model,),\
                                 jac=jacobiSEFull, method=method, \
                                 atol=atol, rtol=rtol, t_eval=t_eval)
   t=sol.t
   sFull=numpy.reshape(numpy.transpose(sol.y),(len(t),model.n,model.m+1))
   s1=sFull[:,:,1:]
   ysol=sFull[:,:,0]
   se=model.calculateUncertainty(s1)
   print('Done simultaneous LSODA SE')
   return t,ysol,se,s1

def solveSequential(model,tmax,atol,rtol,method='LSODA',nt=201,\
      t0=0,y0=numpy.array([]), sIn=numpy.array([])):

   t_eval=numpy.linspace(t0,tmax,nt)

   if y0.size==0:
      y0=numpy.zeros(model.n)

   
   solIVP=scipy.integrate.solve_ivp(dfdy,[t0, tmax],y0, args=(model,), jac=jacobi,
                                  method=method, atol=atol, rtol=rtol, t_eval=t_eval)

   #y is n x nt (odeint nt x n)
   sol=numpy.transpose(solIVP.y)
   t=solIVP.t
   print('shape (y) {}'.format(sol.shape))
   model.setY(t,sol)
   s0=numpy.zeros((model.n,model.m))
   if sIn.size==0:
      s0=sIn
   s0=s0.ravel()
    
   solIVPSE=scipy.integrate.solve_ivp(dfdyS,[0, tmax],s0, args=(model,), jac=jacobiSE,
                               method=method, atol=atol, rtol=rtol,t_eval=t_eval)

   sraw=numpy.reshape(numpy.transpose(solIVPSE.y),(len(solIVPSE.t),model.n,model.m))
   #interpolate on t
   s1=numpy.zeros((len(t),model.n,model.m))
   for i in range(model.n):
      for j in range(model.m):
         tck = scipy.interpolate.splrep(solIVPSE.t, sraw[:,i,j], s=0)
         s1[:,i,j]=scipy.interpolate.splev(t, tck, der=0)

   se=model.calculateUncertainty(s1)
   return t,sol,se,s1

def solveSimultaneousOdeint(model,tmax,nt=201,t0=0,y0=numpy.array([]), sIn=numpy.array([])):
   t = numpy.linspace(t0,tmax, nt)
   if y0.size==0:
      y0=numpy.zeros(model.n)
   
   s0=numpy.zeros((model.n,model.m+1))
   if sIn.size!=0:
      s0[:,1:]=sIn
   #set initial condition
   s0[:,0]=y0
   s0=s0.ravel()
   solSE1=scipy.integrate.odeint(dfdySFull, s0, t, args=(model,),Dfun=jacobiSEFull,tfirst=True)
   sFull=numpy.reshape(solSE1,(len(t),model.n,model.m+1))
   s1=sFull[:,:,1:]
   sol=sFull[:,:,0]
   se=sys.calculateUncertainty(s1)
   print('Done simultaneous SE')
   return t,sol,se,s1

def solveSequentialOdeint(model,tmax,nt=201,t0=0,y0=numpy.array([]),sIn=numpy.array([])):
   t = numpy.linspace(t0,tmax, nt)
   if y0.size==0:
      y0=numpy.zeros(sys.n)
   sol = scipy.integrate.odeint(dfdy, y0=y0, t=t, args=(model,),Dfun=jacobi,tfirst=True)
   print('shape (y) {}'.format(sol.shape))

   model.setY(t,sol)
   s0=numpy.zeros((model.n,model.m))
   if sIn.size==0:
      s0=sIn
   s0=s0.ravel()
   solSE=scipy.integrate.odeint(dfdyS, s0, t, args=(model,),Dfun=jacobiSE,tfirst=True)
   s1=numpy.reshape(solSE,(len(t),model.n,model.m))
   se=model.calculateUncertainty(s1)
   print('Done sequential SE')
   return t,sol,se,s1