import os
import numpy as np
import xarray as xr
import matplotlib.pyplot as plt
import matplotlib.colors as colors

import getpass
import act

# Note - recommended ARM Live token to be set as an environmental variable

# Check for env variables, otherwise set your username and token here!
if "ARM_USERNAME" in os.environ:
    arm_username = os.getenv("ARM_USERNAME")
else:
    arm_username = input("Enter ARM username: ").strip()
if "ARM_TOKEN" in os.environ:
    arm_token = os.getenv("ARM_TOKEN")
else:
    arm_token = getpass.getpass("Enter ARM token (hidden): ").strip()

# Set the datastream and start/enddates
aps_datastream = 'dstaosapsM1.b1'
met_datastream = "dstmetM1.b1"

startdate = '2026-08-04'
enddate = '2026-08-04'

# 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_aps = act.discovery.download_arm_data(arm_username, arm_token, aps_datastream, startdate, enddate)
result_met = act.discovery.download_arm_data(arm_username, arm_token, met_datastream, startdate, enddate)

# Let's read in the data using ACT and check out the data
ds_aps = act.io.read_arm_netcdf(result_aps)
ds_met = act.io.read_arm_netcdf(result_met)

ds_aps

ds_met

# Since ACT is built upon xarray, we can use xarray's built in plotting to quickly visualize the data
ds_aps['total_N_conc'].plot()

# We can also use xarray's built in slicing to quickly subset the data.  
# Let's look at some suspect data points in the temperature data
ds_aps['total_N_conc'].sel(time=slice('2026-08-04T05:45', '2026-08-04T08:30')).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_aps.clean.cleanup()

# Query the ARM DQR Webservice
variable = 'total_N_conc'
ds_aps = act.qc.add_dqr_to_qc(ds_aps, variable=variable)

ds_aps

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

ds_met.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_met, figsize=(15, 10), subplot_shape=(2,))

# Plot up the variable in the first plot
display.plot("wspd_vec_mean", 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("wspd_vec_mean", subplot_index=(1,))
plt.show()

# And plot again!
# Create a plotting display object with 2 plots
# Note we have to create a dictionary of datasets to pass in
display = act.plotting.TimeSeriesDisplay(
    {'APS': ds_aps, 'MET': ds_met}, figsize=(15, 10), subplot_shape=(2,)
)

# LogNorm requires strictly positive values, so mask out the zero/missing
# bins before plotting rather than letting them distort the color scale
ds_aps['dN_dlogDp'] = ds_aps['dN_dlogDp'].where(ds_aps['dN_dlogDp'] > 0)

# Plot up the APS number size distribution (dN_dlogDp) as a 2D PSD.
# Need to specify the dsname so it knows which dataset to use for this
# data. This is helpful when datasets have similar variable names
display.plot(
    'dN_dlogDp',
    dsname='APS',
    subplot_index=(0,),
    cmap='viridis',
    norm=colors.LogNorm(vmin=1, vmax=float(ds_aps['dN_dlogDp'].max())),
    cbar_label=r'dN/dlog$D_p$ (cm$^{-3}$)',
    ylabel='Aerodynamic Diameter (µm)',
)
# Diameter bins are log-spaced, so a log y-axis keeps fine- and
# coarse-mode particles both visible
display.axes[0].set_yscale('log')

# Plot up the MET wind speed for comparison
display.plot('wspd_vec_mean', dsname='MET', subplot_index=(1,))

# Day/night background only makes sense on the 1D MET line plot -
# on the 2D PSD panel above it would be completely hidden behind the mesh
display.day_night_background(dsname='MET', subplot_index=(1,))

plt.show()

# ACT's accumulate_precip turns an instantaneous rain-gauge variable into a
# running total, appended to the dataset as "<variable>_accumulated"
ds_met = act.utils.data_utils.accumulate_precip(ds_met, 'tbrg_precip_total')

# Plot the instantaneous rate on the left axis and the running total on a
# twin axis so both are visible on the same time base
display = act.plotting.TimeSeriesDisplay(ds_met, figsize=(15, 5))
display.plot('tbrg_precip_total', label='TBRG Precip Rate')

ax2 = display.axes[0].twinx()
ax2.plot(ds_met['time'], ds_met['tbrg_precip_total_accumulated'], color='tab:red')
ax2.set_ylabel('Accumulated Precipitation (mm)')

display.day_night_background()
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_met)

# 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('temp_mean', hist_kwargs={'range': [30, 45]})
plt.show()

# Let's directly compare the two datasets using a scatter plot. 
# First we need to clean up the datasets and average to the same time grid
# Once again, we take advantage of the xarray data structures and use xarray's resample function
ds_aps = ds_aps.resample(time='1Min').nearest()
ds_met = ds_met.resample(time='1Min').nearest()

# Combine the datasets together using xarray's merge function.
ds_combined = xr.merge([ds_aps, ds_met], compat='override')

# 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_combined)
display.set_xrng([25, 35])
display.set_yrng([25, 35])

# And then we can plot the data!
display.plot_scatter('box_temperature', 
                     'temp_mean', 
                     m_field='wspd_vec_mean',
                     marker='x',
                     cbar_label=r'Wind Speed ($^\circ$)',
)

# Display the 1:1 ratio line
display.set_ratio_line()

plt.show()

# A heatmap bins both variables and colors by point density - useful once a
# scatter plot like the one above gets crowded with overlapping points
display = act.plotting.DistributionDisplay(ds_combined)

display.plot_heatmap(
    'box_temperature',
    'temp_mean',
    x_bins=20,
    y_bins=20,
    cmap='viridis',
)

plt.show()

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

# Plot out the data rose using the WindRose display object
display = act.plotting.WindRoseDisplay(ds_combined)

# Let's try a different type of data rose type that is using contours amd the polar y axis is wind speed
# depending on wind direction and speed
display.plot_data('wdir_vec_mean', 'wspd_vec_mean', 'total_N_conc', num_dirs=20, plot_type='contour', contour_type='mean')
plt.show()