#%% Plotting EWR Flume Tests

# Alexander Rey, 2022

#%% Import
import pandas as pd
import numpy as np
import math
import matplotlib.pyplot as plt
import matplotlib.cm as cm
import matplotlib.colors as mcolors
import geopandas as gp
gp.options.use_pygeos = True
from shapely import geometry, ops

# Map Plotting
import cartopy.crs as ccrs
import contextily as ctx

# Interpolation
import scipy as sp
from scipy.interpolate import griddata
from scipy.interpolate import LinearNDInterpolator, interp1d

# Lowess interpolation
import statsmodels.api as sm

import pathlib as pl

import datetime

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

#%% Import Observations from database

# Read in combined dataset
obs_IN = pd.read_csv("//srv-ott3.baird.com/Projects/12828.101 English Wabigoon River/05_Analyses/02_Data Analysis/output/combined dataset-SGJ.csv")
# Add Datetime
obs_IN['DateTime'] = pd.to_datetime(obs_IN.Sampledate_x)

# Add months
obs_IN['Month'] = obs_IN['DateTime'].dt.month

# Convert to geodataframe
obs_gdf = gp.GeoDataFrame(obs_IN, geometry=gp.points_from_xy(
                obs_IN.loc[:, 'Longitude'], obs_IN.loc[:, 'Latitude'], crs="EPSG:4326")).to_crs(crs="EPSG:32615")

# Adjust Units
obs_gdf.loc[obs_gdf.Unit == 'NG/G', 'Sample_NG/G'] = obs_gdf.Samplevalue
obs_gdf.loc[obs_gdf.Unit == 'ng/g', 'Sample_NG/G'] = obs_gdf.Samplevalue
obs_gdf.loc[obs_gdf.Unit == 'UG/G', 'Sample_NG/G'] = obs_gdf.Samplevalue * 1000
obs_gdf.loc[obs_gdf.Unit == 'ug/g', 'Sample_NG/G'] = obs_gdf.Samplevalue * 1000
obs_gdf.loc[obs_gdf.Unit == 'MG/KG', 'Sample_NG/G'] = obs_gdf.Samplevalue * 1000
obs_gdf.loc[obs_gdf.Unit == 'mg/kg', 'Sample_NG/G'] = obs_gdf.Samplevalue * 1000
obs_gdf.loc[obs_gdf.Unit == 'NG/L', 'Sample_NG/L'] = obs_gdf.Samplevalue
obs_gdf.loc[obs_gdf.Unit == 'ng/L', 'Sample_NG/L'] = obs_gdf.Samplevalue
obs_gdf.loc[obs_gdf.Unit == 'UG/L', 'Sample_NG/L'] = obs_gdf.Samplevalue * 1000
obs_gdf.loc[obs_gdf.Unit == 'ug/L', 'Sample_NG/L'] = obs_gdf.Samplevalue * 1000
obs_gdf.loc[obs_gdf.Unit == 'MG/L', 'Sample_NG/L'] = obs_gdf.Samplevalue * 1000000
obs_gdf.loc[obs_gdf.Unit == 'mg/L', 'Sample_NG/L'] = obs_gdf.Samplevalue * 1000000

# Add centerline and reset index
obs_gdf = gp.sjoin_nearest(river_centerline_merge_gpd2, obs_gdf, how='right', max_distance=250).reset_index()

# Sediment Hg GeoDataFrame where media is Sediment
obsMask = (((obs_gdf.Media_x == 'SED') | (obs_gdf.Media_x == 'SOIL')) &
           ((obs_gdf.Parameter == 'Mercury') | (obs_gdf.Parameter == 'Mercury (Hg) Total')
            | (obs_gdf.Parameter == 'Total Mercury') | (obs_gdf.Parameter == 'Mercury (Hg)')) &
            (obs_gdf.DateTime > datetime.datetime(2000, 1, 1)) & obs_gdf.RiverKM.notna())

obs_HgSed = obs_gdf.loc[obsMask, :]

# Sediment Hg GeoDataFrame for surface sediment
obsMask = (((obs_gdf.Media_x == 'SED') | (obs_gdf.Media_x == 'SOIL')) &
           ((obs_gdf.Parameter == 'Mercury') | (obs_gdf.Parameter == 'Mercury (Hg) Total')
            | (obs_gdf.Parameter == 'Total Mercury') | (obs_gdf.Parameter == 'Mercury (Hg)')) &
            (obs_gdf.DateTime > datetime.datetime(2000, 1, 1)) & (obs_gdf.RiverKM.notna()) &
           ((obs_gdf.TopDepth_x < 5) | (obs_gdf.TopDepth_x.isna())))

