# Plotting case figures tacking ERA-Interim files as reference
import os
import numpy as np
import generic_tools as gen
import nc_var_tools as ncvar
import drawing_tools as drw
import subprocess as sub

main = 'casefigures.py'

# Starting files from scratch
filescratch = False

# Starting figures from scratch
figscratch = False

# PyNCplot folder
pyHOME = '/home/lluis/PyNCplot'

# CDO folder
cdoHOME = '/usr/bin/cdo'

# GRIB folder
gribfold='/media/lluis/ExtDiskC_ext4/DATA/ECMWF/ERA-Interim/'

# head, middle and tail of grib name files
gribheadf = 'ERAI_'
gribmidf = '200512'
gribtailf = '.grib'

# Does the the code of the variable appears in the name of the grib file
gribcodeInfile = True

# GRIB table
gribtable = '128ECMWF'

# Folder where to keep generated netcdf files and figures
outfold = 'out'

# Dates to plot
idate='20051212120000'
edate='20051216120000'

# Slice netcdf
slicevals = 'lev|0,lat|-1,lon|-1'

# Map values
mapvalues='cyl,l'

# kind of figure output 
kindfig = 'png'

# Reverse of matrices in figure
remap = 'flip@y'

# Map range as ['SWlon','SWlat','NElon','NElat']
maprange=[-10., 25., 40., 55.]

# Axis values
# [dimxyfmt]=[dxs],[dxf],[Ndx],[ordx],[dys],[dyf],[Ndy],[ordx]
dimxyfmt = ' auto'

# Characteristics variables to make the files
#  [CFvarn]@[level (hpa)] level='sfc' for surface values
varlevs = ['z|100000', 'ta|100000', 'ua|100000', 'va|100000', 'hur|100000',          \
  'z|50000', 'ta|50000', 'ua|50000', 'va|50000', 'hur|50000', 'z|30000', 'ta|30000', \
  'ua|30000', 'va|30000', 'hur|30000', 'ps|sfc', 'psl|sfc']

# Figures
#   shdcont: shadow-contour plots 
#     [shadvnl]@[contvnl]@[minshad],[maxshad]@[mincnt],[maxcnt],[Nlev]@[colbar]@[cntype]
#   shd2cont: shadow, 2-contour plots 
#     [shadvnl]@[contvnl1]@[contvnl2]@[minshad],[maxshad]@[mincnt1],[maxcnt1],[Nlev1]@[mincnt2],[maxcnt2],[Nlev2]@[colbar]@[cntype1]@[cntype2]
#   shdvec: shadow-vector plots 
#     [shadvnl]@[xvecnl],[yvecnl]@[minshad],[maxshad]@[colbar]@[vectype]
#   shdcontvec: shadow-contour-vector plots 
#     [shadvnl]@[contvnl]@[xvecnl],[yvecnl]@[minshad],[maxshad]@[mincnt],[maxcnt],[Nlev]@[colbar]@[cntype]@[vectype]
#   shdcontbarb: shadow-contour-barb plots 
#     [shadvnl]@[contvnl]@[xnbarbnl],[ybarbnl]@[minshad],[maxshad]@[mincnt],[maxcnt],[Nlev]@[colbar]@[cntype]@[barbtype]

#### ### ## #
#    [XXXvnl]: [vn]|[lev], [vn] CF-variable name at level [lev]
#    [minXXX],[maxXXX]: minimun and maximum values to plot
#    [colbar]: name of the colormap to use
#    [cntype]=[ckind]|[clabfmt] values for the contour plot
#      ckind]: kind of contours
#        'cmap': as it gets from colorbar
#        'fixc,[colname]': fixed color [colname], all stright lines
#        'fixsigc,[colname]': fixed color [colname], >0 stright, <0 dashed line
#      clabfmt: format of the labels in the contour (None, also possible)
#    [vectype]= [frequency],[color],[length],[windname],[windunits]  values for the vector plot
#    [barbtype]: values for the barb plot

shd2cont = ['hur|100000@z|100000@ta|100000@60.,100.@-700.,4000.,8@265.,300.,8@BuPu@fixc,green|None@fixc,red|None',
'hur|100000@z|30000@psl|sfc@60.,100.@82000.,93000.,8@99000.,104000.,8@BuPu@fixc,green|None@fixc,red|None']
#shdcontbarb = [[shadvnl]@[contvnl]@[xnbarbnl],[ybarbnl]@[minshad],[maxshad]@[mincnt],[maxcnt],[Nlev]@[colbar]@[cntype]@[barbtype]]

