"""
@author: AJMR

Plotting and animation script for NTC sediment tests

April 17, 2023
"""

import os
import pandas as pd
import geopandas as gpd
import netCDF4 as nc
import math
import numpy as np

import matplotlib as mpl
import matplotlib.pyplot as plt
from scipy.interpolate import LinearNDInterpolator, interp1d
import matplotlib.animation as animation

import dfm_tools as dfmt

# Geotiff
import cartopy.crs as ccrs
import contextily as ctx
import pathlib as pl

from shapely.geometry import Point, MultiPoint
import alphashape

import rioxarray
import xarray as xr


import cartopy.crs as ccrs
import contextily as ctx
from dfm_tools.regulargrid import scatter_to_regulargrid

from pathlib import Path

# Colormaps
import cm_xml_to_matplotlib as cm

#%% Read Model Log

runLog = pd.read_excel('C:/Users/arey/files/Projects/Newtown/Model Log NTC.xlsx', 'ModelLog')

dataPath = "FlowFM/dflowfm/output/FlowFM_map.nc"
histPath = "FlowFM/dflowfm/output/FlowFM_his.nc"

#%% Import Model results for static images

# Define which models to import
modelPlot = [144]

d3dfm_DataArray = [None] * (max(modelPlot)+1)
sedMass         = [None] * (max(modelPlot)+1)  # Available Mass of Sediment
erodeSed        = [None] * (max(modelPlot)+1)  # Erosion and deposition
initialBed      = [None] * (max(modelPlot)+1)  # Erosion and deposition
modelCRS        = [None] * (max(modelPlot)+1)  # Model CRS
sedConc         = [None] * (max(modelPlot)+1)  # Model CRS

for i in modelPlot:
    file_nc_map = Path(runLog['Run Location'][i]) / dataPath

    # Open Map file as xarrray Dataset
    d3dfm_DataArray[i] = dfmt.open_partitioned_dataset(file_nc_map.as_posix())

    # Get Var info
    vars_pd = dfmt.get_ncvarproperties(d3dfm_DataArray[i])

    # Get last timestep
    tStep = d3dfm_DataArray[i].time[-1].values

    # Import sand sediment class
    sedMass[i]  = d3dfm_DataArray[i].mesh2d_bodsed[-1, :, 2].compute()
    initialBed[i]= d3dfm_DataArray[i].mesh2d_s1[0, :]  - d3dfm_DataArray[i].mesh2d_waterdepth[0, :]
    finalBed    = d3dfm_DataArray[i].mesh2d_s1[-1, :] - d3dfm_DataArray[i].mesh2d_waterdepth[-1, :]

    erodeSed[i] = finalBed - initialBed[i]
    erodeSed[i] = erodeSed[i].compute()
    initialBed[i] = initialBed[i].compute()

    # Import concentration of sediment 2 at timestep 95 and layer 10
    sedConc[i] = d3dfm_DataArray[i].mesh2d_sedfrac_concentration[94, 2, :, 9].compute()

    modelCRS[i] = vars_pd['EPSG_code']['projected_coordinate_system']

#%% Define axis limits
axLim_Xmin = []
axLim_Xmax = []

axLim_Ymin = []
axLim_Ymax = []

# Set axis limits
# Full domain
axLim_Xmin.append(298500)
axLim_Xmax.append(306800)

axLim_Ymin.append(58500)
axLim_Ymax.append(68000)

# Zoom North River
axLim_Xmin.append(303500)
axLim_Xmax.append(305500)

axLim_Ymin.append(66000)
axLim_Ymax.append(68000)

# Zoom Tribs
axLim_Xmin.append(305500)
axLim_Xmax.append(306800)

axLim_Ymin.append(60000)
axLim_Ymax.append(62500)

# Zoom far trib
axLim_Xmin.append(305800)
axLim_Xmax.append(305900)

axLim_Ymin.append(60150)
axLim_Ymax.append(60400)


#%% Plot using DFM functions

