## Python Script to Process EWR Bathymetry Data
# AJMR: March 2022

import pandas as pd
import geopandas as gp
import matplotlib.pyplot as plt
import contextily as ctx
import numpy as np
from shapely import geometry, ops


#%% Read in and process centerline shapefile
mapbox = 'https://api.mapbox.com/styles/v1/alexander0042/ckemxgtk51fgp19nybfmdcb1e/tiles/256/{z}/{x}/{y}@2x?access_token=pk.eyJ1IjoiYWxleGFuZGVyMDA0MiIsImEiOiJjazVmdG4zbncwMHY4M2VrcThwZGUzZDFhIn0.w6oDHoo1eCeRlSBpwzwVtw'

river_centerline = gp.read_file('//srv-ott3/Projects/12828.101 English Wabigoon River/03_Data/02_Physical/16_Waterline/Centerline_for_Modelling_merge_UTMZ15.shp')

river_centerlineExploded = river_centerline.explode(ignore_index=True)

tempMulti = river_centerlineExploded.iloc[[4, 3, 1, 2], 1]
# Put the sub-line coordinates into a list of sublists
outcoords = [list(i.coords) for i in tempMulti]
# Flatten the list of sublists and use it to make a new line
river_centerline_merge = geometry.LineString([i for sublist in outcoords for i in sublist])
river_centerline_merge_gpd = gp.GeoSeries(river_centerline_merge)

# Interpolate the centerline to every 25 m
distance_delta = 25
distances = np.arange(0, river_centerline_merge.length, distance_delta)
points = [river_centerline_merge.interpolate(distance) for distance in distances] + [river_centerline_merge.boundary[1]]
# multipoint = ops.unary_union(points)
interpolated_centerline = geometry.LineString(points)

# Turn the linestring into a geodataframe and add distance
# river_centerline_merge_gpd2 =\
#     gp.GeoDataFrame(geometry=gp.points_from_xy(
#         river_centerline_merge.xy[0], river_centerline_merge.xy[1], crs="EPSG:32615"))
#
# river_centerline_merge_gpd2['Distance'] = river_centerline_merge_gpd2.distance(river_centerline_merge_gpd2.shift(1))
# river_centerline_merge_gpd2.loc[0, 'Distance'] = 0
#
# river_centerline_merge_gpd2['RiverKM'] = river_centerline_merge_gpd2['Distance'].cumsum()

# Turn interpolated inestring into a geodataframe and add distance
river_centerline_merge_gpd3 = \
    gp.GeoDataFrame(geometry=gp.points_from_xy(
        interpolated_centerline.xy[0], interpolated_centerline.xy[1]), crs="EPSG:32615")

river_centerline_merge_gpd3['Distance'] = river_centerline_merge_gpd3.distance(river_centerline_merge_gpd3.shift(1))
river_centerline_merge_gpd3.loc[0, 'Distance'] = 0

river_centerline_merge_gpd3['RiverKM'] = river_centerline_merge_gpd3['Distance'].cumsum()


#%% Read in the processed waterline shapefile
river_edges = gp.read_file('//srv-ott3/Projects/12828.101 English Wabigoon River/03_Data/02_Physical/16_Waterline/WaterlineToLines_Grass_Simple5m_RevB.shp')

# Interpolate the edges to every 10 m
distance_delta = 10
distances = np.arange(0, river_edges.geometry[0].length, distance_delta)
points = [river_edges.geometry[0].interpolate(distance) for distance in distances] + [river_edges.geometry[0].boundary[1]]
# multipoint = ops.unary_union(points)
interpolated_edges = geometry.LineString(points)

# Create geodataframes for south edge
river_edges_south = \
    gp.GeoDataFrame(geometry=gp.points_from_xy(
        interpolated_edges.xy[0], interpolated_edges.xy[1]), crs="EPSG:32615")
river_edges_south['SouthBoundary'] = river_edges_south['geometry']