obs_HgSurfaceSed = obs_gdf.loc[obsMask, :]

# obs_HgSurfaceSed.to_file('C:/Users/arey/files/Projects/Grassy Narrows/LocalData/obs_HgSurfaceSed.geojson', driver="GeoJSON")
obs_HgSurfaceSed.loc[:, ['Sample_NG/G', 'Samplenumber_x', 'geometry']].to_file('C:/Users/arey/files/Projects/Grassy Narrows/LocalData/obs_HgSurfaceSed2.shp')

# Group sediment core data to take average/min/max
obs_HgSurfaceSed_Avg = obs_HgSurfaceSed.groupby(['Sitecode', 'DateTime']).agg({'Sample_NG/G': ['mean', 'min', 'max'],
                                                                'RiverKM': ['mean'],
                                                                'Longitude': ['mean'],
                                                                'Latitude': ['mean'],
                                                                'DateTime': ['mean']}).reset_index()
# Flatten Multiindex
obs_HgSurfaceSed_Avg.columns = ["_".join(a) for a in obs_HgSurfaceSed_Avg.columns.to_flat_index()]

# Rename for shapefile
obs_HgSurfaceSed_Avg.rename(columns={"Sample_NG/G_mean": "Mean_NG/G",
                                     "Sample_NG/G_max": "Max_NG/G",
                                     "Sample_NG/G_min": "Min_NG/G"}, inplace=True)

# Convert to geodataframe and save
obs_HgSurfaceSed_Avg_gdf = gp.GeoDataFrame(obs_HgSurfaceSed_Avg, geometry=gp.points_from_xy(
                obs_HgSurfaceSed_Avg.loc[:, 'Longitude_mean'], obs_HgSurfaceSed_Avg.loc[:, 'Latitude_mean'],
                crs="EPSG:4326")).to_crs(crs="EPSG:32615")

obs_HgSurfaceSed_Avg_gdf.loc[:, ['Mean_NG/G', 'Max_NG/G', 'Min_NG/G', 'geometry']].to_file('C:/Users/arey/files/Projects/Grassy Narrows/LocalData/obs_HgSurfaceSed_Avg.shp')


# Group sediment core data to take average/min/max
obs_HgSed_Avg = obs_HgSed.groupby(['Sitecode', 'DateTime']).agg({'Sample_NG/G': ['mean', 'min', 'max'],
                                                                'RiverKM': ['mean'],
                                                                'Longitude': ['mean'],
                                                                'Latitude': ['mean'],
                                                                'DateTime': ['mean']}).reset_index()
# Merge Index
obs_HgSed_Avg.columns = obs_HgSed_Avg.columns.map('|'.join).str.strip('|')

# Fix Names
obs_HgSed_Avg.rename(columns={"RiverKM|mean": "RiverKM", "Longitude|mean": "Longitude",
                              "Latitude|mean": "Latitude", "DateTime|mean": "DateTime"}, inplace=True)

# Create new GeoDataFrame
obs_HgSed_Avg = gp.GeoDataFrame(obs_HgSed_Avg, geometry=gp.points_from_xy(
                obs_HgSed_Avg.loc[:, 'Longitude'], obs_HgSed_Avg.loc[:, 'Latitude'], crs="EPSG:4326")).to_crs(crs="EPSG:32615")



# Water Hg GeoDataFrame where media is SurfaceWater (SW)
obsMask = (((obs_gdf.Media_x == 'SW')) &
           ((obs_gdf.Parameter == 'Mercury') | (obs_gdf.Parameter == 'Mercury (Hg)-Total')
            | (obs_gdf.Parameter == 'Total Mercury') | (obs_gdf.Parameter == 'Mercury (Hg)')) &
            (obs_gdf.DateTime > datetime.datetime(2000, 1, 1)) & obs_gdf.RiverKM.notna())

obs_HgWater = obs_gdf.loc[obsMask, :]

# Save
# obs_HgWater.loc[:, ['Sample_NG/L', 'Samplenumber_x', 'geometry']].to_file('C:/Users/arey/files/Projects/Grassy Narrows/LocalData/obs_HgWater.shp')

