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

# Set your username and token 
username = 'yxie'
token = '8bb157033dfd0a5'

# Set the datastream and start/enddates
datastream = 'gucsebsM1.b1'

# Event #1 at January 2, 2022
startdate1 = '2022-01-02'
enddate1 = '2022-01-07'

# Event #2 at January 25, 2022
startdate2 = '2022-01-25'
enddate2 = '2022-01-30'

# Event #3 at April 3, 2023
startdate3 = '2023-04-03'
enddate3 = '2023-04-08'

# We are looking at 5 days after the event
numdate = 6

# 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
result1 = act.discovery.download_arm_data(username, token, datastream, startdate1, enddate1)
result2 = act.discovery.download_arm_data(username, token, datastream, startdate2, enddate2)
result3 = act.discovery.download_arm_data(username, token, datastream, startdate3, enddate3)

# Let's read in the data using ACT and check out the data
ds_rad1 = act.io.read_arm_netcdf(result1)
ds_rad2 = act.io.read_arm_netcdf(result2)
ds_rad3 = act.io.read_arm_netcdf(result3)

# Quality Check the data
# Remove the bad data
ds_rad1.clean.cleanup()
#ds_rad = act.qc.arm.add_dqr_to_qc(ds_rad)
ds_rad1.qcfilter.datafilter(rm_assessments=['Bad'], del_qc_var=False)

ds_rad2.clean.cleanup()
ds_rad2.qcfilter.datafilter(rm_assessments=['Bad'], del_qc_var=False)

ds_rad3.clean.cleanup()
ds_rad3.qcfilter.datafilter(rm_assessments=['Bad'], del_qc_var=False)

# check the data structure
ds_rad1

# Visualize the SEBS albedo measurement
plt.figure(figsize=(5.5,2), dpi=150)
plt.plot(ds_rad1['time'], ds_rad1['albedo'], '.-')
plt.xticks(fontsize=7)
plt.yticks(np.arange(0,1.1,0.2),fontsize=8)
plt.ylabel('albedo')
plt.grid()
plt.title('Control event')
plt.show()

plt.figure(figsize=(5.5,2), dpi=150)
plt.plot(ds_rad2['time'], ds_rad2['albedo'], '.-')
plt.xticks(fontsize=7)
plt.yticks(np.arange(0,1.1,0.2),fontsize=8)
plt.ylabel('albedo')
plt.grid()
plt.title('Black carbon event')
plt.show()

plt.figure(figsize=(5.5,2), dpi=150)
plt.plot(ds_rad3['time'], ds_rad3['albedo'], '.-')
plt.xticks(fontsize=7)
plt.yticks(np.arange(0,1.1,0.2),fontsize=8)
plt.ylabel('albedo')
plt.grid()
plt.title('Dust event')
plt.show()

daily_swdn1 = ds_rad1["down_short_hemisp"].groupby("time.day").mean()
daily_swup1 = ds_rad1["up_short_hemisp"].groupby("time.day").mean()
daily_alb1 = daily_swup1 / daily_swdn1

daily_swdn2 = ds_rad2["down_short_hemisp"].groupby("time.day").mean()
daily_swup2 = ds_rad2["up_short_hemisp"].groupby("time.day").mean()
daily_alb2 = daily_swup2 / daily_swdn2

daily_swdn3 = ds_rad3["down_short_hemisp"].groupby("time.day").mean()
daily_swup3 = ds_rad3["up_short_hemisp"].groupby("time.day").mean()
daily_alb3 = daily_swup3 / daily_swdn3

print("Control event: ")
print(daily_alb1.values)

print("Black carbon event: ")
print(daily_alb2.values)

print("Dust event: ")
print(daily_alb3.values)

plt.figure(figsize=(5,3), dpi=150)

plt.plot(np.arange(1,7), daily_alb1, '.--', label='control event')
plt.plot(np.arange(1,7), daily_alb2, 'o-', label='black carbon event')
plt.plot(np.arange(1,7), daily_alb3, '^-', label='dust event')

plt.xlabel('days')
plt.ylabel('daily albedo')
plt.legend(fontsize=8)

# calculate the mean by hour of each day
ds_rad["albedo"].groupby("time.hour").mean()

# convert xarray to numpy array
albnew = ds_rad["albedo"].to_numpy()
albnew_mean = np.nanmean(np.reshape(albnew, (48, numdate)), axis=1)

albnew_mean

# Remove the diurnal cycle from the surface albedo
alb = albnew - np.repeat(albnew_mean, numdate)
alb.shape

plt.plot(ds_rad['time'], alb, '.-')