availfigk = ['shad2cont']

# dictionary of figures
figures = {}
figures['shad2cont'] = shd2cont

####### ###### ##### #### ### ## #


#######    #######
## MAIN
    #######

cdolonlatS = str(maprange[0]) + ',' + str(maprange[2]) + ',' + str(maprange[1]) +    \
  ',' + str(maprange[3])

cdoidate = gen.datetimeStr_conversion(idate, 'YmdHMS', 'Y-m-d_H:M:S').replace('_','T')
cdoedate = gen.datetimeStr_conversion(edate, 'YmdHMS', 'Y-m-d_H:M:S').replace('_','T')

# Files in folder
gribfiles = gen.files_folder_HMT(folder=gribfold, head=gribheadf, middle=gribmidf,   \
  tail=gribtailf)
print 'grib files to use:', gribfiles

# Mapping variables in file
gribvarfiles = {}
for gribf in gribfiles:
    ins = sub.Popen([cdoHOME,"showcode",gribfold + '/' + gribf], stdout=sub.PIPE)
    codes = ins.communicate()
    codels = codes[0]
    codelist = codels.replace('\n','').split(' ')
    while gen.searchInlist(codelist,''): codelist.remove('')
    gribvarfiles[gribf] = list(np.array(codelist, dtype=int))

# creation of netcdf variable & level file from grib
ncfiles = []
# Time ranges for each file
dimts = {}
for vln in varlevs:
    varn = vln.split('|')[0]
    varlev = vln.split('|')[1]
    varcode = gen.CF_gribequiv(varn, gribtable)[0]
    varGRIBname = gen.CF_gribequiv(varn, gribtable)[1]

    varcodeS = str(varcode)

    ofilen = outfold + '/' + varn + '_' + varcodeS + '_' + varlev + '_' + idate +    \
      '-' + edate + '.nc'
    ncfiles.append(ofilen)

    if filescratch: 
        ins = 'rm -f ' + ofilen #+ ' >& /dev/null'
        print 'ins:', ins
        sub.call(ins, shell=True)

    if not os.path.isfile(ofilen):
        # Looking in which file there is the variable
        for gribfilen in gribvarfiles.keys():
            codesfile = gribvarfiles[gribfilen]
            if gen.searchInlist(codesfile, varcode):
                foundcode = True
                if varlev != 'sfc':
                    ins = cdoHOME + ' -f nc selcode,' + varcodeS + ' -sellevel,' +   \
                      varlev + ' -sellonlatbox,' + cdolonlatS + ' -seldate,' +       \
                      cdoidate + ',' + cdoedate
                else:
                    ins = cdoHOME + ' -f nc selcode,' + varcodeS + ' -sellonlatbox,'+\
                      cdolonlatS + ' -seldate,' + cdoidate + ',' + cdoedate
                try:
                    with gen.Capturing() as output:
                        sout = sub.call(ins + ' ' + gribfold+'/' + gribfilen + ' ' + \
                          ofilen, shell=True)
                except:
                    print gen.errormsg
                    print ins+' '+gribfold+'/' + gribfilen + ' ' + ofilen
                    print sout
                    for s1out in output: print s1out
                    quit(-1)

                ncvar.chvarname(varn, ofilen, varGRIBname)
                print "    created file '" + ofilen + "'"
                break

        if not foundcode:
            print gen.errormsg
            print "  variable: '"+varn+"' with code:",varcode, " not found !!"
            quit(-1)

# Checking time consitency
for ofilen in ncfiles:
    ins = sub.Popen([cdoHOME,"showtimestamp",ofilen], stdout=sub.PIPE)
    dts = ins.communicate()
    Ts = dts[0]
    dt = Ts.replace('\n','').split(' ')
    while gen.searchInlist(dt,''): 
        dt.remove('')
    dimts[ofilen] = dt

