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

# # SparseEdges tests

# ## A series of tests for the SparseEdges package

# 
# Table of content
# 
# * [What is the SparseEdges package?](#What-is-the-SparseEdges-package?) 
# * [Installing](#Installing) 
# * [testing one step](#testing-one-step)
# 

# In[1]:


get_ipython().run_line_magic('load_ext', 'autoreload')
get_ipython().run_line_magic('autoreload', '2')

from __future__ import division, print_function
get_ipython().run_line_magic('matplotlib', 'inline')
import matplotlib.pyplot as plt

import numpy as np
np.set_printoptions(precision=4)#, suppress=True)


# In[2]:


get_ipython().run_line_magic('cd', '-q ../probe')


# In[3]:


from NeuroTools.parameters import ParameterSet
from SLIP import Image
from LogGabor import LogGabor
from SparseEdges import SparseEdges


# In[4]:


#! defining framework
#!-------------------
mp = SparseEdges('https://raw.githubusercontent.com/bicv/SparseEdges/master/default_param.py')
mp.pe.N = 4
mp.pe.MP_alpha=1.
mp.pe.do_whitening = False


# In[5]:


get_ipython().system('ls -l  ../probe/cache_dir/experiment_test_*')

!rm -f  ../probe/cache_dir/experiment_test_*!ls -l  cache_dir/experiment_test_whitening.npy cache_dir/experiment_test_whitening_MSE.npy
# ## Defining units, axis and conventions

# cardinal points in X, Y (clockwise):  E (=right=3 o'clock=horizontal), ESE, S, WSW, W, WNW, N, ENE :

# In[6]:


X = np.array([0, 1./4, 1./2, 1./4, 0, -1./4, -1./2])#, -1./4])\n",
Y = np.array([1./2, np.sqrt(3)/4, 0, -np.sqrt(3)/4, -1./2, -np.sqrt(3)/4, 0.])#, np.sqrt(3)/
Theta = np.arctan2(-X, Y)  # perpendicular to circle
Theta = ((Theta + np.pi/2  - np.pi/mp.pe.N_phi/2 ) % (np.pi)) - np.pi/2  + np.pi/mp.pe.N_phi/2
print(['%0.1f' % (angle*180/np.pi) for angle in Theta])


# We set the [0,0] index of the array in the upper left left corner of the axes.
# 
# The X, Y image axis are defined in the matrix style (origin on the upper left; X going vertically down, Y going right horizontally) - the angle is defined as the angle wrt to the horizontal axis :

# In[7]:


fig, a = mp.show_edges(np.vstack(((X/2. + .5)*mp.pe.N_X, (Y/2. + .5)*mp.pe.N_X, Theta, .1*np.ones(X.shape), np.ones(X.shape), np.zeros(X.shape) )))
t = [a.text(Y_, X_, r'(%0.1f,%0.1f), $\theta=$%0.1f' % (X_, Y_, angle*180/np.pi))  for (X_, Y_, angle)  in zip((X/2. + .5)*mp.pe.N_X, (Y/2. + .5)*mp.pe.N_X, Theta)]
a.set_xlabel('Y (columns in the image matrix - from west to east)')
_ = a.set_ylabel('X (rows in the image matrix - from north to south)')


# Overlaying an image :

# In[8]:


mp = SparseEdges('https://raw.githubusercontent.com/bicv/SparseEdges/master/default_param.py')
mp.pe.N = 128 # number of edges
mp.pe.figsize_edges = 9

