"""
radon.py - Radon and inverse radon transforms

Based on code of Justin K. Romberg 
(http://www.clear.rice.edu/elec431/projects96/DSP/bpanalysis.html)
J. Gillam and Chris Griffin.

References: 
    -B.R. Ramesh, N. Srinivasa, K. Rajgopal, "An Algorithm for Computing
    the Discrete Radon Transform With Some Applications", Proceedings of
    the Fourth IEEE Region 10 International Conference, TENCON '89, 1989.
    -A. C. Kak, Malcolm Slaney, "Principles of Computerized Tomographic
    Imaging", IEEE Press 1988.
"""

import numpy as np
from scipy.fftpack import fftshift, fft, ifft
from ._project import homography

def radon(image, theta=None):
    """
    Calculates the radon transform of an image given specified
    projection angles.

    Parameters
    ----------
    image : array_like, dtype=float
        Input image.
    theta : array_like, dtype=float, optional (default np.arange(180))
        Projection angles (in degrees).
    
    Returns
    -------
    output : ndarray
        Radon transform.

    """
    if image.ndim != 2:
        raise ValueError('The input image must be 2-D')
    if theta == None:
        theta = np.arange(180)    
    height, width = image.shape
    diagonal = np.sqrt(height ** 2 + width ** 2)
    heightpad = np.ceil(diagonal - height) + 2
    widthpad = np.ceil(diagonal - width) + 2
    padded_image = np.zeros((int(height + heightpad),
                             int(width + widthpad)))
    y0, y1 = int(np.ceil(heightpad / 2)), \
             int((np.ceil(heightpad / 2) + height))
    x0, x1 = int((np.ceil(widthpad / 2))), \
             int((np.ceil(widthpad / 2) + width))

    padded_image[y0:y1, x0:x1] = image
    out = np.zeros((max(padded_image.shape), len(theta)))

    h, w = padded_image.shape
    shift0 = np.array([[1, 0, -w/2.],
                       [0, 1, -h/2.],
                       [0, 0, 1]])

    shift1 = np.array([[1, 0, w/2.],
                       [0, 1, h/2.],
                       [0, 0, 1]])

    def build_rotation(theta):
        T = -np.deg2rad(theta)

        R = np.array([[np.cos(T), -np.sin(T), 0],
                      [np.sin(T), np.cos(T), 0],
                      [0, 0, 1]])

        return shift1.dot(R).dot(shift0)

    for i in range(len(theta)):
        rotated = homography(padded_image,
                             build_rotation(-theta[i]))

        out[:,i] = rotated.sum(0)[::-1]

    return out


def iradon(radon_image, theta=None, output_size=None,
           filter="ramp", interpolation="linear"):
    """
    Inverse radon transform.

    Reconstruct an image from the radon transform.
        
    Parameters
    ----------
    radon_image : array_like, dtype=float
        Image containing radon transform (sinogram).
    theta : array_like, dtype=float, optional (default np.arange(180))
        Reconstruction angles (in degrees).
    output_size : int
        Number of rows and columns in the reconstruction.
    filter : str, optional (default ramp)
        Filter used in frequency domain filtering. Ramp filter used by default.
        Filters available: ramp, shepp-logan, cosine, hamming, hann
        Assign None to use no filter.
    interpolation : str, optional (default linear)
        Interpolation method used in reconstruction.
        Methods available: nearest, linear.
    
    Returns
    -------
    output : ndarray
      Reconstructed image.

    Notes
    -----
    It applies the fourier slice theorem to reconstruct an image by
    multiplying the frequency domain of the filter with the FFT of the
    projection data.
    
    """
    if radon_image.ndim != 2:
        raise ValueError('The input image must be 2-D')
    if theta == None:
        theta = np.arange(180)
    th = (np.pi / 180.0) * theta        
    # if output size not specified, estimate from input radon image
    if not output_size:
        output_size = 2 * np.floor(radon_image.shape[0] / (2 * np.sqrt(2)))
    n = radon_image.shape[0]
  
    img = radon_image.copy()
    # resize image to next power of two for fourier analysis
    # speeds up fourier and lessens artifacts
    order = max(64, 2 ** np.ceil(np.log(2 * n) / np.log(2)))
    # zero pad input image
    img.resize((order, img.shape[1]))
    # construct the fourier filter
    freqs = np.zeros((order, 1))
    
    f = fftshift(abs(np.mgrid[-1:1:2 / order])).reshape(-1, 1)
    w = 2 * np.pi * f
    # start from first element to avoid divide by zero
    if filter == "ramp":
        pass
    elif filter == "shepp-logan":
        f[1:] = f[1:] * np.sin(w[1:] / 2) / (w[1:] / 2)
    elif filter == "cosine":
       f[1:] = f[1:] * np.cos(w[1:] / 2)
    elif filter == "hamming":
       f[1:] = f[1:] * (0.54 + 0.46 * np.cos(w[1:]))
    elif filter == "hann":
       f[1:] = f[1:] * (1 + np.cos(w[1:])) / 2
    elif filter == None:
        f[1:] = 1
    else:
        raise ValueError("Unknown filter: %s" % filter)
        
    filter_ft = np.tile(f, (1, len(theta)))        
    # apply filter in fourier domain
    projection = fft(img, axis=0) * filter_ft
    radon_filtered = np.real(ifft(projection, axis=0))
    # resize filtered image back to original size
    radon_filtered = radon_filtered[:radon_image.shape[0], :]
    reconstructed = np.zeros((output_size, output_size))
    mid_index = np.ceil(n/2);
    x = output_size
    y = output_size
    [X, Y] = np.mgrid[0.0:x, 0.0:y]
    xpr = X - (output_size + 1.0) / 2.0
    ypr = Y - (output_size + 1.0) / 2.0
    
    # reconstruct image by interpolation
    if interpolation == "nearest":   
        for i in range(len(theta)):
            k = np.round(mid_index + xpr * np.sin(th[i]) - ypr * np.cos(th[i]))
            reconstructed += radon_filtered[
                ((((k > 0) & (k < n)) * k) - 1).astype(np.int), i]
    elif interpolation == "linear":
        for i in range(len(theta)):
          t = xpr*np.sin(th[i]) - ypr*np.cos(th[i])
          a = np.floor(t)
          b = mid_index + a 
          b0 = ((((b + 1 > 0) & (b + 1 < n)) * (b + 1)) - 1).astype(np.int)
          b1 = ((((b > 0) & (b < n)) * b) - 1).astype(np.int)
          reconstructed += (t - a) * radon_filtered[b0, i] + \
                           (a - t + 1) * radon_filtered[b1, i]
    else:
        raise ValueError("Unknown interpolation: %s" % interpolation)
        
    return reconstructed * np.pi / (2 * len(th))