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,\
      't0':0,\
      'tmax':1,\
      'mode':'IVPSimultaneous',\
      'nt':201,
      'tUnit':'day',\
      'network':'labkey-klimt.json',\
      'startFromRef':'None',\
      'project':'Analysis/Run',\
      'schema':'lists',\
      'query':'runs',\
      'view':'withStrings'}

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
   t0,y0,S1=loadSolutionFromRef(setup)

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

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

   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)

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

def loadSolutionFromRef(setup):
   if setup['startFromRef']=='None':
      return 0,None,None

   nixSuite=os.path.join(os.path.expanduser('~'),'software','src','nixSuite')
   sys.path.append(os.path.join(nixSuite,'wrapper'))
   import nixWrapper
   nixWrapper.loadLibrary('labkeyInterface')
      
   import labkeyInterface
   import labkeyFileBrowser

   net=labkeyInterface.labkeyInterface()
   fconfig=os.path.join(os.path.expanduser('~'),'.labkey',setup['network'])
   net.init(fconfig)
   #debug
   net.getCSRF()

   #db=labkeyDatabaseBrowser.labkeyDB(net)
   fb=labkeyFileBrowser.labkeyFileBrowser(net)

   #refFilter={'variable':'runRef','value':setup['startFromRef'],'oper':'eq'}
   #ds=db.selectRows(setup['project'],setup['schema'],setup['query'],[refFilter],setup['view']'WithStrings')
   #ordering requires the latest to be first
   #ref=ds['rows'][0]['runRef']
   #parFileString=ds['rows'][0]['parameterFileString']
   #parFiles=parFileString.split(';')
   ref=setup['startFromRef']
   remoteDir=fb.formatPathURL(setup['project'],'/'.join(['jobs',ref]))
   localDir=os.path.join(os.path.expanduser('~'),'temp')
   inFiles=['t.txt','sol.txt','se.txt','qt.txt','sOut.txt']
   #inFiles.extend(parFiles)
   for f in inFiles:
       localPath=os.path.join(localDir,f)
       remotePath='/'.join([remoteDir,f])
       fb.readFileToFile(remotePath,localPath)

   t=numpy.loadtxt(os.path.join(localDir,'t.txt'))
   sol=numpy.loadtxt(os.path.join(localDir,'sol.txt'))
   se=numpy.loadtxt(os.path.join(localDir,'se.txt'))
   sOut=read3D(os.path.join(localDir,'sOut.txt'))
   qt=numpy.loadtxt(os.path.join(localDir,'qt.txt'))
   return t[-1],sol[-1],sOut[-1]