# Interpolate the edges to every 10 m North
distances = np.arange(0, river_edges.geometry[1].length, distance_delta)
points = [river_edges.geometry[1].interpolate(distance) for distance in distances] + [river_edges.geometry[1].boundary[1]]
# multipoint = ops.unary_union(points)
interpolated_edges = geometry.LineString(points)

# Create geodataframes for North edge
river_edges_north = \
    gp.GeoDataFrame(geometry=gp.points_from_xy(
        interpolated_edges.xy[0], interpolated_edges.xy[1]), crs="EPSG:32615")
river_edges_north['NorthBoundary'] = river_edges_north['geometry']

# river_edges_south = \
#     gp.GeoDataFrame(geometry=gp.points_from_xy(
#         river_edges.loc[0, 'geometry'].xy[0], river_edges.loc[0, 'geometry'].xy[1]), crs="EPSG:32615")
# river_edges_south['SouthBoundary'] = river_edges_south['geometry']

# river_edges_north = \
#     gp.GeoDataFrame(geometry=gp.points_from_xy(
#         river_edges.loc[1, 'geometry'].xy[0], river_edges.loc[1, 'geometry'].xy[1]), crs="EPSG:32615")
# river_edges_north['NorthBoundary'] = river_edges_north['geometry']

#%% Plot the centerline and the waterline
fig, axes = plt.subplots(nrows=1, ncols=1, figsize=(20, 14))
fig.patch.set_facecolor('white')

axes.set_xlim(350000, 520000)
axes.set_ylim(5510000, 5580000)

axes.set_xlim(505438, 508368)
axes.set_ylim(5518465, 5520152)

ctx.add_basemap(axes, source=mapbox, crs='EPSG:32615')

# river_centerlineExploded.loc[[0], 'geometry'].plot(ax=axes, color='tab:blue')
# river_centerlineExploded.loc[[1], 'geometry'].plot(ax=axes, color='tab:orange')
# river_centerlineExploded.loc[[2], 'geometry'].plot(ax=axes, color='tab:green')
# river_centerlineExploded.loc[[3], 'geometry'].plot(ax=axes, color='tab:red')
# river_centerlineExploded.loc[[4], 'geometry'].plot(ax=axes, color='tab:purple')

river_edges.loc[[0], 'geometry'].plot(ax=axes, color='tab:red', linestyle='-')
river_edges.loc[[1], 'geometry'].plot(ax=axes, color='tab:orange', linestyle='-')
#
# river_centerline_merge_gpd.plot(ax=axes, marker='s')
#
# river_centerline_merge_gpd2.plot(ax=axes, column='RiverKM', marker='o')
#
river_centerline_merge_gpd3.plot(ax=axes, column='RiverKM', marker='^')

plt.show()

#%% Find the closest point on the centerline to the North and South waterlines

river_centerline_Bounds = river_centerline_merge_gpd3.sjoin_nearest(river_edges_south, how='left', distance_col='South_Distance')
# Remove Index Right Column
river_centerline_Bounds = river_centerline_Bounds.drop(['index_right'], axis=1)
river_centerline_Bounds = river_centerline_Bounds.sjoin_nearest(river_edges_north, how='left', distance_col='North_Distance')
river_centerline_Bounds = river_centerline_Bounds.drop(['index_right'], axis=1)

#%% Add points to cenetrline gdf at 10,10,90% of the distance between the two waterline points through the centerline

# Take x and y for south distances from centerline
for t_percent in range(10, 100, 10):
    transect_name = 'transectSouth_' + str(t_percent)
    river_centerline_Bounds[transect_name] = gp.points_from_xy((river_centerline_Bounds.geometry.x -
                                        river_centerline_Bounds['SouthBoundary'].x) * t_percent * 0.01 + river_centerline_Bounds['SouthBoundary'].x,
                                        (river_centerline_Bounds.geometry.y -
                                        river_centerline_Bounds['SouthBoundary'].y) * t_percent * 0.01 + river_centerline_Bounds['SouthBoundary'].y,
                                        crs="EPSG:32615")

