#from netCDF4 import Dataset

import xarray
from scipy.stats import norm
import numpy as np
import matplotlib.pyplot as plt

list_indices = ['tasmax']
indices = list_indices[0]

#test
#files = '/g/data3/w97/dc8106/AMZ_def_EXPs/121GPsc_E0/AMZDEF.daily_tasmin.tasmax.pr.1978_2011_121GPsc_E0.nc'

#remove CTL diff to CTL file from analysis
#del(files[11])

#var = np.zeros((len(files),12418,145,192),dtype=np.float32)
#t = []

print "Open file"
data = xarray.open_dataset('/g/data3/w97/dc8106/AMZ_def_EXPs/121GPsc_E0/AMZDEF.daily_tasmin.tasmax.pr.1978_2011_121GPsc_E0.nc')
#tasmax = data.tasmax
#tasmin = data.tasmin
#lat = data.lat
#lon = data.lon
lons,lats = np.meshgrid(data.lon,data.lat)

print "File opened. Calculate stats"

# To calculate the PDF of the distribution using norm.pdf() we need to know the mean and std. deviation of the distribution.
# First create an array to save the mean and std. deviation for each grid cell. This array needs to have the dimensions:
# (2, nlat, nlon). 2 is because we save 2 stats per grid cell
# We'll put the mean first so stats[0,:,:] will contain the means of all grid cells and stats[1,:,:] will be the std. dev.
stats = np.zeros((2,data.tasmax.shape[1],data.tasmax.shape[2]),dtype=np.float32)

# Let's transform it into an xarray DataArray so we keep the information of what all the axes are with the data:
stats_xr = xarray.DataArray(stats, coords=[[0,1],data.coords['lat'],data.coords['lon']],dims=['stats','lat','lon'])

stats_xr[0,:,:] = data.tasmax.mean(dim='time')  # xarray methods allow you to refer to dimensions by name rather than position
stats_xr[1,:,:] = data.tasmax.std(dim='time')

print "Stats calculated. Plot 1 PDF"
# Now if we want to plot the PDF at a location we simply need to do:
# Define our x axis. Here we'll plot the PDF for values ranging from [min-10, max+10] where min (resp. max) is the minimum
# (resp. maximum) value of tasmax at the location we want to plot
# Location:
ilon=50
ilat=30

# xarray comes with a handy way to select a point using the dimension names instead of the position in the array:
tasmax_t = data.tasmax.isel(lat=ilat,lon=ilon)
mean_loc = stats_xr.isel(stats=0,lat=ilat,lon=ilon)
std_loc  = stats_xr.isel(stats=1,lat=ilat,lon=ilon)

x = np.linspace(tasmax_t.min().values-10,tasmax_t.max().values-10,100)  # The last number is the number of values

#norm.pdf(x,mean,std) calculates the PDF for a normal distribution with the given mean and stats and return the value of the PDF
# at the x point. plt.plot() will plot those PDFs for all the points in the x array as defined above: 100 pts between min-10 and max+10.

pdf=plt.plot(x,norm.pdf(x,mean_loc,std_loc))
plt.show()