################################################################################################
#Contact roktem@lbl.gov for questions
#Date : June 2020
#This script reads PCCP data for a time instance and displays relevant plots
#Usage: python pccp_snapshot.py <input_nc_filename> <output_plot_filename> <time_instance>
#       e.g. python pccp_snapshot.py sgppccpE45.c1.20190512.130000.nc sample_plot 201100
#       will read data from sgppccpE45.c1.20190512.130000.nc and save the output in sample_plot.png for time 20:11:00

import matplotlib as mpl
import numpy as np
import numpy.ma as ma
import sys
from mpl_toolkits.mplot3d import Axes3D
import matplotlib.pyplot as plt
from netCDF4 import Dataset
import time

plt.rc('xtick', labelsize=14)
plt.rc('ytick', labelsize=14)
plt.rc('axes', labelsize=16)
plt.rc('figure', titlesize=16)

###2D plot of x,y variables in subplot(2,2,si)
def plot2D(x,y,fig,si,xlabel,ylabel):
    ax = fig.add_subplot(2,2,si)
    ax.scatter(x, y, s=1, marker='x', c= 'gray')
    ax.set_xlabel(xlabel)
    ax.set_ylabel(ylabel)
    ax.xaxis.labelpad = 5
    ax.yaxis.labelpad = 5

###3D plot of x,y,z variables in subplot(2,2,si)
def plot3D(x,y,z,fig,si):
    #check if data point count is sufficient for display
    if (len(x)>10):
        ax = fig.add_subplot(2,2,si, projection='3d')
        x = [x[0:len(x)]]
        y = [y[0:len(y)]]
        z = [z[0:len(z)]]
        x1 = int(min(min(x)))
        x2 = int(max(max(x)))
        y1 = int(min(min(y)))
        y2 = int(max(max(y)))
      
        ax.scatter(x, y, z, c='gray', marker='o')
        ax.xaxis.set_ticks(np.arange(x1,x2,int((x2-x1+2)/4)+0.5))
        ax.yaxis.set_ticks(np.arange(y1,y2,int((y2-y1+2)/4)+0.5))
        ax.view_init(elev=15, azim=-70)
        ax.set_xlabel('X [km] ')
        ax.set_ylabel('Y [km] ')
        ax.xaxis.labelpad = 15
        ax.yaxis.labelpad = 15
        ax.zaxis.set_ticks(np.arange(0,9,2))
        ax.set_zlabel('Z [km] ')

#filter out values that exceed 50km
def FilterRange(data_in, indnz):
    if (indnz==-1):
        indnz = (abs(data_in) < 50000).nonzero()
    data_out = data_in[indnz]/1000 # convert to km
    return data_out,indnz

##########MAIN####################################################
##################################################################`
def main():
    #check input arguments
    if (len(sys.argv)<4):
        print('Usage: python pccp_snapshot.py in_filename out_plot_name time(HHMMSS)')
        print('e.g. python pccp_snapshot.py sgppccpE44.c1.20180420.130000.nc sample 170000')
        sys.exit();
              
    inFileName = sys.argv[1] #input nc filename
    outFileName = sys.argv[2] #output png filename
    curTimeStr = sys.argv[3] #time to display
    curTime = int(curTimeStr[0:2])*3600+int(curTimeStr[2:4])*60+int(curTimeStr[4:6]) #convert to integer
    
    #read input
    try:
        dataset = Dataset(inFileName, 'r', format='NETCDF4')
    except:
        print('Could not open file:', inFileName)
    
    x_relative_var = dataset.variables['x_relative']
    y_relative_var = dataset.variables['y_relative']
    z_relative_var = dataset.variables['z_relative']
    timeoff = dataset.variables['time'][:]
    
    tind = np.array(range(0,len(timeoff)))
    ti = tind[timeoff == curTime ]
    if (len(ti)==0): #if no data is available at the time entered
        dt = abs(timeoff-curTime)
        tn = tind[dt==min(dt)]
        print('No data available at the entered time, the nearest available time is at ',time.strftime("%H:%M:%S", time.gmtime(timeoff[tn])))
    else:
        x_slice,ind_nonzero = FilterRange(x_relative_var[ti][:][:],-1)
        y_slice,ind_nonzero = FilterRange(y_relative_var[ti][:][:],ind_nonzero)
        z_slice,ind_nonzero = FilterRange(z_relative_var[ti][:][:],ind_nonzero)
        print(len(x_slice),' cloud points are extracted')
        
        fig = plt.figure(figsize=(24,12))
        fig.suptitle('PCCP data from file '+inFileName +' at '+curTimeStr + ' UTC')
        plot2D(x_slice,y_slice,fig,1,'direction eastward [km]','direction northward [km]')
        plot2D(x_slice,z_slice,fig,2,'direction eastward [km]','altitude above the ground [km]')
        plot2D(y_slice,z_slice,fig,3,'direction northward [km]','altitude above the ground [km]')
        plot3D(x_slice,y_slice,z_slice,fig,4)
        
        fig.savefig('{}.png'.format(outFileName))
    dataset.close()


##################################################################`
main()