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

# # Horizon holes filling
# 
# **Author: M. Ravasi, KAUST**

# Welcome to the **Matrix-free inverse problems with PyLops** tutorial!
# 
# The aim of this tutorial is to fill holes in seismic horizons (also known in the computer vision community as **inpainting**). 
# 
# 
# This notebook uses an horizon from Matt's paper [M. Hall (2007). Smooth operator: Smoothing seismic interpretations and attributes. The Leading Edge 26, p16–20](https://library.seg.org/doi/10.1190/1.2431821) and it is heavily inspired by [S. Fomel DWT jupyter notebook](https://github.com/sfomel/ipython/blob/master/DWT.ipynb) As a by-product you will learn to:
# 
# - Familiarize with the `pylops.signalprocessing` submodule, and more specifically with the 2D-FFT (``FFT2D``) and Wavelet transform (``DWT``);
# - Learn how to use sparse solvers within the `pylops.optimization.sparsity` submodule, and more specifically the FISTA solver.

# From a mathematical point of view we write the impainting problem as follows:
# 
# $$
# J = arg min_{\mathbf{p}} ||\mathbf{d} - \mathbf{M} \mathbf{P} \mathbf{p}||_2 + \epsilon ||\mathbf{p}||_1
# $$
# 
# where $\mathbf{M}$ is a masking operator, $\mathbf{P}$ is a sparsyfing transform, $\mathbf{d}$ is the horizon with holes, and $\mathbf{m}=\mathbf{P}\mathbf{p}$ is the filled horizon we wish to obtain.

# Let's first import the libraries we need in this tutorial

# In[1]:


# Run this when using Colab (will install the missing libraries)
# !pip install pylops scooby


# In[2]:


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

import numpy as np
import matplotlib.pyplot as plt
import scooby

from pylops.basicoperators import *
from pylops.optimization.sparsity import *
from pylops.signalprocessing import FFT2D, DWT, DWT2D
from pylops.utils.tapers import taper2d


# ## Data loading

# Let's start by loading a seismic horizon

# In[3]:


# You can download this horizon using Madagascar's command: Fetch('horizon.asc','hall')
# And it is now also provided in our data folder.
f = np.loadtxt('data/horizon.asc')
ig = f[:, -1]-65
ig = ig.reshape(291, 196).T
nyorig, nxorig = ig.shape
y = 33.139 + np.arange(nyorig) * 0.01
x = 35.031 + np.arange(nxorig) * 0.01

# We apply some tapering on the edges
#tap = taper2d(nxorig, nyorig, (25, 25))
#ig = ig * tap

# Pad to closest multiple of 2 (for DWT)
ig = np.pad(ig, ((0, 256-nyorig), (0, 512-nxorig)))
ny, nx = ig.shape

plt.figure(figsize=(15, 12))
plt.imshow(ig[:nyorig, :nxorig], cmap='jet', vmin=-14, vmax=14,
           extent=(x[0], x[-1], y[0], y[1]), origin='lower')
plt.xlabel('y (km)')
plt.ylabel('x (km)')
plt.title('Horizon')
plt.axis('tight');


# Let's now create some holes in the horizon

# In[4]:


mask = np.ones_like(ig)
mask[125:145, 150:170] = 0
mask[150:170, 50:70] = 0
mask[75:95, 175:195] = 0
igholes = ig * mask

plt.figure(figsize=(15, 12))
plt.imshow(igholes[:nyorig, :nxorig], cmap='jet', vmin=-14, vmax=14,
           extent=(x[0], x[-1], y[0], y[1]), origin='lower')
plt.xlabel('y (km)')
plt.ylabel('x (km)')
plt.title('Horizon with holes')
plt.axis('tight');


# ## Reconstruction with DWT

# First of all we need to create the sparsyfing operator

# In[5]:


wavkind = 'sym9'

# Cascaded DWT
Wopy, Wopx = DWT((ny,nx), 1, wavelet=wavkind, level=5), DWT((ny,nx), 0, wavelet=wavkind, level=5)
Wop = Wopy*Wopx
dimswav = Wopy.dimsd

# DWT2D
#Wop = DWT2D((ny, nx), wavelet=wavkind, level=5)
#dimswav = Wop.dimsd

ig_wav = Wop * ig.ravel()
iginv = Wop.H * ig_wav
iginv = iginv.reshape(ny, nx)

fig, axs = plt.subplots(1, 3, figsize=(16, 4))
axs[0].imshow(ig[:nyorig, :nxorig], cmap='jet', vmin=-14, vmax=14,
              extent=(x[0], x[-1], y[0], y[1]), origin='lower')
axs[0].set_title('Image')
axs[0].axis('tight')
axs[1].imshow(ig_wav.reshape(dimswav), cmap='gray', vmin=0, vmax=5e-1)
axs[1].set_title('DWT2 coefficients')
axs[1].axis('tight')
axs[2].imshow(iginv[:nyorig, :nxorig], cmap='jet', vmin=-14, vmax=14,
              extent=(x[0], x[-1], y[0], y[1]), origin='lower')