# Sediment surface MeHg GeoDataFrame
obsMask = (((obs_gdf.Media_x == 'SED') | (obs_gdf.Media_x == 'SOIL')) &
           ((obs_gdf.Parameter == 'Methylmercury (as MeHg)') | (obs_gdf.Parameter == 'Methyl mercury')) &
            (obs_gdf.DateTime > datetime.datetime(2000, 1, 1)) & obs_gdf.RiverKM.notna() &
            ((obs_gdf.TopDepth_x < 5) | (obs_gdf.TopDepth_x.isna())))

obs_MeHgSed = obs_gdf.loc[obsMask, :]

# Save
# obs_MeHgSed.loc[:, ['Sample_NG/G', 'Samplenumber_x', 'geometry']].to_file('C:/Users/arey/files/Projects/Grassy Narrows/LocalData/obs_MeHgSed.shp')


# Group sediment core data to take average/min/max
obs_MeHgSed_Avg = obs_MeHgSed.groupby(['Sitecode', 'DateTime']).agg({'Sample_NG/G': ['mean', 'min', 'max'],
                                                                'RiverKM': ['mean'],
                                                                'Longitude': ['mean'],
                                                                'Latitude': ['mean'],
                                                                'DateTime': ['mean']}).reset_index()
# Flatten Multiindex
obs_MeHgSed_Avg.columns = ["_".join(a) for a in obs_MeHgSed_Avg.columns.to_flat_index()]

# Rename for shapefile
obs_MeHgSed_Avg.rename(columns={"Sample_NG/G_mean": "Mean_NG/G",
                                     "Sample_NG/G_max": "Max_NG/G",
                                     "Sample_NG/G_min": "Min_NG/G"}, inplace=True)

# Create new GeoDataFrame
obs_MeHgSed_Avg_gdf = gp.GeoDataFrame(obs_MeHgSed_Avg, geometry=gp.points_from_xy(
                obs_MeHgSed_Avg.loc[:, 'Longitude_mean'], obs_MeHgSed_Avg.loc[:, 'Latitude_mean'], crs="EPSG:4326")).to_crs(crs="EPSG:32615")

obs_MeHgSed_Avg_gdf.loc[:, ['Mean_NG/G', 'Max_NG/G', 'Min_NG/G', 'geometry']].to_file('C:/Users/arey/files/Projects/Grassy Narrows/LocalData/obs_MeHgSed_Avg.shp')



# MeHg in Water
obsMask = (((obs_gdf.Media_x == 'SW')) &
           ((obs_gdf.Parameter == 'Methylmercury (as MeHg)') | (obs_gdf.Parameter == 'Methyl mercury')) &
            (obs_gdf.DateTime > datetime.datetime(2000, 1, 1)) & obs_gdf.RiverKM.notna())

obs_MeHgWater = obs_gdf.loc[obsMask, :]

# Save
obs_MeHgWater.loc[:, ['Sample_NG/L', 'Samplenumber_x', 'geometry']].to_file('C:/Users/arey/files/Projects/Grassy Narrows/LocalData/obs_MeHgWater.shp')

#%% Plot Observations
fig, axes = plt.subplots(nrows=1, ncols=1, figsize=(30, 15))

axes.set_xlim(494649, 513647)
axes.set_ylim(5514653, 5525256)

# Add basemap
mapbox = 'https://api.mapbox.com/styles/v1/alexander0042/ckemxgtk51fgp19nybfmdcb1e/tiles/256/{z}/{x}/{y}@2x?access_token=pk.eyJ1IjoiYWxleGFuZGVyMDA0MiIsImEiOiJjazVmdG4zbncwMHY4M2VrcThwZGUzZDFhIn0.w6oDHoo1eCeRlSBpwzwVtw'
ctx.add_basemap(axes, source=mapbox, crs='EPSG:26915')

obs_HgSed_Avg.plot(column='Sample_NG/G|mean', axes=axes, vmin=0, vmax=10000)

plt.show()

#%% Plot River Profile
name = "Set1"
cmap = cm.get_cmap(name)
colors = cmap.colors

# Create a HSV colormap sifted by 180 degrees
# create the "hsv" colormap
hsv = plt.get_cmap('hsv')

# shift the colormap so that red is in the middle
shifted_hsv = mcolors.LinearSegmentedColormap.from_list('shifted_hsv',
                    np.roll(hsv(np.linspace(0.0, 1.0, 256)), 128, axis=0))

interp_xvals = np.linspace(0, 100000, num=1000)

# Bin averaging function
def average_in_bins(df, bin_width):
    binned_df = pd.cut(df.index, bins=range(0, int(df.index.max() + bin_width), bin_width), right=False)
    return df.groupby(binned_df).mean()


