%matplotlib inline
import numpy as np

import matplotlib.pyplot as plt         # to generate plots
from mpl_toolkits.basemap import Basemap # plot on map projections
import matplotlib.dates as mdates

import datetime             
from netCDF4 import Dataset # http://unidata.github.io/netcdf4-python/
from netCDF4 import netcdftime
from netcdftime import utime

"""
Explanation: Primary analysis of precipitation data
This notebook only show some basic analysis on precipitation data based on previous notebooks
* read data
* check variables
* visualze spatial map
* calcualte climatology
* calculate no-wegithed global mean time series
* plot time series
* save data
Precipitation data
Precipitation data is downloaded from GPCC Global Precipitation Climatology Centre (https://www.esrl.noaa.gov/psd/data/gridded/data.gpcc.html).
* Temporal Coverage: Monthly values 1901/01 through 2013 (full V7).
* Spatial Coverage: 0.5 degree latitude x 0.5 degree longitude global grid (720x180)
1. Read in precipitation data
Here, let's see how to open netcdf data in python and generate yearly climatology of global precipitation. First, we will import the required libraries, then we will open the monthly average global precipitaton.
1.1 Load basic libraries
End of explanation
"""


#!wget ftp://ftp.cdc.noaa.gov/Datasets/gpcc/full_v7/precip.mon.total.v7.nc

file = 'data\precip.mon.total.v7.nc'
variable = 'precip'

"""
Explanation: 1.2 Set input NetCDF file info
You can download data using wget or aria2c under Linux. In this case, the data is already downloaded and put into the folder of data.
End of explanation
"""


fh     = Dataset(file, mode='r') # file handle, open in read only mode
lon    = fh.variables['lon'][:]
lat    = fh.variables['lat'][:]
nctime = fh.variables['time'][:]
t_unit = fh.variables['time'].units
pr = fh.variables[variable][:]

try :
    t_cal = fh.variables['time'].calendar
except AttributeError : # Attribute doesn't exist
    t_cal = u"gregorian" # or standard

fh.close() # close the file

"""
Explanation: 1.3 Extract variables
End of explanation
"""


#netcdftime.utime?
#netcdftime.utime.num2date?
utime = netcdftime.utime(t_unit, calendar = t_cal)
datevar= utime.num2date(nctime)

datevar[0:5]

"""
Explanation: 1.4 Parse time
If you want to learn more about what a specific function does, you can use the help function: help(function) or function?.
For example, help(np.mean) or netcdftime.utime?
End of explanation
"""


lat

"""
Explanation: 2. Access data
2.1 Output all latitude values
End of explanation
"""


lat[0] # Caution! Python’s indexing starts with zero

lat[-1] # gives the last value of the vector

"""
Explanation: 2.2 Access the first and the last value of lat
End of explanation
"""


pr[0,:,:] # or simply pr[0]

pr[0].shape # (360, 720)

"""
Explanation: 2.3 Get the precipitation field of the first month
End of explanation
"""


plt.imshow(pr[0]) #data for the first time-step (January 1901)
plt.title('Jan 1901 Monthly Rainfall [mm]')

"""
Explanation: 3. Visualize data
3.1 Have a quick shot
The precipitaton at Jan-1901 is visualized using imshow, which looks just like a image without using any geolocation information.
End of explanation
"""


lons, lats = np.meshgrid(lon,lat)

m = Basemap(projection='kav7', lon_0=0)
m.drawmapboundary(fill_color='Gray', zorder=0)
m.drawparallels(np.arange(-90.,99.,30.), labels=[1,0,0,0])
m.drawmeridians(np.arange(-180.,180.,60.), labels=[0,0,0,1])

h = m.pcolormesh(lons,lats,pr[0], shading='flat',latlon=True)
m.colorbar(h, location='bottom', pad="15%", label='Rainfall [mm/month]')

plt.title('Jan 1901 Monthly Rainfall [mm]')

"""
Explanation: 3.2 Visualize on a Map Projection
The matplotlib basemap is used here, which is a library for plotting 2D data on maps in Python. It is similar in functionality to the matlab mapping toolbox, the IDL mapping facilities, GrADS, or the Generic Mapping Tools. PyNGL and CDAT are other libraries that provide similar capabilities in Python (https://matplotlib.org/basemap/users/intro.html).
basemap supports many projecitons:
* Azimuthal Equidistant Projection
* Gnomonic Projection
* Orthographic Projection
* Geostationary Projection
* Near-Sided Perspective Projection
* Mollweide Projection
* ...
End of explanation
"""


#pr # Python automatically created a masked array with the data and the mask

"""
Explanation: 4. Work with masked arrays
In the previous section we noticed that our precipitation data-set only has values over land. This is where
masked arrays come into play.
4.1 Have a look at pr
End of explanation
"""


pr.data # gives you the data array only

pr.mask # gives you True where the data is masked (over the ocean) and False where data is available (over land)

"""
Explanation: 4.2 Have a look at the data and the mask separately
End of explanation
"""


pr.data[~pr.mask] # the tilde inverts the mask
pr.data

"""
Explanation: 4.3 Create a vector of precipitation data where values are not masked
End of explanation
"""


clim = np.mean(pr, axis=0) # because pr is a masked array, all masked values are ignored in the calculation.
clim.shape

"""
Explanation: 4.4 Create a yearly climatology field from the precipitation data
End of explanation
"""


lons, lats = np.meshgrid(lon,lat)
m = Basemap(projection='kav7',lon_0=0)
m.drawmapboundary(fill_color='Gray', zorder=0)
m.drawparallels(np.arange(-90.,99.,30.), labels=[1,0,0,0])
m.drawmeridians(np.arange(-180.,180.,60.), labels=[0,0,0,1])