axs[2].set_title('Reconstructed image')
axs[2].axis('tight');


# Let's now solve the inverse problem

# In[6]:


Op = Diagonal(mask.ravel())

igdata = Op * ig.ravel()
igfilled_wav, niter, cost = FISTA(Op * Wop.H, igdata, niter=100, 
                                  #threshkind='soft', eps=1e1, 
                                  threshkind='soft-percentile', perc=10, 
                                  returninfo=True, show=True)
igfilled = np.real((Wop.H * igfilled_wav).reshape(ny, nx))

plt.figure(figsize=(12, 3))
plt.plot(cost, 'k', lw=2)
plt.title('Cost function');


# In[7]:


fig, axs = plt.subplots(1, 3, figsize=(16, 4))
axs[0].imshow(ig[:nyorig, :nxorig], cmap='jet', vmin=-14, vmax=14,
              extent=(x[0], x[-1], y[0], y[1]), origin='lower')
axs[0].set_title('Image')
axs[0].axis('tight')
axs[1].imshow(igfilled[:nyorig, :nxorig], cmap='jet', vmin=-14, vmax=14,
              extent=(x[0], x[-1], y[0], y[1]), origin='lower')
axs[1].set_title('Reconstructed image')
axs[1].axis('tight');
axs[2].imshow(ig[:nyorig, :nxorig]-igfilled[:nyorig, :nxorig], cmap='jet', vmin=-14, vmax=14,
              extent=(x[0], x[-1], y[0], y[1]), origin='lower')
axs[2].set_title('Error image')
axs[2].axis('tight');


# In[8]:


plt.figure(figsize=(15, 12))
plt.imshow(igfilled[:nyorig, :nxorig], cmap='jet', vmin=-14, vmax=14,
           extent=(x[0], x[-1], y[0], y[1]), origin='lower')
plt.xlabel('y (km)')
plt.ylabel('x (km)')
plt.title('Filled Horizon')
plt.axis('tight');


# At this point you may wonder if the DWT is the best sparsifying transform for this data. 
# 
# I encourage you to experiment with more transforms, for example:
# 
# - ``FFT2D`` (as shown below);
# - ``Curvelet`` transform (see [curvelops](https://github.com/PyLops/curvelops) for a PyLops wrapper);
# - Other transforms such as the Shearlet transform (you try to wrap [pyshearlab](http://na.math.uni-goettingen.de/pyshearlab/) into PyLops).

# In[9]:


Wop = FFT2D((ny, nx))
ig_wav = Wop * ig.ravel()
igholes_wav = Wop * igholes.ravel()
iginv = Wop.H * ig_wav
iginv = np.real(iginv).reshape(ny, nx)

fig, axs = plt.subplots(1, 3, figsize=(16, 4))
axs[0].imshow(ig[:nyorig, :nxorig], cmap='jet', vmin=-14, vmax=14,
              extent=(x[0], x[-1], y[0], y[1]), origin='lower')
axs[0].set_title('Image')
axs[0].axis('tight')
axs[1].imshow(np.fft.fftshift(np.abs(np.reshape(ig_wav, (ny,nx)))), 
              cmap='gray', vmin=0, vmax=5)
axs[1].set_title('DWT2 coefficients')
axs[1].axis('tight')
axs[2].imshow(iginv[:nyorig, :nxorig], cmap='jet', vmin=-14, vmax=14,
              extent=(x[0], x[-1], y[0], y[1]), origin='lower')
axs[2].set_title('Reconstructed image')
axs[2].axis('tight');


# In[10]:


Op = Diagonal(mask.ravel())

igdata = Op * ig.ravel()
igfilled_wav, niter, cost = FISTA(Op * Wop.H, igdata.astype(np.complex), niter=100, 
                                  #threshkind='soft', eps=1e1, 
                                  threshkind='soft-percentile', perc=10, 
                                  returninfo=True, show=True)
igfilled = np.real((Wop.H * igfilled_wav).reshape(ny, nx))

plt.figure(figsize=(12, 3))
plt.plot(cost, 'k', lw=2)
plt.title('Cost function');


# In[11]:


fig, axs = plt.subplots(1, 3, figsize=(16, 4))
axs[0].imshow(ig[:nyorig, :nxorig], cmap='jet', vmin=-14, vmax=14,
              extent=(x[0], x[-1], y[0], y[1]), origin='lower')
axs[0].set_title('Image')
axs[0].axis('tight')
axs[1].imshow(igfilled[:nyorig, :nxorig], cmap='jet', vmin=-14, vmax=14,
              extent=(x[0], x[-1], y[0], y[1]), origin='lower')
axs[1].set_title('Reconstructed image')
axs[1].axis('tight');
axs[2].imshow(ig[:nyorig, :nxorig]-igfilled[:nyorig, :nxorig], cmap='jet', vmin=-14, vmax=14,
              extent=(x[0], x[-1], y[0], y[1]), origin='lower')
axs[2].set_title('Error image')
axs[2].axis('tight');


# In[12]:


scooby.Report(core='pylops')


# In[ ]: