swb2_workflow.py

#%%
import os
import matplotlib.pyplot as plt
import rasterio
import geopandas as gpd
import numpy as np
import numpy.ma as ma
import flopy
import pandas as pd
import datetime
import shutil
import netCDF4 as nc
from scipy.interpolate import griddata
import platform
import time


def quick_fix_tmax():
    files = os.listdir(os.path.join('swb_models', 'wahp', 'climate', 'tmin_v01'))
    files = [f for f in files if '2000_' in f]
    # copy files to new directory
    for f in files:
        src_file = os.path.join('swb_models', 'wahp', 'climate', 'tmin_v01', f)
        dst_file = os.path.join('swb_models', 'wahp', 'climate', 'tmax_v01', f)
        shutil.copy2(src_file, dst_file)

    files = os.listdir(os.path.join('swb_models', 'wahp', 'climate', 'tmax_v01'))
    # change string in filenames
    for f in files:
        new_name = f.replace('tmin', 'tmax')
        src_file = os.path.join('swb_models', 'wahp', 'climate', 'tmax_v01', f)
        dst_file = os.path.join('swb_models', 'wahp', 'climate', 'tmax_v01', new_name)
        os.rename(src_file, dst_file)
        print(f"Renamed {f} to {new_name}")

def update_tmax_nodata_value():
    files = os.listdir(os.path.join('swb_models', 'wahp', 'climate', 'tmax_v01'))
    # files = [f for f in files if '2000_' in f]



    for f in files:
        output_path = os.path.join('swb_models', 'wahp', 'climate', 'tmax_v01', f"{f.split('.asc')[0]}_upd.asc")
        # Open the ASCII raster
        with rasterio.open(os.path.join('swb_models', 'wahp', 'climate', 'tmax_v01',f)) as src:
            data = src.read(1)  # Read the first band
            profile = src.profile

        # Optionally replace np.nan with a new nodata value if saving back to .asc
        # new_nodata_value = 0.0
        data = np.where(data == -9999, new_nodata_value, data)
        data[data>0] = data[data>0] + 10.0

        # Update profile for saving as ASCII
        profile.update(driver='AAIGrid', dtype='float32', nodata=new_nodata_value)

        # Write modified raster
        with rasterio.open(output_path, 'w', **profile) as dst:
            dst.write(data.astype(np.float32), 1)

def update_ppt_nodata_value():
    files = os.listdir(os.path.join('swb_models', 'wahp', 'climate', 'ppt_v01'))
    files = [f for f in files if '2000_' in f]

    for f in files:
        output_path = os.path.join('swb_models', 'wahp', 'climate', 'ppt_v01', f"{f.split('.asc')[0]}_upd.asc")
        # Open the ASCII raster
        with rasterio.open(os.path.join('swb_models', 'wahp', 'climate', 'ppt_v01',f)) as src:
            data = src.read(1)  # Read the first band
            profile = src.profile

        # Optionally replace np.nan with a new nodata value if saving back to .asc
        new_nodata_value = 0.0
        data = np.where(data == -9999, new_nodata_value, data)
        # data[data>0] = data[data>0] + 10.0

        # Update profile for saving as ASCII
        profile.update(driver='AAIGrid', dtype='float32', nodata=new_nodata_value)

        # Write modified raster
        with rasterio.open(output_path, 'w', **profile) as dst:
            dst.write(data.astype(np.float32), 1)

def fill_resolve_raster(grid, raster):
    '''

    Parameters:
    ----------
    raster: DEM

    Returns
    -------
    uses pysheds to clean up raster file to perform flow accumulation function
    '''

    # Fill pits in DEM
    pit_filled_dem = grid.fill_pits(raster, max_iter=10000)

    # Fill depressions in DEM
    flooded_dem = grid.fill_depressions(pit_filled_dem, max_iter=10000)
    # assert not flooded_dem.any()

    # Resolve flats in DEM
    inflated_dem = grid.resolve_flats(flooded_dem, max_iter=10000)
    # assert not inflated_dem.any()

    return inflated_dem

def write_prj_file(ras, dir,):
    # write prj file
    prj = open(os.path.join(dir.split('.asc')[0] + ".prj"), "w")
    # call the function and supply the epsg code
    epsg = ras.crs.to_wkt()
    prj.write(epsg)
    prj.close()

def create_fldir_rasters(rdir, grid, dem, name, ):
    """

    Parameters
    ----------
    dir
    dem
    name

    Returns
    -------
    Takes DEM file, clips to watershed of interest, creates flow acc, soils, landuse rasters for SWB
    """

    inflated_dem = fill_resolve_raster(grid, dem)

    # flow dir
    dirmap = (64, 128, 1, 2, 4, 8, 16, 32)
    fldir = grid.flowdir(inflated_dem, dirmap=dirmap)
    # flow acc
    acc = grid.accumulation(fldir, dirmap=dirmap)

    catch = dem.copy()
    # catch = catch.astype(str)
    # convert catch to boolean
    catch[catch < 100.0] = False
    catch[catch>100.0] = True


    # Crop the dem to the target catchment and regenerate the flow dir
    # grid.clip_to(catch)
    # clipped_catch = grid.view(catch)
    clipped_dem = grid.view(inflated_dem)
    clipped_dem[clipped_dem < 100.0] = 0.

    fldir_clip = grid.flowdir(clipped_dem, dirmap=dirmap)

    temp = grid.accumulation(fldir_clip, dirmap=dirmap)

    # fix incorrect flow direction numbers....
    if ~np.isin(fldir_clip, list(dirmap) + [0]).all():
        "Some flow dir values are not recognized, taking them out ...."
        indices = np.where(~np.isin(fldir_clip, list(dirmap) + [0]))
        fldir_clip[indices] = 0


    # grid.to_ascii(fldir_clip, os.path.join('swb2', name, 'input', f"{name}_flw_dir.asc"))
    grid.to_raster(fldir_clip, os.path.join(rdir, 'clipped', f"flw_dir_{name}.tif"))
    convert_tif_to_asc(os.path.join(rdir, 'clipped', f"flw_dir_{name}.tif"), os.path.join('swb2', name, 'input', f"flw_dir_{name}.asc"))

    # write prj file
    write_prj_file(fldir_clip, os.path.join('swb2', name, 'input', f"flw_dir_{name}.asc"),)

