#######################
##### PLANETOPLOT #####
#######################

### A. Spiga     -- LMD -- 06~09/2011 -- General building and mapping capabilities
### T. Navarro   -- LMD -- 10~11/2011 -- Improved use for GCM and added sections + 1Dplot capabilities 
### A. Colaitis  -- LMD --            -- Mostly minor improvements and inter-plot operation capabilities + zrecast interpolation for gcm

def planetoplot (namefiles,\
           level=0,\
           vertmode=0,\
           proj=None,\
           back=None,\
           target=None,
           stride=3,\
           numplot=None,\
           var=None,\
           colorb="def",\
           winds=True,\
           addchar=None,\
           interv=[0,1],\
           vmin=None,\
           vmax=None,\
           tile=False,\
           zoom=None,\
           display=True,\
           itstep=None,\
           hole=False,\
           save="gui",\
           anomaly=False,\
           var2=None,\
           ndiv=10,\
           first=1,\
           mult=1.,\
           zetitle="fill",\
           slon=None,\
           slat=None,\
           svert=None,\
           stime=None,\
           outputname=None,\
           resolution=200,\
           ope=None,\
           fileref=None,\
           minop=0.,\
           maxop=0.,\
           titleref="fill",\
           invert_y=False,\
           xaxis=[None,None],\
           yaxis=[None,None],\
           ylog=False):


    ####################################################################################################################
    ### Colorbars http://www.scipy.org/Cookbook/Matplotlib/Show_colormaps?action=AttachFile&do=get&target=colormaps3.png

    #################################
    ### Load librairies and functions
    from netCDF4 import Dataset
    from myplot import getcoord2d,define_proj,makeplotres,simplinterv,vectorfield,ptitle,latinterv,getproj,wrfinterv,dumpbdy,\
                       fmtvar,definecolorvec,defcolorb,getprefix,putpoints,calculate_bounds,errormess,definesubplot,\
                       zoomset,getcoorddef,getwinddef,whatkindfile,reducefield,bounds,getstralt,getfield,smooth,nolow,\
                       getname,localtime,polarinterv,getsindex,define_axis,determineplot
    from mymath import deg,max,min,mean
    import matplotlib as mpl
    from matplotlib.pyplot import contour,contourf, subplot, figure, rcParams, savefig, colorbar, pcolor, show, plot, clabel
    from matplotlib.cm import get_cmap
    import numpy as np
    from numpy.core.defchararray import find

    ################################
    ### Which plot needs to be done?
    nlon, nlat, nvert, ntime, mapmode, nslices = determineplot(slon, slat, svert, stime)
    if mapmode == 0: winds=False
    if not isinstance(namefiles, np.ndarray): namefiles = [namefiles]
    zelen = len(namefiles)
    if numplot == None:  numplot = zelen*nslices
    print "len(namefiles), nslices, numplot: ", zelen, nslices, numplot
    if fileref is not None:
        all_var  = [[]]*(zelen+2)
    else:
        all_var   = [[]]*zelen
        all_var2  = [[]]*zelen
        all_title = [[]]*zelen
    
    #########################
    ### Loop over the files initially separated by comas to be plotted on the same figure
    k = 0
    firstfile = True
    for namefile in namefiles:
      print namefile 
      ######################
      ### Load NETCDF object
      nc  = Dataset(namefile)

      ##################################
      ### Initial checks and definitions
      ### ... TYPEFILE
      typefile = whatkindfile(nc)                                  
      if firstfile:
         typefile0 = typefile
      elif typefile != typefile0:
         errormess("Not the same kind of files !", [typefile0, typefile])
      ### ... VAR
      if var not in nc.variables: var = False
      ### ... WINDS
      if winds:                                                    
         [uchar,vchar,metwind] = getwinddef(nc)             
         if uchar == 'not found': winds = False
      if not var and not winds: errormess("please set at least winds or var",printvar=nc.variables)
      ### ... COORDINATES, could be moved below
      [lon2d,lat2d] = getcoorddef(nc)                       
      ### ... PROJECTION
      if proj == None:   proj = getproj(nc)                  

