#! /usr/bin/env python
from    netCDF4               import    Dataset
from	numpy		      import	*
import  numpy                 as        np
import  matplotlib.pyplot     as        mpl
from matplotlib.cm import get_cmap
import pylab
import matplotlib.colors as mcolors
import colorsys
from FV3_utils import *

############################
filename1=name+"_A.nc"
var="v" #variable
tint=[30,31] #Time must be as written in the input file
xarea=[-180,179]

nc1=Dataset(filename1)

lat=getvar(nc1,"latitude")
lon=getvar(nc1,"longitude")
alt=getvar(nc1,"altitude")
tim=getvar(nc1,"Time")
############################

def make_colormap(seq):
    """Return a LinearSegmentedColormap
    seq: a sequence of floats and RGB-tuples. The floats should be increasing
    and in the interval (0,1).
    """

    seq = [(None,) * 3, 0.0] + list(seq) + [1.0, (None,) * 3]
    print(seq)
    cdict = {'red': [], 'green': [], 'blue': []}
    for i, item in enumerate(seq):
        if isinstance(item, float):
            r1, g1, b1 = seq[i - 1]
            r2, g2, b2 = seq[i + 1]
            cdict['red'].append([item, r1, r2])
            cdict['green'].append([item, g1, g2])
            cdict['blue'].append([item, b1, b2])
    print(cdict)
    return mcolors.LinearSegmentedColormap('CustomMap', cdict)

def diverge_map(high, low):
    '''
    low and high are colors that will be used for the two
    ends of the spectrum. they can be either color strings
    or rgb color tuples
    '''
    c = mcolors.ColorConverter().to_rgb
    if isinstance(low, str): low = c(low)      #si low=string (color)
    if isinstance(high, str): high = c(high)
    print(high)
    return make_colormap([low, c('white'), 0.55, c('white'),0.65, c('white'), high])


hh=(255,99,71)
hh=(hh[0]/255,hh[1]/255,hh[2]/255)
print(hh)
h=hh #(0.565, 0.392, 0.173)
l=(0.094, 0.310, 0.635)
rvb1=diverge_map(h,l)

myvar=getvar(nc1,var,tint,tim,xarea,lon,t_mean=True,l_mean=True)

mpl.figure(figsize=(20, 10))

font=26

#pal=rvb1 #get_cmap(name="RdYlBu_r")
#pal=get_cmap(name="Spectral_r")
pal=get_cmap(name="rainbow")
lev=np.linspace(-0.1,0.1,25)

xticks=[-90,-60,-30,0,30,60,90]
#yticks=np.linspace(0,240,9)
alt=alt/1000.

CF=mpl.contourf(lat,alt,myvar,lev,cmap=pal,extend='both')
cbar=mpl.colorbar(CF,shrink=1, format="%1.2f")
#cbar.ax.set_title("[K]",y=1.04,fontsize=font)
for t in cbar.ax.get_yticklabels():
      t.set_fontsize(font)

vect=lev
CS=mpl.contour(lat,alt,myvar,vect,colors='k',linewidths=0.5)
#### inline=1 : values over the line
mpl.clabel(CS, inline=1, fontsize=20, fmt='%1.1f',inline_spacing=1)

#mpl.title('Latitude ='+str(tintstr[i]),fontsize=font)
mpl.ylabel('Altitude (km)',labelpad=10,fontsize=font)
mpl.xlabel('Latitude (deg)',labelpad=10, fontsize=font)
mpl.xticks(xticks,fontsize=font)
#mpl.xticks(fontsize=font)
#mpl.yticks(yticks,fontsize=font)
mpl.yticks(fontsize=font)
pylab.ylim([-4,250])

mpl.savefig('meanmeridwind.eps',dpi=200)
mpl.savefig('meanmeridwind.png',dpi=200)
mpl.show()