#!/usr/bin/env python

### A. Spiga -- LMD -- 30/05/2011
### Thanks to A. Colaitis for the parser trick


####################################
####################################
### The main program to plot vectors
def winds (namefile,\
           nvert,\
           proj=None,\
           back=None,\
           target=None,
           stride=3,\
           numplot=4,\
           var=None,\
           colorb=None):

    #################################
    ### Load librairies and functions
    from netCDF4 import Dataset
    from myplot import getcoord2d,define_proj,makeplotpng,simplinterv,vectorfield,ptitle,latinterv,getproj,wrfinterv,dumpbdy
    from matplotlib.pyplot import contourf, subplot, figure, rcParams, savefig, colorbar
    import numpy as np

    ######################
    ### Load NETCDF object
    nc  = Dataset(namefile)

    ###################################
    ### Recognize predefined file types
    if 'controle' in nc.variables:   typefile = 'gcm'
    elif 'Um' in nc.variables:       typefile = 'mesoapi'
    elif 'U' in nc.variables:        typefile = 'meso'
    else:                            
        print "typefile not supported."
        exit()

    ##############################################################
    ### Try to guess the projection from wrfout if not set by user 
    if typefile in ['mesoapi','meso']:
        if proj == None:       proj = getproj(nc)
                                    ### (il faudrait passer CEN_LON dans la projection ?)
    elif typefile in ['gcm']:
        if proj == None:       proj = "cyl"   
                                    ## pb avec les autres (de trace derriere la sphere ?)

    ############################################
    #### Choose underlying topography by default 
    if not back:
        if not var:                                        back = "mola"    ## if no var:         draw mola
        elif typefile in ['mesoapi','meso'] \
             and proj not in ['merc','lcc']:               back = "molabw"  ## if var but meso:   draw molabw
        else:                                              pass             ## else:              draw None

    ####################################################
    ### Get geographical coordinates and plot boundaries
    if typefile in ['mesoapi','meso']:
        [lon2d,lat2d] = getcoord2d(nc)
        lon2d = dumpbdy(lon2d)
        lat2d = dumpbdy(lat2d)
    elif typefile in ['gcm']:
        [lon2d,lat2d] = getcoord2d(nc,nlat="latitude",nlon="longitude",is1d=True)
    if proj == "npstere":    [wlon,wlat] = latinterv("North_Pole")
    elif proj == "lcc":      [wlon,wlat] = wrfinterv(lon2d,lat2d)
    else:                    [wlon,wlat] = simplinterv(lon2d,lat2d)

    ##################
    ### Get local time
    if typefile in ['mesoapi','meso']:  ltst = int(getattr(nc, 'GMT') + 0.5*(wlon[0]+wlon[1])/15.)
    elif typefile in ['gcm']:           ltst = 0

    ##############################################################################
    ### Get winds and know if those are meteorological winds (ie. zon, mer) or not
    if typefile is 'mesoapi': 
        [u,v] = getwinds(nc)
        metwind = True  ## meteorological (zon/mer)
    elif typefile is 'gcm':
        [u,v] = getwinds(nc,charu='u',charv='v')
        metwind = True  ## meteorological (zon/mer)
    elif typefile is 'meso':
        [u,v] = getwinds(nc,charu='U',charv='V')
        metwind = False ## geometrical (wrt grid)
        print "Beware ! Not using meteorological winds. You trust numerical grid as being (x,y)."

    ####################################################
    ### Load the chosen variables, whether it is 2D or 3D
    if var:
        if var not in nc.variables: 
            print "not found in file:",var
            exit()
        else:   
            dimension = np.array(nc.variables[var]).ndim
            if dimension == 2:     field = nc.variables[var][:,:]
            elif dimension == 3:   field = nc.variables[var][:,:,:]
            elif dimension == 4:   field = nc.variables[var][:,nvert,:,:]  
    nt = len(u[:,0,0,0])

    #########################################
    ### Name for title and graphics save file
    basename = "WINDS"
    if var:
        basename = basename + "_" + var
    if   typefile is 'meso':                    stralt = "_lvl" + str(nvert)
    elif typefile is 'mesoapi':  
        zelevel = int(nc.variables['vert'][nvert])
        if 'altitude_abg'   in nc.dimensions:   stralt = "_"+str(zelevel)+"m-ALS"
        elif 'bottom_top'   in nc.dimensions:   stralt = "_"+str(zelevel)+"m-AMR"
        elif 'pressure'     in nc.dimensions:   stralt = "_"+str(zelevel)+"Pa" 
    else:                                       stralt = ""          
    basename = basename + stralt

    ##################################
    ### Open a figure and set subplots
    fig = figure()
    if   numplot > 0:    
        if   numplot == 4: 
            sub = 221
            fig.subplots_adjust(wspace = 0.1, hspace = 0.3)
            rcParams['font.size'] = int( rcParams['font.size'] * 2. / 3. )
        elif numplot == 2: 
            sub = 121
            fig.subplots_adjust(wspace = 0.3)
            rcParams['font.size'] = int( rcParams['font.size'] * 3. / 4. )
        elif numplot == 3: 
            sub = 131
            fig.subplots_adjust(wspace = 0.3)
            rcParams['font.size'] = int( rcParams['font.size'] * 2. / 3. )
        elif numplot == 6: 
            sub = 231
            fig.subplots_adjust(wspace = 0.4, hspace = 0.0)
            rcParams['font.size'] = int( rcParams['font.size'] * 2. / 3. )
        elif numplot == 8: 
            sub = 331 #241
            fig.subplots_adjust(wspace = 0.1, hspace = 0.3)
            rcParams['font.size'] = int( rcParams['font.size'] * 1. / 2. )
        elif numplot == 9:
            sub = 331 
            fig.subplots_adjust(wspace = 0.1, hspace = 0.3)
            rcParams['font.size'] = int( rcParams['font.size'] * 1. / 2. )
        elif numplot == 1:
            sub = 99999
        else:
            print "supported: 1,2,3,4,6,8"
            exit()
        ### Prepare time loop
        if nt <= numplot or numplot == 1:  
            tabrange = [0]
            numplot = 1
        else:                          
            tabrange = range(0,nt,int(nt/numplot))  #nt-1
            tabrange = tabrange[0:numplot]
    else: 
        tabrange = range(0,nt,1)
        sub = 99999
    print tabrange

    #################################
    ### Time loop for plotting device
    found_lct = False
    for i in tabrange:

       ### General plot settings
       if numplot > 1: 
           subplot(sub)
           found_lct = True
       elif numplot == 1:
           found_lct = True 
            ### If only one local time is requested (numplot < 0)
       elif numplot <= 0: 
           if (ltst+i)%24 + numplot != 0:   continue
           else:                            found_lct = True

       ### Map projection
       m = define_proj(proj,wlon,wlat,back=back)
       x, y = m(lon2d, lat2d)

       #### Contour plot
       if var:
           if typefile in ['mesoapi','meso']:    what_I_plot = dumpbdy(field[i,:,:])
           elif typefile in ['gcm']:             
               if dimension == 2:                what_I_plot = field[:,:]
               elif dimension == 3:              what_I_plot = field[i,:,:]
           contourf(x, y, what_I_plot, 30)
           if colorb:     colorbar(fraction=0.05,pad=0.1)

       ### Vector plot
       if   typefile in ['mesoapi','meso']:    
           [vecx,vecy] = [dumpbdy(u[i,nvert,:,:]), dumpbdy(v[i,nvert,:,:])]
           key = True
       elif typefile in ['gcm']:               
           [vecx,vecy] = [        u[i,nvert,:,:] ,         v[i,nvert,:,:] ]
           key = False
       if metwind:  [vecx,vecy] = m.rotate_vector(vecx, vecy, lon2d, lat2d)
       if var == None and back == "vishires":  colorvec = 'w'
       else:                                   colorvec = 'k'
       vectorfield(vecx, vecy,\
                      x, y, stride=stride, csmooth=stride,\
                      scale=15., factor=300., color=colorvec, key=key)
                                        #200.         ## or csmooth=2
       
       ### Next subplot
       ptitle( basename + "_LT"+str((ltst+i)%24) )
       sub += 1

    ##########################################################################
    ### Save the figure in a file in the data folder or an user-defined folder
    if not target:    zeplot = namefile+"_"+basename
    else:             zeplot = target+"/"+basename
    if numplot <= 0:  zeplot = zeplot + "_LT"+str(abs(numplot))
    if found_lct:     makeplotpng(zeplot,pad_inches_value=0.35)   
    else:             print "Local time not found"



