import act
import numpy as np
import xarray as xr
import matplotlib.pyplot as plt

# Set your username and token here!
username = 'username'
token = 'token'

# Set the datastream and start/enddates
datastream = 'kcgmplpolfsM1.b1'
startdate = '2024-11-10'
enddate = '2024-11-11'

# Use ACT to easily download the data.  Watch for the data citation!  Show some support
# for ARM's instrument experts and cite their data if you use it in a publication
result = act.discovery.download_arm_data(username, token, datastream, startdate, enddate)

# Let's read in the data using ACT and check out the data
ds_mpl = act.io.read_arm_netcdf(result)

ds_mpl

# Let's take a look at the quality control information associated with a variable from the MPL
variable = 'signal_return_co_pol'

# First, for many of the ACT QC features, we need to get the dataset more to CF standard and that
# involves cleaning up some of the attributes and ways that ARM has historically handled QC
ds_mpl.clean.cleanup()

# Apply corrections for the ceilometer, correcting for the vertical height
ds_mpl = act.corrections.correct_mpl(ds_mpl)

# Next, let's take a look at visualizing the quality control information
# Create a plotting display object with 2 plots
display = act.plotting.TimeSeriesDisplay(ds_mpl, figsize=(10, 5), subplot_shape=(1,))

# Plot up the variable in the first plot
display.plot(variable, subplot_index=(0,))

# Plot up the QC variable in the second plot
#display.qc_flag_block_plot(variable, subplot_index=(1,))
plt.show()

# Let's filter out test 5 using ACT.  Yes, it's that simple!
ds_mpl.qcfilter.datafilter(variable, rm_tests=[1, 2], del_qc_var=False)

# There are other ways we can filter data out as well.  Using the
# rm_assessments will filter out by all Bad/Suspect tests that are failing
# ds.qcfilter.datafilter(variable, rm_assessments=['Bad', 'Suspect'], del_qc_var=False)

# Let's check out the attributes of the variable
# Whenever data are filtered out using the datafilter function
# a comment will be added to the variable history for provenance purposes
print(ds_mpl[variable].attrs)

# And plot it all again!
# Create a plotting display object with 2 plots
display = act.plotting.TimeSeriesDisplay(ds_mpl, figsize=(15, 10), subplot_shape=(2,))

# Plot up the variable in the first plot
display.plot(variable, subplot_index=(0,))

# Plot up a day/night background
display.day_night_background(subplot_index=(0,))

# Plot up the QC variable in the second plot
display.qc_flag_block_plot(variable, subplot_index=(1,))
plt.show()

# Query the ARM DQR Webservice
ds_mpl = act.qc.add_dqr_to_qc(ds_mpl, variable=variable)

ds_mpl['qc_' + variable]

# Let's set a longer time period
startdate = '2024-11-07'
enddate = '2024-11-14'

# aps
datastream = 'kcgaosapsS3.b1'
result = act.discovery.download_arm_data(username, token, datastream, startdate, enddate)
ds_aps = act.io.arm.read_arm_netcdf(result)

# AOSMET - Just to get the wind data!
datastream = 'kcgaosmetS2.a1'
result = act.discovery.download_arm_data(username, token, datastream, startdate, enddate)
ds_met = act.io.read_arm_netcdf(result)

# MPL to get the full record
datastream = 'kcg30smplcmask1zwangM1.c1'
result = act.discovery.download_arm_data(username, token, datastream, startdate, enddate)
ds_mpl = act.io.read_arm_netcdf(result)

# Before we proceed to plotting, let's reduce the MPL data down a little bit
# This will remove all data where heights are greater than 5
ds_mpl = ds_mpl.where(ds_mpl.height <= 3, drop=True)

# This will resample to 1 minute
ds_mpl = ds_mpl.resample(time='1min').nearest()

# Let's not forget about QCing the data!
# We can remove all the bad data from each aerosol dataset

