# -*- coding: utf-8 -*-
"""
Python script for bright band identification using the "Combined Algorithm". 

1. Identify the bright band top and bottom using velocity algorithm.
2. Find where Z is max or where the first derivative of Z changes the most 
between the bright band top and bottom marked in step 1. 

@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 Radar 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 decimal time 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 

#############################################################################
# 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)

############################################################################
# =================== Prepare empty arrays for velocity ===================#
 
# Array to store the locations of raindrops                                        
vel_rain = np.zeros((len(Range),len(time)))

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

#############################################################################
# ========== 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] >= 3.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 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)

#############################################################################
# ====================== Prepare empty arrays for Z ========================#

# Convert Z from log to linear
Z = 10**dBZ

# Array to store maxima of linear Z                                         
max_Z = np.zeros((len(Range),len(time))) 
    
# Array to store derivative of Z, i.e. dZ/dH
dZ = np.zeros((len(Range),len(time)))      
   
# Array to store the magnitude of the first derivative (dZ)     
mag_dZ = np.zeros((len(Range),len(time))) 

# Array to mark where the dZ changes most significantly         
max_dZ = np.zeros((len(Range),len(time)))          

##############################################################################
# ================= Calculate the (central) derivative of Z =================#
def derivative():
    for i in range (0,len(time)):
        for j in range (len(Range)-2,-1,-1):
            if j == 0:
                break
            dZ[j,i] = ((Z[j+1,i]-Z[j-1,i]))/(Range[j+1]-Range[j-1])
            # Store the magnitude of how much dZ changes
            if j<(len(Range)-1):
                mag_dZ [j,i] = dZ[j,i]-dZ[j-1,i]
    return()

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

# Mask the variable so that the 0 values are not plotted
mag_dZ = mask_array(mag_dZ)
dZ = mask_array(dZ)

#############################################################################
# ===== Prepare arrays to mark the top and bottom of the melting layer =====#

# Array to store the locations of the bright band top  
mark_top = np.zeros((len(Range),len(time)))
     
# Array to store the locations of the bright band bottom           
mark_bottom = np.zeros((len(Range),len(time)))               

# Function to identify bright band locations using the "Combined Algorithm"
def bb_mark():
    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 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
            
            # Mark the max Z or max changes of dZ in the region
            for j in range(np.max(x),np.min(y)):
                if mag_dZ[j,i] == np.max(mag_dZ[:,i]):
                    max_dZ [j,i] = 1
                if Z[j,i] == np.max(Z[:,i]):
                    max_Z [j,i] = 1
    return()

# Run the function and print out time spent on 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)
max_dZ = mask_array(max_dZ)
max_Z = mask_array(max_Z)

#############################################################################
# Function to give the marks units so that the final plot will show Time in 
# hours and Range in metres. 
# 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)])
    # 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]]
    for j in range(0,len(b)):
        time_mark[0,j] = timehr[b[j]]
    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)]
    return(range_mark,time_mark)

############################################################################
# ============================== Plotting ================================ #
############################################################################

fig1 = plt.figure()
ax = fig1.add_subplot(1, 1, 1)

# Set colormap
cmap = matplotlib.cm.jet 
# Basic radar reflectivity factor profile plot                            
plt.pcolor(timehr,Range,dBZ,cmap=cmap)
plt.title('Copernicus Radar Reflectivity Factor Profile ('+date+')')


# 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(max_Z == 1).T
a,b = Markingunits(a,b,max_Z)
x,y = np.argwhere(max_dZ == 1).T
x,y = Markingunits(x,y,max_dZ)

plt.plot(b,a,'^',ms = 10,color = 'grey',markeredgecolor='black',\
         alpha = 0.5,label = 'Linear Z')
plt.plot(y,x,'kx',alpha = 0.5,label = 'Derivative')

# Annotate plots
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_combined_'+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:
       fig1.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()
       ax1 = plt.subplot(1,4,1)
       plt.plot(dBZ[:,int(index*len(time)/24)],Range,'k-',linewidth = 2)
       plt.xlabel('dB',fontsize = 18)
       plt.ylabel('Range (m)',fontsize = 18)
       plt.ylim([0,6000])
       plt.title('dBZ at '+ stimestr +' on '+ date,fontsize = 14)
       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,4,2)
       plt.plot(Z[:,int(index*len(time)/24)],Range,'b-',linewidth = 2)
       plt.ylim([0,6000])
       plt.title('$Z_{linear}$ at '+ stimestr +' on '+ date,fontsize = 14)
       plt.xlabel('$mm^6 \cdot m^{-3}$')
       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)
       
       ax3 = plt.subplot(1,4,3)
       plt.plot(dZ[:,int(index*len(time)/24)],Range,'r-',linewidth = 2)
       plt.ylim([0,6000])
       plt.title('First Derivative at '+ stimestr +' on '+ date,fontsize = 13)
       major_ticks = np.arange(0, 6000, 500)
       minor_ticks = np.arange(0, 6000, 100)
       
       ax3.set_yticks(major_ticks)
       ax3.set_yticks(minor_ticks, minor=True)
       # Set the width of grid lines
       ax3.grid(which='minor', alpha=0.2)
       ax3.grid(which='major', alpha=0.5)
        
       ax4 = plt.subplot(1,4,4)
       plt.plot(Vel[:,int(index*len(time)/24)],Range,'-',\
                    linewidth = 2,color='brown')
       plt.xlabel('$m \cdot s^{-1}$',fontsize = 18)
       plt.ylim([0,6000])
       plt.title('Doppler Velocity'+ stimestr +' on '+ date,fontsize = 13)
       major_ticks = np.arange(0, 6000, 500)
       minor_ticks = np.arange(0, 6000, 100)
       
       ax4.set_yticks(major_ticks)
       ax4.set_yticks(minor_ticks, minor=True)
       # Set the width of grid lines
       ax4.grid(which='minor', alpha=0.2)
       ax4.grid(which='major', alpha=0.5)
       
       plt.tight_layout()
       return(index)

coords = []

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