#!/usr/bin/env python
# coding: utf-8

# # Developing `make_ww3_wind_file` Worker
# 
# Code experiments and verification for the `make_ww3_wind_file` worker.

# wwatch3 requires the wind forcing file that is input to `ww3_prnc`
# to have a structure similar this ocean currents file:
# 
# ```
# netcdf GSL5km_2009_CUR {
# dimensions:
#     longitude = 150 ;
#     latitude = 236 ;
#     time = UNLIMITED ; // (8760 currently)
# variables:
#     double time(time) ;
#         time:long_name = "Time" ;
#         time:time_origin = "2009-01-01 01:00:00" ;
#         time:delta_t = "0000-00-00 01:00:00" ;
#         time:units = "days since 2009-01-01T01:00:00Z" ;
#     double latitude(latitude, longitude) ;
#         latitude:long_name = "Latitude" ;
#         latitude:units = "degrees_north" ;
#     double longitude(latitude, longitude) ;
#         longitude:long_name = "Longitude" ;
#         longitude:units = "degrees_east" ;
#     double uas(time, latitude, longitude) ;
#         uas:standard_name = "uas" ;
#         uas:long_name = "East component of current" ;
#         uas:units = "m s-1" ;
#     double vas(time, latitude, longitude) ;
#         vas:standard_name = "vas" ;
#         vas:long_name = "North component of current" ;
#         vas:units = "m s-1" ;
# 
# ```
# (Example courtesy of Caroline Sevigny at Université du Québec à Rimouski)
# 
# The worker needs to transform the `time_counter`,
# `nav_lon`,
# `nav_lat`,
# `u_wind`,
# and `v_wind` variable from files like
# 
# `/results/forcing/atmospheric/GEM2.5/operational/ops_y2017m03d17.nc`
# 
# to produce an netCDF4 file with a structure similar to above that `ww3_prnc` will accept.
# 
# It turns out that `xarray` makes it quite easy to produce such a file.

# In[1]:


import arrow
import cmocean
import matplotlib.pyplot as plt
import numpy as np
import numpy.ma as ma
import xarray as xr


# In[2]:


get_ipython().run_line_magic('matplotlib', 'inline')


# In[4]:


hrdps = xr.open_dataset('/results/forcing/atmospheric/GEM2.5/operational/fcst/ops_y2017m04d06.nc')
hrdps


# We need to create a new dataset with:
# 
# * the `time_counter` coordinate renamed to `time` because that is the coordinate/dimension name that `ww3_prnc` expects
# * a `latitude` coordinate derived from the `nav_lat` variable because that is one of the variable names that `ww3_prnc` expects
# * a `longitude` coordinate derived from the `nav_lon` variable because that is one of the variable names that `ww3_prnc` expects
# * a `u_wind` variable with its `time_counter` coordinate renamed to `time`
# * a `v_wind` variable with its `time_counter` coordinate renamed to `time`

# In[5]:


ds = xr.Dataset(
    data_vars={
        'u_wind': hrdps.u_wind.rename({'time_counter': 'time'}),
        'v_wind': hrdps.v_wind.rename({'time_counter': 'time'}),
    },
    coords={
        'time': hrdps.time_counter.rename('time').rename({'time_counter': 'time'}),
        'latitude': hrdps.nav_lat,
        'longitude': hrdps.nav_lon,
    }
)
del ds.u_wind.attrs['coordinates']
del ds.v_wind.attrs['coordinates']
ds


# The `coordinates` attribute has to be deleted from `u_wind` and `v_wind`
# because the `to_netcdf()` method creates it and complains if it already exists.

# In[6]:


ds.to_netcdf('SoG_wind_20170406.nc')


# In[7]:


get_ipython().system('/usr/bin/ncdump -cst SoG_wind_20170406.nc')


# With a `ww3_prnc.inp` file like:
# ```
# $ WAVEWATCH III NETCDF Field preprocessor input ww3_prnc_wind.inp
# $
# $ Forcing type, grid type, time in file, header 
#    'WND' 'LL' T T
# $
# $ Dimension variable names
#   x y
# $
# $ Wind component variable names
#   u_wind v_wind
# $
# $ Forcing source file path/name
# $ File is produced by make_ww3_wind_file worker
#   'wind/SoG_wind_20170121.nc'
# ```
# `ww3_prnc` will process files like the above into `wind.ww3` files,
# and `ww3_shel` will accept them.
# 
# Here's a comparison of the wind speeds from 
# `/results/forcing/atmospheric/GEM2.5/operational/fcst/ops_y2017m01d22.nc`
# and those in the output from a wwatch3 run that used them as forcing:

# In[71]:


hrdps = xr.open_dataset('/results/forcing/atmospheric/GEM2.5/operational/fcst/ops_y2017m01d21.nc')
wwf = xr.open_dataset('../../analysis-doug/notebooks/SoG-waves/SoG_ww3_fields_20170121_20170123_ops_wind.nc')


# In[72]:


wwf_u_wind = ma.masked_invalid(wwf.uwnd)
wwf_v_wind = ma.masked_invalid(wwf.vwnd)

time = 0
fig, axs = plt.subplots(1, 2, figsize=(18, 9))
mesh = axs[0].pcolormesh(
    hrdps.nav_lon, hrdps.nav_lat, np.sqrt(hrdps.u_wind[time]**2 + hrdps.v_wind[time]**2),
    cmap=cmocean.cm.speed,
)
cbar = fig.colorbar(mesh, ax=axs[0])
cbar.set_label('Wind speed [m/s]')
axs[0].set_title('HRDPS wind input to ww3_prnc')