def clip_and_resample_dem():
    import rasterio
    from rasterio.warp import reproject, Resampling

    # Input paths
    reference_raster_path = os.path.join('gis', 'raster', 'awc.asc.tif')
    target_raster_path = os.path.join('gis', 'raster', 'dem_clipped.tif')
    output_raster_path = os.path.join('gis', 'raster', 'dem_aligned.tif')

    # Open reference raster
    with rasterio.open(reference_raster_path) as ref:
        ref_profile = ref.profile
        ref_transform = ref.transform
        ref_shape = (ref.height, ref.width)
        ref_crs = ref.crs

    # Open target raster
    with rasterio.open(target_raster_path) as src:
        target_data = src.read(1)
        target_profile = src.profile

        # Create output array
        aligned_data = np.empty(ref_shape, dtype=target_profile['dtype'])

        # Reproject and resample
        reproject(
            source=target_data,
            destination=aligned_data,
            src_transform=src.transform,
            src_crs=src.crs,
            dst_transform=ref_transform,
            dst_crs=ref_crs,
            dst_resolution=(ref_transform.a, -ref_transform.e),
            resampling=Resampling.nearest  # or bilinear, cubic, etc.
        )
    # Write aligned raster
    with rasterio.open(output_raster_path, 'w', **ref_profile) as dst:
        dst.write(aligned_data, 1)


def plot_spatial_monthly_average(maps, swb_dir, odir, years):
    print('plotting monthly averages')
    for year in years:
        for month in range(1, 13):
            fig, axes = plt.subplots(1, len(maps.keys()), figsize=[14, 6])
            for i, key in enumerate(maps.keys()):
                ax = axes[i]
                arr = np.loadtxt(os.path.join(swb_dir, 'output', 'monthly',
                                              f'swb__{key}_{year}_{str(month).zfill(2)}_MEAN.asc'), skiprows=6)
                arr = arr / 12 / 3.281 * 1000  # convert in to mm
                arr[arr < 0] = np.nan
                arr_map = ax.imshow(arr, cmap='plasma', norm=matplotlib.colors.LogNorm())
                arr_map.set_clim(0.01, 100)
                cb = plt.colorbar(arr_map, shrink=0.7)
                ax.set_title(f'{key.title()} mm')
            fig.suptitle(f'Month {month}')
            fig.savefig(os.path.join(odir, f'spatial_outputs_mean{year}_month{month}.png'))

            plt.close(fig)


def plot_spatial_daily_event(maps, swb_dir, odir, dates):
    import matplotlib
    print('plotting rain event')
    for date in dates:
        fig, axes = plt.subplots(1, len(maps.keys()), figsize=[14, 6])
        m, d, y = date.month, date.day, date.year
        for i, key in enumerate(maps.keys()):
            ax = axes[i]
            arr = np.loadtxt(os.path.join(swb_dir, 'output', 'daily',
                                          f'swb__{key}_{y}_{str(m).zfill(2)}_{str(d).zfill(2)}.asc'), skiprows=6)
            arr = arr / 12 / 3.281 * 1000  # convert in to mm
            arr[arr < 0] = np.nan
            arr_map = ax.imshow(arr, cmap='plasma', norm=matplotlib.colors.LogNorm())
            arr_map.set_clim(0.1, 1000)
            cb = plt.colorbar(arr_map, shrink=0.7)
            ax.set_title(f'{key.title()} mm')
        fig.suptitle(f'Date {m}-{d}-{y}')
        fig.savefig(os.path.join(odir, f'spatial_outputs_{y}_{m}_{d}.png'))

        plt.close(fig)

def plot_spatial_annual_avg_swb2(keys, dict, o_dir):
    import matplotlib
    fig, axes = plt.subplots(1, len(keys.keys()), figsize=[16, 8])

    for i, key in enumerate(keys):
        ax = axes[i]
        # spatial annual average in mm/d
        arr_year = dict[keys[key]].sum(axis=0) / 1.0 / (365.25*24)

        arr_map = ax.imshow(arr_year, cmap='viridis', norm=matplotlib.colors.LogNorm(vmin=0.0000001, vmax=0.005))
        cb = plt.colorbar(arr_map, shrink=0.7)
        ax.set_title(f'{key} mm/d')
    fig.savefig(os.path.join(o_dir,'spaital_annual_avg.png'))
    plt.close(fig)


def plot_spatial_monthly_avg_swb2(keys, dict, o_dir, dates):

    months = dates.month
    for month in np.unique(months):
        fig, axes = plt.subplots(1, len(keys.keys()), figsize=[16, 8])
        for i, key in enumerate(keys):
            ax = axes[i]
            # spatial annual average in mm/d
            arr= dict[keys[key]]
            arr_m = arr[np.where(months==month)].mean(axis=0)
            arr_m[arr_m == 0] = np.nan
            arr_map = ax.imshow(arr_m, cmap='viridis', )
            cb = plt.colorbar(arr_map, shrink=0.7)
            ax.set_title(f'{key} mm/d')

        fig.suptitle(f'Month {month}')
        fig.savefig(os.path.join(o_dir,f'spaital_month_{month}_avg.png'))
        plt.close(fig)


def create_file_reference(component_name, o_dir):
        '''
        This is a simple convenience function that will form a path and filename to a
        given water budget component
        '''
        # specify the prefix, path to SWB2 output, timeframe, and resolution
        output_path = os.path.join(o_dir, )
        file = [f for f in os.listdir(output_path) if 'swb2_gross_precipitation' in f]
        prefix = 'swb2_'
        start_year = file[0].split('_')[4]
        end_year = file[0].split('_')[6]
        ncol = file[0].split('_')[10].split('.')[0]
        nrow = file[0].split('_')[8]
        return (output_path +'\\'+ prefix + component_name + '__' + start_year + '_to_'
                + end_year + '__' + nrow + '_by_' + ncol + '.nc'), start_year, end_year

def read_nc_swb2_output(wb_components, o_dir):
    ''' Reads in SWB netCDF outputs and returns a dictionary of the components converted to mm'''
    wb={}
    for key in wb_components:
        name, start_year, end_year = create_file_reference(key, o_dir)
        data = nc.Dataset(name)
        data_nc = data.variables[key]
        # convert in/d to ft/d
        data_vals = ma.masked_where(np.isnan(data_nc),data_nc) / 12   # Convert inch to ft
        wb[key] = data_vals

    return wb, start_year, end_year