for i in modelPlot:
    # Create figure
    fig, axesFig = plt.subplots(1, 4, figsize=(8, 2))

    axes = axesFig.flatten()

    # Loop through axis and plot with different limits
    for subIDX in range(0, 4):

        # # Plot Sediment Mass
        # smp = sedMass[i].ugrid.plot(ax=axes[subIDX], edgecolors='face', cmap='turbo',
        #                       vmin=0, vmax=2000, add_colorbar=False)

        # # Plot Erosion
        # smp = erodeSed[i].ugrid.plot(ax=axes[subIDX], edgecolors='face', cmap='coolwarm',
        #                       vmin=-1, vmax=1, add_colorbar=False)

        # Plot Bed
        smp = initialBed[i].ugrid.plot(ax=axes[subIDX], edgecolors='face', cmap='nipy_spectral',
                              vmin=-25, vmax=0, add_colorbar=False)


        # Set axis labels
        axes[subIDX].set_xlabel('')
        axes[subIDX].set_ylabel('')
        axes[subIDX].set_title('')
        axes[subIDX].set_xticks([])
        axes[subIDX].set_yticks([])

        axes[subIDX].set_xlim(axLim_Xmin[subIDX], axLim_Xmax[subIDX])
        axes[subIDX].set_ylim(axLim_Ymin[subIDX], axLim_Ymax[subIDX])

        # Add basemap
        # ctx.add_basemap(ax=axes, source=ctx.providers.Esri.WorldImagery, crs=modelCRS[i], attribution=False)
        mapbox = 'https://api.mapbox.com/styles/v1/alexander0042/ckemxgtk51fgp19nybfmdcb1e/tiles/256/{z}/{x}/{y}@2x?access_token=pk.eyJ1IjoiYWxleGFuZGVyMDA0MiIsImEiOiJjazVmdG4zbncwMHY4M2VrcThwZGUzZDFhIn0.w6oDHoo1eCeRlSBpwzwVtw'
        ctx.add_basemap(axes[subIDX], source=mapbox, crs=modelCRS[i])


    # cbar = plt.colorbar(smp, ax=axesFig.ravel().tolist(), label='Sand Mass [kg/m2]')
    # cbar = plt.colorbar(smp, ax=axesFig.ravel().tolist(), label='Erosion/ deposition [m]')
    cbar = plt.colorbar(smp, ax=axesFig.ravel().tolist(), label='Initial Bed [mNAVD88]')

    plt.show()

    # fig.savefig('C:/Users/arey/files/Projects/Newtown/SedFigs2/SedMass_' + runLog['Run'][i] + '.png',
    #             bbox_inches='tight', dpi=300)

    # fig.savefig('C:/Users/arey/files/Projects/Newtown/SedFigs2/ErodeDep_' + runLog['Run'][i] + '.png',
    #             bbox_inches='tight', dpi=300)

    # fig.savefig('C:/Users/arey/files/Projects/Newtown/SedFigs2/InitialBed_' + runLog['Run'][i] + '.png',
    #             bbox_inches='tight', dpi=300)

#%% Plot using DFM functions
#Cmap Path
cmap_path = 'C:/Users/arey/Repo/MATLAB_Q/Downloads/KeyColormaps/'