ds_aps.clean.cleanup()
ds_aps = act.qc.arm.add_dqr_to_qc(ds_ccn)
ds_aps.qcfilter.datafilter(rm_assessments=['Bad'], del_qc_var=False)

ds_met.clean.cleanup()
ds_met = act.qc.arm.add_dqr_to_qc(ds_met)
ds_met.qcfilter.datafilter(rm_assessments=['Bad'], del_qc_var=False)

ds_mpl.clean.cleanup()
ds_mpl = act.qc.arm.add_dqr_to_qc(ds_mpl)
ds_mpl.qcfilter.datafilter(rm_assessments=['Bad'], del_qc_var=False)

ds_aps

# We can pass a dictionary to the display objects with multiple datasets
# So let's plot all this up!
display = act.plotting.TimeSeriesDisplay({'mpl': ds_mpl, 'aps': ds_aps, 'met': ds_met},
                                         subplot_shape=(4,), figsize=(10,16))

# MPL Plot
# Variable names of interest linear_depol_ratio, linear_depol_snr, backscatter_snr
display.plot('linear_depol_ratio', dsname='mpl', subplot_index=(0,), cmap='HomeyerRainbow')
display.set_yrng([0, 3], subplot_index=(0,))

# CNN Plot
display.plot('total_N_conc', dsname='aps', subplot_index=(1,))
display.day_night_background(dsname='aps', subplot_index=(1,))

# CNN-kappa Plot
display.plot('total_V_conc', dsname='aps', subplot_index=(2,))
display.day_night_background(dsname='aps', subplot_index=(2,))

# MET Plot
display.plot('wind_speed', dsname='met', subplot_index=(3,))
display.day_night_background(dsname='met', subplot_index=(3,))

plt.subplots_adjust(hspace=0.3)
plt.savefig('output.png')
plt.show()

ds_combined.total_N_conc.plot()

# We already should have the data loaded up so let's explore with some data roses
# First we need to combine data and to do that, we need to get it on the same time grid
ds_combined = xr.merge([ds_met.resample(time='30min').nearest(), ds_aps.resample(time='30min').nearest()], compat='override')

# Plot out the data rose using the WindRose display object
display = act.plotting.WindRoseDisplay(ds_combined)
display.plot_data('wind_direction', 'wind_speed', 'total_N_conc', num_dirs=15, plot_type='line', line_plot_calc='mean')
plt.show()

# We do the same thing as before but call the DistributionDisplay class
display = act.plotting.DistributionDisplay(ds_ccn)

# And then we can plot the data!  Note that we are passing a range into the
# histogram function to set the min/max range of the data
display.plot_stacked_bar('total_N_conc', bins=20, hist_kwargs={'range': [0, 40]})
plt.show()

# We can create these plots in groups as well but we need to know
# how many there will be ahead of time for the shape
display = act.plotting.DistributionDisplay(ds_aps, figsize=(15, 15), subplot_shape=(6, 4))
groupby = display.group_by('hour')

# And then we can plot the data in groups!  The main issue is that it doesn't automatically
# Annotate the group on the plot.  We're also setting the titile to blank to save space
groupby.plot_group('plot_stacked_bar', None, field='total_N_conc', set_title='', bins=20, hist_kwargs={'range': [0, 40]})

# We want these graphs to have the same axes, so we can easily run through
# each plot and modify the axes.  Right now, we can just hard code these in
for i in range(len(display.axes)):
    for j in range(len(display.axes[i])):
        display.axes[i, j].set_xlim([0, 40])
        
plt.subplots_adjust(wspace=0.35)
        
plt.show()

# Let's merge the aps and ACSM data together and plot out some distribution plots
# First we need to combine data and to do that, we need to get it on the same time grid
ds_combined = xr.merge([ds_met.resample(time='30min').nearest(), ds_aps.resample(time='30min').nearest()], compat='override')