def plot_time_series_swb2(dates, wb, o_dir):
    keys = wb.keys()
    df = pd.DataFrame(columns = keys, index=dates)
    df2 = pd.DataFrame(columns = keys, index=dates)
    fig,ax = plt.subplots(1,1, figsize = [12,8])

    for key in keys:
        arr = wb[key]
        idom = np.where(arr[0] >= 0, 1, 0)
        ts = np.sum(arr, axis=(1, 2)) * 660 * 660 # convert ft/d to ft^3/d
        df.loc[:, key] = ts / 43560  # convert ft^3/d to acre-ft/d
        df2.loc[:, key] = (ts*12) /(np.nansum(idom)*660*660)# total ft3/d to in/day
        ax.plot(dates, ts/43560, label=key)

    plt.legend()
    ax.set_ylabel('Sum acre-ft/d')
    ax.grid()
    fig.savefig(os.path.join(o_dir, 'water_balance_daily.png'))
    plt.close(fig)
    df.to_csv(os.path.join(o_dir, 'water_balance_daily_acreft_day.csv'))
    df2.to_csv(os.path.join(o_dir, 'water_balance_daily_in_day.csv'))
    return df


def preprocess_head_data():
    fname = os.path.join("data","heads","Calibration_Targets_v2.xlsx")
    df = pd.read_excel(fname,sheet_name="All")
    dfinfo = pd.read_excel(fname,sheet_name="Final_data")

    dfhlu = pd.read_excel(fname,sheet_name="hlu_map")
    dfhlu.columns = dfhlu.columns.str.lower()

    df.columns = df.columns.str.lower()
    df=df[['well_id','date','wle_ahd','hlu']]
    df.rename(columns={"wle_ahd":"head","well_id":"bore"},inplace=True)
    df.bore = df.bore.str.lower()
    df.date = pd.to_datetime(df.date, format="%yyyy-%mm-%dd %hh:%MM:%ss")

    dfinfo.columns = dfinfo.columns.str.lower()
    dfinfo.bore = dfinfo.bore.str.lower()
    dfinfo = dfinfo[['bore','x','y']]
    dfinfo.drop_duplicates(subset="bore",inplace=True)

    # map bore locations into df
    df2 = df.merge(dfinfo,on="bore",how="left")
    df2 = df2.loc[df2.hlu!="don't know"]
    df2.dropna(inplace=True)

    df2["zone"] = df2.hlu.apply(lambda x: dfhlu.loc[dfhlu.hlu==x]["hlu zone"].tolist())

    df2.to_csv(os.path.join("data","heads","heads.csv"),index=False)

    return

def plot_water_balance_location(keys_spatial, wb, o_dir, dates, boreholes):
    import geopandas as gpd
    keys = wb.keys()
    keys = ['gross_precipitation', 'net_infiltration', 'runoff', 'actual_et', 'rejected_net_infiltration', 'soil_storage']

    grid = gpd.read_file(os.path.join('data','shp', 'swb_grid.shp'))

    nrows = wb['gross_precipitation'].shape[1]
    ncols = wb['gross_precipitation'].shape[2]
    rows = np.repeat(np.arange(nrows, 0, -1), ncols)
    cols = cols = np.tile(np.arange(1, ncols+1), nrows)
    grid['row'] = rows
    grid['col'] = cols
    wells = gpd.read_file(os.path.join('data', 'shp', 'Bore_Master_Nov_2023_pr.shp'))

    well_grid = gpd.sjoin(wells, grid, how='inner')

    for borehole in boreholes:
        fig, ax = plt.subplots(1, 1, figsize=[16, 12])
        idx = well_grid.loc[well_grid['Bore'] == borehole, ['row', 'col']].values[0]-1
        for key in keys:
            arr = wb[key][:,idx[0],idx[1]]
            if key == 'net_infiltration':
                print(f'Borehole {borehole}: net infil {np.round(arr.sum()/2,2)} mm')
                fig.suptitle(f'Borehole {borehole} Water Balance\nNet Infil {np.round(arr.sum()/2,2)} mm')
            ax.plot(dates, arr, label=key)

        plt.legend()
        ax.set_ylabel('mm/d')
        ax.grid()
        fig.savefig(os.path.join(o_dir, f'water_balance_daily_{borehole}.png'))
        plt.close(fig)


def plot_swb2_results(sdir, o_dir,):
    # from osgeo import gdal


    # make figures directory
    o_dir_fig = os.path.join(o_dir, 'figures')
    if not os.path.exists(o_dir_fig):
            os.mkdir(o_dir_fig)

    wb_components = ['gross_precipitation', 'rainfall', 'net_infiltration', 'runoff', 'actual_et',
                      'interception',  'soil_storage', 'snowfall',
                     'interception', 'snow_storage', 'snowmelt', ]
                     # 'snow_storage','snowmelt',]

    wb_dict, start_year, end_year = read_nc_swb2_output(wb_components, o_dir)
    dates = pd.date_range(start_year, end_year)

    df = plot_time_series_swb2(dates, wb_dict, o_dir_fig)

    keys_spatial = {'Recharge': 'net_infiltration', 'Interception': 'interception',
                    'Actual ET': 'actual_et', 'Runoff': 'runoff'}
    plot_spatial_annual_avg_swb2(keys_spatial, wb_dict, o_dir_fig)
    plot_spatial_monthly_avg_swb2(keys_spatial, wb_dict, o_dir_fig, dates)

    boreholes = ['11','23551','23934','DJ2','JE1','JE2A','JE2B','JE3A','JE3B','OB203C','OB204C','OB206C',
                 'OB25','OB41','OB42','OB42A','OBN227C','OBN228C','RP1N1','RP1N2','WSMB12','WSMB13','BH06','BH07']

    # plot_water_balance_location(keys_spatial, wb_dict, o_dir_fig, dates, boreholes)


def plot_annual_precip(sdir, o_dir,):
    # from osgeo import gdal


    # make figures directory
    o_dir_fig = os.path.join(o_dir, 'figures')
    if not os.path.exists(o_dir_fig):
            os.mkdir(o_dir_fig)

    wb_components = ['gross_precipitation', ]
                     # 'snow_storage','snowmelt',]

    wb_dict, start_year, end_year = read_nc_swb2_output(wb_components, o_dir)
    dates = pd.date_range(start_year, end_year)

    df = plot_time_series_swb2(dates, wb_dict, o_dir_fig)

    keys_spatial = {'Recharge': 'net_infiltration', 'Interception': 'interception',
                    'Actual ET': 'actual_et', 'Runoff': 'runoff'}
    plot_spatial_annual_avg_swb2(keys_spatial, wb_dict, o_dir_fig)
    plot_spatial_monthly_avg_swb2(keys_spatial, wb_dict, o_dir_fig, dates)

    boreholes = ['11','23551','23934','DJ2','JE1','JE2A','JE2B','JE3A','JE3B','OB203C','OB204C','OB206C',
                 'OB25','OB41','OB42','OB42A','OBN227C','OBN228C','RP1N1','RP1N2','WSMB12','WSMB13','BH06','BH07']

    # plot_water_balance_location(keys_spatial, wb_dict, o_dir_fig, dates, boreholes)



