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

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).dot(mS)+system.fSY(y,t)
    return numpy.ravel(mOut)

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

def solveSimultaneous(sys,tmax,atol,rtol,method='LSODA'):
    S1=numpy.zeros((sys.n,sys.m+1))
    y0=numpy.zeros(sys.n)
    #set initial condition
    S1[:,0]=y0
    S1=S1.ravel()
    sol=scipy.integrate.solve_ivp(dfdySFull,[0, tmax],S1, args=(sys,), jac=jacobiSEFull,
                                  method=method, atol=atol, rtol=rtol)
    t=sol.t
    sFull=numpy.reshape(numpy.transpose(sol.y),(len(t),sys.n,sys.m+1))
    s1=sFull[:,:,1:]
    ysol=sFull[:,:,0]
    se=sys.calculateUncertainty(ysol,s1)
    print('Done simultaneous LSODA SE')
    return t,ysol,se

def solveSequential(sys,tmax,atol,rtol,method='LSODA'):
   y0=numpy.zeros(sys.n)
   solIVP=scipy.integrate.solve_ivp(dfdy,[0, tmax],y0, args=(sys,), jac=jacobi,
                                  method=method, atol=atol, rtol=rtol)
#y is n x nt (odeint nt x n)
   sol=numpy.transpose(solIVP.y)
   t=solIVP.t
   print('shape (y) {}'.format(sol.shape))
   sys.setY(t,sol)
   S0=numpy.zeros((sys.n,sys.m))
   S0=S0.ravel()
    
   solIVPSE=scipy.integrate.solve_ivp(dfdyS,[0, tmax],S0, args=(sys,), jac=jacobiSE,
                               method=method, atol=atol, rtol=rtol)
   sraw=numpy.reshape(numpy.transpose(solIVPSE.y),(len(solIVPSE.t),sys.n,sys.m))
#interpolate on t
   s1=numpy.zeros((len(t),sys.n,sys.m))
   for i in range(sys.n):
      for j in range(sys.m):
         tck = scipy.interpolate.splrep(solIVPSE.t, sraw[:,i,j], s=0)
         s1[:,i,j]=scipy.interpolate.splev(t, tck, der=0)

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

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

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

   sys.setY(t,sol)
   S0=numpy.zeros((sys.n,sys.m))
   S0=S0.ravel()
   solSE=scipy.integrate.odeint(dfdyS, S0, t, args=(sys,),Dfun=jacobiSE,tfirst=True)
   s1=numpy.reshape(solSE,(len(t),sys.n,sys.m))
   se=sys.calculateUncertainty(sol,s1)
   print('Done sequential SE')
   return t,sol,se