# -*- coding: utf-8 -*-
"""
Python Script to identify radar bright band using
the Doppler (radial) velocity, i.e. the Velocity Algorithm.

1. If the velocity is between 0.5 and 1.5 m/s, the precipitation is classified 
as snow. The lowest range is the bright band top.

2. If the velocity is over 3.5 m/s, the precipitation is classified as rain.
The lowest range is the bright band bottom.

3. If the distance between the top and the bottom is within 750 m, it is a bright 
band.

@author: Dongqi Lin 2018
Contact: ee17dl@leeds.ac.uk
"""
##########################################################################
# Import packages 
import sys
import timeit
import matplotlib
import matplotlib.pylab as plt
import matplotlib.mlab as mlab
import numpy as np
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)

# Begin to calculate time spent on execution
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'][:]  

# 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) 
    
# Radar reflectivity factor at horizontal polarisation from spectral processing
# Format = [Range,time]
dBZ = np.swapaxes(nc.variables['ZED_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
 
############################################################################
# Prepare empty arrays
 
# 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)))

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

#############################################################################
# Function to classify the velocity for solid and liquid hydrometeors
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)

##############################################################################
# Mark the top and bottom of the bright band 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 bright band bottom                 
mark_bottom = np.zeros((len(Range),len(time)))               

def vel_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 0<=np.min(y)-np.max(x) <= 25:
            mark_top[np.min(y),i] = 1
            mark_bottom[np.max(x),i] = 1
    return()

# Run the function and print out time spent on execution
print('Bright Band Marking Done: ',timeit.timeit(stmt = vel_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)


#############################################################################
# Function to give units to the marks 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 ================================ #
############################################################################
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 coordinates of the variable to plot
a,b = np.argwhere(mark_top == 1).T
a,b = Markingunits(a,b,mark_top)

m,n = np.argwhere(mark_bottom == 1).T
m,n = Markingunits(m,n,mark_bottom)

plt.plot(b,a,'o',color = 'white',markeredgecolor='black',ms = 7,\
         label = 'BB_Vel_top',alpha = 0.5)
plt.plot(n,m,'kx',ms = 7,label = 'BB_Vel_bottom',alpha = 0.5)

# Annotate plot
plt.legend(loc='upper left')
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])
       
# Set grid spacing
# Major ticks every 20, minor ticks every 5
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_Vel_'+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=(6,8))
       ax1 = plt.subplot(1,1,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)
       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)
       
       plt.tight_layout()
       return(index)

coords = []

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