import ivp
import os
import cModel
import sys
import json
import numpy
import time
import scipy.interpolate

defaultValues={\
      'method':'LSODA',\
      'atol':1e-4,\
      'rtol':1e-4,\
      'tmax':1,\
      'mode':'IVPSimultaneous',\
      'nt':201,
      'tUnit':'day'}

def get(setup,par):
   try:
      return setup[par]
   except KeyError:
      pass
   return defaultValues[par]

def parseSetup(setupFile):
   with open(setupFile,'r') as f:
      setup=json.load(f)
   out={}
   for p in defaultValues:
      out[p]=get(setup,p)
   return out

def getScale(setup):
   tUnit=setup['tUnit']
   if tUnit=='min':
      return 1
   if tUnit=='hour':
      return 60
   if tUnit=='day':
      return 24*60
   if tUnit=='month':
      return 24*60*30
   if tUnit=='year':
      return 24*60*30*365
   return 1



def main(parFiles,jobDir):
   model=cModel.model()
   setupFile=parFiles[0]
   modelFile=parFiles[1]
   parameterFile=parFiles[2]
   model.parse(modelFile,parameterFile)
   setup=parseSetup(setupFile)
   scale=getScale(setup)
   tmax=setup['tmax']*scale

   start_time=time.time()
   if setup['mode']=='SequentialOdeint':
      t,sol,se,s1=ivp.solveSequentialOdeint(model,tmax,setup['nt'])
   if setup['mode']=='SimultaneousOdeint':
      t,sol,se,s1=ivp.solveSimultaneousOdeint(model,tmax,setup['nt'])

   if setup['mode']=='IVP':
      t,sol,se,s1=ivp.solveSequential(model,tmax,atol=setup['atol'],
         rtol=setup['rtol'],method=setup['method'])
        
   if setup['mode']=='IVPSimultaneous':
      t,sol,se,s1=ivp.solveSimultaneous(model,tmax,atol=setup['atol'],
         rtol=setup['rtol'],method=setup['method'])

   end_time=time.time()
   print('Time: {:.3f} s'.format(end_time-start_time))
   #store
   #interpolate on s1
   #s1 has shape ft x n x m
   #where ft is LSODE driven number of points
   qt,sOut=interpolate(setup,model,t,s1,tmax)

   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)
   numpy.savetxt(os.path.join(jobDir,'qt.txt'),qt)
   #this is 3D, so new routines
   write3D(os.path.join(jobDir,'sOut.txt'),sOut)

def interpolate(setup,model,t,s1,tmax):
   #interpolate on s1
   #s1 has shape ft x n x m
   #where ft is LSODE driven number of points
   sOut=numpy.zeros((setup['nt'],model.n,model.m))
   qt=numpy.linspace(0,tmax,setup['nt'])
   for i in range(model.n):
      for j in range(model.m):
         y=s1[:,i,j]
         f = scipy.interpolate.interp1d(t, y, kind='cubic')
         sOut[:,i,j]=f(qt)
   return qt,sOut



def write3D(fName,a):
   with open(fName,'w') as f:
      f.write('#Shape {}\n'.format(a.shape))
      i=0
      for data_slice in a:
         f.write('#\t Slice {}\n'.format(i))
         numpy.savetxt(f,data_slice)
         i+=1
      
def read3D(fName):
   new_data=numpy.loadtxt(fName)
   #read and parse first row to get shape
   with open(fName,'r') as f:
      firstLine=f.readline()

   shape=firstLine.replace('#Shape ','').\
      replace('\n','').\
      replace('(','').\
      replace(')','').\
      split(',')
   shape=[int(x) for x in shape]
   return new_data.reshape(shape)