def plot_swb_water_balance(name):
    import plotly.express as px
    import plotly.io as pio
    pio.renderers.default = 'browser'
    pd.options.plotting.backend = "plotly"
    import plotly.graph_objects as go
    import xarray as xr

    wb_components = ['gross_precipitation', 'rainfall', 'net_infiltration', 'runoff',  'actual_et',
                      'interception', 'rejected_net_infiltration', 'soil_storage', 'snowfall',
                     'interception', 'snow_storage', 'snowmelt', ]

    wb_dict, start_year, end_year = read_nc_swb2_output(wb_components, os.path.join('swb2', name, 'output'))

    df_swb = pd.read_csv(os.path.join( 'swb2', name, 'output',
                                    'figures', "water_balance_daily_acreft_day.csv") , index_col=0)

    colors = ['blue', 'green', 'red', 'orange', 'purple', 'brown', 'pink', 'gray', 'darkred', 'cyan',
              'magenta', 'gold', 'black', 'darkorange', 'darkgreen', 'darkblue', 'gold']

    # make plotly figure
    fig = go.Figure()

    for i, col in enumerate(df_swb.columns):

        # plot magela creek
        fig.add_trace(go.Scatter(
            x=df_swb.index,
            y=df_swb[col]/ 43560,
            mode='lines',
            line=dict(color=colors[i], width=1),
            name=col))

    # ylabel
    fig.update_layout(
        title='Daily Water Balance',
        xaxis_title='Date',
        yaxis_title='Water Balance (acre-ft/d)',
        legend_title='Components',
        template='plotly_white',
        height=600,
        width=1000,
    )
    fig.show()
    fig.write_html(os.path.join('swb2', name, 'output', 'figures', f"daily_water_balance.html"))

 # def create_bounding_shapefile():
 #     import rasterio
 #     from shapely.geometry import box, mapping
 #     import geopandas as gpd
 #
 #     # Path to your raster file
 #     raster_path = os.path.join("swb_models", "wahp", "input",  "awc_convert_inft.asc")
 #
 #     # Open the raster
 #     with rasterio.open(raster_path) as src:
 #         bounds = src.bounds  # (left, bottom, right, top)
 #         crs = src.crs  # Coordinate Reference System
 #
 #     # Create a shapely box from bounds
 #     extent_geom = box(bounds.left, bounds.bottom, bounds.right, bounds.top)
 #
 #     # Convert to GeoDataFrame
 #     gdf = gpd.GeoDataFrame([{"geometry": extent_geom}], crs=crs)
 #
 #     # Save to shapefile (or GeoJSON)
 #     gdf.to_file(os.path.join("gis", "shp", "raster_extent_outline.shp"))

def clip_ras_to_watershed(ras, shp, ras_dir):
    with rasterio.open(ras) as src:
        # Reproject shapefile to match DEM
        shp = shp.to_crs(src.crs)

        # Convert to GeoJSON geometry
        geoms = [json.loads(shp.geometry.to_json())["features"][i]["geometry"] for i in range(len(shp))]

        # Clip DEM using rasterio.mask
        dem_data, dem_transform = mask(src, geoms, crop=True)
        dem_meta = src.meta.copy()

    # Update metadata for the new raster
    dem_meta.update({
        "driver": "GTiff",
        "height": dem_data.shape[1],
        "width": dem_data.shape[2],
        "transform": dem_transform
    })

    # Save the clipped raster to a new file
    if 'dem' in ras:
        filename = f'dem_{name}_clipped.tif'
    elif 'lu' in ras:
        filename  = f'lu_{name}_clipped.tif'
    elif 'soils' in ras:
        filename  = f'soils_{name}_clipped.tif'
    elif 'awc' in ras:
        filename = f'awc_{name}_clipped.tif'
    dem_path = os.path.join('gis', 'raster','clipped', filename)
    with rasterio.open(dem_path, "w", **dem_meta) as dest:
        dest.write(dem_data)

    print(f"Clipped raster written to {dem_path}")


def convert_tif_to_asc(input_tif, output_asc):
    """
    Convert a GeoTIFF file to an ASCII grid file.

    Parameters:
    input_tif (str): Path to the input GeoTIFF file.
    output_asc (str): Path to the output ASCII grid file.
    """
    with rasterio.open(input_tif) as src:
        data = src.read(1)  # Read the first band
        transform = src.transform

        # Extract ASCII header information
        ncols = src.width
        nrows = src.height
        xllcorner = transform.c
        yllcorner = transform.f + (nrows * transform.e)  # transform.e is typically negative
        cellsize = 660 # transform.a
        nodata_value = src.nodata if src.nodata is not None else -9999

        # Replace NaNs or masked values
        data = np.where(np.isnan(data), nodata_value, data)

        with open(output_asc, 'w') as f:
            f.write(f"NCOLS {ncols}\n")
            f.write(f"NROWS {nrows}\n")
            f.write(f"XLLCORNER {xllcorner}\n")
            f.write(f"YLLCORNER {yllcorner}\n")
            f.write(f"CELLSIZE {cellsize}\n")
            f.write(f"NODATA_VALUE {nodata_value}\n")

            for row in data:
                f.write(' '.join(map(str, row)) + '\n')
        if 'climate' not in output_asc:
            write_prj_file(src, output_asc,)

    print(f"Converted {input_tif} to {output_asc}")



def resample_raster(d, name):
    """
    Resample a raster to a specified resolution using rasterio.
    Parameters:
    d (dict): Dictionary containing the raster file path and target resolution.
    """

    for key in d.keys():
        input_raster = d[key][0]
        target_resolution =d[key][1]

        with rasterio.open(input_raster) as src:
            # Calculate new dimensions
            new_width = int(src.width * (src.res[0] / target_resolution[0]))
            new_height = int(src.height * (src.res[1] / target_resolution[1]))

            # Resample the raster
            data = src.read(
                out_shape=(src.count, new_height, new_width),
                resampling=Resampling.bilinear
            )

            # Update metadata
            profile = src.profile.copy()
            profile.update({
                'width': new_width,
                'height': new_height,
                'transform': src.transform * src.transform.scale(
                    (src.width / new_width),
                    (src.height / new_height)
                )
            })

            # Write the resampled raster to a new file
            output_raster = os.path.join('gis', 'raster', 'resampled', f"{key}_{name}_resampled.tif")
            with rasterio.open(output_raster, 'w', **profile) as dst:
                dst.write(data)

        print(f"Resampled raster written to {output_raster}")