Nncfiles = len(ncfiles)
dt0 = len(dimts[ncfiles[0]])
for ifn in range(Nncfiles-1):
    gfilen = ncfiles[ifn+1]
    dt1 = len(dimts[gfilen])
    if dt1 != dt0:
        print gen.warnmsg
        print "  file '" + gfilen + "' has different time-steps:", dt1, " than '" +  \
          ncfiles[0] + "' with:", dt0

# Plotting
for figk in figures.keys():
    print figk
    figplot = figures[figk]

    for figp in figplot:
        print "  '" + figp + "' ..."
        figvalues = figp.split('@')

        # shad2cont
        if figk == 'shad2cont':
            shdvarn = figvalues[0].split('|')[0]
            shdlev = figvalues[0].split('|')[1]
            cnt1varn = figvalues[1].split('|')[0]
            cnt1lev = figvalues[1].split('|')[1]
            cnt2varn = figvalues[2].split('|')[0]
            cnt2lev = figvalues[2].split('|')[1]
            shdrange = figvalues[3]
            cnt1range = figvalues[4]
            cnt2range = figvalues[5]
            shdcolbar = figvalues[6]
            cnt1type = figvalues[7].replace('|',':')
            cnt2type = figvalues[8].replace('|',':')
            
            shdcode = gen.CF_gribequiv(shdvarn, gribtable)[0]
            cnt1code = gen.CF_gribequiv(cnt1varn, gribtable)[0]
            cnt2code = gen.CF_gribequiv(cnt2varn, gribtable)[0]

            shdfile= outfold + '/' + shdvarn + '_' + str(shdcode) + '_' + shdlev +   \
              '_' + idate+'-'+edate+'.nc'
            cnt1file= outfold + '/' + cnt1varn + '_' + str(cnt1code) + '_' + cnt1lev+\
              '_' + idate+'-'+edate+'.nc'
            cnt2file= outfold + '/' + cnt2varn + '_' + str(cnt2code) + '_' + cnt2lev+\
               '_' + idate+'-'+edate+'.nc'

            ncfile= shdfile + ',' + cnt1file + ',' + cnt2file
            varfignames = shdvarn + '@' + shdlev + ',' + cnt1varn + '@' + cnt1lev +  \
              ',' +  cnt2varn + '@' + cnt2lev
            varnames = shdvarn + ',' + cnt1varn + ',' +  cnt2varn

            timevals = dimts[shdfile]
            for it in range(len(timevals)):
                slicev = 'time|' + str(it) + ',' + slicevals
                slicesv = slicev + ':' + slicev + ':' + slicev
                timev = timevals[it].replace(':','').replace('-','').replace('T','')
 
                if shdlev == cnt1lev == cnt2lev:
                   titlefig= varnames.replace(',','|')+'|@'+ shdlev + '|on|' + timev
                else:
                   titlefig = varfignames.replace(',', '|') + '|on|' + timev

                values = varfignames + ':' + slicesv + ':lon:lat:auto:' + shdcolbar+ \
                  ',auto,auto:' + cnt1type + ':' + cnt2type + ':' + shdrange + ':' + \
                  cnt1range + ':' + cnt2range + ':' + titlefig + ':' + kindfig +     \
                  ':' + remap + ':' + mapvalues + ':yes'

                ins = pyHOME + '/drawing.py -o draw_2D_shad_2cont -f ' + ncfile +    \
                  " -S '" + values + "' -v " + varnames

                oimgn = outfold + '/' + figk + '_' + varfignames.replace('@','_p').replace(',','_') +\
                  '_' + timev + '.' + kindfig

                print 'Lluis ins:', ins

                if figscratch: sub.call('rm ' + oimgn, shell=True)

                if not os.path.isfile(oimgn):
                    try:
                        with gen.Capturing() as output:
                            sout = sub.call('python ' + ins, shell=True)
                    except:
                        print gen.errormsg
                        print 'python ' + ins
                        for s1out in output: print s1out
                        quit(-1)
                    if sout != 0:
                        print gen.errormsg
                        print '  : $ python ' + ins
                        sub.call('python ' + ins, shell=True)
                        quit(-1)

                    ins = 'mv 2Dfields_shadow-2contour.' + kindfig + ' ' + oimgn
                    sub.call(ins, shell=True)

                    #quit()

        else:
            print gen.errormsg
            print "  figure kind '" + figk + "' not found !!"
            print '    available ones:',  availfigk