fig, axes = plt.subplots(nrows=4, ncols=1, figsize=(8, 12), sharex=True)
axes[0].set_prop_cycle(color=colors)
axes[1].set_prop_cycle(color=colors)
axes[2].set_prop_cycle(color=colors)
axes[3].set_prop_cycle(color=colors)

# Plot surface Hg

# Plot using a different color for each month
obs_HgSurfaceSed.plot.scatter('RiverKM', 'Sample_NG/G', ax=axes[0], c='Month',
                              label='All Samples', vmin=1, vmax=12, cmap=shifted_hsv)

# obs_HgSurfaceSed.plot.scatter('RiverKM', 'Sample_NG/G', ax=axes[0], color='grey', label='All Samples')
# obs_HgSurfaceSed_Avg_gdf.plot.scatter('RiverKM_mean', 'Mean_NG/G', ax=axes[0], color='magenta', label='Site Averaged')

# obs_HgSurfaceSed_Index = obs_HgSurfaceSed.copy()
# obs_HgSurfaceSed_Index['RiverKM2'] = obs_HgSurfaceSed['RiverKM']
# obs_HgSurfaceSed_Index = obs_HgSurfaceSed_Index.set_index('RiverKM')
#
# binwidths = [10, 100, 1000, 10000]
# for binWidthIDX, binWidth in enumerate(binwidths):
#     average_in_bins(obs_HgSurfaceSed_Index, binWidth).dropna(subset=['Sample_NG/G']).\
#         plot.line('RiverKM2', 'Sample_NG/G', ax=axes[0],
#                   label=str(binWidth) + ' m binned Samples')

# Find and plot Lowess regression for each season
SeasonMonths = [[12, 1, 2], [3, 4, 5], [6, 7, 8], [9, 10, 11]]
SeasonNames = ['Winter', 'Spring', 'Summer', 'Fall']
SeasonColors = ['tab:blue', 'tab:green', 'darkred', 'tab:orange']

for seasonIDX, season in enumerate(SeasonMonths):
    # Skip Winter
    if seasonIDX == 0:
        continue

    seasonMask = obs_HgSurfaceSed['DateTime'].dt.month.isin(season)
    lowess = sm.nonparametric.lowess(obs_HgSurfaceSed.loc[seasonMask, 'Sample_NG/G'],
                                     obs_HgSurfaceSed.loc[seasonMask, 'RiverKM'],
                                     frac=0.5, xvals=interp_xvals, it=2)
    axes[0].plot(interp_xvals, lowess, label='Lowess:' + SeasonNames[seasonIDX],
                 linewidth=4, color=SeasonColors[seasonIDX])

# lowess = sm.nonparametric.lowess(obs_HgSurfaceSed['Sample_NG/G'], obs_HgSurfaceSed['RiverKM'],
#                                  frac=0.3, xvals=interp_xvals, it=2)
# axes[0].plot(interp_xvals, lowess, label='Lowess regression', linewidth=4, color='black')


axes[0].set_ylabel('Surface Sediment Total Hg [ng/g]')
axes[0].set_ylim([0, 20000])
axes[0].set_title('Surface Sediment Total Hg')
axes[0].legend()

obs_MeHgSed.plot.scatter('RiverKM', 'Sample_NG/G', ax=axes[1], c='Month',
                              label='All Samples', vmin=1, vmax=12, cmap=shifted_hsv)

# obs_MeHgSed_Avg.plot.scatter('RiverKM_mean', 'Mean_NG/G', ax=axes[1], color='magenta', label='Site Averaged')
axes[1].set_ylabel('Surface Sediment Total MeHg [ng/g]')

# obs_MeHgSed_Index = obs_MeHgSed.copy()
# obs_MeHgSed_Index['RiverKM2'] = obs_MeHgSed_Index['RiverKM']
# obs_MeHgSed_Index = obs_MeHgSed_Index.set_index('RiverKM')
#
# binwidths = [10, 100, 1000, 10000]
# for binWidthIDX, binWidth in enumerate(binwidths):
#     average_in_bins(obs_MeHgSed_Index, binWidth).dropna(subset=['Sample_NG/G']).\
#         plot.line('RiverKM2', 'Sample_NG/G', ax=axes[1],
#                   label=str(binWidth) + ' m binned Samples')

