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

# ### Gage adjustment
# 
# -----------------------------------------------------
# ---

# In[1]:


import os
import math
import copy
import warnings

import numpy as np
import numpy.ma as ma

import matplotlib.pyplot as plt
import cartopy.crs as ccrs
import xarray as xr

import wradlib as wrl

try:
    get_ipython().run_line_magic("matplotlib inline")
except:
    plt.ion()

warnings.filterwarnings('ignore')


# ### **PART 1):** Reading the RADAR file

# In[2]:


filename = "CAPPI-2km-Depth-H&V.nc"


# In[3]:


grid = xr.open_dataset(filename)
display(grid)


# ### **PART 2):** Reading rain gauges

# In[4]:


# Load data from text file
obs = np.loadtxt('obs.txt')
obs.shape


# In[5]:


#obs


# In[6]:


# Load coordinate text file data
obs_coords = np.loadtxt('obs_coords.txt')
obs_coords.shape


# In[7]:


#obs_coords


# ### **PART 3):**Preparing the data

# In[8]:


# select base 2d grid
grid = grid.isel(time=0, z=0)


# ### Radar Horizontal
# 
# - create mask
# - fill nan with zero
# - stack xy into radar_obs

# In[9]:


nanmask_h = xr.where(np.isnan(grid.R_KDP_Corr_h_depth), True, False)
radar_h = grid.R_KDP_Corr_h_depth.fillna(0)
radar_h_stack = radar_h.stack(radar_obs=("x", "y"))
display(radar_h_stack)


# ### Radar Vertical
# - create mask
# - fill nan with zero
# - stack xy into radar_obs

# In[10]:


nanmask_v = xr.where(np.isnan(grid.R_KDP_Corr_v_depth), True, False)
radar_v = grid.R_KDP_Corr_v_depth.fillna(0)
radar_v.plot()
radar_v_stack = radar_v.stack(radar_obs=("x", "y"))
display(radar_v_stack)


# ### **PART 3): **Adjustment of Rain Gauges**

# In[11]:


# get grid coords
grid_coords = xr.concat([radar_h_stack.x, radar_h_stack.y], "xy").transpose(..., "xy")
display(grid_coords)


# ## AdjustADD with OrdinaryKriging

# In[12]:


# Estrutura dos seus dados
print("obs:", obs.shape)
print("obs_coords:", obs_coords.shape)
print("radar_h:", radar_h.shape)
print("grid_coords:", grid_coords.shape)
print("radar_v:", radar_v.shape)


# In[13]:


get_ipython().run_cell_magic('time', '', "# Additive Error Model with OrdinaryKriging Exponential\n# class wradlib.ipol.OrdinaryKriging(src, trg, cov='1.0 Exp(10000.)', nnearest=12)\naddadjuster_ok_exp = wrl.adjust.AdjustAdd(obs_coords, \n                                      grid_coords.values,\n                                      cov = '1.0 Exp(10000.)',\n                                      ipclass=wrl.ipol.OrdinaryKriging)\naddadjusted_ok_exp_values = addadjuster_ok_exp(obs, radar_h_stack.values)\n")


# ### create DataArray, set values, unstack and transpose

# In[14]:


addadjusted_ok_exp = xr.zeros_like(radar_h_stack)
addadjusted_ok_exp.values = addadjusted_ok_exp_values
addadjusted_ok_exp = addadjusted_ok_exp.unstack().transpose("y", "x")
display(addadjusted_ok_exp)


# In[15]:


# open figure
fig = plt.figure(figsize=(10, 3.5))

# "Radar Rainfall: Real"
ax = fig.add_subplot(121)
radar_h.where(~nanmask_h).plot(ax=ax)
ax.set_title("Radar Rainfall: Real")

# "Additive Error Model with Nearest"
ax = fig.add_subplot(122)
addadjusted_ok_exp.where(~nanmask_h).plot(ax=ax)
ax.set_title("Additive Error Model with OK EXPl")

plt.tight_layout()
# salva figura


# ## AdjustADD with ExternalDriftKriging

# Docstring excerpt:
# 
# ```python
# src_drift : :class:`numpy:numpy.ndarray`
#         ndarray of floats, shape (nsrcpoints, )
#         values of the external drift at each source point
# trg_drift : :class:`numpy:numpy.ndarray`
#         ndarray of floats, shape (ntrgpoints, )
#         values of the external drift at each target point
# ```
# 
# Not sure, if the now provided values for `src_drift` and `trg_drift` are the correct ones to provide there as `drift`.

# In[16]:


get_ipython().run_cell_magic('time', '', "addadjuster_kde_exp = wrl.adjust.AdjustAdd(obs_coords, \n                                           grid_coords.values,\n                                           cov = '1.0 Exp(10000.)',\n                                           nnearest=12,\n                                           src_drift=obs, # need obs values here\n                                           trg_drift=radar_v_stack.values, # need radar v values here\n                                           ipclass=wrl.ipol.ExternalDriftKriging)\naddadjusted_kde_exp_values = addadjuster_kde_exp(obs, radar_h_stack.values)\n")


# ### create DataArray, set values, unstack and transpose

# In[17]:


addadjusted_kde_exp = xr.zeros_like(radar_h_stack)
addadjusted_kde_exp.values = addadjusted_kde_exp_values
addadjusted_kde_exp = addadjusted_kde_exp.unstack().transpose("y", "x")
display(addadjusted_kde_exp)


# In[18]:


# open figure
fig = plt.figure(figsize=(10, 3.5))

# "Radar Rainfall: Real"
ax = fig.add_subplot(121)
radar_h.where(~nanmask_h).plot(ax=ax)
ax.set_title("Radar Rainfall: Real")

# "Additive Error Model with Nearest"
ax = fig.add_subplot(122)
addadjusted_kde_exp.where(~nanmask_h).plot(ax=ax)
ax.set_title("Additive Error Model with OK EXPl")

plt.tight_layout()
# salva figura