"""
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.misc import imrotate
from scipy.interpolate import interp1d
from scipy.fftpack import fftshift, fft, ifft
import math


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)))
    for i in range(len(theta)):
        rotated = imrotate(padded_image, -theta[i])
        out[:,i] = rotated.sum(0)[::-1]
    return out


def iradon(radon_image, theta=None, output_size=None, filter="ramp", interpolation="linear"):
    """
    Reconstructs an image from radon transformed data.
        
    Parameters
    ----------
    radon_image : array_like, dtype=float
        Image containing radon transform.
    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 = (math.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 * math.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]
# XXX slow with some artifacts
#    elif interpolation == "spline":
#        axis = np.arange(0, radon_filtered.shape[0]) - mid_index
#        for i in range(len(theta)):
#            print i
#            t = xpr*np.sin(th[i]) - ypr*np.cos(th[i])
#            #f = interp1d(axis, radon_filtered[:, i], kind="cubic", bounds_error=False, fill_value=0) 
#            f = interp1d(axis, radon_filtered[:, i], kind="linear", bounds_error=False, fill_value=0)
#            reconstructed += f(t).reshape(output_size, output_size)
    else:
        raise ValueError("Unknown interpolation: %s" % interpolation)
        
    return reconstructed * math.pi / (2*len(th))