##########################################################
      if typefile == "gcm":
          lat = nc.variables["latitude"][:] 
          lon = nc.variables["longitude"][:]
          if "Time" in nc.variables:
             time = nc.variables["Time"][:]
          elif "time" in nc.variables:
             time = nc.variables["time"][:]
          else:
             print "no time axis found."
             exit()
          vert = nc.variables["altitude"][:]
      elif typefile in ['meso','mesoapi']:
          if mapmode == 0:
              if var in ['PHTOT','W']:    vertdim='BOTTOM-TOP_PATCH_END_STAG'
              else:                       vertdim='BOTTOM-TOP_PATCH_END_UNSTAG'
              if var in ['V']:  latdim='SOUTH-NORTH_PATCH_END_STAG'
              else:             latdim='SOUTH-NORTH_PATCH_END_UNSTAG'
              if var in ['U']:  londim='WEST-EAST_PATCH_END_STAG'
              else:             londim='WEST-EAST_PATCH_END_UNSTAG'
              lon = np.arange(0,getattr(nc,londim),1)
              lat = np.arange(0,getattr(nc,latdim),1)
              if vertmode is None: vertmode=0
              if vertmode == 0:   vert = np.arange(0,getattr(nc,vertdim),1)
              else:               vert = nc.variables["vert"][:]
              time = np.arange(0,len(nc.variables["Times"]),1)
       #if firstfile:
       #   lat0 = lat
       #elif len(lat0) != len(lat):
       #   errormess("Not the same latitude lengths !", [len(lat0), len(lat)])
       #elif sum((lat == lat0) == False) != 0:
       #   errormess("Not the same latitudes !", [lat,lat0])
       ## Faire d'autre checks sur les compatibilites entre fichiers!!
##########################################################

      if firstfile:
         ##########################
         ### Define plot boundaries
         ### todo: possible areas in latinterv in argument (ex: "Far_South_Pole")
         if proj in ["npstere","spstere"]: [wlon,wlat] = polarinterv(lon2d,lat2d)
         elif proj in ["lcc","laea"]:      [wlon,wlat] = wrfinterv(lon2d,lat2d)
         else:                             [wlon,wlat] = simplinterv(lon2d,lat2d)
         if zoom:                          [wlon,wlat] = zoomset(wlon,wlat,zoom) 

         #########################################
         ### Name for title and graphics save file
         basename = getname(var=var,winds=winds,anomaly=anomaly)
         basename = basename + getstralt(nc,level)  ## can be moved elsewhere for a more generic routine
      
      print "var, var2: ", var, var2
      if var: all_var[k] = getfield(nc,var)
      if var2: all_var2[k] = getfield(nc,var2)
      
      print "k", k
      print "all_var[k].shape", all_var[k].shape
      k += 1
      firstfile = False
      #### End of for namefile in namefiles

    ##################################
    ### Load reference file if there is one
    if fileref is not None:
       ncref  = Dataset(fileref)
       ### ... VAR
       if var not in nc.variables: var = False
       if var: 
             all_var[k] = getfield(ncref,var)
             if ope is "-":
                all_var[k+1]= all_var[k-1] - all_var[k]
             elif ope is "+":
                all_var[k+1]= all_var[k-1] + all_var[k]
             elif ope is not None:
                print "non-implemented operation. Available op. are + and -."
             numplot = numplot+2

    ##################################
    ### Open a figure and set subplots
    fig = figure()
    subv,subh = definesubplot( numplot, fig ) 
 
    #################################
    ### Time loop for plotting device
    found_lct = False
    nplot = 1
    itime = first
    error = False
    if itstep is None and numplot > 0: itstep = int(24./numplot)
    elif numplot <= 0:                 itstep = 1 
    
    #for nplot in range(numplot):
    while error is False:
       print "nplot", nplot
       print error
       
       ### Which local time ?
       ltst = localtime ( interv[0]+itime*interv[1], 0.5*(wlon[0]+wlon[1]) )

       ### General plot settings
       #print itime, int(ltst), numplot, nplot
       if numplot >= 1: 
           if nplot > numplot: break
           if numplot > 1:     
               if typefile not in ['geo']:  subplot(subv,subh,nplot)
           found_lct = True
       ### If only one local time is requested (numplot < 0)
       elif numplot <= 0: 
           if int(ltst) + numplot != 0:	
                 itime += 1 
                 if found_lct is True: break     ## because it means LT was found at previous iteration
                 else:                 continue  ## continue to iterate to find the correct LT
           else: 
                 found_lct = True

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

