# Python script to draw trajectories
#   L. Fita, LMD. May 2014
## e.g. # trajectories.py -m /home/lflmd/etudes/domains/WL_HyMeX/geo_em.d01.nc,XLONG_M,XLAT_M -s 'control:cu1:mp1:pblsfc1:rad1:shallow:wlmdza' -f /home/lflmd/etudes/WRF_LMDZ/WL_HyMeX/analysis/superstorm -w wss -M cyl,i -l -4.,33.,16.,44.

from optparse import OptionParser
import numpy as np
from netCDF4 import Dataset as NetCDFFile
import os
import re
import nc_var_tools as ncvar
import drawing_tools as drwpy

main='trajectories.py'

errormsg='ERROR -- error -- ERROR -- error'
warnmsg='WARNING -- warning -- WARNING -- warning'

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

parser = OptionParser()
parser.add_option("-f", "--folder", dest="infold", 
                  help="folder with the data", metavar="FOLD")
parser.add_option("-l", "--lonlatlims", dest="Lllim", 
                  help="OPTIONAL: limits of the map [lonmin, latmin, lonmax, latmax]", metavar="FOLD")
parser.add_option("-M", "--mapKind", dest="mapkind", 
                  help="""drawing coastaline ([proj],[res])
                    [proj]: projection
                      * 'cyl', cilindric
                    [res]: resolution:
                      * 'c', crude
                      * 'l', low
                      * 'i', intermediate
                      * 'h', high
                      * 'f', full      
                  """, metavar="FILE")
parser.add_option("-m", "--map", dest="map", 
                  help="[file],[lonname],[latname] ([file] with the [lon],[lat] matrices", metavar="FILE")
parser.add_option("-s", "--Simulations", dest="sims", 
                  help="':' list separated of the labels of the simulations", metavar="LABEL")
parser.add_option("-t", "--TimeInterval", dest="tint", 
                  help="OPTIONAL: dates interval to use [init],[endt]", metavar="FOLD")
parser.add_option("-w", "--simulated_variable", dest="svarname",
                  help="simulated variable to use", metavar="VAR")

(opts, args) = parser.parse_args()

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

sims= opts.sims.split(':')
leglabels = sims

Nsims = len(sims)

if opts.mapkind is not None:
    geofile = opts.map.split(',')[0]
    lonn = opts.map.split(',')[1]
    latn = opts.map.split(',')[2]

    if not os.path.isfile(geofile):
        print errormsg
        print '  ' + main + ' map file: "' + geofile + '" does not exist !!'
        quit(-1)

    geoobj = NetCDFFile(geofile, 'r')
    lonobj = geoobj.variables[lonn]
    latobj = geoobj.variables[latn]

if opts.tint is not None:
    tini = np.float(opts.tint.split(',')[0])
    tend = np.float(opts.tint.split(',')[0])
else:
    tini = -1
    tend = -1

if opts.Lllim is not None:
    lonlatlims = np.zeros((4), dtype=np.float)
    lonlatlims[0] = np.float(opts.Lllim.split(',')[0])
    lonlatlims[1] = np.float(opts.Lllim.split(',')[1])
    lonlatlims[2] = np.float(opts.Lllim.split(',')[2])
    lonlatlims[3] = np.float(opts.Lllim.split(',')[3])
##    lonlatlims = np.matrix(opts.Lllim.split(','), dtype=np.float)

else:
    lonlatlims = None

lontrjv = []
lattrjv = []

for sim in sims:
    trjfile=opts.infold + '/' + sim + '/tevolboxtraj_' + opts.svarname + '.nc'
    if not os.path.isfile(trjfile):
        print errormsg
        print '  ' + main + ' trajectory file: "' + trjfile + '" does not exist !!'
        quit(-1)

    trjobj = NetCDFFile(trjfile, 'r')
    lontr = trjobj.variables['trlon']
    lattr = trjobj.variables['trlat']

    lontrjv.append(lontr[:])
    lattrjv.append(lattr[:])
    trjobj.close()

drwpy.plot_Trajectories(lontrjv, lattrjv, leglabels, tini, tend, lonobj, latobj,     \
  lonlatlims, 'trajectories', 'pdf', opts.mapkind)