"""
@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

import pytz

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

# modelPlot  = [20, 22]
# modelPlot  = [18, 20, 21, 22]
# modelPlot  = [18, 21]
modelPlot = [146]

d3dfm_DataArray = [None] * (max(modelPlot)+1)
modelVelX  = [None] * (max(modelPlot)+1)
modelVelY  = [None] * (max(modelPlot)+1)
tSteps        = [None] * (max(modelPlot)+1)

# Time steps to import
tStepsImport = range(1000, 7000)

for i in modelPlot:
    # Create variables
    modelVelX[i] = [None] * (max(tStepsImport)+1)
    modelVelY[i] = [None] * (max(tStepsImport)+1)

    file_nc_map = Path(runLog['Run Location'][i]) / dataPath

    d3dfm_DataArray[i] = dfmt.open_partitioned_dataset(file_nc_map.as_posix())

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

    # Get time steps
    tSteps[i] = d3dfm_DataArray[i].time.values

    # Loop through time steps and import data
    for t in tStepsImport:
        # Import 3D velocity data
        modelVelXin = d3dfm_DataArray[i].mesh2d_ucx[t, :, :].compute()
        modelVelYin = d3dfm_DataArray[i].mesh2d_ucy[t, :, :].compute()

        # Average over layers
        modelVelX[i][t]  = modelVelXin.mean(axis=1)
        modelVelY[i][t]  = modelVelYin.mean(axis=1)

        print(t)

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

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

    # Open hist file dataset
    data_xr = xr.open_mfdataset(file_nc_hist, preprocess=dfmt.preprocess_hisnc)

    # Get stations
    stations_pd = data_xr['stations'].to_dataframe()

    # Convert to Pandas
    modelHistWL[i] = data_xr.waterlevel.sel(stations=['West_WL', 'East_WL']).to_pandas()

    # Put in UTC if needed
    if i == 104 or i == 105 or i >= 135:
        modelHistWL[i] = modelHistWL[i].tz_localize(
            pytz.timezone('EST')).tz_convert(pytz.utc)
    else:
        modelHistWL[i] = modelHistWL[i].tz_localize(
            pytz.timezone('UTC'))


#%% Velocity Quiver animations
plt.rcParams['animation.ffmpeg_path'] = \
    'C:/Users/arey/files/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 = [84]
tStepsPlot = tStepsImport
tStepsPlot = range(970, 971)
tStepsPlot = range(955, 990)


vmin        = 0   # Legend min
vmax        = 0.05 # Legend max
scale       = 30  # Size of arrows, smaller is bigger #25
qmax        = 0.002# Scaling threshold. Values above this are set to 1
amin        = 0.1   # Minimum arrow length. A dot is used below this value

for m in modelPlot:
    # Find water level plot range
    wlPlotMinIDX = wlHistTimeIDX = np.argmin(abs(tSteps[m][tStepsPlot[0]] - modelHistWL[m].index.values))
    wlPlotMaxIDX = wlHistTimeIDX = np.argmin(abs(tSteps[m][tStepsPlot[-1]] - modelHistWL[m].index.values))

    cmapName = 'turbo'
    cmap = mpl.cm.turbo

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

    # Corner limits
    xLimits = [306088, 306444]
    yLimits = [61370, 61741]

    # Setup Video
    metadata = dict(title='NTC Velocity', artist='Matplotlib',
                    comment='AJMR June 2, 2023')
    writer = animation.FFMpegWriter(fps=1, metadata=metadata, codec='h264', bitrate=5000)
    fig, axes = plt.subplots(figsize=(10, 10))

    writer.setup(fig, 'C:/Users/arey/files/Projects/Newtown/Vel_Animation4.mp4')

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

    frameCount = 1

    for t in tStepsPlot: #289
        # Clear axes for video
        axes.cla()

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

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


        # Cropt to zoom limits
        croppedGridMask = (d3dfm_DataArray[m].mesh2d_face_x >= xLimits[0]) & \
                          (d3dfm_DataArray[m].mesh2d_face_x <= xLimits[1]) & \
                          (d3dfm_DataArray[m].mesh2d_face_y >= yLimits[0]) & \
                          (d3dfm_DataArray[m].mesh2d_face_y <= yLimits[1])
        croppedGridMask = croppedGridMask.compute()

        # Rasterize Ugrid
        rasterU = dfmt.rasterize_ugrid(d3dfm_DataArray[m]['mesh2d_ucx'].isel(
            time=t, mesh2d_nLayers=4).where(croppedGridMask, drop=True).to_dataset(), resolution=10)
        rasterV = dfmt.rasterize_ugrid(d3dfm_DataArray[m]['mesh2d_ucy'].isel(
            time=t, mesh2d_nLayers=4).where(croppedGridMask, drop=True).to_dataset(), resolution=10)

        # Using fixed grid-> select desiered time step/layer/variable, convert to dataset, rasterize, then to numpy
        overGrid_U_Scale = rasterU.to_array().to_numpy().squeeze()
        overGrid_V_Scale = rasterV.to_array().to_numpy().squeeze()

        r = np.power(np.add(np.power(overGrid_U_Scale, 2), np.power(overGrid_V_Scale, 2)), 0.5)

        curPlotMask = np.sqrt((overGrid_U_Scale) ** 2 + (overGrid_V_Scale) ** 2) > qmax

        overGrid_U_Plot = np.zeros(overGrid_U_Scale.shape)
        overGrid_V_Plot = np.zeros(overGrid_V_Scale.shape)

        # Normalize arrows to 1 where not zero
        overGrid_U_Plot[r != 0] = overGrid_U_Scale[r != 0] / r[r != 0]
        overGrid_V_Plot[r != 0] = overGrid_V_Scale[r != 0] / r[r != 0]

        # Scale arrows above qmax
        overGrid_U_Plot[~curPlotMask] = overGrid_U_Plot[~curPlotMask] * (r[~curPlotMask] / qmax)
        overGrid_V_Plot[~curPlotMask] = overGrid_V_Plot[~curPlotMask] * (r[~curPlotMask] / qmax)

        # Calculate velocity magnitude
        velMag = np.sqrt((overGrid_U_Scale) ** 2 + (overGrid_V_Scale) ** 2)

        # Using fixed grid
        qv = axes.quiver(rasterU['x'].values, rasterU['y'].values, overGrid_U_Plot, overGrid_V_Plot, velMag, scale=scale,
                          color='k', headlength=5, headwidth=5, headaxislength=5, width=0.002, minlength=amin) #200

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

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

        # Set colormap for quiver
        qv.set_cmap(cmap)

        # Add title
        axes.title.set_text('Mid-depth velocity: ' + np.datetime_as_string(tSteps[m][t], unit='m'))

        # Add inset wl plot
        axin1 = axes.inset_axes([0.1, 0.1, 0.4, 0.15])

        # Find matching time step in model water level history
        wlHistTimeIDX = np.argmin(abs(tSteps[m][t] - modelHistWL[m].index.values))

        # Plot model water level history
        axin1.plot(modelHistWL[m].index, modelHistWL[m])
        # Add current point
        axin1.scatter(modelHistWL[m].index[wlHistTimeIDX],
                      modelHistWL[m]['East_WL'][wlHistTimeIDX], 10, 'r')

        axin1.set_xticks([])
        axin1.set_yticks([])

        # Set time axis limits
        axin1.set_xlim([modelHistWL[m].index[wlPlotMinIDX],
                        modelHistWL[m].index[wlPlotMaxIDX]])
        # Set inset axis y limits
        axin1.set_ylim([-0.5, 2])


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

        if frameCount == 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 without designated tick marks
            cb1 = mpl.colorbar.ColorbarBase(cax, cmap=cmap,
                                            norm=norm,
                                            orientation='vertical')

            # Colorbar lavel
            cb1.set_label('Water velocity [m/s]')

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

        # Increase frame count
        frameCount += 1

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



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