#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Fri Sep 17 18:47:06 2021

@author: Camaron.George
"""
import os
import glob
from datetime import datetime

import numpy as np
import netCDF4 as nc

from pytides.tide import Tide
import pytides.constituent as con

from scipy.interpolate import griddata

# path = '/scratch2/NCEPDEV/ohd/Camaron.George/'
# grid = 'schismRuns/EastGulf/hgrid.gr3'
# out = 'schismRuns/EastGulf/tides.nc'
# startPred = datetime(2020,9,2,13)#time zero for the prediction
# hours = np.arange(60.0)#number of hours in the prediction
# predTimes = Tide._times(startPred,hours)


def generate_tidal_levels(params_path, grid_file, output_file, start_time, total_hours):
    lon = []
    lat = []
    elems = []
    bnodes = []
    with open(grid_file) as f:
        next(f)
        line = f.readline()
        ne = int(line.split()[0])
        nn = int(line.split()[1])
        for i in range(nn):
            line = f.readline()
            lon.append(float(line.split()[1]))
            lat.append(float(line.split()[2]))
        for i in range(ne):
            line = f.readline()
            elems.append(line)
        next(f)
        next(f)
        line = f.readline()
        nb = int(line.split()[0])
        for j in range(nb):
            bnodes.append(int(f.readline()))

    lo = [lon[b-1] for b in bnodes]
    la = [lat[b-1] for b in bnodes]

    consfiles = glob.glob(os.path.join(params_path, 'TidalConst/*'))
    consfiles.sort()  # 'k1','k2','m2','n2','o1','p1','q1','s2'
    pha = np.zeros((len(lo), len(consfiles)))
    amp = np.zeros((len(lo), len(consfiles)))

    col = 0
    for file in consfiles:
        data = nc.Dataset(file, 'r')
        if col == 0:
            x = data.variables['lon'][:]
            y = data.variables['lat'][:]
            x, y = np.meshgrid(x, y)
            i = np.where(x > 180.0)
            x[i] = x[i]-360.0
            i = np.where((x < max(lo)+1) & (x > min(lo)-1) & (y < max(la)+1) & (y > min(la)-1))
            x = x[i]
            y = y[i]

        a = data.variables['amplitude'][:]
        a = a[i]
        p = data.variables['phase'][:]
        p = p[i]

        # test = griddata((x,y),a,(lo,la),method='nearest')
        amp[:, col] = griddata((x, y), a, (lo, la), method='nearest')
        pha[:, col] = griddata((x, y), p, (lo, la), method='nearest')
        col += 1

    pred_times = Tide._times(start_time, total_hours)

    wl = np.zeros((len(pred_times), amp.shape[0]))
    cons = [c for c in con.noaa if c != con._Z0]
    n = [3, 34, 0, 2, 5, 29, 25, 1]  # corresponds to position in list in constituent.py
    cons = [cons[c] for c in n]
    model = np.zeros(len(cons), dtype=Tide.dtype)
    model['constituent'] = cons
    for i in range(amp.shape[0]):
        model['amplitude'] = amp[i, :]
        model['phase'] = pha[i, :]

        tide = Tide(model=model, radians=False)
        wl[:, i] = (tide.at(pred_times))/100.0

    if os.path.exists(output_file):
        print(f"Removing existing tidal level file {output_file}")
        os.remove(output_file)

    # open a netCDF file to write
    ncout = nc.Dataset(output_file, 'w', format='NETCDF4')

    # define axis size
    ncout.createDimension('time', None)
    ncout.createDimension('nOpenBndNodes', amp.shape[0])

    # create time axis
    nctime = ncout.createVariable('time', 'f8', ('time',))

    # create water level time series
    ncwl = ncout.createVariable('time_series', 'f8', ('time', 'nOpenBndNodes',))

    # copy axis from original dataset
    nctime[:] = np.arange(651600.0, 865000, 3600.0)
    ncwl[:] = wl

    ncout.close()


if __name__ == '__main__':
    params_path = "/glade/work/rcabell/ecflow/hydro-workflow/coastal/Tides"
    grid_file = "/glade/work/rcabell/ecflow/hydro-workflow/coastal/prvi/hgrid.gr3"
    output_file = "/glade/u/home/rcabell/work/ecflow/hydro-workflow/jobdir/prvi/medium_range_mem1/coastal/prvi/tides.nc"
    start_time = datetime(2020, 9, 9, 0, 0)
    total_hours = range(54)
    generate_tidal_levels(params_path, grid_file, output_file, start_time, total_hours)