#! /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
from mpl_toolkits.basemap import Basemap, shiftgrid
from FV3_utils import *
from input import *

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()
lvls=np.linspace(37,57,21)

altitude=1 # in km

### Data
filename1=name+'_A.nc'
# name=name+'.nc'
filename2=name+'.nc'
print("Plot "+filename1)
nc1=Dataset(filename1)
nc2=Dataset(filename2)
ts=nc2.variables["temperature"][:,0,:,:]
# ts=nc2.variables["tsurf"][:,:,:]
u=nc1.variables["u"][:,:,:,:]
v=nc1.variables["v"][:,:,:,:]
lat=getvar(nc1,"latitude")
alt=getvar(nc1,"altitude")
lon=getvar(nc1,"longitude")
ps=nc2.variables["ps"][:,:]
# ps=nc2.variables["phisinit"][:,:]/0.6169/1000.  # altitude km

numalt=getind(altitude,alt)
print('numalt =',numalt,'altitude=',alt[numalt])
u=u[:,numalt,:,:]
v=v[:,numalt,:,:]
u=np.mean(u,axis=0)
v=np.mean(v,axis=0)
ts=np.mean(ts,axis=0)
ps=np.mean(ps,axis=0)

if lon[0]<0:
    ts=switchlon(ts)
    u=switchlon(u)
    v=switchlon(v)
    ps=switchlon(ps)
    # topo=switchlon(topo)
    lon=lon+180.

### Figure
fig=mpl.figure(figsize=(20, 10))
# CF=mpl.contourf(lon, lat, ts,lvls,cmap=pal,norm=norm,extend='both')
# cbar=mpl.colorbar(CF, shrink=1, ticks=lvls[::2]) #,format="%1.0f")
# cbar.ax.set_title("[K]",y=1.04,fontsize=font)
# for t in cbar.ax.get_yticklabels():
#   t.set_fontsize(font)
# mpl.title('Surface temperatures and winds',fontsize=font)

#vect=lvls
# CS=mpl.contour(lon,lat,topo,np.linspace(-6,6,21),colors='k',linewidths=0.5)
#lab=mpl.clabel(CS, inline=1, fontsize=20, fmt='%1.2e',inline_spacing=1)
#for l in lab:
#  l.set_rotation(0)

CF=mpl.contourf(lon, lat, ps,cmap=pal,norm=norm,extend='both')
cbar=mpl.colorbar(CF, label="Pa")
cbar.set_label("Pa", size=font)
for t in cbar.ax.get_yticklabels():
  t.set_fontsize(font)
mpl.title('Surface pressures and winds',fontsize=font)

getwinds(lon,lat,u,v,1,1,100,0.002,5)

mpl.grid()
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.tight_layout()
mpl.savefig('mapwinds1km.png',bbox_inches='tight',dpi=70)
mpl.show()