# AUTOGENERATED! DO NOT EDIT! File to edit: notebooks/09_raster_process.ipynb (unless otherwise specified).

__all__ = ["query_window_by_polygon", "query_window_by_gdf"]


# Cell
from pathlib import PosixPath
from typing import Union

# hide
# export
import rasterio as rio
import rasterio.mask
from rasterio.io import DatasetReader
from shapely.geometry import Polygon

# Cell


def query_window_by_polygon(
    input_raster: Union[str, DatasetReader, PosixPath],
    output_path: str,
    geometry: Polygon,
    mask=False,
) -> None:
    """Query a cloud-optimized geotiff based on polygon and save to disk
    Returns a subset of a rasterio dataset based on a window
    defined by (left,right,bottom,top) coordinates. This function
    assumes that the CRS of the bounds is based on the input dataset

    Args:
        input_raster (str, DatasetReader, PosixPath): Input rasterio dataset, which can be specified
            by a local filename or an open rasterio dataset
        output_path (str): Path to save output file
        geometry (polygon): polygon to crop
    Returns:
        None

    Cloud optimized geotiffs are also supported, which allows you
    to query a geotiff hosted in an HTTP server.

    """

    # If input_raster is specified as PosixPath, convert to string explicitly
    if isinstance(input_raster, PosixPath):
        input_raster = str(input_raster)

    # # Open new dataset if input is a path to a file
    # if isinstance(input_raster,str):
    #     input_dset = rio.open(input_raster)
    # else:
    #     input_dset = input_raster

    # get the window_bounds for the geometry
    window_bounds = geometry.bounds

    # Check that window_bounds has the correct number of elements
    assert (
        len(window_bounds) == 4
    ), "window_bounds must have exactly four items (left,bottom,right,top)"

    # Unroll window_bounds
    left, bottom, right, top = window_bounds

    with rio.open(input_raster) as input_dst:

        # Get profile of input_dset
        input_profile = input_dst.profile

        # Specify window and query subset
        window = rio.windows.from_bounds(left, bottom, right, top, input_dst.transform)
        subset = input_dst.read(window=window)

        # Get the shape of the output subset
        number_of_bands, height, width = subset.shape

        # Get the transformation of the subset based on the window
        win_transform = input_dst.window_transform(window)

        # Update metadata for output
        output_profile = input_profile.copy()
        update_params = {
            "count": number_of_bands,
            "height": height,
            "width": width,
            "transform": win_transform,
        }
        output_profile.update(update_params)

        # Write image to output_file
        with rio.open(output_path, "w", **output_profile) as output_dst:
            output_dst.write(subset)
            output_dst.colorinterp = input_dst.colorinterp

    if mask:
        with rio.open(output_path) as dst:
            masked_image, masked_transform = rio.mask.mask(dst, [geometry], crop=True)
        with rio.open(output_path, "w", **output_profile) as output_dst:
            update_params = {
                "height": masked_image.shape[1],
                "width": masked_image.shape[2],
                "transform": masked_transform,
            }
            output_profile.update(update_params)
            output_dst.write(masked_image)

    return


# Cell


def query_window_by_gdf(
    input_raster: Union[str, DatasetReader, PosixPath],
    output_folder: str,
    gdf,
    name_col=None,
    mask=False,
) -> None:
    """Query a subset of a geotiff based on a window and save to disk
    Returns a subset of a rasterio dataset based on a window
    defined by (left,right,bottom,top) coordinates. This function
    assumes that the CRS of the bounds is based on the input dataset

    Cloud optimized geotiffs are also supported, which allows you
    to query a geotiff hosted in an HTTP server.

    Args:
        input_raster (str, DatasetReader, PosixPath): Input rasterio dataset, which can be specified
            by a local filename or an open rasterio dataset
        output_folder (str): Path to save output file
        gdf,
        name_col (str): column name to base output filepath on. If none, outputs will be named sequentially
            as `output_0.tif`
    Returns:
        None
    """

    gdf = gdf.copy()

    # Check CRS match
    with rio.open(input_raster) as dst:
        assert dst.meta["crs"] == gdf.crs, "input_raster and gdf CRS must match!"

    if name_col is None:
        name_col = "name"
        gdf[name_col] = "output_" + gdf.reset_index().index.astype(str) + ".tif"
    else:
        gdf[name_col] = gdf[name_col] + ".tif"

    for i, row in gdf.iterrows():
        polygon = row.geometry
        output_name = row[name_col]
        output_path = output_folder / output_name
        print(output_path)
        query_window_by_polygon(input_raster, output_path, polygon, mask)