####################################################
####################################################
### A simple program to get wind vectors' components
def getwinds (nc,charu='Um',charv='Vm'):
    import numpy as np
    u = nc.variables[charu]
    v = nc.variables[charv]
    if charu == 'U': u = u[:, :, :, 0:len(u[0,0,0,:])-1]
    if charv == 'V': v = v[:, :, 0:len(v[0,0,:,0])-1, :]
                     ### ou alors prendre les coordonnees speciales
    return u,v



###########################################################################################
###########################################################################################
### What is below relate to running the file as a command line executable (very convenient)
if __name__ == "__main__":
    import sys
    from optparse import OptionParser    ### to be replaced by argparse
    from api_wrapper import api_onelevel
    parser = OptionParser()
    parser.add_option('-f', action='store', dest='namefile',    type="string",  default=None,  help='[NEEDED] name of WRF file')
    parser.add_option('-l', action='store', dest='nvert',       type="float",   default=0,     help='vertical level (def=0)(-i 2: p,mbar)(-i 3,4: z,km)')
    parser.add_option('-p', action='store', dest='proj',        type="string",  default=None,  help='projection')
    parser.add_option('-b', action='store', dest='back',        type="string",  default=None,  help='background')
    parser.add_option('-t', action='store', dest='target',      type="string",  default=None,  help='destination folder')
    parser.add_option('-s', action='store', dest='stride',      type="int",     default=3,     help='stride vectors (def=3)')
    parser.add_option('-v', action='store', dest='var',         type="string",  default=None,  help='variable contoured')
    parser.add_option('-n', action='store', dest='numplot',     type="int",     default=4,     help='number of plots (def=1)(<0: 1 plot of LT -*numplot*)')
    parser.add_option('-i', action='store', dest='interp',      type="int",     default=None,  help='interpolation method (2: press, 3: z-amr, 4:z-als)')
    parser.add_option('-c', action='store', dest='colorb',      type="string",  default=None,  help='colorbar')
    (opt,args) = parser.parse_args()
    if opt.namefile is None: 
        print "I want to eat one file at least ! Use winds.py -f name_of_my_file. Or type winds.py -h"
        exit()
    print "Options:", opt

    zefile = opt.namefile    
    zelevel = opt.nvert   
    stralt = None
    if opt.interp is not None:
        if opt.nvert is 0 and opt.interp is 4:  zelevel = 0.010
        if   opt.var is None    :  zefields = 'uvmet'
        else                    :  zefields = 'uvmet,'+opt.var
        zefile = api_onelevel (  path_to_input   = '', \
                                 input_name      = zefile, \
                                 path_to_output  = opt.target, \
                                 fields          = zefields, \
                                 interp_method   = opt.interp, \
                                 onelevel        = zelevel )
        zelevel = 0 ## so that zelevel could play again the role of nvert

    winds (zefile,int(zelevel),\
           proj=opt.proj,back=opt.back,target=opt.target,stride=opt.stride,var=opt.var,numplot=opt.numplot,colorb=opt.colorb)
