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

# In[1]:


import numpy as np
get_ipython().run_line_magic('', 'matplotlib inline')
import matplotlib.pyplot as plt


# Make sure we are in the right directory.

# In[2]:


ls data


# In[5]:


data = np.loadtxt('../geocomp-0118/data/Penobscot_xl1155.txt')


# Inspect the first ten rows.

# In[16]:


plt.figure(figsize=(10, 10))
plt.imshow(data.T, cmap='Reds', vmin=-7000, vmax=7000, interpolation='bicubic')


# In[5]:


inlines = 748 - 104 + 1
xlines = 1248 - 302 + 1

print(inlines, 'inlines')
print(xlines, 'crosslines')


# We'll transform the first two columns to normalized versions of inline and crossline.

# In[6]:


data[:,0] -= 103 
data[:,1] -= 301


# In[7]:


data[:10]


# Make an empty array with the same shape. We're going to put our data in here.

# In[8]:


horizon = np.empty((inlines, xlines))
print(horizon)
print(horizon.shape)


# In[9]:


amplitude = np.empty((inlines, xlines))

for sample in data:
    inline = int(sample[0])
    xline = int(sample[1])
    z_value = sample[2]
    amp_value = sample[3]
    
    horizon[inline - 1, xline - 1] = z_value
    amplitude[inline - 1, xline - 1] = amp_value


# In[10]:


plt.imshow(horizon)
plt.show()


# In[11]:


plt.imshow(amplitude)
plt.show()


# In[12]:


np.amin(amplitude), np.amax(amplitude)


# The max value is a problem — that's why we're not seeing anything in the <code>imshow()</code>. So let's clip it.

# In[13]:


amplitude = np.clip(amplitude,-20000,20000)
plt.imshow(amplitude)
plt.show()


# In[14]:


amplitude


# ## Symmetric nearest neighbour

# In[15]:


def snn_old(horizon, kernel_size=5):
 
    # make an empty array the same shape as the input
    output = np.empty_like(horizon) 
    inlines, xlines = horizon.shape
    offset = kernel_size/2
 
    for x in range(int(inlines - 2*offset)):
        x += offset # Correct for offset
        for y in range(int(xlines - 2*offset)):
            y += offset
            centre = horizon[(x),(y)]
            nearest_neighbours = [centre]
            for a in range(kernel_size**2 // 2):
                i, j = divmod(a, kernel_size)
                i -= offset # transform to relative coordinates in kernel
                j -= offset
                value = horizon[int(x+i), int(y+j)]
                opposite = horizon[int(x-i), int(y-j)]
                closest = min(value-centre, opposite-centre)
                nearest_neighbours.append(closest+centre)
            output[x,y] = np.mean(nearest_neighbours)
 
    return output


# In[127]:


get_ipython().run_line_magic('timeit', 'x = snn_old(amplitude)')


# In[20]:


x = snn_old(amplitude)


# In[196]:


import scipy.ndimage

def snn(horizon, size=5, include=True):
    
    # TODO: See how it handles Nans,
    # Consider removing them, interpolate over them,
    # and put them back at end.
    
    def nearest(a, num):
        return a.flat[np.abs(a - num).argmin()]
    
    def func(this, pairs, include):
        centre = this[this.size // 2]
        select = [nearest(this[p], centre) for p in pairs]
        if include:
            select += [centre]
        return np.mean(select)

    pairs = [[i, size**2-1 - i] for i in range(size**2 // 2)]
    return scipy.ndimage.generic_filter(horizon, func, size=size,
                                        extra_keywords={'pairs': pairs,
                                                        'include': include}
                                       )


# In[55]:


get_ipython().run_line_magic('timeit', 'snn(amplitude)')


# In[72]:


y_ = snn(amplitude)


# In[82]:


plt.imshow(amplitude[:200, 600:], cmap='viridis')


# In[83]:


plt.imshow(x[:200, 600:], cmap='viridis')


# In[84]:


plt.imshow(y[:200, 600:], cmap='viridis')


# In[87]:


plt.imshow(z[:200, 600:], cmap='viridis')


# In[85]:


plt.imshow(y[:200, 600:] - y_[:200, 600:], cmap='viridis')


# In[86]:


plt.imshow(x[:200, 600:] - y_[:200, 600:], cmap='viridis')


# ## Kuwahara

# In[100]:


def kuwahara(horizon, size=5):

    # TODO: See how it handles Nans,
    # Consider removing them, interpolate over them,
    # and put them back at end.

    def func(this, s, k):
        t = this.reshape((s, s))
        sub = np.array([t[:k, :k].flatten(),
                        t[:k, k-1:].flatten(),
                        t[k-1:, :k].flatten(),
                        t[k-1:, k-1:].flatten()]
                       )
        select = sub[np.argmin(np.var(sub, axis=1))]
        return np.mean(select)

    if not size // 2:
        size += 1
        
    k = int(np.ceil(size / 2))

    return scipy.ndimage.generic_filter(horizon,
                                        func,
                                        size=size,
                                        extra_keywords={'s':size,
                                                        'k':k,
                                                       }
                                       )


# In[101]:


get_ipython().run_line_magic('timeit', 'kuwahara(amplitude)')


# In[ ]:


z = kuwahara(amplitude)


# In[87]:


plt.imshow(z[:200, 600:], cmap='viridis')


# ## Conservative mean

# In[194]:


def conservative(horizon, size=5, supercon=False):

    # TODO: See how it handles Nans,
    # Consider removing them, interpolate over them,
    # and put them back at end.

    def func(this, k, supercon):
        centre = this[k]
        rest = [this[:k], this[-k:]]
        mi, ma = np.nanmin(rest), np.nanmax(rest)
        if centre < mi:
            return mi if supercon else np.mean(rest)
        elif centre > ma:
            return ma if supercon else np.mean(rest)
        else:
            return centre

    if not size // 2:
        size += 1

    k = int(np.floor(size**2 / 2))

    return scipy.ndimage.generic_filter(horizon,
                                        func,
                                        size=size,
                                        mode='nearest',
                                        extra_keywords={'k':k,
                                                        'supercon': supercon,
                                                        }
                                        )


# In[185]:


supercon = False
this = np.arange(25)
this[12] = -1
s = int(np.sqrt(this.size))
k = int(np.floor(this.size/2))
centre = this[k]
rest = [this[:k], this[-k:]]
mi, ma = np.nanmin(rest), np.nanmax(rest)
if centre < mi:
    centre = mi if supercon else np.mean(rest)
if centre > ma:
    centre = ma if supercon else np.mean(rest)
this = np.insert(rest, k, centre)
this


# In[184]:


np.floor(5**2/2)


# In[166]:


#w = conservative(amplitude)


# In[122]:


#%timeit conservative(amplitude)


# In[152]:


plt.imshow(w[:200, 600:], cmap='viridis')


# In[192]:


this = np.arange(25).reshape((5,5))
this[2,2] = -1
this


# In[ ]:





# In[195]:


conservative(this, size=3, supercon=True)


# In[176]:


snn(this, size=3)


# In[171]:


kuwahara(this, size=3)


# In[173]:


scipy.ndimage.generic_filter(this, np.mean, size=3, mode='nearest')


# In[189]:


a = np.testing.assert_array_equal(conservative(this, size=3, supercon=True), conservative(this, size=3, supercon=True))


# In[190]:


a


# In[ ]: