import sys, os

from optparse import OptionParser

import glob

import shutil

#Creates a shell script that can be used to run stage 2 and 3 of the climate layer scripts. It sets up all the 

#input arguments for master_script based on provided arguments.

#Inputs:

#    inLatLonList: A text file that includes lat_lon of the tiles you want to process,one per line in format 10.0_-120.0 

#    outputFolder:  The path to the dir where you want to place the output files (e.g. /nobackupp4/climateLayers/output20Deg/reg4/

#    outDirForShFiles: The directory where you want to place the shell script serial files.

#    filePerDir: How many sh files to put in each subdirectory, this is useful if you want to run a subset.

#    lowStatCount: What's the minimum station count you want to use

#    higStatCount: The max station count you want to use

#    scriptFn: The filename of the master script you want to use to run stage4. 

#              For example 'Rscript /climateCode/environmental-layers/climate/research/world/interpolation/master_script_09012014_normal.R'

#Output:

#    Places number of files into directory.

def main():

  usage = "usage: %prog [options] <inLatLonList> <outputFolder> <outDirForShFiles> <filesPerDir> <lowStatCount> <highStatCount> <scriptFn>\n"+\

            "Prepare input for stage 4 "

  parser = OptionParser(usage=usage)

  opts,args=parser.parse_args()

  if len(args)<2:

     sys.exit(parser.print_help())

  inFn = args[0]

  outputFolder = args[1]

  outDir = args[2]

  fPer = int(args[3])

  lowTresh = int(args[4])

  hiTresh = int(args[5])

  scriptFn = args[6]

  inPtr = open(inFn,"r")

  if inPtr is None:

    print "Error: %s doesnt' exist" % inFn

  lines = inPtr.readlines()

  inPtr.close()

  found = 0

  notFound = 0

  if os.path.exists(outDir) is False:

    print "Warning: %s doesn't exists, creating" % outDir

    os.mkdir(outDir)

  wrap = "%s/wrapper.sh" % outDir

  if os.path.exists(wrap) is False:

    print "Warning: %s doesn't exists, creating" % wrap

    wrapFn = os.path.abspath(os.path.dirname(sys.argv[0]))+"/wrapper.sh"

    shutil.copyfile(wrapFn,wrap)

    os.chmod(wrap,0750)

  qs = "%s/serial_TEMPLATE.qsub" % outDir

  if os.path.exists(qs) is False:

    print "Warning: %s doesn't exists, creating" % qs

    qsFn = os.path.abspath(os.path.dirname(sys.argv[0]))+"/serial_TEMPLATE.qsub"

    shutil.copyfile(qsFn,qs)

    os.chmod(wrap,0750)

  rankCount = 0

  subDirNum = (len(lines)/fPer) + 2

  for i in range(0,subDirNum):

    subDirName = "%s/dirSub_%d" % (outDir,i)

    if os.path.exists(subDirName) is False:

      print "Creating %s" % subDirName

      os.mkdir(subDirName)

  dirCounter = 0

  shpFiles = []

  for l in lines:

    l = l.rstrip().split(",")

    print l[0]

    searchStr = "%s/subset/shapefiles/shp_%s*.shp" % (outputFolder,l[0])

    shpFn = glob.glob(searchStr)

    if len(shpFn) > 0:

      print shpFn

      print shpFn[0]

      shpFn = shpFn[0].replace('.shp','')

      print shpFn

      print shpFn.split('_')[3]

      numStat = 1

      shpFiles.append([shpFn,numStat])

    else:

       notFound += 1

       print "Error finding %s" % searchStr

  sortedShpFiles = sorted(shpFiles,key=lambda x: x[1])

  for s in sortedShpFiles:

    numStat = s[1]

    if (numStat < lowTresh):

      print s[0],numStat

      print "Not enough stations",numStat

      continue 

    if (numStat > hiTresh):

      print s[0],numStat

      print "Too many stations"

      continue 

    baseFn = os.path.basename(s[0])

    print baseFn

    baseSpl = baseFn.split('_')

    ll = "%s_%s" % (baseSpl[1],baseSpl[2]) 

    metTest = "%s/%s/met_stations_outfiles_obj_gam_CAI_%s.RData" % (outputFolder,ll,ll)

    if os.path.exists(metTest) is False:

       print "No met object"

       continue     

    methodTest =  "%s/%s/method_mod_obj_gam_CAI_dailyTmax%s.RData" % (outputFolder,ll,ll)

    if os.path.exists(methodTest) is True:

       print "Method object exists"

       continue 

    outStr = "%s %s wgs84Grid %s %s %s %s/subset/mean_LST_%s_jun_6_wgs84_10deg.tif FALSE %s/%s/covar_obj_%s.RData %s/%s/met_stations_outfiles_obj_gam_CAI_%s.RData 10 4800 > %s/outLogs/%s_stage4.log 2>  %s/outLogs/%s_stage4_err.log" % (scriptFn,ll,ll,outputFolder,s[0],outputFolder,ll,outputFolder,ll,ll,outputFolder,ll,ll,outputFolder,ll,outputFolder,ll)

    print outStr

    outFnSt = "%s/dirSub_%d/input_%d.in" % (outDir,dirCounter,rankCount)

    outPtr = open(outFnSt,"w+")

    if outPtr is None:

       print "Error: Opening %s" % outFnSt

    else:

       outPtr.write("#!/bin/bash"+"\n")

       outPtr.write(outStr+"\n")

       outPtr.close()

       if rankCount > fPer:

         dirCounter += 1

         rankCount = 0

       else:

         rankCount += 1

    outPtr.close()

    found += 1

  print "Found %d" % found

  print "Not found %d" % notFound

  print "End"

if __name__ == '__main__':

   main()

#PBS -S /bin/bash

#PBS -l select=1:ncpus=14:model=ivy

#PBS -l walltime=3:00:00

#PBS -j n

#PBS -m be

#PBS -N name

#PBS -V

#PBS -q normal

#!!!!!MUST SET OPTIONS BELOW!!!!

CORES=1

SUBDIRNUM=0

HDIR=

source /u/aguzman4/sharedModules/etc/environ.sh

module load mpi-mvapich2/1.4.1/intel

#$PBS_O_WORKDIR

date +"%T"

LOGSTDOUT=$HDIR/climate_ivy_serial_0b.stdout

LOGSTDERR=$HDIR/climate_ivy_serial_0b.stderr

mpiexec -pernode -comm none -np $CORES $HDIR/wrapper.sh $SUBDIRNUM $HDIR 1> $LOGSTDOUT 2> $LOGSTDERR

date +"%T"

#!/bin/bash 

rank=$MPIEXEC_RANK

indx=$1

hdir=$2

sh < $hdir/dirSub_$indx\/input_${rank}.in