# Plot out the data rose using the Distribution display object
display = act.plotting.DistributionDisplay(ds_combined)
display.plot_scatter('total_N_conc', 'wind_speed', m_field='time')

plt.show()

# Let's try a heatmap with this as well!
display = act.plotting.DistributionDisplay(ds_combined, figsize=(12, 5), subplot_shape=(1, 2))

display.plot_scatter('total_N_conc', 'wind_speed', m_field='time', subplot_index=(0, 0))
display.plot_heatmap('total_N_conc', 'wind_speed', subplot_index=(0, 1), x_bins=50, y_bins=50, threshold=0)

plt.show()

for diameter in ds_aps.diameter_aerodynamic.values:
    ds_aps[f"dN_dlogDp_{diameter}"] = ds_aps["dN_dlogDp"].sel(diameter_aerodynamic=diameter)

# Let's try one last plot type with this dataset
# Violin plots!
display = act.plotting.DistributionDisplay(ds_aps, figsize=(14, 5))

diameters_to_use = ds_aps.diameter_aerodynamic.values[::4]
for diameter in diameters_to_use:
    display.plot_violin(f"dN_dlogDp_{diameter}", positions=[diameter])

# Let's add some more information to the plots
# Update the tick information
display.axes[0].set_xticks(diameters_to_use)
display.axes[0].set_xticklabels([f"{int(value)} \n um" for value in list(diameters_to_use)])

plt.show()

# Let's set a longer time period
startdate = '2024-11-07'
enddate = '2024-11-14'

# SONDE
datastream = 'kcgsondewnpnM1.b1'
result = act.discovery.download_arm_data(username, token, datastream, startdate, enddate)
result.sort()
ds_sonde = act.io.read_arm_netcdf(result[-1])

# Plot enhanced Skew-T plot
display = act.plotting.SkewTDisplay(ds_sonde)
display.plot_enhanced_skewt(color_field='alt')

plt.show()

# Now we can plot up a wind rose of that entire month's worth of data
windrose = act.plotting.WindRoseDisplay(ds_met, figsize=(10,8))
windrose.plot('wind_direction', 'wind_speed', spd_bins=np.linspace(0, 10, 5))
windrose.axes[0].legend()
plt.show()

ds_met

# Let's accumulate the precipitation data from the three different sensors in the MET System
# These instruments include a tipping bucket rain gauge, optical rain gauge, and a present weather detector
ds_met = act.utils.data_utils.accumulate_precip(ds_met, "rain_intensity")

# We can plot them out easily in a loop.  Note, they have _accumulated on the end of the name
display = act.plotting.TimeSeriesDisplay(ds_met, figsize=(8, 6))
display.plot('rain_intensity_accumulated', label="rain_intensity")

# Add a day/night background
display.day_night_background()
plt.subplots_adjust(bottom=0.20)
plt.legend()
plt.show()

# We're going to use some test data that already exists within ACT
# Let's set a longer time period
startdate = '2024-11-10T00:00:00'
enddate = '2024-11-10T06:00:00'

# SONDE
datastream = 'kcgdlppiM1.b1'
result = act.discovery.download_arm_data(username, token, datastream, startdate, enddate)
result.sort()

ds = act.io.read_arm_netcdf(result)
ds
# Returns the wind retrieval information in a new object by default
# Note that the default snr_threshold of 0.008 was too high for the first profile
# Reducing it to 0.002 makes it show up but the quality of the data is likely suspect.
ds_wind = act.retrievals.compute_winds_from_ppi(ds, snr_threshold=0.01)

# Plot it up
display = act.plotting.TimeSeriesDisplay(ds_wind)
display.plot_barbs_from_spd_dir('wind_speed', 'wind_direction', invert_y_axis=False)

#Update the x-limits to make sure both wind profiles are shown
display.axes[0].set_xlim([np.datetime64(startdate), np.datetime64(enddate)])

plt.ylim(0, 2000)

plt.show()