# Find and plot Lowess regression
# lowess = sm.nonparametric.lowess(obs_MeHgSed['Sample_NG/G'], obs_MeHgSed['RiverKM'],
#                                  frac=0.2, xvals=interp_xvals, it=2)
# axes[1].plot(interp_xvals, lowess, label='Lowess regression', linewidth=4)

# Find and plot Lowess regression for each season
for seasonIDX, season in enumerate(SeasonMonths):
    # Skip Winter
    if seasonIDX == 0:
        continue

    seasonMask = obs_MeHgSed['DateTime'].dt.month.isin(season)
    lowess = sm.nonparametric.lowess(obs_MeHgSed.loc[seasonMask, 'Sample_NG/G'],
                                     obs_MeHgSed.loc[seasonMask, 'RiverKM'],
                                     frac=0.5, xvals=interp_xvals, it=2)
    axes[1].plot(interp_xvals, lowess, label='Lowess:' + SeasonNames[seasonIDX],
                 linewidth=4, color=SeasonColors[seasonIDX])

# axes[0].set_ylim([0, 20000])
axes[1].set_title('Surface Sediment Total MeHg')
axes[1].legend()


obs_HgWater.plot.scatter('RiverKM', 'Sample_NG/L', ax=axes[2], c='Month',
                         label='All Samples', vmin=1, vmax=12,
                         cmap=shifted_hsv)
# Find and plot Lowess regression
# lowess = sm.nonparametric.lowess(obs_HgWater['Sample_NG/L'], obs_HgWater['RiverKM'],
#                                  frac=0.2, xvals=interp_xvals, it=2)
# axes[2].plot(interp_xvals, lowess, label='Lowess regression', linewidth=4)
# Find and plot Lowess regression for each season
for seasonIDX, season in enumerate(SeasonMonths):
    # Skip Winter
    if seasonIDX == 0:
        continue

    seasonMask = obs_HgWater['DateTime'].dt.month.isin(season)
    lowess = sm.nonparametric.lowess(obs_HgWater.loc[seasonMask, 'Sample_NG/L'],
                                     obs_HgWater.loc[seasonMask, 'RiverKM'],
                                     frac=0.4, xvals=interp_xvals, it=2)
    axes[2].plot(interp_xvals, lowess, label='Lowess:' + SeasonNames[seasonIDX],
                 linewidth=4, color=SeasonColors[seasonIDX])

axes[2].set_ylabel('Water Total Hg [ng/L]')
axes[2].set_ylim([0, 50])
axes[2].set_title('Water Total Hg')
axes[2].legend()

obs_MeHgWater.plot.scatter('RiverKM', 'Sample_NG/L', ax=axes[3], c='Month',
                           label='All Samples', vmin=1, vmax=12, cmap=shifted_hsv)

# Find and plot Lowess regression
# lowess = sm.nonparametric.lowess(obs_MeHgWater['Sample_NG/L'], obs_MeHgWater['RiverKM'],
#                                  frac=0.2, xvals=interp_xvals, it=2)
# axes[3].plot(interp_xvals, lowess, label='Lowess regression', linewidth=4)
for seasonIDX, season in enumerate(SeasonMonths):
    # Skip Winter
    if seasonIDX == 0:
        continue

    seasonMask = obs_MeHgWater['DateTime'].dt.month.isin(season)
    lowess = sm.nonparametric.lowess(obs_MeHgWater.loc[seasonMask, 'Sample_NG/L'],
                                     obs_MeHgWater.loc[seasonMask, 'RiverKM'],
                                     frac=0.25, xvals=interp_xvals, it=2)
    axes[3].plot(interp_xvals, lowess, label='Lowess:' + SeasonNames[seasonIDX],
                 linewidth=4, color=SeasonColors[seasonIDX])

axes[3].set_ylabel('Water Total MeHg [ng/L]')
axes[3].set_ylim([0, 5])
axes[3].set_title('Water Total MeHg')
axes[3].legend()


axes[3].set_xlim([0, 100000])
axes[3].set_xlabel('Distance Along River [m]')


plt.show()
fig.savefig("//srv-ott3.baird.com/Projects/12828.101 English Wabigoon River/05_Analyses/02_Data Analysis/Figures/RiverDataProfiles_Lowess_Months_RevB.png",
            bbox_inches='tight', dpi=200)


#%% Plot Profiles
fig, axes = plt.subplots(nrows=1, ncols=1, figsize=(7, 7))
obs_HgSed.plot.scatter('RiverKM', 'TopDepth_x', 12, 'Sample_NG/G', ax=axes, vmax=20000)

