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

# # Dfsu - 3D sigma-z
# This notebook demonstrates, reading from a sigma-z dfsu file, top- and bottom layers, extracting a profile, save selected layers to new dfsu file and save to mesh. 
# 
# It also shows how to read a sigma-z vertical slice file. 

# In[1]:


import mikeio


# In[2]:


filename = "../tests/testdata/oresund_sigma_z.dfsu"
dfs = mikeio.open(filename)
dfs


# In[3]:


dfs.geometry.plot.mesh();


# ## Save geometry to new mesh file
# Note that for sigma-z files, the mesh will be not match the original mesh in areas where z-layers are present! 

# In[4]:


outmesh = "mesh_oresund_extracted.mesh"
dfs.geometry.to_mesh(outmesh)


# ## Evaluate top layer

# In[5]:


ds = dfs.read(layers="top")
print(ds)
max_t = ds['Temperature'].to_numpy().max()
print(f'Maximum temperature in top layer: {max_t:.1f}')


# ## Find position of max temperature and plot
# 
# Find position of max temperature in first time step 

# In[6]:


timestep = 0
elem = ds['Temperature'][timestep].to_numpy().argmax()
max_x, max_y = ds.geometry.element_coordinates[elem,:2]
print(f'Position of maximum temperature: (x,y) = ({max_x:.1f}, {max_y:.1f})')


# In[7]:


ax = ds['Temperature'].isel(time=timestep).plot()
ax.plot(max_x, max_y, marker='*', markersize=20);


# # Read 1D profile from 3D file
# Find water column which has highest temperature and plot profile for all 3 time steps.

# In[8]:


dsp = dfs.read(x=max_x, y=max_y) # select vertical column from dfsu-3d 
dsp["Temperature"].plot();


# Note that the vertical column data is extrapolated to the bottom and surface! 
# 
# The extrapolation can avoided using "extrapolate=False":

# In[9]:


dsp["Temperature"].plot(extrapolate=False, marker='o');


# If the data has more than a few timesteps, it can be more convenient to plot as 2d pcolormesh. We will simulate this by interpolating to 30min data. 
# 
# Note that pcolormesh will plot using the static z information!

# In[10]:


dspi = dsp.Salinity.interp_time(dt=1800)
dspi.plot.pcolormesh();


# # Read top layer of a smaller area

# In[11]:


bbox = [310000, 6192000, 380000, 6198000]
ds_sub = dfs.read(area=bbox, layers="top")
ds_sub


# In[12]:


ds_sub.Temperature.plot();


# Plot subset inside original model domain

# In[13]:


ax = ds_sub.Temperature.plot(figsize=(6,7))
dfs.geometry.plot.outline(ax=ax, title=None);


# In[14]:


ds_sub.to_dfs("oresund_data_extracted.dfsu")


# # Select top 2 layers and write to new file
# get_layer_elements() can take a list of layers. Layers are counted positive from the bottom starting at 0 or alternatively counted negative from the top starting at -1. Here we take layers -1 and -2, i.e., the two top layers. 
# 
# Next data is read from source file and finally written to a new dfsu file (which is now sigma-only dfsu file).

# In[15]:


ds_top2 = dfs.read(layers=[-2, -1])
ds_top2


# In[16]:


outfile = "oresund_top2_layers.dfsu"
ds_top2.to_dfs(outfile)


# # Read vertical slice (transect)

# In[17]:


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


# In[18]:


print(ds.geometry.bottom_elements[:9])
print(ds.geometry.n_layers_per_column[:9])
print(ds.geometry.top_elements[:9])


# In[19]:


ds.Temperature.plot();


# # Clean up

# In[20]:


import os
os.remove("mesh_oresund_extracted.mesh")
os.remove("oresund_data_extracted.dfsu")
os.remove("oresund_top2_layers.dfsu")