# Take x and y for north distances from centerline
for t_percent in range(10, 100, 10):
    transect_name = 'transectNorth_' + str(t_percent)
    river_centerline_Bounds[transect_name] = gp.points_from_xy((river_centerline_Bounds.geometry.x -
                                        river_centerline_Bounds['NorthBoundary'].x) * t_percent * 0.01 + river_centerline_Bounds['NorthBoundary'].x,
                                        (river_centerline_Bounds.geometry.y -
                                        river_centerline_Bounds['NorthBoundary'].y) * t_percent * 0.01 + river_centerline_Bounds['NorthBoundary'].y,
                                        crs="EPSG:32615")

#%% Plotting Check
fig, axes = plt.subplots(nrows=1, ncols=1, figsize=(20, 14))
fig.patch.set_facecolor('white')

axes.set_xlim(505438, 508368)
axes.set_ylim(5518465, 5520152)
ctx.add_basemap(axes, source=mapbox, crs='EPSG:32615')


for t_percent in range(10, 100, 10):
    transect_name = 'transectSouth_' + str(t_percent)
    river_centerline_Bounds.loc[1:500, transect_name].plot(ax=axes)

    transect_name = 'transectNorth_' + str(t_percent)
    river_centerline_Bounds.loc[1:500, transect_name].plot(ax=axes)


river_edges.loc[[0], 'geometry'].plot(ax=axes, color='k', linestyle='-')
river_edges.loc[[1], 'geometry'].plot(ax=axes, color='k', linestyle='-')

river_centerline_merge_gpd.plot(ax=axes, color='k')
river_centerline_merge_gpd3.plot(ax=axes, column='RiverKM', marker='^', color='k')

plt.show()

#%% Read in bathymetry observations
bathy_obs = pd.read_csv('//oak-inundation/D/12828.101 EWR_Delft3FM_MTB/Bathy/combined_ZeroRemoved_RevC.xyz',
                        delim_whitespace=True, header=None, names=['x', 'y', 'z'])

# Create geodataframe
bathy_obs_gdf = gp.GeoDataFrame(bathy_obs, geometry=gp.points_from_xy(
                bathy_obs.x, bathy_obs.y), crs="EPSG:32615")

bathy_obs_gdf = bathy_obs_gdf.drop(['x', 'y'], axis=1)


#%% Find the closest point (to a 50 m max) on the transects to observations
river_transect_bathy = dict()

# Does not include 0% or 100%
for t_percent in range(10, 100, 10):
    # Set geometry to specific transectfor merge
    transect_name = 'transectSouth_' + str(t_percent)
    transect_merge_tmp = river_centerline_Bounds.loc[:, ['RiverKM', transect_name]]
    transect_merge_tmp.set_geometry(transect_name, inplace=True)

    # Create buffered polygons around transect lines
    transect_buffer = geometry.LineString(transect_merge_tmp.geometry).buffer(2.5, resolution=8)
    transect_buffer_gdf = gp.GeoDataFrame(geometry=[transect_buffer], crs="EPSG:32615")

    # Find observations within buffered polygons
    bathy_obs_gdf_tmp = bathy_obs_gdf.sjoin(transect_buffer_gdf, how='inner', predicate='within')
    bathy_obs_gdf_tmp = bathy_obs_gdf_tmp.drop(['index_right'], axis=1)

    river_transect_bathy[transect_name] = transect_merge_tmp.sjoin_nearest(bathy_obs_gdf_tmp, how='left', max_distance=50)
    river_transect_bathy[transect_name] = river_transect_bathy[transect_name].drop(['index_right'], axis=1)

    bathy_obs_gdf_tmp = []
    # # Rename bathy column and drop index
    # river_centerline_Bounds.rename(columns={'z': transect_name + '_z'}, inplace=True)
    # river_centerline_Bounds = river_centerline_Bounds.drop(['index_right'], axis=1)

    # Repeat for North transect
    transect_name = 'transectNorth_' + str(t_percent)
    transect_merge_tmp = river_centerline_Bounds.loc[:, ['RiverKM', transect_name]]
    transect_merge_tmp.set_geometry(transect_name, inplace=True)

    transect_buffer = geometry.LineString(transect_merge_tmp.geometry).buffer(2.5, resolution=8)
    transect_buffer_gdf = gp.GeoDataFrame(geometry=[transect_buffer], crs="EPSG:32615")
    bathy_obs_gdf_tmp = bathy_obs_gdf.sjoin(transect_buffer_gdf, how='inner', predicate='within')
    bathy_obs_gdf_tmp = bathy_obs_gdf_tmp.drop(['index_right'], axis=1)

    river_transect_bathy[transect_name] = transect_merge_tmp.sjoin_nearest(bathy_obs_gdf_tmp, how='left', max_distance=50)
    river_transect_bathy[transect_name] = river_transect_bathy[transect_name].drop(['index_right'], axis=1)
    bathy_obs_gdf_tmp = []
    # # Rename bathy column and drop index
    # river_centerline_Bounds.rename(columns={'z': transect_name + '_z'}, inplace=True)
    # river_centerline_Bounds = river_centerline_Bounds.drop(['index_right'], axis=1)

    print(t_percent)


