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

# # Dfsu - Speed and direction

# In[1]:


import numpy as np
import mikeio


# In[2]:


ds = mikeio.read("../tests/testdata/HD2D.dfsu")
ds


# This file is missing current direction :-(
# 
# Lets'fix that!

# ## Calculate speed & direction

# In[3]:


ds.U_velocity


# In order to use Numpy functions on a DataArray, we first convert the DataArrays (U, V) to standard NumPy ndarrays.

# In[4]:


u = ds.U_velocity.to_numpy()
v = ds.V_velocity.to_numpy()


# In[5]:


direction = np.mod(90 -np.rad2deg(np.arctan2(v,u)),360)


# ## Write new file

# In[6]:


from mikeio.eum import ItemInfo, EUMUnit, EUMType

ds["Current direction"] = mikeio.DataArray(direction, time= ds.time, item = ItemInfo("Current direction", EUMType.Current_Direction, EUMUnit.degree), geometry=ds.geometry)
ds


# In[7]:


ds.to_dfs("speed_direction.dfsu")


# In[8]:


nds = mikeio.read("speed_direction.dfsu")
nds


# # Plot

# In[9]:


step = 1

ax = ds.Current_speed[step].plot(figsize=(10,10))
ax.set_ylim([None, 6903000])
ax.set_xlim([607000, None])
ec = ds.geometry.element_coordinates
x = ec[:,0]
y = ec[:,1]
u = ds.U_velocity.to_numpy()
v = ds.V_velocity.to_numpy()
ax.quiver(x, y, u[step], v[step], scale=6, minshaft=3);


# ## Plot quiver on Cartesian overlay instead 
# Create overset grid and interpolate data on to this

# In[10]:


g = ds.geometry.get_overset_grid(dx=50)
g


# In[11]:


g.projection


# In[12]:


ui = ds.U_velocity.interp_like(g)
vi = ds.V_velocity.interp_like(g)


# In[13]:


ui.plot();


# In[14]:


ax = ds.Current_speed.plot(figsize=(10,10))

u = ui.to_numpy()
v = vi.to_numpy()
ax.quiver(g.x, g.y, u[step], v[step], scale=8, minshaft=5)

ax.set_ylim([None, 6903000])
ax.set_xlim([607000, None])
ax.set_title(f'Current speed with overset grid; {ds.time[step]}')
ax.set_xlabel("Easting (m)")
ax.set_ylabel("Northing (m)");


# ## Clean up

# In[15]:


import os
os.remove("speed_direction.dfsu")