import numpy as np
import matplotlib.pyplot as plt
from netCDF4 import Dataset
import cartopy.crs as ccrs
import datetime as dt
import glob, os, sys
import scipy as sp


# Read in the text file of MJO phases ----------------------------------------
# Skip two lines of header
text_data = np.genfromtxt("/nfs/a240/earceb/Teaching/SOEE2710_Python_module/course_material/level05/MJO_phases.txt",skip_header=2)
MJO_phase = text_data[:,5]
MJO_amp = text_data[:,6]


# Create a datetime array for every day 1st Jan 2008 to 31st Dec 2017
# State with the first datetime
base = dt.date(2008, 1, 1)
# Iterate through the days, adding 1 day each time
MJO_dates = np.array([base + dt.timedelta(days=x) for x in range(0, np.size(MJO_phase))])

# Locate positions of DJF data and extract time, phase and amp data
# Extract month data
MJO_months = [MJO_dates[x].month for x in range(0, np.size(MJO_dates))]
MJO_months = np.array(MJO_months)
DJF_days = np.where(((MJO_months==1) | (MJO_months==2) | (MJO_months==12)))    
# Extract time, phase and amplitude data for DJF days 2008-2017
MJO_dates_DJF = MJO_dates[DJF_days[0]]
MJO_amp_DJF = MJO_amp[DJF_days[0]]
MJO_phase_DJF = MJO_phase[DJF_days[0]]


# Read in the rainfall data -------------------------------------------------
# Get the rainfall data file names, sorted by name
filenames = sorted(glob.glob('/nfs/a240/earceb/Teaching/SOEE2710_Python_module/course_material/level05/TRMM_data/*.nc4'))
nc_fid = Dataset(filenames[0],'r') # Open first file
#print(nc_fid.variables) # Have a look what is in file - check the units
#print(nc_fid.variables.keys()) # Have a look what is in file - variable names only

# Extract lat and lon from 1st file
lat = nc_fid.variables['lat'][:]
lon = nc_fid.variables['lon'][:]

# Set up an empty array to fill with precip data
ppt = np.empty((np.size(filenames),np.size(lon),np.size(lat)))
# Set up an empty array to fill with times
rain_dates = [] 

# Loop through all the files and extract the rainfall data
count = 0 # Count to place data in ppt array
for f in filenames:
    
    #print(f) # Print current file name
    nc_fid = Dataset(f,'r') # Open file
    
    # Read data and add to ppt array
    ppt[count,:,:] = nc_fid.variables['precipitation'][:] # Read data and add to ppt array 
   
    # Retrieve time data from file name and add to dictionary
    temp_date = dt.datetime.strptime(f[94:101+1], "%Y%m%d")
    rain_dates.append(temp_date)
    count = count+1 # Increment the count


# Loop through each MJO phase and extract rainfall data for each phase--------------------------------
# Create empty array of ppt data by phase
ppt_by_phase = np.empty((8,np.size(lon),np.size(lat)))
for mp in np.arange(1,8+1):
    print(mp)

    # Work out how many days fit criteria of MJOphase=mp and MJOamp>1
    num_days = 0
    for dd in np.arange(0,len(MJO_dates_DJF)):
        if MJO_phase_DJF[dd]==mp and MJO_amp_DJF[dd]>1:
            num_days = num_days+1

    # Create temporary array to hold rainfall data
    rainfall_temp = np.empty((num_days,np.size(lon),np.size(lat))) 

    # Locate days in Dec, Jan or Feb when MJOphase=mp and MJOamp>1
    count = 0
    for dd in np.arange(0,len(MJO_dates_DJF)):
        
        # Add rainfall data to temporary array 
        if MJO_phase_DJF[dd]==mp and MJO_amp_DJF[dd]>1: 
            rainfall_temp[count,:,:] = ppt[dd,:,:] 
            count = count + 1

    # Average rainfall data in temporary array and add to ppt_by_phase array
    ppt_by_phase[mp-1,:,:] = np.nanmean(rainfall_temp, axis=0) 
        

# Compute mean rainfall and rainfall anomally by MJO phase
ppt_mean = np.mean(ppt, axis=0)



# Create plots of rainfall anomally, 1 for each phase-------------------------- 
for mp in np.arange(1,8+1):
    
    # Initiate new Figure
    fig = plt.figure()

    # Set up axes and select map projection
    ax1 = plt.subplot(1,1,1,projection=ccrs.PlateCarree())
    
    # Set axis limits to show Africa only
    ax1.set_extent([25, 180, -30, 30], crs=ccrs.PlateCarree())

    # Add features
    ax1.coastlines()

    # Find anomally from long-term mean
    data2plot =  ppt_by_phase[mp-1,:,:]-ppt_mean
    
    # Apply gaussian filter
    data2plot = np.transpose(ppt_by_phase[mp-1,:,:]-ppt_mean)
    data2plot_smoothed = sp.ndimage.filters.gaussian_filter(data2plot, [3.0, 3.0], mode='constant')
    
    # Change values outside contour limits to min and max contours limits
    data2plot_smoothed[data2plot_smoothed > 6] = 6
    data2plot_smoothed[data2plot_smoothed < -6] = -6
    
    # Plot data
    im1 = ax1.contourf(lon, lat, data2plot_smoothed, np.arange(-6,6+1,2), transform=ccrs.PlateCarree(), cmap=plt.cm.RdBu)
    
    # Add a title
    ax1.set_title('Rainfall anomally: MJO phase:' + str(mp))

    # Add colorbar
    # Label the colorbar axes to make further changes
    cbar = plt.colorbar(im1, orientation='horizontal')
    # Add colorbar label
    cbar.ax.set_xlabel('mm/day')


    # Make the plot visible
    #plt.show()
    
    # Save plot
    filename_start = "/nfs/a240/earceb/Teaching/SOEE2710_Python_module/course_material/level05/plots/MJO_rainfall_"
    filename_end = ".png"
    plt.savefig(filename_start.strip() + str(mp).strip() + filename_end.strip(), dpi=300, format='png')