####################################################################
       if typefile in ['meso','mesoapi'] and mapmode == 1: 
               indextime = itime
               indexlon = None
               indexlat = None
               indexvert = level  ## ou svert ??? 
               nlon = 1
               nlat = 1
               nvert = 1
               ntime = 1 
       else:
           ## get all indexes to be taken into account for this subplot and then reduce field
           ## We plot 1) all lon slices 2) all lat slices 3) all vert slices 4) all time slices and then go to the next slice
           indexlon  = getsindex(slon,(nplot-1)%nlon,lon)
           indexlat  = getsindex(slat,((nplot-1)//nlon)%nlat,lat)
           indexvert = getsindex(svert,((nplot-1)//(nlon*nlat))%nvert,vert)  
           indextime = getsindex(stime,((nplot-1)//(nlon*nlat*nvert))%ntime,time)  

       if fileref is not None:
           index_f = ((nplot-1)//(nlon*nlat*nvert*ntime))%(len(namefiles)+2)
       else:
           index_f = ((nplot-1)//(nlon*nlat*nvert*ntime))%len(namefiles)

       print nlon, nlat, nvert, ntime  
       print slon, slat, svert, stime       
       print index_f
       print "index lon,lat,vert,time", max(indexlon), max(indexlat), max(indexvert), max(indextime)
####################################################################

       #### Contour plot
       if var2:
           what_I_plot, error = reducefield(all_var2[index_f], d4=indextime, d1=indexlon, d2=indexlat , d3=indexvert )
           #what_I_plot = what_I_plot*mult
           if not error:
              if mapmode == 1:
                  if typefile in ['mesoapi','meso']:    what_I_plot = dumpbdy(what_I_plot,6)
              zevmin, zevmax = calculate_bounds(what_I_plot)
              zelevels = np.linspace(zevmin,zevmax,num=ndiv+1) #20)
              if var2 == 'HGT':  zelevels = np.arange(-10000.,30000.,2000.)
              if mapmode == 0:
                  x, y = define_axis(lon,lat,vert,time,indexlon,indexlat,indexvert,\
                        indextime,what_I_plot.shape, len(all_var2[index_f].shape),vertmode)
              ### If we plot a 2-D field
              if len(what_I_plot.shape) is 2:
                  cs = contour(x,y,what_I_plot, zelevels, colors='k', linewidths = 1 ) #0.33 colors='w' )# , alpha=0.5)
                  if typefile in ['gcm']: clabel(cs,fmt = '%1.2e')
              ### If we plot a 1-D field
              elif len(what_I_plot.shape) is 1:
                  plot(what_I_plot,x) 
           else:
              errormess("There is an error in reducing field !")

       #### Shaded plot
       if var:
           what_I_plot, error = reducefield(all_var[index_f], d4=indextime, d1=indexlon, d2=indexlat , d3=indexvert )
           what_I_plot = what_I_plot*mult
           if not error: 
               fvar = var
               ###
               if anomaly:
                   what_I_plot = 100. * ((what_I_plot / smooth(what_I_plot,12)) - 1.)
                   fvar = 'anomaly'
               #if mult != 1:     
               #    fvar = str(mult) + "*" + var
               ###
               if mapmode == 1:
                   if typefile in ['mesoapi','meso']:    what_I_plot = dumpbdy(what_I_plot,6)
               elif mapmode == 0:
                   what_I_plot, x, y = define_axis(lon,lat,vert,time,indexlon,indexlat,indexvert,\
                         indextime,what_I_plot, len(all_var[index_f].shape),vertmode)
               if (fileref is not None) and (index_f is numplot-1):
                  zevmin, zevmax = calculate_bounds(what_I_plot,vmin=minop,vmax=maxop)
               else:
                  zevmin, zevmax = calculate_bounds(what_I_plot,vmin=vmin,vmax=vmax)
               if colorb in ["def","nobar"]:   palette = get_cmap(name=defcolorb(fvar.upper()))
               else:                           palette = get_cmap(name=colorb)
               ##### 2D field
               if len(what_I_plot.shape) is 2:
                 if not tile:
                     if not hole: what_I_plot = bounds(what_I_plot,zevmin,zevmax)
                     #zelevels = np.linspace(zevmin*(1. + 1.e-7),zevmax*(1. - 1.e-7)) #,num=20)
                     zelevels = np.linspace(zevmin,zevmax)#,num=ndiv+1)
                     #contourf(what_I_plot, zelevels, cmap = palette )

                     contourf( x, y, what_I_plot, zelevels, cmap = palette)

                     zxmin, zxmax = xaxis
                     zymin, zymax = yaxis
                     if zxmin is not None:
                        mpl.pyplot.xlim(xmin=zxmin)
                     if zxmax is not None:
                        mpl.pyplot.xlim(xmax=zxmax)
                     if zymin is not None:
                        mpl.pyplot.ylim(ymin=zymin)
                     if zymax is not None:
                        mpl.pyplot.ylim(ymax=zymax)

                     if invert_y:
                        lima,limb = mpl.pyplot.ylim()
                        mpl.pyplot.ylim(limb,lima)
                     if ylog:
                        mpl.pyplot.semilogy()

                 else:
                     if hole:  what_I_plot = nolow(what_I_plot)
                     pcolor( x, y, what_I_plot, cmap = palette, vmin=zevmin, vmax=zevmax )
                 if colorb != 'nobar' and var != 'HGT' :        
                     if (fileref is not None) and (index_f is numplot-1):
                        colorbar(fraction=0.05,pad=0.03,format="%.2f",\
                                           ticks=np.linspace(zevmin,zevmax,min([ndiv+1,20])),\
                                           extend='neither',spacing='proportional')
                     else:
                        colorbar(fraction=0.05,pad=0.03,format=fmtvar(fvar.upper()),\
                                           ticks=np.linspace(zevmin,zevmax,min([ndiv+1,20])),\
                                           extend='neither',spacing='proportional')

                                           # both min max neither
               ##### 1D field
               elif len(what_I_plot.shape) is 1:
                 plot(x,what_I_plot)
               #### Other cases: (maybe plot 3-D field one day or movie ??)
               else:
                 print "WARNING!!! ",len(what_I_plot.shape),"-D PLOT NOT SUPPORTED !!!"
                 print "field dimensions: ", what_I_plot.shape
                 exit()
               
           else:
               errormess("There is an error in reducing field !")

       ### Vector plot --- a adapter
       if winds:
           vecx, error = reducefield( getfield(nc,uchar), d4=indextime, d3=indexvert )
           vecy, error = reducefield( getfield(nc,vchar), d4=indextime, d3=indexvert )
           #what_I_plot, error = reducefield(all_var[index_f], d4=indextime, d1=indexlon, d2=indexlat , d3=indexvert )
           if not error:
               if typefile in ['mesoapi','meso']:    
                   [vecx,vecy] = [dumpbdy(vecx,6,stag=uchar), dumpbdy(vecy,6,stag=vchar)]
                   key = True
               elif typefile in ['gcm']:            
                   key = False
               if metwind:  [vecx,vecy] = m.rotate_vector(vecx, vecy, lon2d, lat2d)
               if var == False:       colorvec = definecolorvec(back)
               else:                  colorvec = definecolorvec(colorb)
               vectorfield(vecx, vecy,\
                          x, y, stride=stride, csmooth=2,\
                          #scale=15., factor=300., color=colorvec, key=key)
                          scale=20., factor=250., color=colorvec, key=key)
                                            #200.         ## or csmooth=stride
    
       ### Next subplot
       if typefile in ['mesoapi','meso']:
            plottitle = basename
            if addchar:  plottitle = plottitle + addchar + "_LT"+str(ltst)
            else:        plottitle = plottitle + "_LT"+str(ltst)
       else:
            if fileref is not None:
                if index_f is numplot-1:
                   plottitle = basename+' '+"fig(1) "+ope+" fig(2)"
                elif index_f is numplot-2:
                   plottitle = basename+' '+fileref
                else:
                   plottitle = basename+' '+namefiles[index_f]
            else:
                plottitle = basename+' '+namefiles[index_f]
       if mult != 1:           plottitle = str(mult) + "*" + plottitle
       if zetitle != "fill":   
          plottitle = zetitle
          if titleref is "fill":
             titleref=zetitle
          if fileref is not None:
             if index_f is numplot-2:
                plottitle = titleref
             if index_f is numplot-1:
                plottitle = "fig(1) "+ope+" fig(2)"
#       if indexlon is not None:
#         plottitle = plottitle + " lon: " + str(min(lon[indexlon])) +" "+ str(max(lon[indexlon]))
#       if indexlat is not None:
#         plottitle = plottitle + " lat: " + str(min(lat[indexlat])) +" "+ str(max(lat[indexlat]))
#       if indexvert is not None:
#         plottitle = plottitle + " vert: " + str(min(vert[indexvert])) +" "+ str(max(vert[indexvert]))
#       if indextime is not None:
#         plottitle = plottitle + " time: " + str(min(time[indextime])) +" "+ str(max(time[indextime]))
       ptitle( plottitle )
       itime += itstep
       if nplot >= numplot:
          error = True
       nplot += 1

  





      
    ##########################################################################
    ### Save the figure in a file in the data folder or an user-defined folder
    if outputname is None:
       if typefile in ['meso','mesoapi']:   prefix = getprefix(nc)
       elif typefile in ['gcm']:            prefix = 'LMD_GCM_'
       else:                                prefix = ''
    ###
       zeplot = prefix + basename 
       if addchar:         zeplot = zeplot + addchar
       if numplot <= 0:    zeplot = zeplot + "_LT"+str(abs(numplot))
    ###
       if not target:      zeplot = namefile[0:find(namefile,'wrfout')] + zeplot
       else:               zeplot = target + "/" + zeplot  
    ###
    else:
       zeplot=outputname

    if found_lct:     
        pad_inches_value = 0.35
        print "save", save
        if save == 'png': 
            if display: makeplotres(zeplot,res=100.,pad_inches_value=pad_inches_value) #,erase=True)  ## a miniature
            makeplotres(zeplot,res=resolution,pad_inches_value=pad_inches_value,disp=False)
        elif save in ['eps','svg','pdf']:
            makeplotres(zeplot,         pad_inches_value=pad_inches_value,disp=False,ext=save)
        elif save == 'gui':
            show()
        else: 
            print "save mode not supported. using gui instead."
            show()
    else:   print "Local time not found"

    ###############
    ### Now the end
    return zeplot
