import ivp
import os
import cModel
import sys
import json
import numpy
import time

def get(setup,par):
   defaultValues={\
      'method':'LSODA',\
      'atol':1e-4,\
      'rtol':1e-4,\
      'tmax':1,\
      'mode':'IVPSimultaneous',\
      'nt':201,
      'tUnit':'day'}
   try:
      return setup[par]
   except KeyError:
      pass
   return defaultValues[par]

def main(parFiles,jobDir):
    print('runSolver')
    sys=cModel.model()
    modelFile=parFiles[1]
    parameterFile=parFiles[2]
    sys.parse(modelFile,parameterFile)
    with open(parFiles[0],'r') as f:
        setup=json.load(f)
    tmax=get(setup,'tmax')
    atol=get(setup,'atol')
    rtol=get(setup,'rtol')
    mode=get(setup,'mode')
    method=get(setup,'method')
    nt=get(setup,'nt')
    tUnit=get(setup,'tUnit')
    print('Running with tmax={} {}'.format(tmax,tUnit))

    if tUnit=='min':
        pass
    if tUnit=='hour':
        tmax*=60
    if tUnit=='day':
        tmax*=24*60
    if tUnit=='month':
        tmax*=24*60*30
    if tUnit=='year':
        tmax*=24*60*30*365

    start_time=time.time()
    if mode=='SequentialOdeint':
        t,sol,se=ivp.solveSequentialOdeint(sys,tmax,nt)
    if mode=='SimultaneousOdeint':
        t,sol,se=ivp.solveSimultaneousOdeint(sys,tmax,nt)

    if mode=='IVP':
        t,sol,se=ivp.solveSequential(sys,tmax,atol=atol,rtol=rtol,method=method)
        
    if mode=='IVPSimultaneous':
        t,sol,se=ivp.solveSimultaneous(sys,tmax,atol=atol,rtol=rtol,method=method)

    end_time=time.time()
    print('Time: {:.3f} s'.format(end_time-start_time))
    #store
    numpy.savetxt(os.path.join(jobDir,'t.txt'),t)
    numpy.savetxt(os.path.join(jobDir,'sol.txt'),sol)
    numpy.savetxt(os.path.join(jobDir,'se.txt'),se)

if __name__=="__main__":
    parFiles=sys.argv[1].split(';')
    jobDir=sys.argv[2]
    main(parFiles,jobDir)