# Uncomment if running on CoLab
##!pip install act-atmos>=2.2.23

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
cf_datastream = 'sgpmetE13.b1'
ef_datastream = "sgpmetE39.b1"

startdate = '2026-04-01'
enddate = '2026-04-30'

# 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_cf_met = act.discovery.download_arm_data(arm_username, arm_token, cf_datastream, startdate, enddate)
result_ef_met = act.discovery.download_arm_data(arm_username, arm_token, ef_datastream, startdate, enddate)

# Let's read in the data using ACT and check out the data
ds_cf_met = act.io.read_arm_netcdf(result_cf_met)
ds_ef_met = act.io.read_arm_netcdf(result_ef_met)

ds_cf_met

# Since ACT is built upon xarray, we can use xarray's built in plotting to quickly visualize the data
ds_cf_met['temp_mean'].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_cf_met['temp_mean'].sel(time=slice('2026-04-15T00:00', '2026-04-17T00:00')).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_cf_met.clean.cleanup()

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

ds_cf_met

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

# 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_cf_met.sel(time=slice('2026-04-15T00:00', '2026-04-17T00:00')), 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()

#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({'CF': ds_cf_met, 'EF': ds_ef_met}, figsize=(15, 10), subplot_shape=(2,))

# Plot up the variable in the first plot
# 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('temp_mean', dsname='CF', subplot_index=(0,))
ax1 = display.axes[0].twinx()
ax1.plot(ds_cf_met.time, ds_cf_met.tbrg_precip_total, color='orange')

# Plot up the EF-MET temperature and precipitation
display.plot('temp_mean', dsname='EF', subplot_index=(1,))
ax2 = display.axes[1].twinx()
ax2.plot(ds_ef_met.time, ds_ef_met.tbrg_precip_total, color='orange')

# Plot up a day/night background
display.day_night_background(dsname='CF', subplot_index=(0,))
display.day_night_background(dsname='EF', subplot_index=(1,))

#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({'CF': ds_cf_met, 'EF': ds_ef_met}, figsize=(15, 10), subplot_shape=(2,))

# Let's accumulate the precipitation using ACT to make this easier to read
ds_cf_met = act.utils.data_utils.accumulate_precip(ds_cf_met, 'tbrg_precip_total')
ds_ef_met = act.utils.data_utils.accumulate_precip(ds_ef_met, 'tbrg_precip_total')

# Let's add in titles and more to the SMPS plot
title = 'MET Temperature and Precipitation at SGP CF (EF13)'
display.plot('temp_mean', dsname='CF', subplot_index=(0,), set_title=title,
             ylabel='Temperature (C)', color='black')
display.set_yrng([0, 35], subplot_index=(0,))
ax1 = display.axes[0].twinx()
ax1.plot(ds_cf_met.time, ds_cf_met.tbrg_precip_total_accumulated, color='red')
ax1.set_ylabel('Precipitation (mm)')
ax1.set_ylim([0, 140.0])

# Plot up the MET btemperature and precipitation
title = 'MET Temperature and Cumulative Precipitation at SGP EF39'
display.plot('temp_mean', dsname='EF', subplot_index=(1,), set_title=title, color='black')
display.set_yrng([0, 35], subplot_index=(1,))
ax2 = display.axes[1].twinx()
ax2.plot(ds_ef_met.time, ds_ef_met.tbrg_precip_total_accumulated, color='red')
ax2.set_ylabel('Precipitation (mm)')
ax2.set_ylim([0, 140.0])

# Plot up a day/night background
display.day_night_background(dsname='CF', subplot_index=(0,))
display.day_night_background(dsname='EF', subplot_index=(1,))

# 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_cf_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': [4, 20]})
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_cf_qc = ds_cf_met.where(ds_cf_met.temp_mean > 0).resample(time='10min').nearest()
ds_ef_qc = ds_ef_met.where(ds_ef_met.temp_mean > 0).resample(time='10min').nearest()

# Note - plenty of ways to do this, but adding a suffix to distinguish between two sites
suffix = "_ef"
exclude = ["base_time", "time_bounds", "time_offset", "lat", "lon", "alt"]
ds_ef_qc = ds_ef_qc.rename_vars({
    var: f"{var}{suffix}"
    for var in ds_ef_qc.data_vars
    if var not in exclude
})

# quick check to make sure we've eliminated the questionable data
ds_cf_qc.temp_mean.plot(label='CF-MET')

# Combine the datasets together using xarray's merge function.
ds_combined = xr.merge([ds_cf_qc, ds_ef_qc], 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([0, 40])
display.set_yrng([0, 40])

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

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

plt.show()

# Plot out the data rose using the WindRose display object
display = act.plotting.WindRoseDisplay(ds_combined)
display.plot_data('wdir_vec_mean_ef', 'wspd_vec_mean_ef', 'tbrg_precip_total_accumulated_ef', 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', 'tbrg_precip_total_accumulated', num_dirs=20, plot_type='contour', contour_type='mean')
plt.show()