%load_ext autoreload
%autoreload 2

import numpy as np
import xarray as xr
import matplotlib.pyplot as plt

"""
Explanation: Navigating between various vertical coordinates in the ocean
In oceanography it is often required to transform the vertical coordinate system of data. The large scale circulation for example, follows surfaces of constant density rather than constant depth levels, thus certain analyses require a coordinate transformation.
The general process of vertical coordinates transformation as discussed here, consists of two steps: Regridding and remapping.
Regridding: The process of definining the target depth coordinates (e.g. the depth defined by certain density levels or e.g. just a different depth spacing)
Remapping: Converting a data profile from the source depth profile to the target depth profile.
There are several methods to perform both regridding and remapping implemented in xarrayutils.vertical_coordinates. In the following we will see how to implement different combinations of method for both processes. For general applications the conservative conservative_remap function is recommended for remapping, due to the property of conserving total tracer mass (see examples below).

At the moment these functions require very explicit input (dimension names etc), which will be simplified with a wrapper function in the future.
End of explanation
"""


import fsspec
ds = xr.open_zarr(fsspec.get_mapper('gcs://pangeo-data/balwada/channel_ridge_resolutions/20km/tracer_10day_snap'), consolidated=True)
ds = ds.isel(time=slice(0, 20))
# ds.to_zarr('offline_backup.zarr')
# ds = xr.open_zarr('offline_backup.zarr')

"""
Explanation: Dont forget to install
conda install gcsfs zarr fsspec
for the docs
We are going to investigate the coordinate transformations using data from an indealized high resolution model run produced by Dhruv Balwada(to read more about the data see his paper).
End of explanation
"""


for dim in ['Z', 'Zp1', 'Zl', 'Zu']:
    ds.coords[dim] = -ds[dim]

ds.Z

"""
Explanation: Adjusting the vertical coordinate orientation
All the following steps crucially assume that the depth values are increasing with the depth dimension. The values can be negative but then have progress towards less negative values as you follow the logical depth index. 
This dataset has negative depth values that decrease. The simplest fix is to flip the sign of all depth values.
End of explanation
"""


# check monotonicity for all profiles
ds = ds.mean('time')
ds = ds.isel(YC=slice(10,-10), Z=slice(28,60), Zp1=slice(28,61))
# since we operate on the depth dimension we convert it to a single chunk to avoid doing this all the time later
ds = ds.chunk({'Z':-1})

assert (ds.T.diff('Z') < 0).all(['Z']).all()

"""
Explanation: Another issue with this dataset is that in the surface, the profile is sometimes unstable (density does not strictly increase with depth). This will cause issues when interpolating, and thus for now we use a time mean of the dataset between ~200-1500m depth.
End of explanation
"""


# define a new depth array
z_new = np.arange(10,2000, 20)

ds_z_new = ds.interp(Z=z_new)
ds_z_new.Z

"""
Explanation: Transforming to a different (spatially uniform) depth coordinate system
In this case we will transform the dataset to different depth coordinates (we leave out the regridding step and manually provide a new depth grid). This example might be useful when you want to convert different depth grids (e.g. from several observational products) into a uniform grid.
Using linear interpolation as remapping
For the simplest case (using linear interpolation for both the regridding and remapping), xarray offers all the necessary tools built in.
End of explanation
"""


ds.PTRACER01.isel(XC=40, YC=40).plot(y='Z', yincrease=False)
ds_z_new.PTRACER01.isel(XC=40, YC=40).plot(ls='--', y='Z', yincrease=False)

"""
Explanation: The depth coordinate Z is now regularly spaced instead of the surface refined resolution of the original dataset, we 'remapped' all data onto a new vertical grid using linear interpolation. Lets compare the actual data.
End of explanation
"""


dz_original =  ds.drF #vertical cell thickness of the model grid
tracer_intz_original = (ds.PTRACER01 * dz_original).sum('Z')
dz_new = 20 #This is easy to infer since the grid is uniformly spaced
tracer_intz_new = (ds_z_new.PTRACER01 * dz_new).sum('Z')