for i in modelPlot:
    fig, axes = plt.subplots(nrows=1, ncols=1, subplot_kw={'projection': ccrs.epsg(32118)}, figsize=(6, 6))
    fig.patch.set_facecolor('white')

    fig.tight_layout(pad=3.0)

    # Load cmaps
    cmap = 'AJMR_Sed5_RevE.xml'
    plotcmap = cm.make_cmap(cmap_path + cmap)

    pc = plot_netmapdata(ugrid_all[i].verts, values=modelMaxShear[i][0, :],
                         ax=axes, linewidth=0.5, cmap=plotcmap)

    axes.set_xlim(305900, 306400)
    axes.set_ylim(61500, 62200)

    # axes.quiver(modelX[i], modelX[i], U[i], V[i], color='w', scale=1.00)

    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:32118')

    extent = axes.get_extent()
    axes.set_xticks(np.linspace(extent[0], extent[1], 4))
    axes.set_yticks(np.linspace(extent[2], extent[3], 4))

    pc.set_clim([0, 0.145])
    cbarTicks = [0, 0.0378, 0.063, 0.0826, 0.11, 0.145]
    cbar = fig.colorbar(pc, ax=axes, shrink=0.95,
                        ticks=cbarTicks)
    cbar.set_label('Total Max Shear Stress [N/m^2]')

    # Add label on colorbar
    cbar.ax.text(0.08, (cbarTicks[0]+cbarTicks[1])/2, 'CLAY', ha='center', va='center',
                 rotation=90, fontweight='bold')
    cbar.ax.text(0.08, (cbarTicks[1]+cbarTicks[2])/2, 'FSILT', ha='center', va='center',
                 rotation=90, fontweight='bold')
    cbar.ax.text(0.08, (cbarTicks[2]+cbarTicks[3])/2, 'MSILT', ha='center', va='center',
                 rotation=90, fontweight='bold')
    cbar.ax.text(0.08, (cbarTicks[3]+cbarTicks[4])/2, 'CSILT', ha='center', va='center',
                 rotation=90, fontweight='bold')
    cbar.ax.text(0.08, (cbarTicks[4]+cbarTicks[5])/2, 'FSAND', ha='center', va='center',
                 rotation=90, fontweight='bold')

    plt.show()

    fig.savefig('C:/Users/arey/files/Projects/Newtown/SedFigs/Shear_Class_' + runLog['Run'][i] + '.png',
                bbox_inches='tight', dpi=300)

#%% Import Model results for animations

# modelPlot  = [20, 22]
# modelPlot  = [18, 20, 21, 22]
# modelPlot  = [18, 21]
modelPlot = [25, 26]

modelSedConc  = [None] * (max(modelPlot)+1)
modelMaxShear_animate  = [None] * (max(modelPlot)+1)
ugrid_all     = [None] * (max(modelPlot)+1)
tSteps        = [None] * (max(modelPlot)+1)

for i in modelPlot:
    file_nc_map = Path(runLog['Run Location'][i]) / dataPath

    # Get Var info
    vars_pd, dims_pd = get_ncvardimlist(file_nc=file_nc_map.as_posix())

    # Import
    tSteps[i] = get_timesfromnc(file_nc=file_nc_map.as_posix(), varname='time')

    ugrid_all[i]        = get_netdata(file_nc=file_nc_map.as_posix())
    modelSedConc[i] = get_ncmodeldata(file_nc=file_nc_map.as_posix(), varname='mesh2d_sedfrac_concentration',
                                          timestep='all', station=0, layer=0) #timestep='all' OR timestep=range(0, 30)
    modelMaxShear_animate[i]     = get_ncmodeldata(file_nc=file_nc_map.as_posix(), varname='mesh2d_tausmax',
                                          timestep='all')


# %% Import water levels from hist
modelHistWL   = [None] * (max(modelPlot)+1)
modelHistTime  = [None] * (max(modelPlot)+1)
# tSteps        = [None] * (max(modelPlot)+1)

for i in modelPlot:
    file_nc_hist = Path(runLog['Run Location'][i]) / histPath

    vars_pd, dims_pd = get_ncvardimlist(file_nc=file_nc_hist.as_posix())

    # Get station names
    histStations = get_hisstationlist(file_nc=file_nc_hist.as_posix(), varname='waterlevel')

    # Import
    modelHistTime[i] = get_timesfromnc(file_nc=file_nc_hist.as_posix(), varname='time')
    modelHistWL[i]   = get_ncmodeldata(file_nc=file_nc_hist.as_posix(), varname='waterlevel',
                                       timestep='all', station=1)


#%% Sediment animations
plt.rcParams['animation.ffmpeg_path'] = \
    'C:/Users/arey/Local/ffmpeg-2022-02-14-git-59c647bcf3-full_build/bin/ffmpeg.exe'

mapbox = 'https://api.mapbox.com/styles/v1/alexander0042/ckex9vtri0o6619p55sl5qiyv/tiles/256/{z}/{x}/{y}@2x?access_token=pk.eyJ1IjoiYWxleGFuZGVyMDA0MiIsImEiOiJjazVmdG4zbncwMHY4M2VrcThwZGUzZDFhIn0.w6oDHoo1eCeRlSBpwzwVtw'
x, y, arrow_length = 0.93, 0.95, 0.12
# fontprops = fm.FontProperties(size=12)

