import numpy as np
from netCDF4 import Dataset
import matplotlib.pyplot as plt
import matplotlib.colors as colors
import pandas as pd
from mpl_toolkits.basemap import Basemap


#m.fillcontinents(color='0.8', lake_color='white', zorder=0)
def PlotMaps(data,lat,lon,title,yax,cmap):
    plt.figure()
    map = Basemap(llcrnrlon=89,llcrnrlat=-15,urcrnrlon=150.,urcrnrlat=27.,
                resolution='c', projection='merc', lat_0 = 39.5, lon_0 = 1)

    lons,lats= np.meshgrid(lon,lat)
    x,y = map(lons,lats)
    map.drawmapboundary(fill_color='white', zorder=-1)
    parallels = np.arange(-15,30,10.) 
    meridians = np.arange(90,151,10.) 
    map.drawparallels(parallels,labels=[1,0,0,0],fontsize=10)
    map.drawmeridians(meridians,labels=[0,0,0,1],fontsize=10)
    map.readshapefile('./coastline/ne_10m_coastline', 'ne_10m_coastline')
    
    pdata = map.contourf(x,y,data,cmap=cmap)
    cb = map.colorbar(pdata,"bottom", size="3%", pad="10%")
    plt.title(title)
    cb.set_label(yax)
    plt.savefig(title+'.png')


def Drought_characteristics(pdsi, land, drought_threshold):
    print(land.shape)
    mean_Duration = np.nan * np.ones(land.shape)
    mean_Duration[land==-99] = np.nan

    mean_Num_Event = np.nan * np.ones(land.shape)
    mean_Num_Event[land==-99] = np.nan

    mean_Intensity = np.nan * np.ones(land.shape)
    mean_Intensity[land==-99] = np.nan

    mean_Interarrival = np.nan * np.ones(land.shape)
    mean_Interarrival[land==-99] = np.nan

    for ig, v in np.ndenumerate(land.data):        
        if v == 1:
            # IDENFIY NUMBER OF DROUGHT
            df_pdsi = pd.DataFrame(pdsi[:,ig[0],ig[1]],columns=['pdsi'])
            np_pdsi = pdsi[:,ig[0],ig[1]]

            # DROUGHT PERIOD DETECTION
            # tag rows based on the threshold
            df_pdsi['tag'] = df_pdsi['pdsi'] <= drought_threshold

            # first row is a True preceded by a False (first month of an event)
            fst = df_pdsi.index[df_pdsi['tag'] & ~ df_pdsi['tag'].shift(1).fillna(False)]
            
            # last row is a True followed by a False (last month of an event)
            lst = df_pdsi.index[df_pdsi['tag'] & ~ df_pdsi['tag'].shift(-1).fillna(False)]

            # drought event information
            total_month_droughts = len(np_pdsi[np_pdsi <= drought_threshold])
            duration = [(i, j) for i, j in zip(fst, lst)]
            num_events = len(duration)
            D_ind = np.zeros(num_events)		# Duration index
            I_ind = np.zeros(num_events)		# Intensity index
            for j in range(0,num_events):
                istart, iend = duration[j]
                D_ind[j] = iend - istart + 1 
                I_ind[j] = np.mean(np_pdsi[istart:iend+1])


            # NON-DROUGHT PERIOD DETECTION
            # tag rows based on the threshold
            df_pdsi['tag'] = df_pdsi['pdsi'] > drought_threshold

            # first row is a True preceded by a False (first month of an event)
            fst = df_pdsi.index[df_pdsi['tag'] & ~ df_pdsi['tag'].shift(1).fillna(False)]
            
            # last row is a True followed by a False (last month of an event)
            lst = df_pdsi.index[df_pdsi['tag'] & ~ df_pdsi['tag'].shift(-1).fillna(False)]

            # No drought event information
            total_month_nd = len(np_pdsi[np_pdsi > drought_threshold])
            duration_nd = [(i, j) for i, j in zip(fst, lst)]
            num_events_nd = len(duration_nd)
            D_nd_ind = np.zeros(num_events_nd)		# Duration index
            for j in range(0,num_events_nd):
                istart, iend = duration_nd[j]
                D_nd_ind[j] = iend - istart + 1 

            mean_Duration[ig[0],ig[1]] = np.mean(D_ind)
            mean_Num_Event[ig[0],ig[1]] = num_events
            mean_Intensity[ig[0],ig[1]] = np.mean(I_ind)
            mean_Interarrival[ig[0],ig[1]] = np.mean(D_ind)+np.mean(D_nd_ind)
            mean_Intensity[mean_Intensity<-20] = np.nan

    return mean_Duration,mean_Num_Event,mean_Intensity,mean_Interarrival



nc_fid = Dataset('/work/projects/G09/SEACLID/Drought_SEA/Dr_PDSI_SPEI_SPI.nc','r')
ori_pdsi = nc_fid.variables['ori_pdsi'][:]
lat = nc_fid.variables['lat'][:]
lon = nc_fid.variables['lon'][:]
month = nc_fid.variables['month'][:]
land = nc_fid.variables['land'][:]

mean_Duration,mean_Num_Event,mean_Intensity,Mean_Interarrival = Drought_characteristics(ori_pdsi, land, -2)

cmap=plt.get_cmap('jet',11)

PlotMaps(mean_Duration,lat,lon,'Mean Duration','[month]',cmap)
PlotMaps(mean_Num_Event,lat,lon,'Number of Events','[-]',cmap)
PlotMaps(mean_Intensity,lat,lon,'Mean Intensity','[-]',cmap)
PlotMaps(Mean_Interarrival,lat,lon,'Mean Interarrival','[month]',cmap)
plt.show()