def clip_project_daily_prism():
    '''
    1. unzip daily prism data
    2. clip it to ND and part of Minnesota
    3. project to SP CS

    '''
    import geopandas as gpd
    import zipfile
    import json
    from rasterio.mask import mask
    import json
    from rasterio.enums import Resampling
    from rasterio.warp import calculate_default_transform, reproject
    from rasterio.transform import Affine
    import tempfile

    shp = gpd.read_file(os.path.join('..', 'gis', 'shp', 'project_area.shp'))
    target_crs = shp.crs.to_string()

    prism_dir = os.path.join(r"S:\IND\NDDWR.C001.ASR-GW-MOD\4000.GIS\zips\PRISM_daily") #os.path.join('..', 'data', 'prism_daily',)

    types = ['ppt', 'tmax', 'tmin']
    for t in types:
        years = os.listdir(os.path.join(prism_dir, t))
        years = years[21:]
        for year in years:
            print(f'clipping {t} data for year {year}')
            files = os.listdir(os.path.join(prism_dir, t, year))
            for day in files:
                zip_path = os.path.join(prism_dir, t, year, day)

                # Create a temporary folder
                temp_dir = tempfile.mkdtemp()

                try:
                    # Unzip contents into the temp folder
                    with zipfile.ZipFile(zip_path, 'r') as zip_ref:
                        zip_ref.extractall(temp_dir)

                    # Find the raster file
                    raster_path = None
                    for root, _, files in os.walk(temp_dir):
                        for file in files:
                            if file.endswith('.bil'):
                                raster_path = os.path.join(root, file)
                                break

                    if raster_path is None:
                        raise FileNotFoundError("No .bil file found in ZIP archive.")

                    # Load the raster and clip to ND and part of Minnesota
                    with rasterio.open(raster_path) as src:
                        # Reproject shapefile to match DEM
                        shp = shp.to_crs(src.crs)

                        # Convert to GeoJSON geometry
                        geoms = [json.loads(shp.geometry.to_json())["features"][i]["geometry"] for i in range(len(shp))]

                        # Clip DEM using rasterio.mask
                        dem_data, dem_transform = mask(src, geoms, crop=True)
                        dem_meta = src.meta.copy()

                    # Update metadata for the new raster
                    dem_meta.update({
                        "driver": "GTiff",
                        "height": dem_data.shape[1],
                        "width": dem_data.shape[2],
                        "transform": dem_transform
                    })

                    filename = f'{day.split('.zip')[0]}_clip.tif'
                    dem_path = os.path.join('..','data', 'prism_daily_clipped', t,  filename )
                    with rasterio.open(dem_path, "w", **dem_meta) as dest:
                        dest.write(dem_data)

                    # then project raster to SP cs
                    # Open clipped raster
                    with rasterio.open(dem_path) as src:
                        # Compute transform and shape for the target CRS
                        transform, width, height = calculate_default_transform(
                            src.crs, target_crs, src.width, src.height, *src.bounds
                        )

                        # Update metadata for output raster
                        kwargs = src.meta.copy()
                        kwargs.update({
                            'crs': target_crs,
                            'transform': transform,
                            'width': width,
                            'height': height
                        })

                        # Reproject and write new raster
                        project_path = os.path.join('..', 'data', 'prism_daily_project', t, filename)
                        with rasterio.open(project_path, 'w', **kwargs) as dst:
                            for i in range(1, src.count + 1):
                                reproject(
                                    source=rasterio.band(src, i),
                                    destination=rasterio.band(dst, i),
                                    src_transform=src.transform,
                                    src_crs=src.crs,
                                    dst_transform=transform,
                                    dst_crs=target_crs,
                                    resampling=Resampling.bilinear  # Choose method based on data
                                )


                    # then resample to 660m resolution
                    # ToDo: don't do this until you clip to SWB domain size
                    # target_resolution = (660, 660)  # Define target resolution
                    # with rasterio.open(project_path) as src:
                    #     # Calculate the new dimensions
                    #     scale_x = src.res[0] / target_resolution[0]
                    #     scale_y = src.res[1] / target_resolution[1]
                    #     new_width = int(src.width * scale_x)
                    #     new_height = int(src.height * scale_y)
                    #
                    #     # Calculate new transform
                    #     transform = src.transform
                    #     new_transform = Affine(
                    #         target_resolution[0], transform.b, transform.c,
                    #         transform.d, -target_resolution[0], transform.f
                    #     )
                    #
                    #     # Prepare output metadata
                    #     kwargs = src.meta.copy()
                    #     kwargs.update({
                    #         'height': new_height,
                    #         'width': new_width,
                    #         'transform': new_transform
                    #     })
                    #
                    #     # Resample and write
                    #     resample_path = os.path.join('..', 'data', 'prism_daily_resampled', t, filename)
                    #     with rasterio.open(resample_path, 'w', **kwargs) as dst:
                    #         for i in range(1, src.count + 1):
                    #             data = src.read(
                    #                 i,
                    #                 out_shape=(new_height, new_width),
                    #                 resampling=Resampling.bilinear  # Options: nearest, bilinear, cubic, etc.
                    #             )
                    #             dst.write(data, i)

                except:
                    continue

                finally:
                    # Clean up: delete the temporary unzipped folder
                    shutil.rmtree(temp_dir)



def clip_climate_to_watershed(ras, shp, o_dir):
    with rasterio.open(ras) as src:
        # Reproject shapefile to match DEM
        shp = shp.to_crs(src.crs)

        # Convert to GeoJSON geometry
        geoms = [json.loads(shp.geometry.to_json())["features"][i]["geometry"] for i in range(len(shp))]

        # Clip DEM using rasterio.mask
        dem_data, dem_transform = mask(src, geoms, crop=True)
        dem_meta = src.meta.copy()



    # Update metadata for the new raster
    dem_meta.update({
        "driver": "GTiff",
        "height": dem_data.shape[1],
        "width": dem_data.shape[2],
        "transform": dem_transform
    })

    filename =ras.split('\\')[-1]
    out_path = os.path.join(o_dir, filename)
    with rasterio.open(out_path, "w", **dem_meta) as dest:
        dest.write(dem_data)

    print(f"Clipped raster written to {out_path}")
    return out_path

