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

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

# Set the datastream and start/enddates
datastream = 'bnfstampS20.b1'
startdate = '2025-03-31'
enddate = '2025-04-30T23:59:59'

# 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_stamp = act.discovery.download_arm_data(username, token, datastream, startdate, enddate)

# datastream = 'bnfaossmpsM1.b1'
# result_smps = act.discovery.download_arm_data(username, token, datastream, startdate, enddate)

# Let's read in the data using ACT and check out the data
ds_stamp = act.io.read_arm_netcdf(result_stamp)
#ds_smps = act.io.read_arm_netcdf(result_smps)

ds_stamp

#this will plot each of the three soil sensors (east, west and south) at one depth 
ds_stamp.clean.cleanup()

ds1=ds_stamp.sel(depth=5)  #pick depth here

# Query the ARM DQR Webservice
south_temp = 'soil_temperature_south'
west_temp = 'soil_temperature_west'
east_temp = 'soil_temperature_east'

south_wc = 'soil_specific_water_content_south'
west_wc = 'soil_specific_water_content_west'
east_wc = 'soil_specific_water_content_east'

display = act.plotting.TimeSeriesDisplay(ds1, figsize=(15, 10), subplot_shape=(2,))

# Plot up the variable in the first plot
display.plot(south_temp, subplot_index=(0,), label='South Temp')
display.plot(west_temp, subplot_index=(0,), label='West Temp')
display.plot(east_temp, subplot_index=(0,), label='East Temp')
display.day_night_background(subplot_index=(0,))
display.axes[0].legend()

display.plot(south_wc, subplot_index=(1,), label='South WC')
display.plot(west_wc, subplot_index=(1,), label='West WC')
display.plot(east_wc, subplot_index=(1,), label='East WC')
display.day_night_background(subplot_index=(1,))
display.axes[1].legend()

plt.show()

#this will plot one sensor (east, west or south) for all 5 depths for the time range

display = act.plotting.TimeSeriesDisplay({'STAMP': ds_stamp}, figsize=(15, 10), subplot_shape=(1,))
display.plot('soil_temperature_south', dsname='STAMP', subplot_index=(0,), force_line_plot=True, labels = True)