import os
from pathlib import Path
import glob
import numpy as np
import xarray as xr
import matplotlib.pyplot as plt
from matplotlib.colors import Normalize
import getpass

# Workarounds for OpenMP + Matplotlib in non-interactive runs
os.environ.setdefault("KMP_DUPLICATE_LIB_OK", "TRUE")
os.environ.setdefault("OMP_NUM_THREADS", "1")
os.environ.setdefault("MPLCONFIGDIR", str(Path("/tmp/matplotlib")))
Path(os.environ["MPLCONFIGDIR"]).mkdir(parents=True, exist_ok=True)

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader, TensorDataset
from sklearn.preprocessing import StandardScaler
from sklearn.metrics import mean_absolute_error, mean_squared_error, r2_score

import act

print("✓ All packages imported successfully")

# 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
los_datastream = 'sgpmwrlosC1.b1'
mfrs_datastream = 'sgpmfrsr7nchC1.b1'

startdate = '2025-01-01'
enddate = '2026-01-01'

# 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_los_met = act.discovery.download_arm_data(arm_username, arm_token, los_datastream, startdate, enddate)
result_mfrs_met = act.discovery.download_arm_data(arm_username, arm_token, mfrs_datastream, startdate, enddate)

#load the datafiles into a single nc file. 
ds_los = act.io.read_arm_netcdf(result_los_met)
ds_mfrs = act.io.read_arm_netcdf(result_mfrs_met)

ds_mfrs

ds_los.to_netcdf('radiometer.nc')
ds_mfrs.to_netcdf('mfrs.nc')

ds_mfrs

ds_mfrs[['direct_normal_narrowband_filter6', 'qc_direct_normal_narrowband_filter6']].to_netcdf('mfrs.nc')

qc_vap_idx = np.where(ds_los['qc_vap'].values == 0)
qc_direct_idx = 

idx = np.where(ds_cf_met['qc_vap'].values == 0)
ds_cf_met['vap'][:10000].plot()
ds_cf_met['vap'][idx][:10000].plot()

ds_cf_met['direct_normal_narrowband_filter6'].values.shape