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

# # Exercise 17.1
# ## Speckle removal with denoising autoencoders (DAEs)
# Small-angle scattering of X-rays or neutrons enables insights into properties of nanostructured materials.
# In the case of X-rays from undulators at synchrotron radiation sources, extreme beam focusing can result in a high degree of coherence. Interference effects called speckles then appear naturally in the recorded data.
# This may be an unwanted effect that makes the measurements appear noisy.  
# 
# Try to remove the speckles from the given test samples:
#  1. Set up and train a deep convolutional autoencoder.
#  2. State the test loss and comment on the reconstruction results for some test images.
#  3. Apply the autoencoder to experimental data from partially coherent illumination and describe your observations.
# 
# Training DAEs on the provided data can be computationally demanding, thus, we recommend to use a GPU for this task.

# In[2]:


import numpy as np
import h5py
import matplotlib.pyplot as plt
from tensorflow import keras

layers = keras.layers



# In[2]:


import gdown
import os
url = "https://drive.google.com/u/0/uc?export=download&confirm=HgGH&id=1CJVbIkLmzDTfC6ftaIUPiuqjoC0ay-wr"
output = 'speckles.npz'

if os.path.exists(output) == False:
    gdown.download(url, output, quiet=True)


# ### Preprocess data

# In[3]:


f = np.load(output)
image_normal = f['target_images']
image_speckle = f['speckle_images']

# logarithmic intensity values
image_normal = np.log10(image_normal + 1.)                       
image_speckle = np.log10(image_speckle + 1.)

# norm input data to max. value of distorted scattering pattern
max_val = np.max(image_speckle, axis=(1, 2, 3), keepdims=True)

image_normal = image_normal / max_val
image_normal = np.clip(image_normal, 0., 1.1)  # limit maximum intensity values

image_speckle = image_speckle / max_val
image_speckle = np.clip(image_speckle, 0., 1.1)  # limit maximum intensity values


# In[4]:


n_train = 20000
n = image_normal.shape[0]

x_train_noisy, x_test_noisy = image_speckle[0:n_train], image_speckle[n_train:]
x_train, x_test = image_normal[0:n_train], image_normal[n_train:]


# ### Plot example data 

# In[5]:


plots = 10
plt.figure(1, (15, 3.5))
idx = np.random.choice(n, 10)

for i in range(plots):
    plt.subplot(2, plots, i+1)
    plt.imshow(image_speckle[idx[i],:,:,0])
    plt.xticks([])
    plt.yticks([])
    plt.title("Noisy")

    plt.subplot(2, plots, plots+i+1)
    plt.imshow(image_normal[idx[i],:,:,0])
    plt.xticks([])
    plt.yticks([])
    plt.title("True")
    
plt.show()


# ## Model building

# In[ ]:


input_img = layers.Input(shape=(64, 64, 1))
# build your model here

autoencoder = keras.models.Model(input_img, decoded)

print(autoencoder.summary())


# #### Plot model to visualize the shortcuts

# In[ ]:


keras.utils.plot_model(autoencoder, show_shapes=True, dpi=60)


# #### Compile and train the DAE

# In[ ]:


autoencoder.compile(...)


# In[ ]:


results = autoencoder.fit(...)


# ### Plot training history

# In[ ]:


plt.figure(1, (12, 4))
plt.subplot(1, 2, 1)
plt.plot(results.history['loss'])
plt.plot(results.history['val_loss'])
plt.ylabel('loss')
plt.xlabel('epoch')
plt.legend(['train', 'val'], loc='upper right')


# ### Investigate the DAE performance using the test data

# In[ ]:


preds = autoencoder.predict(x_test_noisy, verbose=1)


# In[ ]:


plots = 10
n_test = x_test.shape[0]
plt.figure(1, (15, 5.5))
idx = np.random.choice(n_test, 10)

for i in range(plots):
    plt.subplot(4, plots, i+1)
    plt.imshow(x_test_noisy[idx[i],:,:,0])
    plt.xticks([])
    plt.yticks([])
    plt.title("noisy")

    plt.subplot(4, plots, plots+i+1)
    plt.imshow(preds[idx[i],:,:,0])
    plt.xticks([])
    plt.yticks([])
    plt.title("denoised")

    plt.subplot(4, plots, 2*plots+i+1)
    plt.imshow(x_test[idx[i],:,:,0] - preds[idx[i],:,:,0])
    plt.xticks([])
    plt.yticks([])
    plt.title("true")
    
    plt.subplot(4, plots, 3*plots+i+1)
    plt.imshow(x_test[idx[i],:,:,0] - preds[idx[i],:,:,0])
    plt.xticks([])
    plt.yticks([])
    plt.title("residuals")

plt.show()