# -*- coding: utf-8 -*-
"""
Python script to identify the radar bright band using the first derivative of
the Doppler velocity. Bright band strength is measured by analysing "image contrast".
Please note that this is still an ongoing algorithm. How to express the bright band strength
properly is still a problem.

Where the first derivative (with a 360-m interval) is the smallest, it will be marked
as the bright band location. Note that the interval can be changed. 

@author: Dongqi Lin - 2018
Contact: ee17dl@leeds.ac.uk
"""

#############################################################################
import sys
import timeit
import matplotlib
import matplotlib.pylab as plt
import matplotlib.mlab as mlab
import numpy as np
import scipy as sp
from netCDF4 import Dataset as ncfile
from datetime import datetime

# Plot figures in an external window
# This may not work in Jupyther Notebook. Comment out if it is not working.
from IPython import get_ipython
get_ipython().run_line_magic('matplotlib','qt')

# Set figure parameters 
params={     
    'axes.labelsize'  : '16',   
    'axes.titlesize'  : '20',  
    'xtick.labelsize' :'12',
    'ytick.labelsize' :'12',    
    'lines.linewidth' : '1' ,    
    'legend.fontsize' : '14', 
    'figure.figsize'   : '16, 9'    
}
plt.rcParams.update(params)

# Calculate time spent
startTime = datetime.now()


#############################################################################
# ============================== Import Data ===============================#
#############################################################################

# Input date - format yyyymmdd
date = str(sys.argv[1])     
filename ='cfarr-radar-copernicus_chilbolton_'+date+"_fix.nc";
nc = ncfile("../data/radar-copernicus/"+filename, mode='r')

time = nc.variables['time'][:]                       

# Distance from the antenna to the middle of each range gate in m
Range = nc.variables['range'][:]                     

# Radar reflectivity factor at horizontal polarisation from spectral processing
# Format = [Range,time]
dBZ = np.swapaxes(nc.variables['ZED_HC'][:,:],0,1)     

# Radial velocity of scatterers away from instrument at horizontal polarisation 
# from spectral processing - Format = [Range,time]                  
VEL_HC = np.swapaxes(nc.variables['VEL_HC'][:,:],0,1)   

nc.close()

# Switch time in seconds into hrs
timehr = time/3600.0 

# The initial velocity is the velocity fall to radar, and hence is negative.
# For convenience of understanding, change its sign.
Vel = -VEL_HC                                                                      

############################################################################
# ========================= Prepare arrays for velocity ===================#
############################################################################

# Array to store the locations of raindrops
vel_rain = np.zeros((len(Range),len(time)))

# Array to store the locations of snowflakes      
vel_snow = np.zeros((len(Range),len(time)))           

#############################################################################
# Function to mask the varible so that the 0 values are not plotted
def mask_array(X):
    X[X==0]=float('nan')
    X = np.ma.array (X, mask=np.isnan(X))
    return(X)

#############################################################################
# ========== Classify the solid and liquid hydrometeor velocity ============#

def vel_classify():
    for i in range (0,len(time)):
        for j in range (0,len(Range)):
            # Find velocity for raindrops
            if Vel[j,i] >= 2.5:
                vel_rain[j,i] = 1 
            # Find velocity for snowflakes
            if 0.5 <= Vel[j,i] <= 1.5 :
                vel_snow[j,i] = 1 
    return()

# Run the function and print out time spent on its execution
print('Velocity Marking Done: ',timeit.timeit(stmt = vel_classify,number = 1))

# Mask the variables so that the 0 values are not plotted
vel_rain = mask_array(vel_rain)
vel_snow = mask_array(vel_snow)

##############################################################################
# ===================== Calculate (Central) Derivative ===================== #

# Array to store the derivative of velocity
dVel = np.zeros((len(Range),len(time)))  

# Derivative calculation interval in meters
width = 360 

# Derivative interval grid/pixel
dh = int((width/30)/2)

def derivative():
    for i in range (0,len(time)):
       for j in range (len(Range)-(dh+1),dh,-1):
         dVel[j,i] = 1000*((Vel[j+dh,i]-Vel[j-dh,i]))/(Range[j+dh]-Range[j-dh])
    return()

# Run the function and print out time spent on its execution
print('Derivative Calculation Done: ',timeit.timeit(stmt=derivative,number=1))

