#! /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
from matplotlib import ticker
import matplotlib.colors as colors
import datetime
from mpl_toolkits.basemap import Basemap, shiftgrid
from matplotlib.cm import get_cmap
from FV3_utils import * # import name

############################
fa='sans-serif'
hfont = {'fontname':'Arial'}
mpl.rc('font',family=fa)
mpl.rc('pdf',fonttype=42)
font=30
cc=['k']
pal=get_cmap(name="rainbow")
norm=colors.LogNorm()
lvls=np.logspace(-6,-4,21)
norm=None #colors.LogNorm()

### Data
# name='../diagfi2015.nc' # read from FV3_util
print(name)
try:
    nc1=Dataset(name+"_A.nc")
except:
    nc1=Dataset(name+".nc")
alt=nc1.variables["altitude"][:]
lat=nc1.variables["lat"][:]
lon=nc1.variables["lon"][:]
# temp=switchlon(temp)
lon=lon+180.

def plot_alt(altitude = 1000):
    temp=nc1.variables["temperature"][:,:,:,:]
    numalt=getind(altitude,alt)
    temp=temp[:,numalt,:,:]
    temp=np.mean(temp,axis=0)

    min_t=temp.min()
    max_t=temp.max()
    lvls=np.linspace(min_t,max_t,21)

    ### Figure
    fig=mpl.figure(figsize=(15, 10))
    CF=mpl.contourf(lon, lat, temp,lvls,cmap=pal,norm=norm)
    cbar=mpl.colorbar(CF, shrink=1, ticks=lvls[::2],format="%1.1f")
    cbar.ax.set_title("[kg m$^{-2}$]",y=1.04,fontsize=font)
    for t in cbar.ax.get_yticklabels():
        t.set_fontsize(font)

    vect=lvls
    CS=mpl.contour(lon,lat,temp,lvls[:],colors='k',linewidths=0.5)
    lab=mpl.clabel(CS, inline=1, fontsize=20, fmt='%1.1f',inline_spacing=1)
    for l in lab:
        l.set_rotation(0)

    mpl.grid()
    mpl.title(f"Temperatures @ z={altitude/1000}km",fontsize=font)
    mpl.ylabel(r'Latitude',labelpad=10,fontsize=font)
    mpl.xlabel('Longitude',labelpad=10, fontsize=font)
    pylab.ylim([-90,90])
    yticks=np.linspace(-90,90,13)
    pylab.xlim([0,360])
    xticks=np.linspace(0,360,7)
    mpl.yticks(yticks,fontsize=font)
    mpl.xticks(xticks,fontsize=font)
    mpl.savefig(f"maptemp{altitude}",bbox_inches='tight',dpi=70)
    #mpl.show()


plot_alt(1000)
plot_alt(5000)
plot_alt(20000)
plot_alt(50000)
plot_alt(100000)