#! /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 FV3_utils import *
from input import *

############################
# folder="../"
filename1=name+"_A.nc"
filename2=name+".nc"
#filename3="../simu_pole/diagfi2015_S.nc"
var="temperature" #variable
tint=[30,35] #Time must be as written in the input file
print("Reading ", filename1, filename2)

nc1=Dataset(filename1)
nc2=Dataset(filename2)
#nc3=Dataset(filename3)


lat=getvar(nc1,"latitude")
lon=getvar(nc1,"longitude")
alt=getvar(nc1,"altitude")
tim=getvar(nc1,"Time")
#alt2=nc3.variables["altitude"][:]
############################

def zetotarea(dat):

    totarea=0.
    for i in range(size(dat[:,0])):
        for j in range(size(dat[0,:])):
              totarea=totarea+dat[i,j]
    return totarea

def meanarea(dat1,dat2,totarea):
    meandat=np.zeros(dat1.shape[0])
    for t in range(dat1.shape[0]):
        for i in range(dat1.shape[1]):
            for j in range(dat1.shape[2]):
                meandat[t]=meandat[t]+ma.getdata(dat1[t,i,j])*dat2[i,j]/totarea
    return meandat


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

myvar=getvar(nc1,var,tint,t_mean=True)
#myvar2=getvar(nc1,var,tint)
aire = getvar(nc2,"aire")
totarea=zetotarea(aire)
meantemp=meanarea(myvar,aire,totarea)
#meantemp2=meanarea(myvar2,aire,totarea)

font=23

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

mpl.plot(meantemp,alt,'r')
#mpl.plot(meantemp2,alt2,'b--')

mpl.title('Global mean temperature',fontsize=font)
mpl.ylabel('Altitude (km)',labelpad=10,fontsize=font)
mpl.xlabel('Temperature (K)',labelpad=10, fontsize=font)
#mpl.xticks(xticks,fontsize=font)
mpl.xticks(fontsize=font)
#mpl.yticks(yticks,fontsize=font)
mpl.yticks(fontsize=font)
mpl.grid()
# mpl.legend(["Ref","Alt"],prop={'size':27},loc='upper left')

left  = 0.2  # the left side of the subplots of the figure
right = None   # the right side of the subplots of the figure
bottom = None  # the bottom of the subplots of the figure
top = None     # the top of the subplots of the figure
wspace = None  # the amount of width reserved for blank space between subplots
hspace = None  # the amount of height reserved for white space between subplots
mpl.subplots_adjust(left, bottom, right, top, wspace, hspace)
#mpl.subplots_adjust(hspace = .1)

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