print(tracer_intz_original.isel(XC=40, YC=40).load())
print(tracer_intz_new.isel(XC=40, YC=40).load())

"""
Explanation: Visually that looks pretty good and these results might be sufficient for certain applications. The biggest downside of this approach is that the total amount of tracer is not conserved.
End of explanation
"""


from xarrayutils.vertical_coordinates import conservative_remap

# the conservative remapping needs information about the upper and lower bounds of the source and target cells.
bounds_original = ds.Zp1 # depth position of vertical bounding surface
bounds_new = xr.DataArray(np.arange(0,2020, 20), dims=['new_bounds'])

# now we can remap the tracer data again (at the moment the depth dimensions have to be explicitly defined).
ds_z_cons_new = conservative_remap(ds.PTRACER01,bounds_original, bounds_new,
                                   z_dim='Z', z_bnd_dim='Zp1', z_bnd_dim_target='new_bounds', mask=True) # the associated depth dimensions for each array
# replace the new depth dimension values with the appropriate depth 
ds_z_cons_new.coords['remapped'] = xr.DataArray(z_new, coords=[('remapped', z_new)])

ds_z_cons_new

ds.PTRACER01.isel(XC=40, YC=40).plot(y='Z', yincrease=False)
ds_z_cons_new.isel(XC=40, YC=40).plot(ls='--', y='remapped', yincrease=False)

"""
Explanation: The difference might seem small but for certain applications (e.g. budget reconstruction), this is not acceptable.
However,  we can do better by using the conservative_remap function:
Using conservative remapping
For this method we need not only the cell centers, but instead need to provide the depths of the vertical bounding surfaces.
End of explanation
"""


tracer_intz_cons_new = (ds_z_cons_new * dz_new).sum('remapped')
np.isclose(tracer_intz_original.isel(XC=40, YC=40).load(), tracer_intz_cons_new.isel(XC=40, YC=40).load())

"""
Explanation: conservative_remap takes into account every overlap between source and target cells. See for instance the uppermost value of ds_z_cons_new, the low value is due to the fact that there is a traceramount in the upper half of the first source cell, which is then distributed over the much larger targer cell. This ensures that the full tracer amount is conserved to floating point precision.
End of explanation
"""


np.isclose(tracer_intz_original, tracer_intz_cons_new).all()

"""
Explanation: This is in fact true for every grid position:
End of explanation
"""


from xarrayutils.vertical_coordinates import linear_interpolation_regrid

t_vals = np.arange(0.6,3, 0.01)
temperature_values = xr.DataArray(t_vals, coords=[('t', t_vals)]) # define the new temperature grid

z_temp_coord = linear_interpolation_regrid(ds.Z, ds.T, temperature_values, target_value_dim='t') 

plt.subplot(2,1,1)
ds.T.isel(XC=40).plot(x='YC', yincrease=False)
ds.T.isel(XC=40).plot.contour(levels=[2], colors='w',x='YC', yincrease=False)
plt.subplot(2,1,2)
z_temp_coord.isel(XC=40).plot(x='YC')
plt.axhline(2, color='w')

"""
Explanation: Ok this is nice, but it really gets interesting when we define our new depth coordinates 
Switching to potential temperature coordinates using only linear interpolation
Regridding using linear interpolation
The xarray internals can only help us if we want to interpolate on values of a dimension (1D), e.g. if we want to know the temperature as a function of a new depth. For the case of temperature coordinates, we aim to do the opposite: Find the depth for a given temperature value. This is not possible with xarray at the moment, but can be achieved using linear_interpolation_regrid.

The values used to find the new cell depths have to be monotonic! Currently there is no check implemented for this but non-monotonic fields can lead to undesired behaviour. See above for how to check for monotonicity.
End of explanation
"""


from xarrayutils.vertical_coordinates import linear_interpolation_remap

# we cant have nans, so just fill them with an out of bounds value
# z_temp_coord = z_temp_coord.fillna(1e8) # this should be fixed now?