# Set model to plot
# modelPlot = [20, 22]
# modelPlot  = [18, 20, 21, 22]
modelPlot = [25, 26]
cmap_path = 'C:/Users/arey/Repo/MATLAB_Q/Downloads/KeyColormaps/'

for m in modelPlot:
    vmax = 0.75
    vmin = 0
    # vmax = 0.145
    # vmin = 0

    cbarTicks = [0, 0.0378, 0.063, 0.0826, 0.11, 0.145]

    cmapName = 'turbo'
    cmap = mpl.cm.turbo
    # cmapName = 'AJMR_Sed5_RevE.xml'
    # cmap = cm.make_cmap(cmap_path + cmapName)


    # Zoom limits
    # xLimits = [305900, 306500]
    # yLimits = [61400, 62200]

    # Wide Limits
    xLimits = [302719.4, 306934.2]
    yLimits = [60263.7, 64194.8]

    # Setup Video
    metadata = dict(title='NTC Sediment Animation', artist='Matplotlib',
                    comment='AJMR June 28, 2022')
    writer = animation.FFMpegWriter(fps=2, metadata=metadata, codec='h264', bitrate=5000)
    fig, axes = plt.subplots(figsize=(6, 6))

    writer.setup(fig, '//srv-ott3/Projects/11934.201 Newtown Creek TPP – Privileged and Confidential/' +
                 '06_Models/04_Delft3D/Figures/SedConc/HQ3_Wide_RevC_Long_Sed_Cond_'
                 + runLog['Run'][m] + '.mp4')

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

    for i in np.arange(1, len(modelSedConc[m])-1, 1): #289
        # Clear axes for video
        axes.cla()

        # Create new figure for stills
        # fig, axes = plt.subplots(figsize=(6, 6))

        # Clear inset WL Plot
        # if i > 1:
        #     axin1.cla()

        # Plot sediment
        pc = plot_netmapdata(ugrid_all[m].verts, values=modelSedConc[m][i, 0, :, 0],
                             ax=axes, cmap=cmapName,
                             antialiaseds=False)

        # Plot shear
        # pc = plot_netmapdata(ugrid_all[m].verts, values=modelMaxShear_animate[m][i,:],
        #                      ax=axes, cmap=cmap,
        #                      antialiaseds=False)

        # Set map limits
        axes.set_xlim(xLimits[0], xLimits[1])
        axes.set_ylim(yLimits[0], yLimits[1])

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

        # Set map ticks
        axes.set_xticks(np.linspace(xLimits[0], xLimits[1], 4))
        axes.set_yticks(np.linspace(yLimits[0], yLimits[1], 4))

        # Color Limits
        pc.set_clim([vmin, vmax])

        # Add title
        axes.title.set_text('Bottom Sediment Concentration at: ' + tSteps[m][i].strftime("%Y/%m/%d %H:%M"))
        # axes.title.set_text('Total Max Shear Stress at: ' + str(tSteps[m][i]))

        # Add inset wl plot
        axin1 = axes.inset_axes([0.1, 0.1, 0.4, 0.15])
        axin1.plot(modelHistTime[m], modelHistWL[m])
        # Add current point
        # Find closest time
        abs_deltas_from_target_date = np.absolute(modelHistTime[m] - tSteps[m][i])
        axin1.scatter(modelHistTime[m][np.argmin(abs_deltas_from_target_date)], modelHistWL[m][np.argmin(abs_deltas_from_target_date)], 10, 'r')

        # axin1.set_xticks([modelHistTime[m][0], modelHistTime[m][len(modelHistTime[m])-1]])
        axin1.set_xticks([])
        axin1.set_yticks([])

        plt.setp(axin1.get_xticklabels(), backgroundcolor="white")
        plt.setp(axin1.get_yticklabels(), backgroundcolor="white")

        if i == 1: #i < 1000
            fig.subplots_adjust(right=0.8)
            norm = mpl.colors.Normalize(vmin=vmin, vmax=vmax)

            # Create colorbar axis
            # cax = plt.axes([0.82, 0.25, 0.04, 0.5])
            cax = plt.axes([0.85, 0.15, 0.04, 0.7])

            # Add colorbar with designated tick marks
            # cb1 = mpl.colorbar.ColorbarBase(cax, cmap=cmap,
            #                                 norm=norm,
            #                                 orientation='vertical',
            #                                 ticks=cbarTicks)

            # Add colorbar without designated tick marks
            cb1 = mpl.colorbar.ColorbarBase(cax, cmap=cmap,
                                            norm=norm,
                                            orientation='vertical')

            # Colorbar lavel
            cb1.set_label('Sediment Concentration [$kg/m^3$]')
            # cb1.set_label('Total Max Shear Stress [N/m^2]')

            # Add label on colorbar
            # cb1.ax.text(0.08, (cbarTicks[0] + cbarTicks[1]) / 2, 'CLAY', ha='center', va='center',
            #              rotation=90, fontweight='bold')
            # cb1.ax.text(0.08, (cbarTicks[1] + cbarTicks[2]) / 2, 'FSILT', ha='center', va='center',
            #              rotation=90, fontweight='bold')
            # cb1.ax.text(0.08, (cbarTicks[2] + cbarTicks[3]) / 2, 'MSILT', ha='center', va='center',
            #              rotation=90, fontweight='bold')
            # cb1.ax.text(0.08, (cbarTicks[3] + cbarTicks[4]) / 2, 'CSILT', ha='center', va='center',
            #              rotation=90, fontweight='bold')
            # cb1.ax.text(0.08, (cbarTicks[4] + cbarTicks[5]) / 2, 'FSAND', ha='center', va='center',
            #              rotation=90, fontweight='bold')


        # Change label font size for only the primary axis
        for item in ([axes.xaxis.label, axes.yaxis.label] +
                     axes.get_xticklabels() + axes.get_yticklabels()):
            item.set_fontsize(4)


        # Save video frame
        writer.grab_frame()
        plt.show()
        print(i)

        # Save stills
        fig.savefig('//srv-ott3/Projects/11934.201 Newtown Creek TPP – Privileged and Confidential/' +
                    '06_Models/04_Delft3D/Figures/SedConc/Wide_RevC_SedConc_' +
                     runLog['Run'][m] + '_Frame_' + str(i) + '.png',
                    bbox_inches='tight', dpi=300)
        # fig.savefig('//srv-ott3/Projects/11934.201 Newtown Creek TPP – Privileged and Confidential/' +
        #             '06_Models/04_Delft3D/Figures/ShearStress/Wide_ShearStress_' +
        #              runLog['Run'][m] + '_Frame_' + str(i) + '.png',
        #             bbox_inches='tight', dpi=300)

    # Save video
    writer.finish()



