import os
from WhiteboxTools_win_amd64.WBT.whitebox_tools import WhiteboxTools
import wbd_c

whitebox_tools = WhiteboxTools()
ESRI_PNTR = True

# TODO remove 1st order streams
# needs to be done before burning them in
# out_streamorder = os.path.join(out_folder, wbd_c.STRORDER  + '.tif')
# whitebox_tools.strahler_stream_order(fdr, out_streamras, out_streamorder)
# whitebox_tools.set_nodata_value(out_streamorder, out_streamorder, back_value=1.0)


def fix_dem(dem: str, streams: str, out_folder: str, walls: str = None) -> str:
    """Adds walls to dem (optional), burns in streams, fills sinks.  Filled DEM is used for additional steps.

    Args:
        dem (str): Filepath to dem raster.
        streams (str): Filepath to streams shapefile.
        out_folder (str): Output folder.
        walls (str, optional): Filepath to walls shapefile. Defaults to None.

    Returns:
        str: Filepath to filled dem.
    """
    # TODO breach single cells first?
    # TODO change to breach depressions?
    out_walls = os.path.join(out_folder, wbd_c.WALLEDDEM + ".tif")
    out_fill = os.path.join(out_folder, wbd_c.FIL + ".tif")
    if walls:
        whitebox_tools.raise_walls(walls, dem, out_walls)
        whitebox_tools.fill_depressions(out_walls, out_fill)
    else:
        whitebox_tools.fill_depressions(dem, out_fill)
    return out_fill


def stream_ras(dem: str, out_folder: str, threshold_area: int) -> str:
    """Creates stream raster using threshold area from dem.  Stream raster is only used to create catchments.

    Args:
        dem (str): File path to filled dem.
        out_folder (str): Output folder.
        threshold_area (int): Number of upstream cells to initiate a stream.

    Returns:
        str: File path to stream raster.
    """
    out_fac = os.path.join(out_folder, wbd_c.FAC + ".tif")
    out_streamras = os.path.join(out_folder, wbd_c.STR + ".tif")
    whitebox_tools.d8_flow_accumulation(dem, out_fac)
    whitebox_tools.extract_streams(out_fac, out_streamras, threshold_area, zero_background=True)
    return out_streamras


def catchments(fdr: str, streams: str, out_folder: str):
    """Creates catchment raster.  One catchment is created for each reach in the stream raster.

    Args:
        fdr (str): Filepath to flow direction raster.
        streams (str): Filepath to streams raster (value of 1 = stream cell value of 0 = non-stream cell).
        out_folder (str): Output folder.
    """
    out_catchmentsras = os.path.join(out_folder, wbd_c.CAT + ".tif")
    whitebox_tools.subbasins(fdr, streams, out_catchmentsras, esri_pntr=ESRI_PNTR)
    # TODO clip catchment (no data area is messed up >>> problems) not working in whitebox
    # TODO Convert to lines
    # converting to polygons >>> memory error
    # raster to vector lines either >>> strange results or doesn't draw right in pro
    # out_catchmentshape = os.path.join(out_folder, wbd_c.CATCHMENT + '.shp')
    # whitebox_tools.raster_to_vector_polygons(out_catchmentsras, out_catchmentshape)                   # TODO memory error


def execute(dem: str, streams: str, out_folder: str, threshold_area: str, walls: str = None):
    """Executes the above functions to perform terrain preprocessing tools using Whitebox Tools.

    Args:
        dem (str): Filepath to dem raster.
        streams (str): Filepath to streams shapefile.
        out_folder (str): Output folder.
        threshold_area (str): Number of upstream cells to initiate a stream.
        walls (str, optional): Filepath to walls shapefile. Defaults to None.
    """
    out_fdr = os.path.join(out_folder, wbd_c.FDR + ".tif")
    filled_dem = fix_dem(dem, streams, out_folder, walls)
    whitebox_tools.d8_pointer(filled_dem, out_fdr, esri_pntr=ESRI_PNTR)
    streamras = stream_ras(filled_dem, out_folder, threshold_area)
    catchments(out_fdr, streamras, out_folder)


if __name__ == "__main__":
    dem = r"D:\OneDrive - DOI\WBD_Collaboration\AK\Work\IfSAR_Updates\download_0528\19060501_prep\DEM_19060501\dem_19060501.tif"
    stream_shapefile = r"D:\OneDrive - DOI\WBD_Collaboration\AK\Work\IfSAR_Updates\download_0528\19060501_prep\DEM_19060501\NHDFlowline_19060501.shp"
    out_folder = os.path.dirname(dem)
    threshhold_area = 1000
    execute(dem, stream_shapefile, out_folder, threshhold_area)

    # whitebox_tools.subbasins(
    #     r"D:\OneDrive - DOI\WBD_Collaboration\AK\Work\IfSAR_Updates\download_0528\19060501_prep\DEM_19060501\Fdr.tif",
    #     r"D:\OneDrive - DOI\WBD_Collaboration\AK\Work\IfSAR_Updates\download_0528\19060501_prep\DEM_19060501\Str.tif",
    #     r"D:\OneDrive - DOI\WBD_Collaboration\AK\Work\IfSAR_Updates\download_0528\19060501_prep\DEM_19060501\Cat.tif",
    #     False
    # )
