"""Update local copies of NOAA WW3 global wave forecast.

D. Howe
2019-12-13
"""
import os
import netCDF4
import pandas as pd
import xarray as xr

LON = 151.25
LAT = -34
VAR_NAMES = ['htsgwsfc', 'perpwsfc', 'dirpwsfc']
INPUT_DIR = 'csv'
TIMEZONE = 'Australia/Sydney'
DT = pd.Timedelta('6h')  # Time between forecasts


def get_ww3_forecast(lon, lat, var_names, t):
    """Get NOAA WaveWatch III (WW3) global wave model forecast.

    Args:
        lon (float): longitude in degrees.
        lat (float): latitude in degrees.
        var_names (list): variable names.
        t (datetime): timezone-aware local datetime.

    Returns:
        df (dataframe):  forecast time series of selected variables.
        filename (str):  name of forecast file (e.g. 'nww320191212_18z').


    Available variables:
        dirpwsfc   primary wave direction [deg]
        dirswsfc   secondary wave direction [deg]
        htsgwsfc   significant height of combined wind waves and swell [m]
        perpwsfc   primary wave mean period [s]
        perswsfc   secondary wave mean period [s]
        ugrdsfc    u-component of wind [m/s]
        vgrdsfc    v-component of wind [m/s]
        wdirsfc    wind direction (from which blowing) [deg]
        windsfc    wind speed [m/s]
        wvdirsfc   direction of wind waves [deg]
        wvpersfc   mean period of wind waves [s]
    """

    base_url = 'https://nomads.ncep.noaa.gov:9090/dods/wave/nww3'

    # Convert to UTC time, and round to nearest 6 hours
    t_utc = t.tz_convert('UTC').floor('6h')

    dirname = t_utc.strftime('nww3%Y%m%d')
    filename = t_utc.strftime('nww3%Y%m%d_%Hz')
    url = f'{base_url}/{dirname}/{filename}'

    # Try to load dataset
    ds = xr.open_dataset(url)

    # Extract variables from dataset
    ds = ds.sel(lon=lon, lat=lat, method='nearest')
    df = ds[var_names].to_dataframe()

    # Drop lat and lon columns
    df = df.drop(columns=['lat', 'lon'])

    # Convert timestamps to local timezone
    df = df.tz_localize('UTC')
    df = df.tz_convert(t.tz)

    return df, filename


# Get current time
t_now = pd.Timestamp.now(tz=TIMEZONE)

for v in VAR_NAMES:
    csv_name = v + '.csv'
    try:
        # Load local copy of variable time series
        master = pd.read_csv(os.path.join(INPUT_DIR, csv_name),
                             index_col=[0, 1],
                             parse_dates=[1])

        master.columns = master.columns.astype(int)

        # Get start time of most recent forecast
        t = master.index.get_level_values(level=1)[-1]

    except FileNotFoundError:
        index = pd.MultiIndex(levels=[[], []],
                              codes=[[], []],
                              names=['filename', 'start_time'])
        master = pd.DataFrame(index=index)

        # Use start time of 5 days hours before present
        t = t_now - pd.Timedelta('5d')

    while t < t_now:
        t += DT
        try:
            # Get next forecast
            df, filename = get_ww3_forecast(LON, LAT, [v], t)
        except OSError:
            # Stop if next forecast is not available
            break

        # Use filename as column name
        df = df.rename(columns={v: filename})

        # Convert timestamps to relative hours from start of model run
        start_time = df.index[0]
        df.index = ((df.index - start_time).total_seconds() / 3600).astype(int)
        df.index.name = None

        # Transpose to run time series along rows
        df.columns = [[filename], [start_time]]
        df = df.T

        # Add to master dataframe
        df.index.names = ['filename', 'start_time']
        master = master.append(df)

    if master.shape[0] > 0:
        master.to_csv(os.path.join(INPUT_DIR, csv_name), float_format='%g')