"""
Script to read and plot EFDC sediment conditions
Alexander Rey, April 3, 2022
"""
import pandas as pd
import geopandas as gp
import numpy as np
import matplotlib.pyplot as plt
import cartopy.crs as ccrs
import contextily as ctx

# %% Read in EFDC grid
efdc_grid_df = pd.read_csv('//ott-athena.baird.com/E/INPUT_FILES/GRID_FILES/NC_ER_lxly_20140129.inp',
                           delim_whitespace=True, skiprows=4,
                           names=['I', 'J', 'X', 'Y', 'CUE', 'CVE', 'CUN', 'CVN'])

efdc_grid_gdf = gp.GeoDataFrame(
    efdc_grid_df, geometry=gp.points_from_xy(efdc_grid_df.X, efdc_grid_df.Y), crs="EPSG:32118")


ncols = efdc_grid_df.I.max()
nrows = efdc_grid_df.J.max()

# Assign grid cell x values to a numpy grid
efdc_grid_x = np.zeros((nrows+1, ncols+1))
efdc_grid_y = np.zeros((nrows+1, ncols+1))
efdc_grid_x[efdc_grid_df.J.values-1, efdc_grid_df.I.values-1] = efdc_grid_df.X.values
efdc_grid_y[efdc_grid_df.J.values-1, efdc_grid_df.I.values-1] = efdc_grid_df.Y.values


# %% Read in EFDC sediment conditions

# Dry density of each layer
efdc_drydens_df = pd.read_csv('//ott-athena.baird.com/E/INPUT_FILES/DRY_DENSITY_FILES' +
                            '/NC_bed_drydensity_5Layers_20160601',
                            delim_whitespace=True, header=None)
# D90 of skin fraction
efdc_bed90_df = pd.read_csv('//ott-athena.baird.com/E/INPUT_FILES/BED_D90_FILES' +
                            '/NC_BED_D90_20160311_all_1400um',
                            delim_whitespace=True, header=None)

ncols = efdc_grid_df.I.max()
nrows = efdc_grid_df.J.max()

# Flatten the dry density dataframe
drydens_flat = efdc_drydens_df.values.flatten()
drydens_flat = drydens_flat[~np.isnan(drydens_flat)]

efdc_drydens_shaped = drydens_flat.reshape(nrows+1, ncols+1, 5)

# Flatten the bed dataframe
efdc_bed90_flat = efdc_bed90_df.values.flatten()
efdc_bed90_flat = efdc_bed90_flat[~np.isnan(efdc_bed90_flat)]
efdc_bded90_shaped = efdc_bed90_flat.reshape(ncols+1, nrows+1).T

# %% Read in, plot, and save EFDC Fractions
# 5 layers, 7 ftactions
# From "\\ott-athena.baird.com\E\SIMULATION_TEMPLATES\Sedtran\NC_sedtran_7C_2010-01_nonpropwash_1999_2012\Yr1999\spin\run_month.bat"

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

# sedNames = ['bed_content_c1_inorg', 'bed_content_c2_inorg', 'bed_content_c3_inorg',
#             'bed_content_c4_inorg', 'bed_content_c5_inorg', 'bed_content_c6_inorg',
#             'bed_content_c7_inorg']
sedNames = ['bed_content_c5_inorg', 'bed_content_c6_inorg',
            'bed_content_c7_inorg']

# fracNames = ['NC_bed_zero_5layers_20160819', 'NC_bed_zero_5layers_20160819', 'NC_bed_zero_5layers_20160819',
#             'NC_bed_zero_5layers_20160819', 'NC_bed_f1A_5layers_20160819', 'NC_bed_f2_5layers_20160819',
#             'NC_bed_f3_5layers_20160819']
fracNames = ['NC_bed_f1A_5layers_20160819', 'NC_bed_f2_5layers_20160819',
            'NC_bed_f3_5layers_20160819']

for sedIDX, sed in enumerate(sedNames):
    inputFormat = np.ones(20, dtype=int) * 6
    efdc_frac_df = pd.read_fwf('//ott-athena.baird.com/E/INPUT_FILES/BED_FRACTION_FILES' +
                                '/20160819/' + fracNames[sedIDX],
                                widths=inputFormat.tolist(), header=None)

    ncols = efdc_grid_df.I.max()
    nrows = efdc_grid_df.J.max()

    # Flatten the fractions dataframe
    efdc_frac_flat = efdc_frac_df.values.flatten()
    efdc_frac_flat = efdc_frac_flat[~np.isnan(efdc_frac_flat)]
    efdc_frac_shaped = efdc_frac_flat.reshape((nrows+1, ncols+1, 5), order='F')

    fig, axes = plt.subplots(nrows=2, ncols=3, subplot_kw={'projection': ccrs.epsg(32118)}, figsize=(16, 8))
    flat_ax = axes.flatten()
    for sedLayer in range(0, 5):

        # Map Setup
        flat_ax[sedLayer].set_xlim(297579, 309365)
        flat_ax[sedLayer].set_ylim(58211, 67565)
        ctx.add_basemap(flat_ax[sedLayer], source=mapbox, crs='EPSG:32118')


        # Initialize the output file with zeros
        sedLayerOut = np.zeros(((nrows+1) * (ncols+1), 3))
        sedLayerOut[:, 0] = efdc_grid_x.flatten()
        sedLayerOut[:, 1] = efdc_grid_y.flatten()
        sedLayerOut[:, 2] = efdc_frac_shaped[:, :, sedLayer].flatten() / 100

        sedLayersOut_df = pd.DataFrame(sedLayerOut)
        # Remove missing cells
        sedLayersOut_df = sedLayersOut_df.loc[sedLayersOut_df.iloc[:, 0] != 0, :]

        # Save as xyz (space seperated)
        # sedLayersOut_df.to_csv('//srv-ott3.baird.com/Projects/11934.201 Newtown Creek TPP – Privileged and Confidential/' +
        #                         '06_Models/02_EFDC/07_Sediment/Converted_Files/' + sed +
        #                          '_Layer_' + str(sedLayer) + '.xyz', index=False, sep=' ', header=False)

        # Plot
        flat_ax[sedLayer].scatter(sedLayerOut[:, 0], sedLayerOut[:, 1], c=sedLayerOut[:, 2],
                                  vmin=0, vmax=1, s=5)
        flat_ax[sedLayer].set_title('Layer ' + str(sedLayer))

    plt.show()
    plt.title(sed)