## 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 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

#%% Read in water depth to get wet and dry points

depth_xyz = pd.read_csv('//srv-ott3/Projects/12828.101 English Wabigoon River/06_Models/00_Delft3D/04_Bathymetry/Water depth at pressure points - mesh2d_nFaces_ mean.csv',
                        names=['x', 'y', 'z', 'depth'], header=1)

# Make geodataframe
depth_xyz_gp = gp.GeoDataFrame(depth_xyz, geometry=gp.points_from_xy(depth_xyz.x, depth_xyz.y), crs="EPSG:32615")

# Merge with centerline to get river km
depth_xyz_gp = gp.sjoin_nearest(river_centerline_merge_gpd2, depth_xyz_gp, how='right')

#%% Create synthetic Hg data following Reed's equation
# These need to be multiplied by the total mass of dry matter in each cell to get the total Hg/m2 in each layer
CstartHg = 10
CfinalHg = 1
kHg	   = 0.08

depth_xyz_gp['Hg'] = (CstartHg - CfinalHg) * np.exp(-kHg * depth_xyz_gp['RiverKM']/1000) + CfinalHg

CstartMeHg = 0.05
CfinalMeHg = 0.005
kMeHg	   = 0.08

depth_xyz_gp['MeHg'] = (CstartMeHg - CfinalMeHg) * np.exp(-kMeHg * depth_xyz_gp['RiverKM']/1000) + CfinalMeHg


# Set points where depth is zero (no water) to 0
depth_xyz_gp.loc[depth_xyz_gp['depth'] == 0, 'Hg'] = 0
depth_xyz_gp.loc[depth_xyz_gp['depth'] == 0, 'MeHg'] = 0

#%% Hg Test plot
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')

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

# Plot river distance using geopandas
# depth_xyz_gp.plot(ax=axes, column='RiverKM', cmap='viridis')
depth_xyz_gp.loc[depth_xyz_gp['depth'] != 0, :].plot(ax=axes, column='Hg', cmap='viridis')

plt.show()

#%% Save to csv
np.savetxt('//srv-ott3/Projects/12828.101 English Wabigoon River/06_Models/00_Delft3D/99_Setups/SynHg_ug_g.xyz',
           np.array([depth_xyz_gp.geometry.x.values,  depth_xyz_gp.geometry.y.values,  depth_xyz_gp['Hg'].values]).T,
           delimiter=" ")

np.savetxt('//srv-ott3/Projects/12828.101 English Wabigoon River/06_Models/00_Delft3D/99_Setups/SynMeHg_ug_g.xyz',
           np.array([depth_xyz_gp.geometry.x.values, depth_xyz_gp.geometry.y.values, depth_xyz_gp['MeHg'].values]).T,
           delimiter=" ")