import xarray as xr
from glob import glob
import pandas as pd
import os
import numpy as np

# Find all netcdf files
data = sorted(glob("/home/disk/rocinante/DATA/temp/TNC_stormwater/prism/data/4km/2.raw_nc/*.nc"))
odir = "/home/disk/rocinante/DATA/temp/TNC_stormwater/prism/data/4km/"

ds_list = []
ds = xr.Dataset()
print(ds)
# Loop through and process each one
for d in data:
    print('\t', d)
    dx = xr.open_dataset(d)

    ## Crop file
    dx = dx.sel(lat=slice(46.3, 49.4), lon=slice(-125.1, -120.4))
    
    # Add timestep
    dx['time'] = pd.to_datetime(os.path.basename(d), format='PRISM_%Y%m%d.nc')

    # PRISM_201205.nc has an issue where the lat/lon are rounded differently and don't line up
    dx['lat'] = np.round(dx.lat, 6)
    dx['lon'] = np.round(dx.lon, 6)
    
    # Add to list to merge
    ds_list = ds_list + [dx]

# Merge and save data
ds = xr.concat(ds_list, dim='time')
ds = ds.rename({'Band1':'PREC'})
styr = data[0].split('_')[-1].replace('.nc', '')
edyr = data[-1].split('_')[-1].replace('.nc', '')
out = "{}/PRISM_{}-{}.nc".format(odir, styr, edyr)
print(out)
ds.to_netcdf(out)