#! defining a reference test image (see test_Image)
image = np.zeros((mp.pe.N_X, mp.pe.N_Y))
image[mp.pe.N_X//2:mp.pe.N_X//2+mp.pe.N_X//4, mp.pe.N_X//2:mp.pe.N_X//2+mp.pe.N_X//4] = 1
image[mp.pe.N_X//2:mp.pe.N_X//2+mp.pe.N_X//4, mp.pe.N_X//4:mp.pe.N_X//2] = -1

edges = np.vstack(((X/2. + .5)*mp.pe.N_X, (Y/2. + .5)*mp.pe.N_X, Theta, .1*np.ones(X.shape), np.ones(X.shape), np.zeros(X.shape) ))

fig, a = mp.show_edges(edges, image=image)
#fig, a = mp.show_edges(edges, image=im.whitening(image))
t = [a.text(Y_, X_, r'(%0.1f,%0.1f), $\theta=$%0.1f' % (X_, Y_, angle*180/np.pi))  for (X_, Y_, angle)  in zip((X/2. + .5)*mp.pe.N_X, (Y/2. + .5)*mp.pe.N_X, Theta)]
a.set_xlabel('Y (columns in the image matrix - from west to east)')
_ = a.set_ylabel('X (lines in the image matrix - from north to south)')


# In[9]:


mp = SparseEdges('https://raw.githubusercontent.com/bicv/SparseEdges/master/default_param.py')
mp.set_size((32, 64))
mp.pe.N = 128 # number of edges
mp.pe.figsize_edges = 4.5

#! defining a reference test image (see test_Image)
image = np.zeros((mp.pe.N_X, mp.pe.N_Y))
image[mp.pe.N_X//2:mp.pe.N_X//2+mp.pe.N_X//4, mp.pe.N_X//2:mp.pe.N_X//2+mp.pe.N_X//4] = 1
image[mp.pe.N_X//2:mp.pe.N_X//2+mp.pe.N_X//4, mp.pe.N_X//4:mp.pe.N_X//2] = -1

edges = np.vstack(((X/2. + .5)*mp.pe.N_X-1, (Y/2. + .5)*mp.pe.N_X-1, Theta, .5*np.ones(X.shape), np.ones(X.shape), np.zeros(X.shape) ))

fig, a = mp.show_edges(edges, image=mp.whitening(image))


# In[10]:


mp.set_size((256, 256))
#image = mp.imread('../../SLIP/database/lena256.png')
image = mp.imread('https://raw.githubusercontent.com/bicv/SLIP/master/database/lena256.png')

mp.pe.N = 128 # number of edges
mp.pe.figsize_edges = 9

edges = np.vstack(((X/2. + .5)*mp.pe.N_X-1, (Y/2. + .5)*mp.pe.N_X-1, Theta, .1*np.ones(X.shape), np.ones(X.shape), np.zeros(X.shape) ))

fig, a = mp.show_edges(edges, image=mp.whitening(image))


# ## Matching Pursuit

# ### checking basic edge detection

# In[11]:


mp.set_size((128, 128))
#! defining a reference test image (see test_Image)
image = np.zeros((mp.pe.N_X, mp.pe.N_Y))
image[mp.pe.N_X//2:mp.pe.N_X//2+mp.pe.N_X//4, mp.pe.N_X//2:mp.pe.N_X//2+mp.pe.N_X//4] = 1
image[mp.pe.N_X//2:mp.pe.N_X//2+mp.pe.N_X//4, mp.pe.N_X//4:mp.pe.N_X//2] = -1
mp.imshow(image);


# In[12]:


get_ipython().run_cell_magic('writefile', 'experiment_test_MP.py', "import numpy as np\nfrom SparseEdges import SparseEdges\n\nmp = SparseEdges('https://raw.githubusercontent.com/bicv/SparseEdges/master/default_param.py')\nmp.pe.N = 128 # number of edges\nmp.pe.figsize_edges = 9\n\n#! defining a reference test image (see test_Image)\nimage = np.zeros((mp.pe.N_X, mp.pe.N_Y))\nimage[mp.pe.N_X//2:mp.pe.N_X//2+mp.pe.N_X//4, mp.pe.N_X//2:mp.pe.N_X//2+mp.pe.N_X//4] = 1\nimage[mp.pe.N_X//2:mp.pe.N_X//2+mp.pe.N_X//4, mp.pe.N_X//4:mp.pe.N_X//2] = -1\n\n\nwhite = mp.pipeline(image, center=True, do_whitening=True)\n\nimport os\nmatname = os.path.join(mp.pe.matpath, 'experiment_test_MP.npy')\ntry:\n    edges = np.load(matname)\nexcept Exception:\n    edges, C_res = mp.run_mp(white, verbose=True)\n    try:\n        os.mkdir(mp.pe.matpath)\n    except Exception:\n        pass\n    np.save(matname, edges)   \n")


# In[13]:


get_ipython().run_line_magic('run', 'experiment_test_MP.py')


# In[14]:


#fig, a = mp.show_edges(edges, image=mp.whitening(image))
fig, a = mp.show_edges(edges, image=mp.dewhitening(mp.reconstruct(edges)))


# In[15]:


get_ipython().run_cell_magic('writefile', 'experiment_test_whitening.py', "import numpy as np\nfrom SparseEdges import SparseEdges\nmp = SparseEdges('https://raw.githubusercontent.com/bicv/SparseEdges/master/default_param.py')\nmp.pe.N = 2**9\nmp.pe.MP_alpha = 1\nmp.init()\n\n\nimage = mp.imread('https://raw.githubusercontent.com/bicv/SLIP/master/database/serre07_targets/B_N107001.jpg')\n\nwhite = mp.pipeline(image, do_whitening=True)\n    \nimport os\nmatname = os.path.join(mp.pe.matpath, 'experiment_test_whitening.npy')\ntry:\n    edges = np.load(matname)\nexcept Exception:\n    edges, C_res = mp.run_mp(white, verbose=True)\n    np.save(matname, edges)    \n\n\nmatname_MSE = os.path.join(mp.pe.matpath, 'experiment_test_whitening_MSE.npy')\ntry:\n    MSE = np.load(matname_MSE)\nexcept Exception:\n    MSE = np.ones(mp.pe.N)\n    image_rec = np.zeros_like(image)\n    for i_N in range(mp.pe.N):\n        MSE[i_N] =  ((white-image_rec)**2).sum()\n        image_rec += mp.reconstruct(edges[:, i_N][:, np.newaxis])\n\n    np.save(matname_MSE, MSE)     \n")


# In[16]:


get_ipython().run_line_magic('run', 'experiment_test_whitening.py')


# In[17]:


#edges[0, :] = pe.N_X - edges[0,:] # HACK
#edges[2, :] = - edges[2,:] # HACK
#edges[:, 119] = edges[:, 118] # HACK
fig_width_pt = 318.670  # Get this from LaTeX using \showthe\columnwidth
inches_per_pt = 1.0/72.27               # Convert pt to inches
fig_width = fig_width_pt*inches_per_pt  # width in inches
mp.pe.figsize_edges = .382 * fig_width
fig, a = mp.show_edges(edges, image=image, color='brown', show_phase=False) #, pedestal=.0, gamma=.3
#import os
#fig.savefig(os.path.join(mp.pe.figpath, 'SparseLets_A.pdf'))


# In[18]:


#! recontructing the image from the edges
image_rec = mp.reconstruct(edges)
print('MSE-W = ', ((white-image_rec)**2).sum()/((white)**2).sum())


# In[19]:


fig, a = mp.show_edges(edges, image=image_rec*1., color='brown')


# Recontructing the original image by dewhitening the reconstructed image:

# In[20]:


fig, a = mp.show_edges(edges, image=mp.dewhitening(image_rec)*1., color='brown')


# In[21]:


MSE = np.load(os.path.join(mp.pe.matpath, 'experiment_test_whitening_MSE.npy'))
print(MSE.shape, mp.pe.N)


# In[22]:


# checking the quick computation of the MSE in MP
MSE_0 = (white**2).sum()
print('stats on white:', white.mean(), white.std(), np.sqrt(MSE_0), np.sqrt(MSE[0]))
print ('mp.pe.MP_alpha=', mp.pe.MP_alpha)    

MSE_MP = np.ones(mp.pe.N)
MSE_MP[1:] = 1. - np.cumsum(edges[4, :-1]**2) * (2 -  mp.pe.MP_alpha)/mp.pe.MP_alpha / MSE_0

plt.figure(figsize=(12,6))
plt.subplot(111)
plt.plot(np.sqrt(MSE/MSE[0]), 'ro', label='true', alpha=.2)
plt.plot(np.sqrt(MSE_MP), 'g--', label='MP-vec', lw=2)
plt.xlim([0, mp.pe.N])
plt.ylim([0, 1])
plt.xlabel('# atoms')
plt.ylabel('MSE')
_ = plt.legend()


# In[23]:


get_ipython().run_cell_magic('writefile', 'experiment_test_nowhitening.py', "import numpy as np\nfrom SparseEdges import SparseEdges\nmp = SparseEdges('https://raw.githubusercontent.com/bicv/SparseEdges/master/default_param.py')\nmp.pe.N = 256\nmp.pe.do_whitening = False\n\n#image = mp.imread('https://raw.githubusercontent.com/bicv/SLIP/master/database/serre07_targets/B_N107001.jpg')\nimage = mp.imread('../../SLIP/database/serre07_targets/B_N107001.jpg')\n\n\nmp.pe.figsize_edges = 9\nimage = mp.pipeline(image, do_whitening=False)\n\n\nimport os\nmatname = os.path.join(mp.pe.matpath, 'experiment_test_nowhitening.npy')\ntry:\n    edges = np.load(matname)\nexcept Exception:\n    edges, C_res = mp.run_mp(image, verbose=True)\n    np.save(matname, edges)    \n    \n")


# In[24]:


get_ipython().run_line_magic('run', 'experiment_test_nowhitening.py')


# In[25]:


#! recontructing the image from the edges
image_ = mp.reconstruct(edges)
print('MSE = ', ((image-image_)**2).sum()/((image)**2).sum())


# In[26]:


#! recontructing the image from the edges
image_ = mp.reconstruct(edges)
print('WMSE = ', ((mp.whitening(image-image_))**2).sum()/((mp.whitening(image))**2).sum())


# In[27]:


fig, a = mp.show_edges(edges, image=image_*1., show_mask=True, color='brown')


# ## some book keeping for the notebook

# In[28]:


get_ipython().run_line_magic('load_ext', 'watermark')
get_ipython().run_line_magic('watermark', '')


# In[29]:


get_ipython().run_line_magic('load_ext', 'version_information')
get_ipython().run_line_magic('version_information', 'numpy, scipy, matplotlib, sympy, pillow, imageio')


# In[30]:


get_ipython().run_line_magic('cd', '-q ../notebooks')