def resample_climate_to_watershed(ras, o_dir, t):

    target_resolution = (660, 660)  # Define target resolution
    with rasterio.open(ras) as src:
        # Calculate the new dimensions
        scale_x = src.res[0] / target_resolution[0]
        scale_y = src.res[1] / target_resolution[1]
        new_width = int(src.width * scale_x)
        new_height = int(src.height * scale_y)

        # Calculate new transform
        transform = src.transform
        new_transform = Affine(
            target_resolution[0], transform.b, transform.c,
            transform.d, -target_resolution[0], transform.f
        )

        # Prepare output metadata
        kwargs = src.meta.copy()
        kwargs.update({
            'height': new_height,
            'width': new_width,
            'transform': new_transform
        })

        # Resample and write
        filename = ras.split('\\')[-1]
        resample_path = os.path.join(o_dir, filename)
        with rasterio.open(resample_path, 'w', **kwargs) as dst:
            for i in range(1, src.count + 1):
                data = src.read(
                    i,
                    out_shape=(new_height, new_width),
                    resampling=Resampling.bilinear  # Options: nearest, bilinear, cubic, etc.
                )
                dst.write(data, i)

    return resample_path


def match_raster(source_path, reference_path, output_path):
    from rasterio.enums import Resampling

    # Open the reference raster (target dimensions, transform, CRS)
    with rasterio.open(reference_path) as ref:
        ref_meta = ref.meta.copy()
        dst_transform = ref.transform
        dst_crs = ref.crs
        dst_shape = (ref.count, ref.height, ref.width)  # bands, rows, cols

    # Open the source raster (to be aligned)
    with rasterio.open(source_path) as src:
        src_data = src.read(1)  # Read first band
        if 'aws' in source_path:
            # conver integer to float
            src_data[src_data<0] = 0
            src_data = src_data.astype(np.float32)
            src_data = src_data / 150.0 / 10 * 12 #convert aws to awc (1 cm/150 cm, to in/ft)

        if 'ppt' in source_path:
            src_data = src_data * 0.03937008 # convert mm to inches
        if 'tmin' in source_path or 'tmax' in source_path:
            src_data = src_data * 1.8 + 32  # Convert Celsius to Fahrenheit


        src_transform = src.transform
        src_crs = src.crs
        src_nodata = src.nodata if src.nodata is not None else 0
        if src_nodata != 0:
            src_data[src_data == src_nodata] = 0  # Handle nodata values
            src_nodata = 0  # Set nodata to 0 for consistency

        # Prepare empty destination array
        if 'aws' in source_path:
            dest_data = np.full((ref.height, ref.width), src_nodata, dtype=np.float32)
        else:
            dest_data = np.full((ref.height, ref.width), src_nodata, dtype=src.dtypes[0])

        # Reproject and resample
        reproject(
            source=src_data,
            destination=dest_data,
            src_transform=src_transform,
            src_crs=src_crs,
            dst_transform=dst_transform,
            dst_crs=dst_crs,
            resampling=Resampling.nearest,  # or bilinear/cubic as needed
            src_nodata=src_nodata,
            dst_nodata=src_nodata
        )

    # Save the aligned raster
    ref_meta.update({
        "driver": "GTiff",
        "height": ref.height,
        "width": ref.width,
        "transform": dst_transform,
        "crs": dst_crs,
        "nodata": src_nodata
    })

    with rasterio.open(output_path, "w", **ref_meta) as dst:
        dst.write(dest_data, 1)



def create_lookup_table(dir, lookup_table, wshed):
    lookup_table['First_day_of_growing_season'] = 121
    lookup_table['Last_day_of_growing_season'] = 273
    lookup_table.columns = ['LU_code', 'Description', 'ass_impervious', 'CN_1', 'CN_2', 'CN_3', 'CN_4',
                            'max_net_infil_1', 'max_net_infil_2','max_net_infil_3','max_net_infil_4',
                            'Interception_Growing','Interception_Nongrowing','RZ_1','RZ_2','RZ_3','RZ_4','Reference',
                            'First_day_of_growing_season','Last_day_of_growing_season']
    try:
        os.mkdir(os.path.join(dir, 'lookup_table'))
        ldir = os.path.join(dir, 'lookup_table')
    except:
        ldir = os.path.join(dir, 'lookup_table')

    with open(os.path.join(ldir, f'look_up_{wshed}.txt'), 'w') as f:
        # f.write(f'# SWB lookup table created {datetime.datetime.today().strftime("%m/%d/%y - %H/%M")}\n')
        #f.write(f'NUM_LANDUSE_TYPES\t{len(lu_types)}\n')
        #f.write(f'NUM_SOIL_TYPES\t{max(soil_types)}\n')
        f.write('\t'.join(lookup_table.columns) + '\n')
        for line in range(lookup_table.shape[0]):
            f.write('\t'.join(lookup_table.loc[line, :].astype(str)) + '\n')


