"""
=======================================================
Denoise images using Local PCA via empirical thresholds
=======================================================

PCA-based denoising algorithms are effective denoising methods because they
explore the redundancy of the multi-dimensional information of
diffusion-weighted datasets. In this example, we will show how to
perform the PCA-based denoising using the algorithm proposed by Manjon et al.
[Manjon2013]_.

This algorithm involves the following steps:

* First, we estimate the local noise variance at each voxel.

* Then, we apply PCA in local patches around each voxel over the gradient
  directions.

* Finally, we threshold the eigenvalues based on the local estimate of sigma
  and then do a PCA reconstruction


To perform PCA denoising without a prior noise standard deviation estimate
please see the following example instead: :ref:`denoise_mppca`

Let's load the necessary modules
"""

import numpy as np
import nibabel as nib
import matplotlib.pyplot as plt
from time import time
from dipy.core.gradients import gradient_table
from dipy.denoise.localpca import localpca
from dipy.denoise.pca_noise_estimate import pca_noise_estimate
from dipy.data import get_fnames
from dipy.io.image import load_nifti
from dipy.io.gradients import read_bvals_bvecs

"""

Load one of the datasets. These data were acquired with 63 gradients and 1
non-diffusion (b=0) image.

"""

dwi_fname, dwi_bval_fname, dwi_bvec_fname = get_fnames('isbi2013_2shell')
data, affine = load_nifti(dwi_fname)
bvals, bvecs = read_bvals_bvecs(dwi_bval_fname, dwi_bvec_fname)
gtab = gradient_table(bvals, bvecs)

print("Input Volume", data.shape)

"""
Estimate the noise standard deviation
=====================================

We use the ``pca_noise_estimate`` method to estimate the value of sigma to be
used in the local PCA algorithm proposed by Manjon et al. [Manjon2013]_.
It takes both data and the gradient table object as input and returns an
estimate of local noise standard deviation as a 3D array. We return a smoothed
version, where a Gaussian filter with radius 3 voxels has been applied to the
estimate of the noise before returning it.

We correct for the bias due to Rician noise, based on an equation developed by
Koay and Basser [Koay2006]_.

"""

t = time()
sigma = pca_noise_estimate(data, gtab, correct_bias=True, smooth=3)
print("Sigma estimation time", time() - t)

"""
Perform the localPCA using the function `localpca`
==================================================

The localpca algorithm takes into account the multi-dimensional information of
the diffusion MR data. It performs PCA on a local 4D patch and
then removes the noise components by thresholding the lowest eigenvalues.
The eigenvalue threshold will be computed from the local variance estimate
performed by the ``pca_noise_estimate`` function, if this is inputted in the
main ``localpca`` function. The relationship between the noise variance
estimate and the eigenvalue threshold can be adjusted using the input parameter
``tau_factor``. According to Manjon et al. [Manjon2013]_, this parameter is set
to 2.3.
"""

t = time()

denoised_arr = localpca(data, sigma, tau_factor=2.3, patch_radius=2)

print("Time taken for local PCA (slow)", -t + time())

"""
The ``localpca`` function returns the denoised data which is plotted below
(middle panel) together with the original version of the data (left panel) and
the algorithm residual image (right panel) .
"""

sli = data.shape[2] // 2
gra = data.shape[3] // 2
orig = data[:, :, sli, gra]
den = denoised_arr[:, :, sli, gra]
rms_diff = np.sqrt((orig - den) ** 2)

fig, ax = plt.subplots(1, 3)
ax[0].imshow(orig, cmap='gray', origin='lower', interpolation='none')
ax[0].set_title('Original')
ax[0].set_axis_off()
ax[1].imshow(den, cmap='gray', origin='lower', interpolation='none')
ax[1].set_title('Denoised Output')
ax[1].set_axis_off()
ax[2].imshow(rms_diff, cmap='gray', origin='lower', interpolation='none')
ax[2].set_title('Residual')
ax[2].set_axis_off()
plt.savefig('denoised_localpca.png', bbox_inches='tight')

print("The result saved in denoised_localpca.png")

"""
.. figure:: denoised_localpca.png
   :align: center

Below we show how the denoised data can be saved.
"""

nib.save(nib.Nifti1Image(denoised_arr,
                         affine), 'denoised_localpca.nii.gz')

print("Entire denoised data saved in denoised_localpca.nii.gz")

"""
.. [Manjon2013] Manjon JV, Coupe P, Concha L, Buades A, Collins DL "Diffusion
                Weighted Image Denoising Using Overcomplete Local PCA" (2013).
                PLoS ONE 8(9): e73021. doi:10.1371/journal.pone.0073021.

.. [Koay2006]  Koay CG, Basser PJ (2006). "Analytically exact correction scheme
               for signal extraction from noisy magnitude MR signals". JMR 179:
               317-322.

.. include:: ../links_names.inc
"""