"""
Script to read and save EFDC porpwash data
Alexander Rey, April 5, 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

import datetime as dt

# %% 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 ship data (from input_sedtran_sedzlj)
efdc_ships_df = pd.read_csv('//ott-athena.baird.com/E/INPUT_FILES/INPUT_SEDTRAN_FILES/' +
                           '20200825_15C/NC_input_sedtran_sedzlj_15C_20200930_hot.inp',
                           delim_whitespace=True, skiprows=141,
                           names=['SHIP_PROP_DEPTH', 'SHIP_PROP_DIAM', 'SHIP_PROP_HUB_DIAM', 'SHIP_POWER_MAX'])

# Delete all strings
efdc_ships_df = efdc_ships_df.apply(pd.to_numeric, errors='coerce')

# Remove empty rows
efdc_ships_df = efdc_ships_df[efdc_ships_df['SHIP_PROP_DIAM'].notna()]
efdc_ships_df = efdc_ships_df.iloc[:-1, :]

# Reindex to count ships
efdc_ships_df = efdc_ships_df.reset_index(drop=False)

# %% Read in EFDC ship tracks
# Loop through years and months to read in all files
for year in range(1999, 2016):
    for month in range(1, 13):
        # Read in file
        efdc_tracks_df = pd.read_csv('//ott-athena.baird.com/E/INPUT_FILES/PROPWASH_INPUT_FILES/' +
                                     '20201026_1999_to_2015/YR' + str(year) + '/' +
                                     'ntc_ship_propwash_input_20201026_yr_' + str(year) + '_mon_' + str(month) + '.inp',
                                     delim_whitespace=True, skiprows=1,
                                     names=['ShipID', 'I', 'J', 'Speed', 'Power'])

        # Loop through and create new columns with ship details
        efdc_tracks_data = np.zeros((len(efdc_tracks_df)-sum(efdc_tracks_df['J'].isna()), 6))

        # If there's a NaN, then it's a row with time data and not a ship track
        # Set time based on nan rows, use that for the next ship tracks
        trackCount = 0
        for trackIDX in range(0, len(efdc_tracks_df)):
            if pd.isna(efdc_tracks_df['J'][trackIDX]):
                trackTime = dt.datetime(year=year, month=month, day=1) + dt.timedelta(seconds=float(efdc_tracks_df['ShipID'][trackIDX]))
            else:
                efdc_tracks_data[trackCount, 0]  = trackTime.timestamp()
                efdc_tracks_data[trackCount, 1:] = efdc_tracks_df.iloc[trackIDX, :]
                trackCount += 1

        # Add to dataframe
        efdc_tracks_data_df = pd.DataFrame(efdc_tracks_data, columns=['Time', 'ShipID', 'I', 'J', 'Speed', 'Power'])
        efdc_tracks_data_df.Time = pd.to_datetime(efdc_tracks_data_df.Time, unit='s')

        # Set shipID to start from 0
        efdc_tracks_data_df.ShipID = efdc_tracks_data_df.ShipID - 1

        # Merge the dataframes

        # Merge ship details
        efdc_tracks_data_df = efdc_tracks_data_df.join(efdc_ships_df, on='ShipID', how='left')

        # Merge with grid
        efdc_tracks_data_df = pd.merge(efdc_tracks_data_df, efdc_grid_df,  how='left', left_on=['I', 'J'],
                            right_on = ['I', 'J']).drop(columns=[
                            'I', 'J', 'X', 'Y', 'CUE', 'CVE', 'CUN', 'CVN'])

        # Save as csv
        efdc_tracks_data_df.to_csv('//srv-ott3.baird.com/Projects/11934.201 Newtown Creek TPP – Privileged and Confidential/' +
                                 '06_Models/02_EFDC/08_Propwash/Converted_Files/Tracks_' + str(year) +
                                  '_' + str(month) + '.csv', index=False, sep=',', header=True)
        # Save as shapefile
        efdc_tracks_data_gdf = gp.GeoDataFrame(
            efdc_tracks_data_df, geometry=efdc_tracks_data_df.geometry, crs="EPSG:32118")

        # Fix datetime
        efdc_tracks_data_gdf.Time = efdc_tracks_data_gdf.Time.dt.strftime('%Y-%m-%d %H:%M:%S')
        efdc_tracks_data_gdf.to_file('//srv-ott3.baird.com/Projects/11934.201 Newtown Creek TPP – Privileged and Confidential/' +
                                 '06_Models/02_EFDC/08_Propwash/Converted_Files/Shapefiles/Tracks_' + str(year) +
                                  '_' + str(month) + '.shp')