h = m.pcolormesh(lons,lats,clim, shading='flat',latlon=True)
m.colorbar(h,location='bottom',pad="15%", label='Rainfall [mm/month]')

plt.title('Yearly Rainfall Climatology')

"""
Explanation: 'clim' can be plotted onto a map projection the same way you plotted 'pr[0]' above.
End of explanation
"""


wgtmat = np.cos(np.tile(abs(lat[:,None])*np.pi/180, (1,len(lon)))) # Dimension: 72x144

"""
Explanation: 5. Calculate and plot the global monthly mean precipitation rate
5.1 Define weighting matrix
End of explanation
"""


mean_pr = np.zeros(pr.shape[0]) # Preallocation
for i in range(pr.shape[0]): # Don’t forget the ‘:’
    mean_pr[i] = np.sum(pr[i] * wgtmat * ~pr.mask[i])/np.sum(wgtmat * ~pr.mask[i])

"""
Explanation: 5.2 Calculate the global monthly mean precipitation rate
End of explanation
"""


prunit = 'mm/month'

"""
Explanation: 5.3 Plot a time series of monthly rainfall
5.3.1 Give a unit to the global monthly mean precipitation rate
End of explanation
"""


fig = plt.figure(figsize=(15, 6))
ax = fig.add_subplot(111)

plt.plot(range(1,pr.shape[0]+1), mean_pr, 'b-', label='precip',linewidth=2)
plt.axis([1,pr.shape[0]+1, np.min(mean_pr)-0.05, np.max(mean_pr)+0.05])

plt.xlabel('Month')
plt.ylabel('Precipitation Rate [' + prunit + ']')
plt.title('Monthly Mean Precipitation Rate Over Land (1901-2013)',fontsize=12)
plt.legend(loc=2, borderaxespad=1., frameon=False)

#plt.grid(True)
ax.autoscale_view()

plt.savefig('image/monthly_global_precip.png')

"""
Explanation: 5.3.2 Have a quick look
End of explanation
"""


fig, ax = plt.subplots(1)
fig.set_figwidth(15)
fig.set_figheight(6)

# plot mean precipitation
ax.plot(datevar, mean_pr)

# rotate and align the tick labels so they look better
fig.autofmt_xdate()

# use a more precise date string for the x axis locations in the toolbar
ax.fmt_xdata = mdates.DateFormatter('%Y-%m-%d')

# add some docoration
plt.xlabel('Month')
plt.ylabel('Precipitation Rate [' + prunit + ']')
plt.title('Monthly Mean Precipitation Rate Over Land (1901-2013)',fontsize=12)

ax.autoscale_view()

"""
Explanation: 5.3.3 Use dates as x-axis to make image a little pretty
End of explanation
"""


np.savez('data/monthlylmeanpr.npz', prunit=prunit, dates=datevar, mean_pr=mean_pr)

"""
Explanation: 5.3.4 Save the monthly data for future use
End of explanation
"""


ntimes = mean_pr.shape[0]
nyr = ntimes/12
nyr

"""
Explanation: 5.4 Calculate annual mean precipitation rates
Firstly, convert mean_pr data from [months] to [nyr|12|], where months=nyrx12 months/year. This can be done using the  reshape function.
Secondly, calculate means of all months whithin a year, which can be done using np.mean along the axis=1.
5.4.1 Calculate the numbler of years (should be 113 years)
The months should be divided by 12.
End of explanation
"""


mean_pr_reshaped = mean_pr[0:ntimes].reshape((nyr,12))
mean_pr_yr = np.mean(mean_pr_reshaped, axis=1)
mean_pr_yr  # print the results to the screen

"""
Explanation: 5.4.2 Calculate annual mean
End of explanation
"""


years = [idx.year for idx in datevar]
years = years[1::12]
print(years[-1])
years[0:13]

"""
Explanation: 5.4.3 Get unique years
End of explanation
"""


fig, ax = plt.subplots(1)
fig.set_figwidth(15)
fig.set_figheight(6)

# plot mean precipitation
ax.plot(years, mean_pr_yr)

# add some docoration
plt.xlabel('Year')
plt.ylabel('Precipitation Rate [' + prunit + ']')
plt.title('Yearly Mean Precipitation Rate Over Land (1901-2013)',fontsize=12)
#plt.grid(True)

ax.autoscale_view()

"""
Explanation: 5.4.4 Have a quick plot
End of explanation
"""


np.savez('data/annualmeanpr.npz', prunit=prunit, years=years, mean_pr_yr=mean_pr_yr)

"""
Explanation: 6. Save and load annual mean precipitation rates
we can save and load several arrays into a single file in uncompressed .npz format using np.savez.
The .npz file format is a zipped archive of files named after the variables they contain. The archive is not compressed and each file in the archive contains one variable in .npy format. For a description of the .npy format, see numpy.lib.format or the NumPy Enhancement Proposal http://docs.scipy.org/doc/numpy/neps/npy-format.html
When opening the saved .npz file with load a NpzFile object is returned. This is a dictionary-like object which can be queried for its list of arrays (with the .files attribute), and for the arrays themselves.
6.1 Save annual mean precipitation rates
End of explanation
"""


npzfile = np.load('data/annualmeanpr.npz')

"""
Explanation: 6.2 Load annual mean precipitation rates
End of explanation
"""


npzfile.files

"""
Explanation: 6.2.1 Check variables in the file
End of explanation
"""


npzfile['years'][0:13]

"""
Explanation: 6.2.2 Check the variable of years
Should the same as 6.4.3
End of explanation
"""


