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

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

# Set the datastream and start/enddates
datastream = 'houaospsap3w1mM1.b1'
startdate = '2022-07-16'
enddate = '2022-07-16'

# 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 = act.io.read_arm_netcdf(result)
ds

# We're going to be focusing on the following variable, so let's get some more information about it
# We can do this by looking at it's attributes
variable = 'Ba_B_Weiss'
ds[variable].attrs

# There's a lot of great functionality in ACT, but there's also a lot in the base xarray Dataset!
ds[variable].plot()

# We can see that there's some missing data in the plot above so let's take a look at the embedded QC!
# 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.clean.cleanup()

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

# Let's filter out test 5 using ACT.  Yes, it's that simple!
ds.qcfilter.datafilter(variable, rm_tests=[5], 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[variable].attrs['history'])

# And plot it all again!
# Create a plotting display object with 2 plots
display = act.plotting.TimeSeriesDisplay(ds, 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 = act.qc.add_dqr_to_qc(ds, variable=variable)

#And plot again!
# Create a plotting display object with 2 plots
display = act.plotting.TimeSeriesDisplay(ds, 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()

# First, let's plot up a histogram of the data
# All the ACT plotting is very similar to what we
# did earlier, first we create a display object
display = act.plotting.DistributionDisplay(ds)

# 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(variable, hist_kwargs={'range': [0, 10]})
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, figsize=(12, 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=variable, set_title='', hist_kwargs={'range': [0, 10]})

# 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, 10])
        display.axes[i, j].set_ylim([0, 50])
        
plt.show()

# Next up, let's do some scatter plots to compare some variables
# Scatter plots are also found in the DistributionDisplay module
display = act.plotting.DistributionDisplay(ds)

# And then we can plot the data!
display.plot_scatter(variable, 'transmittance_blue', m_field='time')

# You can adjust the x-range as you need
# display.set_xrng([0, 20])
plt.show()

# Sometimes these scatter plots hide the number of points there actually
# are in some areas so let's try a heatmap as well
display = act.plotting.DistributionDisplay(ds, figsize=(12, 5), subplot_shape=(1, 2))

# And then we can plot the data!
display.plot_scatter(variable, 'transmittance_blue', m_field='time', subplot_index=(0, 0))
# This can be used to adjust the axes limits
# display.set_xrng([0, 20], subplot_index=(0, 0))

# we can also pass in an array of values for the bins using np.arange(start, stop, step)
display.plot_heatmap(variable, 'transmittance_blue', subplot_index=(0, 1), x_bins=25, y_bins=25)
plt.show()

# Let's try one last plot type with this dataset
# Violin plots!
display = act.plotting.DistributionDisplay(ds)

# And then we can plot the data!
display.plot_violin(variable, positions=[1.0])

# And we can add more variables to it as well!
display.plot_violin('Ba_R_Weiss', positions=[2.0])
display.plot_violin('Ba_G_Weiss', positions=[3.0])

# Let's add some more information to the plots
# Update the tick information
display.axes[0].set_xticks([0.5, 1, 2, 3, 3.5])
display.axes[0].set_xticklabels(['',
                                 'Blue Channel\nAbsoprtion',
                                 'Red Channel\nAbsorption',
                                 'Green Channel\nAbsorption',
                                 '']
                                )

# Update the y-axis label
display.axes[0].set_ylabel('Aerosol Light Absorption Coefficient')
plt.show()