#! /usr/bin/env python
from ppclass import pp
from    netCDF4               import    Dataset
from	numpy		      import	*
import  numpy                 as        np
import  matplotlib.pyplot     as        mpl
from matplotlib.cm import get_cmap
import pylab
from matplotlib import ticker
import matplotlib.colors as mcolors
from mpl_toolkits.basemap import Basemap

############################
filename="diagfi2015.nc"
nb_dataset=1  #change vect as well 
tint=["0,0.001"] #Time must be as written in the input file
#zint=["0"] #alt in km
#xarea="-180,179"
#yarea="-90,90"
#step=5  #step between each bin
var="albedo" #variable
var2="phisinit"
############################

nc=Dataset(filename)
lat=nc.variables["lat"][:]
lon=nc.variables["lon"][:]

myvar = pp(file=filename,var=var,t=tint,compute="mean").getf()  # get data to be changed according to selected variable
myvar2 = pp(file=filename,var=var2,t=tint,compute="mean").getf()  # get data to be changed according to selected variable
print(('shape var1 =',shape(myvar)))
print(('shape lat = ',shape(lat)))
print(('lat = ',lat))


myvar=myvar
myvar2=myvar2/0.6192/1000.  # geopotential ==> altitude

################
# colorbar
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)
######
c = mcolors.ColorConverter().to_rgb
#rvb1 = make_colormap([c('red'), c('violet'), 0.33, c('violet'), c('blue'), 0.66, c('blue')])
rvb1 = make_colormap([c('SaddleBrown'), c('violet'), 0.33, c('violet'), c('white'), 0.66, c('white')])

pp=(0.5450980392156862-0.15, 0.27058823529411763-0.05*3, 0.07450980392156863-0.01)
rvb1 = make_colormap([pp,0.33, c('indianred'),0.71, c('lightcyan'),0.99,c('lightsalmon')])
#rvb1 = make_colormap([c('saddlebrown'),0.33, c('indianred'),0.66, c('white')])
######

div=3.      
nbticksbar=3.
#pal=get_cmap(name="Blues") 	
pal=get_cmap(name="gist_earth_r")
yticks=[-90,-60,-30,0,30,60,90]
#xticks=[-180,-120,-60,0,60,120,180]
xticks=[0,60,120,180,240,300,360]

# changer les longitudes pour mettre TR au centre
vec=shape(myvar)
myvarbis=np.zeros(vec,dtype='f')
myvar2bis=np.zeros(vec,dtype='f')
# i lat : pas de changement
# j lon :
for i in range(vec[0]):
    for j in range(vec[1]):
        if j < int(vec[1]/2.) :
           myvarbis[i,j]=myvar[i,j+int(vec[1]/2)]
           myvar2bis[i,j]=myvar2[i,j+int(vec[1]/2)]
        else: 
           myvarbis[i,j]=myvar[i,j-int(vec[1]/2)]
           myvar2bis[i,j]=myvar2[i,j-int(vec[1]/2)]

lon=lon+180.


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

c=mpl.contour(lon, lat, myvar2bis, 10,levels=[-3,-2,-1,-0.5,-0.2], colors = 'k', linewidths = 0.5)
mpl.clabel(c, fmt='%2.1f', colors='k', fontsize=14)

mpl.imshow(myvarbis,extent=[0,360,-90,90],cmap=rvb1) #,extend='both')

#mpl.colorbar()
mpl.ylabel('Latitude (deg)',labelpad=20,fontsize=29)
mpl.xlabel('Longitude (deg)',labelpad=20, fontsize=29)
mpl.xticks(xticks,fontsize=26)
mpl.yticks(yticks,fontsize=26)
mpl.grid()
mpl.savefig('mapalbini.eps',dpi=200)
mpl.savefig('mapalbini.png',dpi=200)
mpl.show()