"""
@author: AJMR

Plotting and animation script for NTC

June 2, 2022
"""

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

from dfm_tools.get_nc import get_netdata, get_ncmodeldata, plot_netmapdata
from dfm_tools.get_nc_helpers import get_timesfromnc, get_ncvardimlist, get_hisstationlist

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', 'Propwash')

dataPath = "FlowFM_map.nc"
histPath = "FlowFM_his.nc"

#%% Import Model results for static images

# Define which models to import
#modelPlot = [3, 5, 6, 7, 8]
#modelPlot  = [2, 4, 5, 6, 7, 8]
# modelPlot  = [18, 20, 21, 22]
modelPlot = [25, 26]

ugrid_all     = [None] * (max(modelPlot)+1)
tSteps        = [None] * (max(modelPlot)+1)
modelX        = [None] * (max(modelPlot)+1)
modelY        = [None] * (max(modelPlot)+1)
modelZ        = [None] * (max(modelPlot)+1)
modelMaxShear = [None] * (max(modelPlot)+1)


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

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

    # Use final timestep
    tStep = len(tSteps[i])-1

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

    # Import
    ugrid_all[i]        = get_netdata(file_nc=file_nc_map.as_posix())
    modelX[i]            = get_ncmodeldata(file_nc=file_nc_map.as_posix(), varname='mesh2d_face_x')
    modelY[i]            = get_ncmodeldata(file_nc=file_nc_map.as_posix(), varname='mesh2d_face_y')
    modelMaxShear[i]     = get_ncmodeldata(file_nc=file_nc_map.as_posix(), varname='mesh2d_tausmax',
                                          timestep=tStep)

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