#%% Plotting Check

fig, axes = plt.subplots(nrows=1, ncols=1, figsize=(20, 14))
fig.patch.set_facecolor('white')

axes.set_xlim(505438, 508368)
axes.set_ylim(5518465, 5520152)
ctx.add_basemap(axes, source=mapbox, crs='EPSG:32615')

bathy_obs_gdf.loc[:, :].plot(ax=axes, column='z', vmin=320, vmax=350, marker='^')

for t_percent in range(10, 100, 10):
    transect_name = 'transectSouth_' + str(t_percent)
    river_transect_bathy[transect_name].loc[1:500, :].plot(ax=axes, column='z', vmin=320, vmax=350)

    transect_name = 'transectNorth_' + str(t_percent)
    river_transect_bathy[transect_name].loc[1:500, :].plot(ax=axes, column='z', vmin=320, vmax=350)



plt.show()

#%% Merge the lines to 1D profile and save as csv
bathy_FileName = 'C:/Users/arey/files/Projects/Grassy Narrows/Bathymetry/AJMR_Interpolated_bathy_25m.xyz'
shape_FileName = 'C:/Users/arey/files/Projects/Grassy Narrows/Bathymetry/'

# bathy_FileName = '//oak-inundation/D/12828.101 EWR_Delft3FM_MTB/Bathy/AJMR_Interpolated_bathy_10m.xyz'

for t_percent in range(10, 100, 10):
    transect_name = 'transectSouth_' + str(t_percent)

    # Convert to numpy array
    bathyNP = np.vstack((river_transect_bathy[transect_name].geometry.x.values,
                         river_transect_bathy[transect_name].geometry.y.values,
                         river_transect_bathy[transect_name].z.values))
    # Transpose
    bathyNP = bathyNP.T

    # Save using pandas
    pd.DataFrame(bathyNP[~np.isnan(bathyNP[:, 2]), :]).to_csv(bathy_FileName, sep=' ', header=False, index=False, mode='a')

    # Write to shape
    # river_transect_bathy[transect_name].to_file(shape_FileName + transect_name + '.shp')

    # Repeat for North transect
    transect_name = 'transectNorth_' + str(t_percent)
    bathyNP = np.vstack((river_transect_bathy[transect_name].geometry.x.values,
                         river_transect_bathy[transect_name].geometry.y.values,
                         river_transect_bathy[transect_name].z.values))
    bathyNP = bathyNP.T
    pd.DataFrame(bathyNP[~np.isnan(bathyNP[:, 2]), :]).to_csv(bathy_FileName, sep=' ', header=False, index=False, mode='a')
    # river_transect_bathy[transect_name].to_file(shape_FileName + transect_name)