axes.invert_yaxis()
axes.set_xlabel('Distance Along River [m]')
axes.set_ylabel('Sample Depth [m]')
axes.set_xlim([0, 100000])

plt.show()

fig.savefig("//srv-ott3.baird.com/Projects/12828.101 English Wabigoon River/05_Analyses/02_Data Analysis/Figures/HgDepthMonth.png",
            bbox_inches='tight', dpi=200)


#%% Interpolate lowess to grid for hg
# Read in model points and roughness
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)


gridInterpNearest = griddata(np.array([obs_HgSurfaceSed.geometry.x, obs_HgSurfaceSed.geometry.y]).T,
                      obs_HgSurfaceSed['Sample_NG/G'], np.array([bathy_xyz.x, bathy_xyz.y]).T,
                      method='nearest')

gridInterpOut = np.vstack((bathy_xyz.x, bathy_xyz.y, gridInterpNearest))
gridInterpOut = np.transpose(gridInterpOut)

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

gridInterpOut[RoughInterp == 0.05, 2] = 0

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

# Linear interpolation
gridInterpNearest = griddata(np.array([obs_HgSurfaceSed.geometry.x, obs_HgSurfaceSed.geometry.y]).T,
                      obs_HgSurfaceSed['Sample_NG/G'], np.array([bathy_xyz.x, bathy_xyz.y]).T,
                      method='linear')

gridInterpOut = np.vstack((bathy_xyz.x, bathy_xyz.y, gridInterpNearest))
gridInterpOut = np.transpose(gridInterpOut)

# Set Wetlands to zero
gridInterpOut[RoughInterp == 0.05, 2] = 0

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

# Find and plot Lowess regression
lowess = sm.nonparametric.lowess(obs_HgSurfaceSed['Sample_NG/G'], obs_HgSurfaceSed['RiverKM'],
                                 frac=0.3, it=2)

gridInterpLowess= griddata(np.array([obs_HgSurfaceSed.geometry.x, obs_HgSurfaceSed.geometry.y]).T,
                      lowess[:, 1], np.array([bathy_xyz.x, bathy_xyz.y]).T,
                      method='nearest')

gridInterpOut = np.vstack((bathy_xyz.x, bathy_xyz.y, gridInterpLowess))
gridInterpOut = np.transpose(gridInterpOut)

# Set Wetlands to zero
gridInterpOut[RoughInterp == 0.05, 2] = 0

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

#%% Interpolate Hg to grid

# Read in model points and roughness
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)

# Loop through parameters
for i in range(0, 5):
    if i == 0:
        dataOUT  = obs_HgSed_Avg['Sample_NG/G|mean']
        dataOutx = obs_HgSed_Avg.geometry.x
        dataOuty = obs_HgSed_Avg.geometry.y
        dateFileName = 'HgSed_meanDB'
    elif i == 1:
        dataOUT  = obs_MeHgSed_Avg['Sample_NG/G|mean']
        dataOutx = obs_MeHgSed_Avg.geometry.x
        dataOuty = obs_MeHgSed_Avg.geometry.y
        dateFileName = 'MeHgSed_meanDB'
    elif i == 2:
        dataOUT  = obs_HgWater['Sample_NG/L']
        dataOutx = obs_HgWater.geometry.x
        dataOuty = obs_HgWater.geometry.y
        dateFileName = 'HgWater_allDB'
    elif i == 3:
        dataOUT  = obs_MeHgWater['Sample_NG/L']
        dataOutx = obs_MeHgWater.geometry.x
        dataOuty = obs_MeHgWater.geometry.y
        dateFileName = 'MeHgWater_allDB'
    elif i == 4:
        # Synthetic from Reed's Sheet
        dataOUT  = ((10 - 1) * np.exp(-0.08 * np.array(river_centerline_merge_gpd2.loc[:, 'RiverKM']) / 1000) + 1) * 1000
        dataOutx = river_centerline_merge_gpd2.geometry.x
        dataOuty = river_centerline_merge_gpd2.geometry.y
        dateFileName = 'ReedHgSed'

    gridInterp = griddata(np.array([dataOutx, dataOuty]).T, dataOUT, np.array([bathy_xyz.x, bathy_xyz.y]).T,
                                        method='nearest')

    gridInterpOut = np.vstack((bathy_xyz.x, bathy_xyz.y, gridInterp))
    gridInterpOut = np.transpose(gridInterpOut)

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

    gridInterpOut[RoughInterp==0.05, 2] = 0

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