

#example usage https://eroslab.cr.usgs.gov/-/snippets/110

#-----------------------------------------------

import rasterio
import numpy as np

#https://stackoverflow.com/questions/62075847/using-qgis-and-shaply-error-geosgeom-createlinearring-r-returned-a-null-pointer
from shapely import speedups
speedups.disable()

import shapely

def get_common_extent_from_raster_paths(raster_paths_list):
    #pass in a list of raster file paths and get a shapely polygon that shows where all of the raster extents intersect
    #currently assumes matching crs
    #find intersection of raster bboxes by converting them to shapely polygons
    #https://shapely.readthedocs.io/en/stable/manual.html#polygons
    common_polygon_intersection = None
    for file_path in raster_paths_list:
        dataset = rasterio.open(file_path)
        bbox = dataset.bounds
        dataset.close()
        polygon_format = [
            [bbox.left, bbox.bottom],
            [bbox.left, bbox.top],
            [bbox.right, bbox.top],
            [bbox.right, bbox.bottom]
        ]
        a_polygon = shapely.geometry.Polygon(polygon_format)
        if common_polygon_intersection is None:
            common_polygon_intersection = a_polygon
        else:
            common_polygon_intersection = common_polygon_intersection.intersection(a_polygon)
        assert (not common_polygon_intersection.is_empty),(file_path+" extent doesn't intersect with the other rasters")
    return common_polygon_intersection


def subdivide_bbox(shapely_shape, num_pieces_x_axis,num_pieces_y_axis):
    # returns num_pieces^2 boxes because it splits by num_pieces
    # on the x and y coordinate dimension
    # takes a bbox in shapely polygon form, and subdivides it
    # into num_pieces bboxes
    # modded from https://www.matecdev.com/posts/shapely-polygon-gridding.html
    minx, miny, maxx, maxy = shapely_shape.bounds
    #add 1 to num_pieces, bc we are generating intervals
    #so to get num_pieces in one dimension you actually want +1 intervals
    num_pieces_x_axis = num_pieces_x_axis + 1
    num_pieces_y_axis = num_pieces_y_axis + 1
    x_steps = np.linspace(start=minx, stop=maxx, num=num_pieces_x_axis)
    y_steps = np.linspace(start=miny, stop=maxy, num=num_pieces_y_axis)
    grid = []
    for i in range(len(x_steps)-1):
        for j in range(len(y_steps)-1):
            #define a rectangle with clockwise coords
            lower_left = [x_steps[i], y_steps[j]]
            upper_left = [x_steps[i], y_steps[j+1]]
            upper_right = [x_steps[i+1], y_steps[j+1]]
            lower_right = [x_steps[i+1], y_steps[j]]
            poly_ij = shapely.geometry.Polygon([lower_left, upper_left, upper_right, lower_right])
            grid.append(poly_ij)
    return grid


def buffer_from_ul_coordinate_to_lr_by_pixels(origin_x_coord,origin_y_coord,num_pixels_to_buffer,riods):
    # pass in a coordinate,
    # returns a shapely bbox polygon, which is that coordinate buffered by n pixels in the lr direction
    # so it buffers from ul to lr
    # this returns a bbox that has coordinates in the center of the pixel
    # to retreive pixels from a rasterds using this new bbox you can use
    # window = rasterio.features.geometry_window(riods, [buffered_bbox])
    # w = riods.read(1, window=window)
    #1 find the pixel grid coordinates of the spatial coordinates
    origin_row,origin_col = riods.index(origin_x_coord,origin_y_coord)
    #2 buffer the grid space coordinates by num_pixels_to_buffer
    u_row = origin_row
    l_row = origin_row + num_pixels_to_buffer
    l_col = origin_col
    r_col = origin_col + num_pixels_to_buffer
    #3 convert grid space coordinates back to spatial coordinates
    ul_coord = riods.xy(row=u_row, col=l_col)
    ur_coord = riods.xy(row=u_row, col=r_col)
    lr_coord = riods.xy(row=l_row, col=r_col)
    ll_coord = riods.xy(row=l_row, col=l_col)
    buffered_bbox_polygon = shapely.geometry.Polygon([ll_coord, ul_coord, ur_coord, lr_coord])
    # 5 make sure that the buffered bbox is inside the rasteriods
    #https://gis.stackexchange.com/questions/352445/make-shapefile-from-raster-bounds-in-python
    riods_bbox_polygon = shapely.geometry.box(riods.bounds.left,riods.bounds.bottom,riods.bounds.right,riods.bounds.top)
    assert riods_bbox_polygon.contains(buffered_bbox_polygon)#this should probably be an Exception not Assertion
    return buffered_bbox_polygon


def make_bbox_from_coordinate_by_buffering_on_all_sides_by_num_pixels(origin_x_coord,origin_y_coord,num_pixels_to_buffer,riods):
    #pass in a coordinate
    #returns a shapely bbox polygon, which is that coordinate buffered in all directions by num_pixels_to_buffer
    #this returns a bbox that has coordinates in the center of the pixel
    # to retreive pixels from a rasterds using this new bbox you can use
    # window = rasterio.features.geometry_window(riods, [buffered_bbox])
    # w = riods.read(1, window=window)
    #1 get the coordinates in grid space from spatial coordinates
    origin_row,origin_col = riods.index(origin_x_coord,origin_y_coord)
    #2 buffer the grid space coordinates by num_pixels_to_buffer
    u_row = origin_row - num_pixels_to_buffer
    l_row = origin_row + num_pixels_to_buffer
    l_col = origin_col - num_pixels_to_buffer
    r_col = origin_col + num_pixels_to_buffer
    #3 convert grid space coordinates back to spatial coordinates
    ul_coord = riods.xy(row=u_row, col=l_col)
    ur_coord = riods.xy(row=u_row, col=r_col)
    lr_coord = riods.xy(row=l_row, col=r_col)
    ll_coord = riods.xy(row=l_row, col=l_col)
    buffered_bbox_polygon = shapely.geometry.Polygon([ll_coord, ul_coord, ur_coord, lr_coord])
    # 5 make sure that the buffered bbox is inside the rasteriods
    #https://gis.stackexchange.com/questions/352445/make-shapefile-from-raster-bounds-in-python
    riods_bbox_polygon = shapely.geometry.box(riods.bounds.left,riods.bounds.bottom,riods.bounds.right,riods.bounds.top)
    assert riods_bbox_polygon.contains(buffered_bbox_polygon)#this should probably be an Exception not Assertion
    return buffered_bbox_polygon
