"""generate_recession_tables.py

Generate ZSA and ZRFC recession tables (storm demand vs. chainage) based on
Nielsen et al. (1992) for lidar-derived beach profiles.

Reads:
    Profiles 1 to 12 2019 DEM.xlsx

Writes:
    recession_results_zsa.csv
    recession_results_zrfc.csv

D Howe
d.howe@wrl.unsw.edu.au
2022-05-06
"""

import os
import re
import json
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from nielsen import Gridder, BeachProfile

BEACH = 'Roches Beach'
COLUMNS = ['Chainage', 'Easting', 'Northing', 'Elevation']  # Workbook headers
VOLUME = 200  # Storm demand volume for plotting (m^3/m)
FIGURE_DIR = 'png'  # Save profile plots here

# Define block and profile IDs
with open('settings.json', 'r') as f:
    DATA = json.loads(f.read())

print('.', end='', flush=True)  # Show progress

# Load data
xlsx_path = 'Profiles 1 to 12 2019 DEM.xlsx'
workbook = pd.ExcelFile(xlsx_path)
sheets = [s for s in workbook.sheet_names if s[0].isdigit()][::-1]

# Create output dataframs
zsa = pd.DataFrame(index=[[], [], []])
zrfc = pd.DataFrame(index=[[], [], []])
zsa.index = zsa.index.set_names(['beach', 'block', 'profile'])
zrfc.index = zrfc.index.set_names(['beach', 'block', 'profile'])

for s in sheets:
    print('.', end='', flush=True)  # Show progress

    # Get block and profile IDs
    block = DATA[s].pop('block')
    profile = DATA[s].pop('profile')

    # Get additional keyword arguments for Nielsen calculations
    kwargs = DATA[s]

    # Read current sheet in workbook
    df = workbook.parse(s, names=COLUMNS)

    # Extract profile coordinates
    g = Gridder(chainage=df['Chainage'],
                elevation=df['Elevation'],
                easting=df['Easting'],
                northing=df['Northing'],
                **kwargs)

    # Update recession tables
    zsa.loc[(BEACH, block, profile), g.volume] = g.chainage['zsa']
    zrfc.loc[(BEACH, block, profile), g.volume] = g.chainage['zrfc']

    # Plot profiles
    fig, ax = plt.subplots(1, 1, figsize=(12, 4))

    p = BeachProfile(df['Chainage'], df['Elevation'])
    p.plot(v=VOLUME, ax=ax, title=f'Roches Beach, P{s}')

    ax.spines['top'].set_visible(False)
    ax.spines['right'].set_visible(False)

    png_path = os.path.join(FIGURE_DIR, f'P{s}')
    plt.savefig(png_path, bbox_inches='tight', dpi=300)

# Write tables
zsa.to_csv('recession_results_zsa.csv', float_format='%0.3f')
zrfc.to_csv('recession_results_zrfc.csv', float_format='%0.3f')

print('\nFinished.')