# AUTOGENERATED! DO NOT EDIT! File to edit: data.ipynb (unless otherwise specified).

__all__ = ['download_era5_land', 'process_raw_era5_land', 'temperature_processor', 'precipitation_processor',
           'load_usgs_obs', 'load_basins', 'load_era5_forcings', 'load_sacsma']

# Cell
import calendar
import cdsapi
import concurrent.futures
import geopandas as gpd
import os
import pandas as pd
import regionmask
import xarray as xr

# Cell
def download_era5_land(variable, area, year, month, data_dir, ext='grib'):
    """
    Download ERA5 Land data from Copernicus using the cdsapi. If
    the file already exists in data_dir it won't be downloaded. Data is downloaded
    in grib format.

    https://cds.climate.copernicus.eu/cdsapp#!/dataset/reanalysis-era5-land?tab=overview
    """

    m_str = f'{month}'.zfill(2)
    filename = f'{data_dir}/{year}{m_str}_{variable}.{ext}'

    # if the file has already been downloaded return early
    if os.path.isfile(filename):
        print(f'{filename} already exists..skipping.')
        return

    days = [str(i+1) for i in range(calendar.monthrange(year, month)[1])]
    times = [f'{str(i).zfill(2)}:00' for i in range(24)]

    api_options = {
        'variable': variable,
        'area': area,
        'year': f'{year}',
        'month': m_str,
        'day': days,
        'time': times,
        'format': ext,
    }

    client = cdsapi.Client()
    try:
        print(f'Retrieving with the following options and filename: {api_options}, {filename}')
        client.retrieve('reanalysis-era5-land', api_options, filename)
    except Exception as e:
        print(f'{e}')

    return

# Cell
def process_raw_era5_land(ds, dates, variable, processor, output_dir, workers=30):
    """
    Write processed grid to netcdf.

    """

    with concurrent.futures.ThreadPoolExecutor(max_workers=workers) as executor:
        for d in dates:
            nc_file = f'{output_dir}/{d.strftime("%Y%m%d")}_{variable}.nc'

            if os.path.isfile(nc_file):
                continue

            sel_ds = ds.sel(time=d.strftime('%Y-%m-%d'))

            if 'time' not in sel_ds.dims:
                sel_ds = sel_ds.expand_dims('time')
                sel_ds.attrs = {}
            else:
                times = [sel_ds.isel(time=0), sel_ds.isel(time=1)]
                sel_ds = xr.merge(times, compat='no_conflicts').expand_dims('time')

            if 'time' not in sel_ds.dims:
                continue

            sel_ds = processor(sel_ds)

            executor.submit(sel_ds.to_netcdf, nc_file)

# Cell

def temperature_processor(ds):
    """
    Process ERA5-Land temperatures.

    Combine time and step dimensions into a singular timestep dimension. Convert
    Kelvin to Fahrenheit.
    """
    ds = (ds - 273.15) * (9. / 5.) + 32. # kelvin to fahrenheit
    ds = ds.stack(timestep=('time', 'step'))
    ds = ds.assign_coords(timestep=(ds.timestep.time + ds.timestep.step))

    return ds

# Cell
def precipitation_processor(ds):
    """
    Process ERA5-Land precipitation.

    Raw precipitation is a cumulative sum that resets each day. Disaggreate the cumulative
    sum into hourly totals. Combine time and step dimensions into a singular timestep dimension.
    Convert meters to inches.
    """
    ds_shifted = ds.shift(step=1)
    ds_shifted = ds_shifted.fillna(0.0)
    ds = ds - ds_shifted
    ds = ds * 39.2701 # meters to inches
    ds = ds.stack(timestep=('time', 'step'))
    ds = ds.assign_coords(timestep=(ds.timestep.time + ds.timestep.step))

    return ds

# Cell

def load_usgs_obs(f, value_col):
    """Load usgs obs from csv file"""

    df = pd.read_csv(f, parse_dates=['datetime'])
    df['DST'] = df['tz_cd'].replace({'PST': False, 'PDT': True})
    df['dt_local'] = df['datetime'].dt.tz_localize(
        tz='US/Pacific',
        ambiguous=df['DST'].values,
        nonexistent='NaT'
    )
    df['dt_utc'] = df['dt_local'].dt.tz_convert('UTC')
    df = df.dropna()
    df = df[df['dt_utc'].dt.minute == 0]

    df = df.set_index('dt_utc')
    dt_index = pd.date_range(start=df.index[0], end=df.index[-1], freq='1H', tz='UTC')

    df = df.reindex(dt_index)

    df = df[[value_col]]
    df.columns = ['obs']

    return df.interpolate()

# Cell

def load_basins(f, name):
    shapefile = gpd.read_file(f)
    basin_names = shapefile.LID.values
    basin_shapes = shapefile.geometry.values

    return regionmask.Regions(basin_shapes, names=basin_names, abbrevs=basin_names, name=name)

# Cell

def load_era5_forcings(data_dir, file_dict, region_id):
    dfs = []
    for filename, [var, col] in file_dict.items():
        ds = xr.open_dataset(f'{data_dir}/{filename}')
        df = ds.sel(region=region_id).to_dataframe()[[var]].dropna()
        df.columns = [col]
        dfs.append(df)

    df = pd.concat(dfs, axis=1)
    df.index = df.index.tz_localize('UTC')

    return df

# Cell

def load_sacsma(f):
    df = pd.read_csv(f, parse_dates=['datetime'], index_col='datetime')
    df.index = df.index.tz_localize(tz='UTC')

    return df