axs[1].pcolormesh(
    wwf.longitude, wwf.latitude, np.sqrt(wwf_u_wind[time]**2 + wwf_v_wind[time]**2),
    cmap=cmocean.cm.speed,
)
cbar = fig.colorbar(mesh, ax=axs[1])
cbar.set_label('Wind speed [m/s]')
axs[1].set_title('Wind output from ww3_shel')

axs[0].set_xlim(axs[1].get_xlim())
axs[0].set_ylim(axs[1].get_ylim())


# For a more analytical test,
# let's construct a wind field with 2 constant speed patches:
# 
# * a patch of 1 m/s westward (`v` component = 0) at the Sentry Shoal buoy
# * a patch of 1 m/s northward (`u` component = 0) at the Halibut Bank buoy

# In[3]:


hrdps = xr.open_dataset('/results/forcing/atmospheric/GEM2.5/operational/fcst/ops_y2017m01d20.nc')


# Near Sentry Shoal:

# In[4]:


u_patch_data = np.zeros_like(hrdps.u_wind)
u_patch_data[:, 181:185, 105:109] = 1
u_patch = xr.DataArray(
    data=u_patch_data,
    dims=hrdps.u_wind.dims,
    name=hrdps.u_wind.name,
    attrs=hrdps.u_wind.attrs,
)
print(u_patch)
print(u_patch[12, 180:186, 104:110])


# Near Halibut Bank:

# In[5]:


v_patch_data = np.zeros_like(hrdps.v_wind)
v_patch_data[:, 147:151, 139:143] = 1
v_patch = xr.DataArray(
    data=v_patch_data,
    dims=hrdps.v_wind.dims,
    name=hrdps.v_wind.name,
    attrs=hrdps.v_wind.attrs,
)
print(v_patch)
print(v_patch[12, 146:152, 138:144])


# In[7]:


ds = xr.Dataset(
    data_vars={
        'u_wind': u_patch.rename({'time_counter': 'time'}),
        'v_wind': v_patch.rename({'time_counter': 'time'}),
    },
    coords={
        'time': hrdps.time_counter.rename('time').rename({'time_counter': 'time'}),
        'latitude': hrdps.nav_lat,
        'longitude': hrdps.nav_lon,
    }
)
del ds.u_wind.attrs['coordinates']
del ds.v_wind.attrs['coordinates']
ds


# In[88]:


ds.to_netcdf('SoG_wind_20170120_patches.nc')


# In[8]:


time = 0
fig, ax = plt.subplots(1, 1, figsize=(9, 9))
mesh = ax.pcolormesh(
    ds.longitude, ds.latitude, np.sqrt(ds.u_wind[time]**2 + ds.v_wind[time]**2),
    cmap=cmocean.cm.speed,
)
cbar = fig.colorbar(mesh, ax=ax)
cbar.set_label('Wind speed [m/s]')
ax.set_title('HRDPS wind input to ww3_prnc')

ax.set_xlim(234.8, 236.4)
ax.set_ylim(49.2, 50)


# In[3]:


wwf = xr.open_dataset('SoG_ww3_fields_20170120_patches.nc')


# In[4]:


ds = xr.open_dataset('SoG_wind_20170120_patches.nc')


# In[5]:


wwf_u_wind = ma.masked_invalid(wwf.uwnd)
wwf_v_wind = ma.masked_invalid(wwf.vwnd)

time = 0
fig, axs = plt.subplots(1, 2, figsize=(18, 9))
mesh = axs[0].pcolormesh(
    ds.longitude, ds.latitude, np.sqrt(ds.u_wind[time]**2 + ds.v_wind[time]**2),
    cmap=cmocean.cm.speed,
)
cbar = fig.colorbar(mesh, ax=axs[0])
cbar.set_label('Wind speed [m/s]')
axs[0].set_title('HRDPS wind input to ww3_prnc')

axs[1].pcolormesh(
    wwf.longitude, wwf.latitude, np.sqrt(wwf_u_wind[time]**2 + wwf_v_wind[time]**2),
    cmap=cmocean.cm.speed,
)
cbar = fig.colorbar(mesh, ax=axs[1])
cbar.set_label('Wind speed [m/s]')
axs[1].set_title('Wind output from ww3_shel')

for ax in axs:
    ax.set_xlim(234.8, 236.4)
    ax.set_ylim(49.2, 50)


# In[37]:


time = 12

fig, axs = plt.subplots(1, 2, figsize=(18, 9))
mesh = axs[0].pcolormesh(
    ds.longitude, ds.latitude, np.rad2deg(np.arctan2(ds.v_wind[time], ds.u_wind[time])),
    cmap=cmocean.cm.amp,
)
cbar = fig.colorbar(mesh, ax=axs[0])
cbar.set_label('Wind direction [°]')
axs[0].set_title('HRDPS wind input to ww3_prnc')

axs[1].pcolormesh(
    wwf.longitude, wwf.latitude, np.rad2deg(np.arctan2(wwf_v_wind[time], wwf_u_wind[time])),
    cmap=cmocean.cm.amp,
)
cbar = fig.colorbar(mesh, ax=axs[1])
cbar.set_label('Wind direction [°]')
axs[1].set_title('Wind output from ww3_shel')

for ax in axs:
    ax.set_xlim(234.8, 236.4)
    ax.set_ylim(49.2, 50)


# Evolution of significant wave height:

# In[15]:


longitude=slice(234.8, 236.4)
latitude=slice(49.2, 50)

(wwf.hs
    .isel(time=slice(0, 8))
    .sel(longitude=longitude, latitude=latitude)
    .plot(col='time', col_wrap=4, cmap=cmocean.cm.amp)
)


# In[ ]: