#!/usr/bin/env python

#-------------
# Load modules
#-------------
from netCDF4 import Dataset
import numpy as np                         # Numpy import
import scipy.io as spio
import scipy.linalg as linalg
import scipy.fftpack as fft
import scipy.stats as stats
import sys

from bmContourMod import *
from subLatLonMod import *


#----------------------
# Provide the file name
#----------------------
dirName  = '/discover/nobackup/jkouatch/gmiMetFields/MERRA/2x2.5/'
fileName = dirName+'MERRA300.prod.assim.20080316.2x2.5x72.nc'

#------------------
# Opening the file
#------------------
ncFid = Dataset(fileName, mode='r')

#----------------------
# Extracting dimensions
#----------------------
time = ncFid.variables['time'][:]
lev  = ncFid.variables['lev'][:]
lat  = ncFid.variables['lat'][:]
lon  = ncFid.variables['lon'][:]

#------------------------------------
# Extracting the temperature variable
#------------------------------------
temp = ncFid.variables['T'][:]

#-------------
# Closing file
#-------------
ncFid.close()

#--------------------------
# Set the level of interest
#--------------------------
level500 = 29

#---------------
# Slice the data
#---------------
indexLat, indexLon = sliceLatLon(lat, lon, (-45.0,45.0), (-90.0,90.0))

nLat = lat[indexLat[0]:indexLat[-1]+1]
nLon = lon[indexLon[0]:indexLon[-1]+1]

T500 = temp[:,level500,indexLat[0]:indexLat[-1]+1,indexLon[0]:indexLon[-1]+1]
#T500 = temp[:,level500,:,:]  # time, lat, lon

#------------------------------
# Compute the mean and variance
#------------------------------

T500mean = np.mean(T500,0)
T500var  = np.var (T500,0)


#nT500mean = T500mean[indexLat[0]:indexLat[-1]+1, indexLon[0]:indexLon[-1]+1]
#nT500var  =  T500var[indexLat[0]:indexLat[-1]+1, indexLon[0]:indexLon[-1]+1]

#---------------------------
# Plot the mean and variance
#---------------------------

bmContourPlot(T500mean, nLat, nLon, 'fig_TempMeanSlice', 'Spatial Temperature Mean')
bmContourPlot(T500var,  nLat, nLon, 'fig_TempVarianceSlice', 'Spatial Temperature Variance')