import h5py
from matplotlib import pyplot as plt
import numpy as np

radName = input('CRS, EXRAD_NADIR, HIWRAP_KU, HIWRAP_KA?: ')
start   = input('Best guess start index (Choose 0 if unsure): ')
stop    = input('Best guess stop index (Choose -1 if unsure): ')
channel = input('Channel#: ')

start = int(start)
stop  = int(stop)

##Declare input files
f  = h5py.File('/data10/ppantina/IMPACTS-2022/Level1A_test0/'    + radName + '_20220222_L1A.h5', 'r')
f0 = h5py.File('/data10/IMPACTS-2022/Level0/20220222/' + radName + '_20220222_nav.hdf', 'r')

##Find dBZe and mask NANs. Find argmax for all times.
Z_masked = np.ma.masked_invalid(10*np.log10(f['Products/Channel' + channel +'/Data/Zuncal'])) 
Z_idx    = np.argmax(Z_masked, axis = 1)

##Find Vel. Then find vel for all times at dBZe argmax index.
vel = f['Products/Channel' + channel + '/Data/V']
vel_surf = np.zeros(vel.shape[0])
for i, j in enumerate (vel_surf):
  vel_surf[i] = vel[i, Z_idx[i]]

##Calculate platform speed
spd = np.abs(f['Navigation/Data/EastVelocity'][:] + f['Navigation/Data/NorthVelocity'][:]*1j)

##Make a plot
fig = plt.figure()
ax = fig.add_subplot(311)
ax.pcolormesh(Z_masked.T, vmin = -60, vmax = 20,  cmap = 'jet')
ax.set_label('DV')
ax = fig.add_subplot(312)
ax.plot(vel_surf)
ax.set_ylim(-4, 4)
ax.set_ylabel('Vel Error')
ax = fig.add_subplot(313)
ax.plot(spd)
ax.set_ylabel('Vel Platform')
print ('Check this plot for start/stop times')
plt.show()

##Update the times if needed
loop = input ('Update start/stop times? Y or N: ')
if loop =='Y':
  start   = input('Best guess start index: ')
  stop    = input('Best guess stop index: ')
  start   = int(start)
  stop    = int(stop)

  print ('times updated')
else: print ('Continuing')

#Calculate averages
pitch_avg = np.nanmean(f['Navigation/Data/Pitch'][start:stop])
vel_avg = np.nanmean(vel_surf[start:stop])
spd_avg = np.nanmean(spd[start:stop])
rho = vel_avg/spd_avg

offset = rho/3.14159*180

##Find L0 pitch data
g_time       = f0['gps/cpusec'] +  f0['gps/cpunsec'][:]/1e9
g_time_start = f['Time/Data/TimeUTC'][start]
g_time_stop  = f['Time/Data/TimeUTC'][stop ]
g_start      = np.argmin(np.abs(g_time - g_time_start))
g_stop       = np.argmin(np.abs(g_time - g_time_stop))
g_pitch      = np.nanmean(f0['gps/INSPVA_pitch'][g_start:g_stop])

n_time       = f0['nav/cpusec'] +  f0['nav/cpunsec'][:]/1e9
n_time_start = f['Time/Data/TimeUTC'][start]
n_time_stop  = f['Time/Data/TimeUTC'][stop ]
n_start      = np.argmin(np.abs(n_time - n_time_start))
n_stop       = np.argmin(np.abs(n_time - n_time_stop))
n_pitch      = np.nanmean(f0['nav/IWG_pitch'][n_start:n_stop])

print ('vel_error               ', vel_avg)
print ('platform vel            ', spd_avg)
print ('offset(rho)             ', offset)
print ('L1A hardcoded pitch corr', f['Navigation/Information/PitchCorrection'][:])
print ('L1A Pitch               ', pitch_avg)
print ('L0 NAV pitch            ', n_pitch)
print ('L0 GPS pitch            ', g_pitch)