def write_control_file(name, dir, grid, proj, et, runoff, rch_method, climate_files, rooting_method,
                       initial_soil_moisture, initial_snow_cover, growing_season_start, growing_season_end,
                       soil_moisture, ia_method, start_date, end_date, soil_max_method):
    with open(os.path.join('swb2', f'swb_control_file_{name}.ctl'), 'w') as file:
        file.write(f'# SWB Model created {datetime.datetime.today().strftime("%m/%d/%y - %H/%M")}\n')
        file.write('\n',)
        file.write(f'GRID {grid.xcenters.shape[1]} {grid.xcenters.shape[0]} {grid.bounds[0]} {grid.bounds[2]} {660}\n')
        file.write(f'BASE_PROJECTION_DEFINITION {proj}\n')
        file.write('GRID_LENGTH_UNITS FEET\n')
        file.write('SUPPRESS_SCREEN_OUTPUT\n')
        file.write('\n')

        file.write(f'SOIL_MOISTURE_METHOD {soil_moisture}\n')
        file.write(f'INITIAL_PERCENT_SOIL_MOISTURE CONSTANT {initial_soil_moisture}\n')
        file.write(f'INITIAL_SNOW_COVER CONSTANT {initial_snow_cover}\n')
        file.write(f'RUNOFF_METHOD {runoff}\n')
        file.write(f'EVAPOTRANSPIRATION_METHOD {et}\n')
        file.write(f'ROOTING_DEPTH_METHOD {rooting_method}\n')
        file.write(f'DIRECT_RECHARGE_METHOD {rch_method}\n')
        file.write(f'SOIL_STORAGE_MAX_METHOD {soil_max_method}\n')
        file.write(f'AVAILABLE_WATER_CONTENT_METHOD GRIDDED\n')
        file.write('\n')
        file.write('PRECIPITATION_METHOD GRIDDED\n')
        file.write('TEMPERATURE_METHOD GRIDDED\n')
        file.write('\n')
        file.write(f'GROWING_SEASON {growing_season_start} {growing_season_end} FALSE\n')
        file.write('\n')
        file.write(f'PRECIPITATION ARC_GRID {climate_files}/ppt/PRISM_ppt_stable_4kmD2_%y%0m%0d_bil_clip.asc\n')
        file.write(f'PRECIPITATION_GRID_PROJECTION_DEFINITION  {proj}\n')
        file.write('\n')
        file.write(f'TMAX ARC_GRID {climate_files}/tmax/PRISM_tmax_stable_4kmD2_%y%0m%0d_bil_clip.asc\n')
        file.write(f'TMAX_GRID_PROJECTION_DEFINITION  {proj}\n')
        file.write('\n')
        file.write(f'TMIN ARC_GRID {climate_files}/tmin/PRISM_tmin_stable_4kmD2_%y%0m%0d_bil_clip.asc\n')
        file.write(f'TMIN_GRID_PROJECTION_DEFINITION  {proj}\n')
        file.write('\n')
        if et == 'MONTHLY_GRID':
            file.write(f'REFERENCE_ET0 ARC_GRID {climate_files}/refET_%b_%y.asc\n')
            file.write(f'REFERENCE_ET0_PROJECTION_DEFINITION  {proj}\n')
            file.write('\n')

        file.write('OUTPUT_GRID_SUFFIX asc\n')
        file.write('\n')
        file.write(f'INTERCEPTION_METHOD {ia_method}\n')
        file.write('\n')
        file.write('INITIAL_FROZEN_GROUND_INDEX CONSTANT 100.0\n')
        file.write('UPPER_LIMIT_CFGI 83.\n')
        file.write('LOWER_LIMIT_CFGI 55.\n')
        file.write('\n')
        file.write(f'FLOW_DIRECTION ARC_GRID {name}/input/flw_dir_{name}.asc\n')
        file.write(f'FLOW_DIRECTION_PROJECTION_DEFINITION  {proj}\n')
        file.write('\n')
        file.write(f'HYDROLOGIC_SOILS_GROUP ARC_GRID {name}/input/soils_{name}.asc\n')
        file.write(f'HYDROLOGIC_SOILS_GROUP_PROJECTION_DEFINITION  {proj}\n')
        file.write('\n')
        file.write(f'LAND_USE ARC_GRID {name}/input/lu_{name}.asc\n')
        file.write(f'LANDUSE_PROJECTION_DEFINITION  {proj}\n')
        file.write('\n')
        file.write(f'LAND_USE_LOOKUP_TABLE {name}/lookup_table/look_up_{name}.txt\n')
        file.write('\n')
        file.write(f'AVAILABLE_WATER_CONTENT ARC_GRID {name}/input/awc_{name}.asc\n')
        file.write(f'AVAILABLE_WATER_CONTENT_PROJECTION_DEFINITION {proj}\n')
        file.write('\n')
        # file.write(f'SOIL_STORAGE_MAX ARC_GRID input_{wshed}/{name}_max_soil_storage_{wshed}.asc\n')
        # file.write(f'SOIL_STORAGE_MAX_PROJECTION_DEFINITION {proj}\n')
        file.write('\n')
        file.write('IGNORE_MISSING_CLIMATE_DATA\n')
        file.write('\n')
        file.write('# OUTPUT OPTIONS\n')
        file.write('\n')
        file.write(f'OUTPUT ENABLE gross_precipitation\n')
        file.write(f'OUTPUT ENABLE rainfall\n')
        file.write(f'OUTPUT ENABLE interception\n')
        file.write(f'OUTPUT ENABLE recharge\n')
        file.write(f'OUTPUT ENABLE runon\n')
        file.write(f'OUTPUT ENABLE runoff\n')
        file.write(f'OUTPUT ENABLE actual_et\n')
        file.write(f'OUTPUT ENABLE net_infiltration\n')
        file.write(f'OUTPUT ENABLE rejected_net_infiltration\n')
        file.write(f'OUTPUT ENABLE soil_storage\n')
        file.write(f'OUTPUT ENABLE snow_storage\n')
        file.write(f'OUTPUT ENABLE snowmelt\n')

        file.write('\n')
        file.write(f'START_DATE {start_date}\n')
        file.write(f'END_DATE {end_date}\n')
        file.write('\n')
        file.write('EOJ\n')

    file.close()




#%%
if __name__ == "__main__":

    root_dir = os.getcwd()

    build_input_rasters = True
    create_climate_files = True
    create_tables       = True
    build_control_file  = True
    run_swb             = True
    postprocess         = True

    name = 'elk'

    swb_dir = os.path.join("swb2", name)
    # create swb folder if it doesn't exist
    if not os.path.exists(swb_dir):
        os.makedirs(swb_dir)

    # Initial DEM file SWB is built off of
    ras_dir = os.path.join("gis", 'raster')
    DEM_file = os.path.join(ras_dir, "EV_DEM_pr.tif")

    raster = flopy.utils.Raster.load(DEM_file)