#%% Save Sediment Geotiff

# Create mask within bounds
boundSelect = 2
croppedGridMask =  (d3dfm_DataArray[i].mesh2d_face_x >= axLim_Xmin[boundSelect]) &\
                   (d3dfm_DataArray[i].mesh2d_face_x <= axLim_Xmax[boundSelect]) &\
                   (d3dfm_DataArray[i].mesh2d_face_y >= axLim_Ymin[boundSelect]) &\
                   (d3dfm_DataArray[i].mesh2d_face_y <= axLim_Ymax[boundSelect])
croppedGridMask = croppedGridMask.compute()

# Step through time steps
for sedTstep in range(105, 120):
    # Rasterize sediment concentration within bounds
    sedConcTif = dfmt.rasterize_ugrid(d3dfm_DataArray[i]['mesh2d_sedfrac_concentration'].isel(
                    time=sedTstep, mesh2d_nLayers=4, nSedSus=0).where(croppedGridMask, drop=True).to_dataset(),
                                      resolution=2)

    # Add CRS
    sedConcTif.rio.write_crs("epsg:32118", inplace=True)
    sedConcTif.rio.to_raster('C:/Users/arey/files/Projects/Newtown/SedFigs2/GeoTiff_RevB_'
                             + 'Sed_Conc4_Time_' + str(sedTstep) + '.tif')