ds_temp_linear = linear_interpolation_remap(ds.Z, ds.T, z_temp_coord)
# this requires me to downgrade dask, should be solved soon (see here https://github.com/pydata/xarray/pull/3660)
ds_temp_linear

ds_temp_linear.isel(XC=40).plot(x='YC', robust=True)

"""
Explanation: As you can see in this example, the line of constant temperature is moving deeper with increasing y, and that is reflected in the depth values along a constant depth coordinate in the regridded values. Now we can remap other data values on the corresponding depths. The simplest method is again linear interpolation:
Using linear interpolation as remapping
End of explanation
"""


ds_temp_linear = linear_interpolation_remap(ds.Z, ds.PTRACER01, z_temp_coord)
ds_temp_linear.sel(remapped=2, method='nearest').plot(robust=True)

"""
Explanation: As expected, when we remap the temperature field itself, we get a matching horizontal stratification. Now lets do something more interesting and look at the tracer field on a constant temperature surface of 2 deg:
End of explanation
"""


# we need to capture all the tracer cells, so if we dont specify our temperature range covering the Tracer min and max, the total tracer amount will not be conserved.
t_vals = np.hstack([ds.T.min().load().data[np.newaxis], np.arange(0.6,3, 0.01), ds.T.max().load().data[np.newaxis]])
t_vals
# for now the results are ordered by the bin values. Since the temperature is decreasing, we need to flip the logical axis.
# We could also do that after the regridding, but this seems more elegant.
t_vals = np.flip(t_vals)

temperature_values = xr.DataArray(t_vals, coords=[('t', t_vals)]) # define the new temperature grid
temperature_values

# # additionally to covering the full range of `ds.T`, we need to provide the bounds of `z` to the function.
z_temp_bounds = linear_interpolation_regrid(ds.Z, ds.T, temperature_values, z_bounds=ds.Zp1 ,target_value_dim='t', z_bounds_dim='Zp1')
z_temp_bounds.isel(XC=40, YC=40).load()

"""
Explanation: Pretty neat, but again there is no guarantee that the total tracer content is preserved with linear interpolation, but with a few modifications we can implement the conservative remapping here as well.
Using conservative remapping
End of explanation
"""


# now we can remap the tracer data again (at the moment the depth dimensions have to be explicitly defined).
ds_temp_cons = conservative_remap(ds.PTRACER01,bounds_original, z_temp_bounds,
                                   z_dim='Z', z_bnd_dim='Zp1', z_bnd_dim_target='regridded', mask=True) # the associated depth dimensions for each array
# # replace the new depth dimension values with the appropriate depth (here the middle of the temperature cell bounds)
t_vals = z_temp_bounds.coords['regridded'].data
t_vals = 0.5 * (t_vals[1:] + t_vals[0:-1])

ds_temp_cons.coords['remapped'] = xr.DataArray(t_vals, coords=[('remapped', t_vals)])

plt.figure(figsize=[20,5])
plt.subplot(1,3,1)
ds_temp_linear.sel(remapped=2, method='nearest').plot(robust=True)
plt.subplot(1,3,2)
ds_temp_cons.sel(remapped=2, method='nearest').plot(robust=True)
plt.subplot(1,3,3)
(ds_temp_cons.sel(remapped=2, method='nearest')-ds_temp_linear.sel(remapped=2, method='nearest')).plot(robust=True)

"""
Explanation: And now we can use these cell bounding values exactly like we did before.
End of explanation
"""


dz_remapped = z_temp_bounds.diff('regridded').rename({'regridded':'remapped'})
dz_remapped.coords['remapped'] = ds_temp_cons.coords['remapped']

tracer_intz_remapped_temp = (ds_temp_cons*dz_remapped).sum('remapped')

xr.testing.assert_allclose(tracer_intz_original, tracer_intz_remapped_temp)

"""
Explanation: And most importantly, the vertical tracer content is again conserved:
End of explanation
"""


