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

# # Energy auto-encoder: comparison with Xavier Primal-Dual matlab implementation

# ## Setup

# In[ ]:


import numpy as np
import numpy.linalg as la
import scipy.io
import matplotlib.pyplot as plt
get_ipython().run_line_magic('matplotlib', 'inline')
from pyunlocbox import functions, solvers
import time
import scipy, matplotlib, pyunlocbox  # For versions only.

# Import auto-encoder definition.
get_ipython().run_line_magic('run', 'auto_encoder.ipynb')

print('Software versions:')
for pkg in [np, matplotlib, scipy, pyunlocbox]:
    print('  %s: %s' % (pkg.__name__, pkg.__version__))


# ## Hyper-parameters

# * The $\lambda$ are the relative importance of each term in the composite objective function.

# In[ ]:


ld = 10  # Xavier sets the weight of the L1 regularization to 1e-1.


# ## Data

# * The set of data vectors $X \in R^{n \times N}$ is given by patches extracted from a grayscale image.
# * There is as many patches as pixels in the image.
# * The saved patches already have zero mean.

# In[ ]:


mat = scipy.io.loadmat('data/xavier_X.mat')
X = mat['X'].T

N, n = X.shape
Np = np.sqrt(n)
print('N = %d samples with dimensionality n = %d (patches of %dx%d).' % (N, n, Np, Np))

plt.figure(figsize=(8,5))
patches = [24, 1000, 2004, 10782]
for k in range(len(patches)):
    patch = patches[k]
    img = np.reshape(X[patch, :], (Np, Np))
    plt.subplot(1, 4, k+1)
    plt.imshow(img, cmap='gray')
    plt.title('Patch %d' % patch)
plt.show()


# ## Initial conditions

# * $Z$ is drawn from a uniform distribution in ]0,1[.
# * Same for $D$. Its columns were then normalized to unit L2 norm.
# * The sparse code dimensionality $m$ should be greater than $n$ for an overcomplete representation but much smaller than $N$ to avoid over-fitting.
# * The dictionary and sparse codes are initialized this way in the auto_encoder class.

# In[ ]:


if False:
    mat = scipy.io.loadmat('data/xavier_initZD.mat')
    Zinit = mat['Zinit']
    Dinit = mat['Dinit']

    m, N = Zinit.shape
    n, m = Dinit.shape
    print('Sparse code dimensionality m = %d --> %s dictionary' % (m, 'overcomplete' if m > n else 'undercomplete'))

    print('mean(Z) = %f' % np.mean(Zinit))

    d = np.sqrt(np.sum(Dinit*Dinit, axis=0))
    print('Constraints on D: %s' % np.alltrue(d <= 1+1e-15))


# ## Algorithm

# Given $X \in R^{n \times N}$, solve $\min\limits_{Z \in R^{m \times N}, D \in R^{n \times m}} \frac{\lambda_d}{2} \|X - DZ\|_F^2 + \|Z\|_1$ s.t. $\|d_i\|_2 \leq 1$, $i = 1, \ldots, m$

# In[ ]:


ae = auto_encoder(m=64, ld=ld, rtol=1e-5, xtol=None, N_inner=100, N_outer=30)
#ae = auto_encoder(m=64, ld=ld, rtol=None, xtol=1e-6, N_inner=100, N_outer=15)
tstart = time.time()
Z = ae.fit_transform(X)
print('Elapsed time: {:.0f} seconds'.format(time.time() - tstart))


# In[ ]:


if False:
    tstart = time.time()
    Z = ae.transform(X)
    print('Elapsed time: {:.0f} seconds'.format(time.time() - tstart))


# ### Convergence

# In[ ]:


ae.plot_convergence()


# ## Solution analysis

# ### Solution from Xavier

# In[ ]:


mat = scipy.io.loadmat('data/xavier_ZD.mat')
Zxavier = mat['Z'].T
Dxavier = mat['D'].T
print('Elapsed time: %d seconds' % mat['exectime'])


# ### Objective

# In[ ]:


objective(X, Zxavier, Dxavier, ld)
objective(X, Z, ae.D, ld)


# ### Sparse codes

# In[ ]:


sparse_codes(Zxavier)
sparse_codes(Z)


# ### Dictionary

# In[ ]:


dictionary(Dxavier)
dictionary(ae.D)


# In[ ]:


atoms(Dxavier, Np)
atoms(ae.D, Np)