# Mask the variables so that the 0 values are not plotted
dVel = mask_array(dVel)


#############################################################################
# ============== Prepare arrays for bright band identification ============ #

# Array to store the top of the melting layer 
mark_top = np.zeros((len(Range),len(time)))     

# Array to store the bottom of the melting layer           
mark_bottom = np.zeros((len(Range),len(time)))               

# Array to store the maximum of derivative
mark_dv = np.zeros((len(Range),len(time)))                

def bb_mark():
    # Find the melting layer
    for i in range (0,len(time)):
        x = np.argwhere(vel_rain[:,i] == 1).T
        y = np.argwhere(vel_snow[:,i] == 1).T
        
        # Remove data if there is no melting. 
        # i.e. only the velocity for simply snow or rain exists
        if x.size == 0 or y.size == 0:
            continue
        
        # Quality control
        # The melting layer width should be approximately 750 m
        # The interval of 'Range' is ~30 m
        if 0<= np.min(y)-np.max(x) <= 25:
            # Mark the top of the melting layer 
            mark_top[np.min(y),i] = 1
            
            # Mark the bottom of the melting layer
            mark_bottom[np.max(x),i] = 1
            
            # If the bright band is detected, find the location of the minimum
            # first derivative and mark as the bright band location.
            m = np.argwhere(dVel[:,i] == np.min(dVel[:,i])).T            
            mark_dv[m,i] = 1 
    return()

# Run the function and print out time spent on its execution
print('Bright Band Marking Done: ',timeit.timeit(stmt = bb_mark,number = 1))

# Mask the variable so that the 0 values are not plotted
mark_top = mask_array(mark_top)
mark_bottom = mask_array(mark_bottom)

#############################################################################
# Identify BB  strength by measurring the image contrast in the 750-m region
# with the bright band location marked by mark_dv as its midpoint

def Imagecontrast():
    # Prepare arrays to store image contrast and BB location
    Imax = np.zeros([len(Range),len(time)])
    Imin = np.zeros([len(Range),len(time)])
    img_contrast = np.zeros([len(Range),len(time)])

    for i in range (0,len(time)):

        # Calculate contrast in a fixed region
        x = np.argwhere(mark_dv[:,i]==1).T+12
        y = np.argwhere(mark_dv[:,i]==1).T-12          
        a = np.argwhere(mark_dv[:,i]==1).T
        if a.size == 0:  
            continue
        else:
            # Measure image contrast
            Imax[a,i] = np.max(dBZ[np.max(y):np.min(x),i])
            Imin[a,i] = np.min(dBZ[np.max(y):np.min(x),i])
            img_contrast[a,i] = ((Imax[a,i]-Imin[a,i])/(Imax[a,i]+Imin[a,i]))*100    
    return(img_contrast)

# Get the bright band strength calculated by the function
img_bb = Imagecontrast()

#############################################################################
# Function to (1) give the marks units so that the final plot will show Time 
# in hours and Range in metres. (2) Assign bright band strength to
# corresponding bright band locations.
# a - marks at Range axis
# b - marks at Time axis
def Markingunits(a,b,marking_array):
    # Prepare two arrays to store the marks with units
    range_mark = np.zeros([len(a),1])
    time_mark = np.zeros([1,len(b)])
    bb_strength = np.zeros(len(b))
    # Assign units to these marks with the given Range and Time
    for i in range(0,len(a)):
        range_mark[i,0] = Range[a[i]]
        time_mark[0,i] = timehr[b[i]]
    range_mark = mask_array(range_mark)
    time_mark = mask_array(time_mark)
    # Remove NAN values 
    range_mark = range_mark[~np.isnan(range_mark)]
    time_mark = time_mark[~np.isnan(time_mark)]
    
    # Store BB strength for plotting
    for j in range(0,len(a)):
        bb_strength[j] = np.abs(marking_array[a[j],b[j]])
    bb_strength = mask_array(bb_strength)
    return(range_mark,time_mark,bb_strength)
############################################################################
# ================================= Plotting ============================= #
############################################################################
fig = plt.figure()
ax = fig.add_subplot(1, 1, 1)

# Set colormap
cmap = matplotlib.cm.jet

