# Copyright 2014 Diamond Light Source Ltd.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""
.. module:: paganin_filter
   :platform: Unix
   :synopsis: A plugin to apply the Paganin filter.

.. moduleauthor:: Nghia Vo <scientificsoftware@diamond.ac.uk>

"""
import math
import logging
import numpy as np
import pyfftw.interfaces.scipy_fftpack as fft

from savu.plugins.filters.base_filter import BaseFilter
from savu.plugins.driver.cpu_plugin import CpuPlugin
from savu.plugins.utils import register_plugin, dawn_compatible
from savu.data.plugin_list import CitationInformation


@dawn_compatible
@register_plugin
class PaganinFilter(BaseFilter, CpuPlugin):
    """
    A plugin to apply Paganin filter (contrast enhancement) on projections.

    :param Energy: Given X-ray energy in keV. Default: 53.0.
    :param Distance: Distance from sample to detection - Unit is \
        metre. Default: 1.0.
    :param Resolution: Pixel size - Unit is micron. Default: 1.28.
    :u*param Ratio: ratio of delta/beta. Default: 250.0.
    :param Padtopbottom: Pad to the top and bottom of projection. Default: 10.
    :param Padleftright: Pad to the left and right of projection. Default: 10.
    :param Padmethod: Numpy pad method. Default: 'edge'.
    :param increment: Increment all values by this amount before taking the \
        log. Default: 0.0.

    :config_warn: The 'log' parameter in the reconstruction should be set to \
    FALSE.
    :config_warn: Previewing a subset of sinograms will alter the result, due \
    to the global nature of this filter. If this is necessary, ensure they \
    are consecutive.

    """

    def __init__(self):
        logging.debug("initialising Paganin Filter")
        logging.debug("Calling super to make sure that all superclases are " +
                      " initialised")
        super(PaganinFilter, self).__init__("PaganinFilter")
        self.filtercomplex = None
        self.count = 0

    def set_filter_padding(self, in_pData, out_pData):
        in_data = self.get_in_datasets()[0]
        det_x = in_data.get_data_dimension_by_axis_label('detector_x')
        det_y = in_data.get_data_dimension_by_axis_label('detector_y')
        pad_det_y = '%s.%s' % (det_y, self.parameters['Padtopbottom'])
        pad_det_x = '%s.%s' % (det_x, self.parameters['Padleftright'])
        mode = self.parameters['Padmethod']
        pad_dict = {'pad_directions': [pad_det_x, pad_det_y], 'pad_mode': mode}
        in_pData[0].padding = pad_dict
        out_pData[0].padding = pad_dict

    def pre_process(self):
        self._setup_paganin(*self.get_plugin_in_datasets()[0].get_shape())

    def _setup_paganin(self, height, width):
        micron = 10**(-6)
        keV = 1000.0
        distance = self.parameters['Distance']
        energy = self.parameters['Energy'] * keV
        resolution = self.parameters['Resolution'] * micron
        wavelength = (1240.0 / energy) * 10.0**(-9)
        ratio = self.parameters['Ratio']

        height1 = height + 2 * self.parameters['Padtopbottom']
        width1 = width + 2 * self.parameters['Padleftright']
        centery = np.ceil(height1 / 2.0) - 1.0
        centerx = np.ceil(width1 / 2.0) - 1.0

        # Define the paganin filter
        dpx = 1.0 / (width1 * resolution)
        dpy = 1.0 / (height1 * resolution)
        pxlist = (np.arange(width1) - centerx) * dpx
        pylist = (np.arange(height1) - centery) * dpy
        pxx = np.zeros((height1, width1), dtype=np.float32)
        pxx[:, 0:width1] = pxlist
        pyy = np.zeros((height1, width1), dtype=np.float32)
        pyy[0:height1, :] = np.reshape(pylist, (height1, 1))
        pd = (pxx * pxx + pyy * pyy) * wavelength * distance * math.pi

        filter1 = 1.0 + ratio * pd
        self.filtercomplex = filter1 + filter1 * 1j

    def _paganin(self, data):
        pci1 = fft.fft2(np.float32(data))
        pci2 = fft.fftshift(pci1) / self.filtercomplex
        fpci = np.abs(fft.ifft2(pci2))
        result = -0.5 * self.parameters['Ratio'] * np.log(
            fpci + self.parameters['increment'])
        return result

    def process_frames(self, data):
        proj = np.nan_to_num(data[0])  # Noted performance
        proj[proj == 0] = 1.0
        return self._paganin(proj)

    def get_max_frames(self):
        return 'single'

    def get_citation_information(self):
        cite_info = CitationInformation()
        cite_info.description = \
            ("The contrast enhancement used in this processing chain is taken\
             from this work.")
        cite_info.bibtex = \
            ("@article{paganin2002simultaneous,\n" +
             "title={Simultaneous phase and amplitude extraction from a single\
             defocused image of a homogeneous object},\n" +
             "author={Paganin, David and Mayo, SC and Gureyev, Tim E and \
             Miller, Peter R and Wilkins, Steve W},\n" +
             "journal={Journal of microscopy},\n" +
             "volume={206},\n" +
             "number={1},\n" +
             "pages={33--40},\n" +
             "year={2002},\n" +
             "publisher={Wiley Online Library}" +
             "}")
        cite_info.endnote = \
            ("%0 Journal Article\n" +
             "%T Simultaneous phase and amplitude extraction from a single \
             defocused image of a homogeneous object\n" +
             "%A Paganin, David\n" +
             "%A Mayo, SC\n" +
             "%A Gureyev, Tim E\n" +
             "%A Miller, Peter R\n" +
             "%A Wilkins, Steve W\n" +
             "%J Journal of microscopy\n" +
             "%V 206\n" +
             "%N 1\n" +
             "%P 33-40\n" +
             "%@ 1365-2818\n" +
             "%D 2002\n" +
             "%I Wiley Online Library\n")
        cite_info.doi = "doi: DOI: 10.1046/j.1365-2818.2002.01010.x"
        return cite_info
