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)
   #q=[x if not x in out for x in setup]
   q=[x for x in setup if not x in out]
   for x in q:
      out[x]=setup[x]
   return out

def getScale(setup):
   tUnit=setup['tUnit']
   return convertUnit(tUnit)

def convertUnit(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*365
   return 1

def getStrides(setup):
   tmax=setup['tmax']*convertUnit(setup['tUnit'])
   t0=setup['t0']
   try:
      stride=setup['stride']
   except KeyError:
      return [{'label':'','length':tmax-t0}]

   #cover tmax with steps of length stride['length']
   strides=[]
   strideLength=stride['length']*convertUnit(stride['tUnit'])
   i=0
   while t0<tmax:
      strides.append({'length':strideLength,'label':'_step{:03d}'.format(i)})
      t0+=strideLength
      i+=1
   return strides

def getFiles(setup, getAll=False):
   strides=getStrides(setup)
   readable=['t','sol','se','qt','sOut']
   #get last stride
   steps=[strides[-1]]
   if getAll:
      steps=strides
   inFiles=[]
   for s in steps:
      l=s['label']
      inFiles.append({x:'{}{}.txt'.format(x,l) for x in readable})
   return inFiles


def main(parFiles,jobDir,startDir='NONE'):
   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
   print('Using {}'.format(jobDir))
   if startDir!='NONE':
      print('Using solution from {}'.format(startDir))
      t0,y0,s0,lut,lutSE=getStartPointFromDir(startDir)
      print('t0={}'.format(t0))
   else:
      t0,y0,s0,lut,lutSE=getStartPoint(setup)

   #we should reorder S1 to match model.lutSE
   #S1 is nxm, so we should reorder columns
   sIn=numpy.zeros(s0.shape)
   if s0.size>0:
      for x in model.lutSE:
         #this could be modified for different ordering of lut as well
         sIn[:,model.lutSE[x]]=s0[:,lutSE[x]]
   #copy t0 to setup to help get the strides right
   setup['t0']=t0 
   strides=getStrides(setup)
   for step in strides:
      t1=t0+step['length']
      start_time=time.time()
      if setup['mode']=='SequentialOdeint':
         t,sol,se,s1=ivp.solveSequentialOdeint(model,t1,setup['nt'],t0,y0,sIn)
      if setup['mode']=='SimultaneousOdeint':
         t,sol,se,s1=ivp.solveSimultaneousOdeint(model,t1,setup['nt'],t0,y0,sIn)

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

      end_time=time.time()
      #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,t0,t1)
      print('Time: {:.3f} s'.format(end_time-start_time))
      y0=sol[-1]
      t0=t[-1]
      sIn=sOut[-1]

      #store
      writeable={'t':t,'sol':sol,'se':se,'qt':qt,'sOut':sOut}
      l=step['label']
      writeable={os.path.join(jobDir,'{}{}.txt'.format(x,l)):writeable[x]\
         for x in writeable}
      for x in writeable:
         if x.find('sOut')>-1:
            write3D(x,writeable[x],model.lut,model.lutSE)
         else:
            numpy.savetxt(x,writeable[x])
      print('Completed step {}'.format(l))
   print('Completed run')
   #update setup with new content (particularly t0)
   setupOut=os.path.join(jobDir,'setup.json')
   with open(setupOut,'w+') as f:
      f.write(json.dumps(setup))
      #this is 3D, so new routines

def interpolate(setup,model,t,s1,t0,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(t0,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 invertLUT(lut):
   fcode=[None]*len(lut)
   for x in lut:
      fcode[lut[x]]=x
   return fcode

def write3D(fName,a,lut,lutSE):
   with open(fName,'w') as f:
      f.write('#Shape {}\n'.format(a.shape))
      f.write('#Lut {}\n'.format(invertLUT(lut)))
      f.write('#Lut (SE) {}\n'.format(invertLUT(lutSE)))

      i=0
      for data_slice in a:
         f.write('#\t Slice {}\n'.format(i))
         numpy.savetxt(f,data_slice)
         i+=1
      
def parseLUT(line,isSE=False):
   SE=''
   if isSE:
      SE='(SE)'
   header='#Lut {}'.format(SE)
   names=line.replace(header,'').\
      replace('[','').\
      replace(']','').\
      replace(' ','').\
      replace('\n','').\
      replace("'",'').\
      split(',')
   lut={names[i]:i for i in range(len(names))}
   return lut

def parseShape(line):
   shape=line.replace('#Shape ','').\
      replace('\n','').\
      replace('(','').\
      replace(')','').\
      split(',')
   return [int(x) for x in shape]

      
def read3D(fName):
   new_data=numpy.loadtxt(fName)
   #read and parse first three lines to get shape,lut and lutSE
   nLines=3
   with open(fName,'r') as f:
      lines=[f.readline() for i in range(nLines)]

   shape=parseShape(lines[0])
   lut=parseLUT(lines[1])
   lutSE=parseLUT(lines[2],isSE=True)

   return new_data.reshape(shape),lut,lutSE


def loadSolutionFromRef(setup, loadAll=False):
   
   if setup['startFromRef']=='None':
      return 0,numpy.array([]),numpy.array([]),numpy.array([]),numpy.array([])

   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
   #import labkeyDatabaseBrowser

   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)

   ref=setup['startFromRef']
   #refFilter={'variable':'runRef','value':ref,'oper':'eq'}
   #ds=db.selectRows(setup['project'],setup['schema'],setup['query'],[refFilter],setup['view'])

   #parFileString=ds['rows'][0]['parameterFileString']
   #parFiles=parFileString.split(';')
   remoteDir=fb.formatPathURL(setup['project'],'/'.join(['jobs',ref]))

   try:
      localDir=setup['localDir']
   except KeyError:
      localDir=os.path.join(os.path.expanduser('~'),'temp','jobDir')
   if not os.path.isdir(localDir):
      os.mkdir(localDir)

   #download setup
   f='setup.json'
   localPath=os.path.join(localDir,f)
   remotePath='/'.join([remoteDir,f])
   fb.readFileToFile(remotePath,localPath)
   setupFile=os.path.join(localDir,f)

   setupOld=parseSetup(setupFile)
   inFiles=getFiles(setupOld,loadAll)

   #inFiles.extend(parFiles)
   i=0
   n=len(inFiles)
   for fSet in inFiles:
      for x in fSet:
         f=fSet[x]
         localPath=os.path.join(localDir,f)
         remotePath='/'.join([remoteDir,f])
         fb.readFileToFile(remotePath,localPath)
      print('Loaded {}/{} {}'.format(i+1,n,localDir))
      i+=1
   return loadSolutionFromDir(localDir,loadAll)

def merge(a1,a2):
   try: 
      return numpy.concatenate(a1,a2)
   except ValueError:
      return a2

def loadSolutionFromDir(jobDir,loadAll=False):
   setupFile=os.path.join(jobDir,'setup.json')
   setup=parseSetup(setupFile)
   inFiles=getFiles(setup,loadAll)
   isFirst=True 
   operators={x:numpy.loadtxt for x in inFiles[0]}
   operators['sOut']=read3D
   data={}
   i=0
   n=len(inFiles)
   for fSet in inFiles:
      idata={}
      print('Parsing [{}/{}]'.format(i+1,n))
      for f in fSet:
         path=os.path.join(jobDir,fSet[f])
         idata=operators[f](os.path.join(jobDir,fSet[f]))
         if f=='sOut':
            lut=idata[1]
            lutSE=idata[2]
            idata=idata[0]
         try:
            data[f]=numpy.concatenate((data[f],idata))
         except KeyError:
            data[f]=idata
      i+=1
   return data['t'],data['sol'],data['se'],data['sOut'],data['qt'],lut,lutSE,setup

def getStartPointFromDir(jobDir):
   t,sol,se,sOut,qt,lut,lutSE,setup=loadSolutionFromDir(jobDir)
   return t[-1],sol[-1],sOut[-1],lut,lutSE

def getStartPoint(setup):
   if setup['startFromRef']=='None':
      return 0,numpy.array([]),numpy.array([]),[],[]
   t,sol,se,sOut,qt,lut,lutSE,oldSetup=loadSolutionFromRef(setup)
   return t[-1],sol[-1],sOut[-1],lut,lutSE

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