# Load the needed packages
import os
import matplotlib.pyplot as plt
import numpy as np
from netCDF4 import Dataset

import awot
from awot.graph import FlightLevel, RadarVerticalPlot, MicrophysicalVerticalPlot

%matplotlib inline

"""
Explanation: <h2> How to add a field to an existing AWOT instance
End of explanation
"""


file1 = "WCR.OWLES13.20131215.225944_234806.up-down.nc"
#file = "WCR.TEST14.20140618.200302_201559.up.nc"
wcrf1 = os.path.join("/Users/guy/data/king_air/owles2013/wcr", file1)

file2 = "W-CORRECTED.WCR.OWLES13.20131215.225944_234806.up-down.nc"
wcrf2 = os.path.join("/Users/guy/data/king_air/owles2013/wcr/", file2)

start_time = "2013-12-16 01:25:00"
end_time = "2013-12-16 01:30:00"

# Set the project name
Project="OWLeS"

#OWLES
altmin, altmax = 0., 3000.

refmin, refmax = -50., 30.
#velmin, velmax = -16., 16. # Nyquist
velmin, velmax = -3., 3. # Nyquist

"""
Explanation: <b>Supply input data and plotting characteristics. In this case we'll use a file from the OWLeS project and corrected field of velocity data using sounding data.</b>
End of explanation
"""


wcr = awot.io.read_wcr2(fname=wcrf1)

"""
Explanation: <b>Read in the radar data</b>
End of explanation
"""


nc = Dataset(wcrf2)
velcor = nc.variables['Velocity_cor_2']
awot.util.add_dict_to_awot_fields(wcr, 'velocity_corrected', data=velcor[:],
                                 units=velcor.units, longname=velcor.long_name, stdname="Corrected velocity")

print(wcr['fields']['velocity']['data'].shape, wcr['fields']['velocity_corrected']['data'].shape)
print(np.ma.min(wcr['fields']['velocity_corrected']['data']), np.ma.max(wcr['fields']['velocity_corrected']['data']))

print(wcr.keys())
print("Reflectivity Min/Max = %f / %f \n"
      "Velocity Min/Max = %f / %f \n"
      "Altitude Min/Max = %f / %f \n"
      "Height Min/Max = %f / %f \n"
      "Surface Min/Max = %f / %f \n"%(
        wcr['fields']['reflectivity']['data'].min(),wcr['fields']['reflectivity']['data'].max(),
        wcr['fields']['velocity']['data'].min(),wcr['fields']['velocity']['data'].max(),
        wcr['altitude']['data'].min(), wcr['altitude']['data'].max(),
        wcr['height']['data'].min(), wcr['height']['data'].max(),
        wcr['surface']['data'].min(), wcr['surface']['data'].max()))

"""
Explanation: <b>Read a variable from another file and add it to the AWOT dictionary. A mask of invalid data is automatically applied. Additional masking can be accomplished by setting the <i>mask_value</i> keyword.</b>
End of explanation
"""


fig, (ax1, ax2, ax3) = plt.subplots(3, 1, figsize=(7.5,10), sharex=True)


wcr_vp= RadarVerticalPlot(wcr)

wcr_vp.time_height_image('reflectivity', ax=ax1, plot_log10_var=False,
#              start_time=start_time, end_time=end_time,
                         vmin=refmin, vmax=refmax,
                         fill_surface=True,
                         cb_label=r'Reflectivity (dBZ)',
                         height_min=altmin, height_max=altmax, title=file1,
                         ylab=r'Altitude (m)', ylabFontSize=12)
              
wcr_vp.time_height_image('velocity', ax=ax2, plot_log10_var=False,
#              start_time=start_time, end_time=end_time,
                         vmin=velmin, vmax=velmax,
                         fill_surface=True,
                         cmap="PuOr_r",
                         cb_label=r'Doppler Velocity (m s$^{-1}$)',
                         height_min=altmin, height_max=altmax,
                         ylab=r'Altitude (m)', ylabFontSize=12,
                         xlab='UTC Time', xlabFontSize=12)
              
wcr_vp.time_height_image('velocity_corrected', ax=ax3, plot_log10_var=False,
#              start_time=start_time, end_time=end_time,
                         vmin=velmin, vmax=velmax,
                         fill_surface=True,
                         cmap="PuOr_r",
                         cb_label=r'Doppler Velocity (m s$^{-1}$)',
                         height_min=altmin, height_max=altmax,
                         ylab=r'Altitude (m)', ylabFontSize=12,
                         xlab='UTC Time', xlabFontSize=12)

"""
Explanation: <b>Make a vertical time series plot of reflectivity and velocity fields</b>
End of explanation
"""