# Basic radar reflectivity factor profile plot                             
plt.pcolor(timehr,Range,dBZ,cmap=cmap)

# Annotate colorbar
cbar = plt.colorbar(fraction=0.046,pad=0.01,aspect=35)
cbar.ax.set_ylabel('dBZ',rotation=90)
ticklabs = cbar.ax.get_yticklabels()
cbar.ax.set_yticklabels(ticklabs,ha='right')
cbar.ax.yaxis.set_tick_params(pad=40)

# Plot the bright band identified
# Get the index of the variable to plot
a,b = np.argwhere(img_bb != 0).T
a,b,bb_strength = Markingunits(a,b,img_bb)

# Plot strength as errorbar - weaker BB with larger errorbar (currently)
plt.errorbar(b,a,yerr = bb_strength*3,fmt ='o', color = 'white',ms = 7, \
 markeredgecolor='black',label = 'Bright Band',alpha = 0.5,ecolor = 'b',lw = 3)

# Annotate the plot
plt.title('Copernicus Radar Reflectivity Factor Profile ('+date+')')
plt.ylabel('Range (m)')
plt.xlabel('Time (hr)')
plt.xlim([0,24])
plt.xticks(np.arange(0,25,1))
plt.ylim([0,9000])
plt.legend(loc='upper left')       

# Set grid spacing
# Major ticks every 500 m, minor ticks every 100 m
major_ticks = np.arange(0, 9000, 500)
minor_ticks = np.arange(0, 9000, 100)

ax.set_yticks(major_ticks)
ax.set_yticks(minor_ticks, minor=True)
# Set the width of grid lines
ax.grid(which='minor', alpha=0.2)
ax.grid(which='major', alpha=0.5)

plt.tight_layout()
plt.show()

############################################################################
# Save figure
savename = 'BB_strength_'+date+'.pdf';
plt.savefig(savename, format='pdf', dpi=900)

# Print time spent to complete excution
print('Time for the total script:',datetime.now() - startTime)

############################################################################
"""
Oneclick event/ vertical profile 
"""
def oneclick(event):
    # Handle the image clicking system
    global ix, iy
    ix, iy = event.xdata, event.ydata

    print ('x = %d, y = %d'%(ix, iy))

    # Assign global variable to access outside of function
    global coords
    coords.append((ix, iy))

    # Disconnect after clicks
    if len(coords) == 1:
       fig.canvas.mpl_disconnect(cid)
       print('Points Chosen')
                
       index=coords[0][0]
       # Work out what time these were in normal times, rather than decimal time
       stimestr='%02d:%02d'%(np.floor(index),round(60*(index-np.floor(index))))
       print(stimestr)
        
        # Plot various lines
       plt.figure(figsize=(10,8))
       ax1 = plt.subplot(1,2,1)
       plt.plot(Vel[:,int(index*len(time)/24)],Range,'b-',linewidth = 2)
       plt.xlabel('$m \cdot s^{-1}$',fontsize = 18)
       plt.ylabel('Range (m)',fontsize = 18)
       plt.ylim([0,6000])
       plt.title('Doppler Velocity at '+ stimestr +' on '+ date,fontsize = 12)
       major_ticks = np.arange(0, 6000, 500)
       minor_ticks = np.arange(0, 6000, 100)
       
       ax1.set_yticks(major_ticks)
       ax1.set_yticks(minor_ticks, minor=True)
       # Set the width of grid lines
       ax1.grid(which='minor', alpha=0.2)
       ax1.grid(which='major', alpha=0.5)
       
       
       ax2 = plt.subplot(1,2,2)
       plt.plot(dVel[:,int(index*len(time)/24)],Range,'r-',linewidth = 2)
       plt.ylim([0,6000])
       plt.title('First Derivative at '+stimestr+' on '+date,fontsize = 12)
       major_ticks = np.arange(0, 6000, 500)
       minor_ticks = np.arange(0, 6000, 100)

       ax2.set_yticks(major_ticks)
       ax2.set_yticks(minor_ticks, minor=True)
       # Set the width of grid lines
       ax2.grid(which='minor', alpha=0.2)
       ax2.grid(which='major', alpha=0.5)
       
       plt.tight_layout()
       return(index)

coords = []

# Call click functioin
cid = fig.canvas.mpl_connect('button_press_event', oneclick)