## Process EWR Hg Data into a surface
# AJMR: September 2022

#%% Setup

import pandas as pd
import geopandas as gp
import gpxpy
import gpxpy.gpx
import numpy as np
import matplotlib.pyplot as plt
import contextily as ctx
import datetime

import scipy as sp
from scipy.interpolate import griddata

from shapely import geometry, ops


#%% Read in consolidated dataset
df_in = pd.read_csv('//srv-ott3.baird.com/Projects/12828.101 English Wabigoon River/05_Analyses/02_Data Analysis/output/mercury sediment subset-SGJ.csv')

# Convert to geopandas
dataIN_gp = gp.GeoDataFrame(df_in, geometry=gp.points_from_xy(df_in.Longitude, df_in.Latitude, crs="EPSG:4326"))
# Convert to UTM
dataIN_gp.to_crs(crs='EPSG:32615', inplace=True)



#%% Read in centerline shapefile
river_centerline = gp.read_file('//srv-ott3.baird.com/Projects/12828.101 English Wabigoon River/03_Data/02_Physical/16_Waterline/Centerline_for_Modelling_UTMZ15.shp')

river_centerlineExploded = river_centerline.explode(ignore_index=True)
river_centerlineExploded.reset_index(inplace=True)

tempMulti = river_centerlineExploded.iloc[[5,0,1,2,3,4,6,7,9], 4]
# 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)

river_centerline_merge_gpd2 =\
    gp.GeoDataFrame(geometry=gp.points_from_xy(
        river_centerline_merge.xy[0], river_centerline_merge.xy[1], crs="EPSG:32615"))

# Add distance along centerline
river_centerline_merge_gpd2['DistanceFromPrevious'] = river_centerline_merge_gpd2.distance(river_centerline_merge_gpd2.shift(1))
river_centerline_merge_gpd2['RiverKM'] = river_centerline_merge_gpd2['DistanceFromPrevious'].cumsum()
river_centerline_merge_gpd2.iloc[0, 1] = 0
river_centerline_merge_gpd2.iloc[0, 2] = 0

#%% Add riverkm distances to dataset
dataOUT = gp.sjoin_nearest(river_centerline_merge_gpd2, dataIN_gp, how='left', max_distance=100)

#%% Take only surface values
dataOUTmax = dataOUT.groupby(['Samplesiteid'], as_index=False)['Samplevalue'].max()
# Merge geometry and river KM
dataOUTmax = pd.merge(dataOUTmax, dataOUT[['Samplesiteid', 'RiverKM']], on='Samplesiteid', how='left')
dataOUTmax = pd.merge(dataOUTmax, dataOUT[['Samplesiteid', 'geometry']], on='Samplesiteid', how='left')

#%% Set River KM as index
dataOUT.set_index('RiverKM', inplace=True)
dataOUTmax.set_index('RiverKM', inplace=True)


#%% River Profile Plotting
fig, axes = plt.subplots(nrows=1, ncols=1, figsize=(6, 6))

axes.scatter(dataOUT.index, dataOUT.Samplevalue)
axes.scatter(dataOUTmax.index, dataOUTmax.Samplevalue)

plt.show()


#%% Hg Map Plotting
mapbox = 'https://api.mapbox.com/styles/v1/alexander0042/ckemxgtk51fgp19nybfmdcb1e/tiles/256/{z}/{x}/{y}@2x?access_token=pk.eyJ1IjoiYWxleGFuZGVyMDA0MiIsImEiOiJjazVmdG4zbncwMHY4M2VrcThwZGUzZDFhIn0.w6oDHoo1eCeRlSBpwzwVtw'

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

# Full system limits
# axes.set_xlim(350000, 520000)
# axes.set_ylim(5510000, 5580000)

# #  Clay Lake Limits
# axes.set_xlim(466418, 515846)
# axes.set_ylim(5513437, 5545362)

#  Dryden Limits
axes.set_xlim(466418, 515846)
axes.set_ylim(5513437, 5545362)

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

# Plot Hg
# dataOUTmax.plot(ax=axes, column='Samplevalue', legend=True, vmin=0, vmax=20000)

# Plot Hg At surface
dataOUT.loc[((dataOUT.TopDepth == 0) | (dataOUT.TopDepth.isna()))].plot(
    ax=axes, column='Samplevalue', legend=True, vmin=0, vmax=10000)

plt.show()


#%% Interpolate Hg
bathy_xyz = pd.read_csv('C:/Users/arey/files/Projects/Grassy Narrows/LocalData/Bed_Level.xyz',
                        names=['x', 'y', 'z'], header=0, delim_whitespace=True)

rough_xyz = pd.read_csv('C:/Users/arey/files/Projects/Grassy Narrows/LocalData/Roughness.xyz',
                        names=['x', 'y', 'z'], header=0, delim_whitespace=True)

HgInterp  = griddata(np.array([dataOUTmax.geometry.x, dataOUTmax.geometry.y]).T,
                                    np.array(dataOUTmax.Samplevalue).T, np.array([bathy_xyz.x, bathy_xyz.y]).T,
                                    method='nearest')

HgInterpOut = np.vstack((bathy_xyz.x, bathy_xyz.y, HgInterp))
HgInterpOut = np.transpose(HgInterpOut)

# Set Wetlands to zero
# Interpolate Roughness
RoughInterp = sp.interpolate.griddata(np.transpose(np.array([rough_xyz.x, rough_xyz.y])), rough_xyz.z, np.transpose(np.array([bathy_xyz.x, bathy_xyz.y])),
                                    method='nearest')

HgInterpOut[RoughInterp==0.05, 2] = 0

np.savetxt('C:/Users/arey/files/Projects/Grassy Narrows/LocalData/SurfaceInterFiltDB.xyz',
           HgInterpOut, delimiter=" ")


#%% Hg Interp Plotting
mapbox = 'https://api.mapbox.com/styles/v1/alexander0042/ckemxgtk51fgp19nybfmdcb1e/tiles/256/{z}/{x}/{y}@2x?access_token=pk.eyJ1IjoiYWxleGFuZGVyMDA0MiIsImEiOiJjazVmdG4zbncwMHY4M2VrcThwZGUzZDFhIn0.w6oDHoo1eCeRlSBpwzwVtw'

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

#  Clay Lake Limits
axes.set_xlim(466418, 515846)
axes.set_ylim(5513437, 5545362)
# Add Basemap
ctx.add_basemap(axes, source=mapbox, crs='EPSG:32615')
# Plot Hg
axes.scatter(HgInterpOut[:, 0], HgInterpOut[:, 1], 5, HgInterpOut[:, 2],
             vmin=0, vmax=5000)

plt.show()