#%%
    if build_input_rasters:

        print('building input rasters...')
        from pysheds.grid import Grid
        from rasterio.mask import mask
        import json
        from rasterio.warp import reproject, Resampling

        # resample rasters
        raster_dict = {'dem': [DEM_file, (660, 660)],
                       'LU': [os.path.join(ras_dir, 'elk_valley', 'nlcd_lc_2021_ev_pr.tif'), (660, 660)],
                       'soils': [os.path.join(ras_dir, 'elk_valley', 'ssurgo_shg_ev_pr2.tif'), (660, 660)],
                       'awc': [os.path.join(ras_dir, 'elk_valley',  'ssurgo_aws_ev_pr.tif'), (660, 660)]}

        # resample and clip dem,
        resample_raster({'dem': [DEM_file, (660, 660)],}, name)
        # clip DEM tif to watershed extent - ToDo: change to incorporate 'name' i.e., wahp
        gdf = gpd.read_file(os.path.join('gis','shp',f'swb_{name}_domain.shp'))
        clip_ras_to_watershed(os.path.join('gis','raster','resampled',f'dem_{name}_resampled.tif'), gdf, ras_dir)

        # resample and align LU, soils, awc to clipped dem
        ref_path = os.path.join('gis', 'raster', 'clipped', f'dem_{name}_clipped.tif')
        match_raster(raster_dict['LU'][0],ref_path, os.path.join(ras_dir, 'resampled', f'lu_{name}_resampled.tif'))
        match_raster(raster_dict['soils'][0], ref_path, os.path.join(ras_dir, 'resampled', f'soils_{name}_resampled.tif'))
        match_raster(raster_dict['awc'][0], ref_path,os.path.join(ras_dir, 'resampled', f'awc_{name}_resampled.tif'))

        # make swb input directory
        os.makedirs(os.path.join(swb_dir, 'input'), exist_ok=True)

        # clip to watershed and convert to asc
        # Landuse
        lu_file = os.path.join(ras_dir,'resampled',f'lu_{name}_resampled.tif')
        clip_ras_to_watershed(lu_file, gdf, ras_dir)
        out_asc = os.path.join(swb_dir, 'input', f'lu_{name}.asc')
        convert_tif_to_asc(os.path.join(ras_dir,'clipped',f'lu_{name}_clipped.tif'), out_asc)

        # Soils
        soils_file = os.path.join(ras_dir,'resampled',f'soils_{name}_resampled.tif')
        clip_ras_to_watershed(soils_file, gdf, ras_dir)
        out_asc = os.path.join(swb_dir, 'input', f'soils_{name}.asc')
        convert_tif_to_asc(os.path.join(ras_dir,'clipped',f'soils_{name}_clipped.tif'), out_asc)

        # Available Water Content
        awc_file = os.path.join(ras_dir,'resampled',f'awc_{name}_resampled.tif')
        clip_ras_to_watershed(awc_file, gdf, ras_dir)
        out_asc = os.path.join(swb_dir, 'input', f'awc_{name}.asc')
        convert_tif_to_asc(os.path.join(ras_dir,'clipped',f'awc_{name}_clipped.tif'), out_asc)

        # Create flow direction and flow accumulation rasters
        grid = Grid.from_raster(ref_path)
        dem = grid.read_raster(ref_path)
        create_fldir_rasters(ras_dir, grid, dem, name,)

    if create_climate_files:

        print('creating climate files...')
        # Pull clipped and projected prism files - clipped to ND and part of Minnesota
        # clip_project_daily_prism() - clips and projects CONUS prism data (very big)
        climate_dir = os.path.join('data', 'prism_daily_project')
        assert(os.path.exists(climate_dir)), f'Need to make clipped prism data from CONUS prism data'
        types = ['ppt', 'tmax', 'tmin']

        # clip to swb domain - this will be dependent upon model grid, i.e., wahp
        gdf = gpd.read_file(os.path.join('gis','shp',f'swb_{name}_domain.shp'))
        ref_path = os.path.join(ras_dir, 'clipped', f'dem_{name}_clipped.tif')

        for type in types:
            os.makedirs(os.path.join('data', name, 'prism_daily_clip', type), exist_ok=True)
            os.makedirs(os.path.join('data', name, 'prism_daily_resample', type), exist_ok=True)
            os.makedirs(os.path.join(swb_dir, 'climate', type), exist_ok=True)

            clip_path = os.path.join('..', 'data', 'prism_daily_clip', type)
            resamp_path = os.path.join('..', 'data', 'prism_daily_resample', type)

            for file in os.listdir(os.path.join(climate_dir, type)):
                ras_path = os.path.join(climate_dir, type, file)
                # Have to resample, align and then clip
                match_raster(ras_path,ref_path, os.path.join(resamp_path, file))
                clipped_path = clip_climate_to_watershed(os.path.join(resamp_path, file), gdf, clip_path)

                # convert clipped, resampled file from tif to asc and save in climate file
                convert_tif_to_asc(clipped_path, os.path.join(swb_dir, 'climate', type, f'{file.split(".tif")[0]}.asc'))

    if create_tables:
        print('creating lookup tables...')
        # --------- now the look_up tables ---------#
        # Lookup Table Consists of:
        # Total Number of LU Types
        # Total Number of Soil Types
        # Then tab delimited df of
        # 1. LU Code, 2. Description, 3. Assumed Imperviousness, 4. CN/soil type [-], 5. Max Infil/soil type [in/d],
        # 6. Interception growing and dormant season, 7. Root Zone depth/soil type [ft],

        lookup = pd.read_csv(os.path.join('data','swb', f'swb_lookup_table_{name}.csv'))
        create_lookup_table(swb_dir,lookup, name)

    if build_control_file:
        print('building control file...')

        # Create Control File
        # variables - assumptions
        et_method =  'HARGREAVES 46.0 49.0'  # Method, southerly latitude, northerly latidtude 'MONTHLY_GRID'
        runoff = 'CURVE_NUMBER' # 'C-N DOWNHILL'
        climate_files = os.path.join(name, f'climate')
        initial_soil_moisture = 60 # percentage saturation
        initial_snow_cover = 20  # water-equivalent value (inches of water)
        growing_season_start = 135  # May 1
        growing_season_end = 260  # September 30
        interception_method = 'BUCKET'
        soil_moisture = 'THORNTHWAITE-MATHER' # 'T-M'
        rooting_depth = 'STATIC' # 'DYNAMIC', 'NONE'
        direct_recharge_method = 'NONE'
        soil_max_method ='CALCULATED' # 'GRIDDED'

        # climate_file = [f for f in os.listdir(os.path.join(swb_dir, f'climate_{wshed}')) if '.dat' in f]
        start_date = f'01/01/2000'
        end_date = f'12/31/2023'

        # if et_method == 'MONTHLY_GRID':
        #     # calculate monthly sums in inches of SILO potential ET
        #     silo_file = os.path.join('data', 'SILO', 'SILO_station14198_20100101_250401.csv')
        #     et_column = 'et_morton_actual'
        #     create_gridded_monthly_refET(silo_file, et_column, swb_dir, wshed, name, start_date, end_date)

        from pyproj import CRS
        asc = flopy.utils.Raster.load(os.path.join('gis', 'raster', 'clipped', f'dem_{name}_clipped.tif'))

        wkt = asc.crs.to_wkt()
        crs = CRS.from_wkt(wkt)
        proj4 = crs.to_proj4()

        write_control_file(name, swb_dir, asc, proj4,   # grid specifics
                           et_method, runoff, direct_recharge_method, climate_files, rooting_depth,  # variables
                           initial_soil_moisture, initial_snow_cover, growing_season_start,
                           growing_season_end, soil_moisture, interception_method, start_date, end_date,
                           soil_max_method)  # variables

    if run_swb:
        print('running swb2...')
        # prep_deps(swb_dir)

        os.makedirs(os.path.join(swb_dir, 'output'), exist_ok=True)

        os.chdir('swb2')
        os.system(f'swb2.exe swb_control_file_{name}.ctl --output_prefix=swb2_ --output_dir={name}/output')
        os.chdir('..')


    if postprocess:
        print('postprocessing swb2 results...')

        o_dir = os.path.join(swb_dir, 'output')
        if not os.path.exists(o_dir):
            os.mkdir(o_dir)

        plot_swb2_results(swb_dir, o_dir,)
        plot_swb_water_balance(name)