import numpy as np
import matplotlib.pyplot as plt
%matplotlib inline
from IPython.display import HTML 
HTML('../style/course.css') #apply general CSS

import matplotlib
from scipy import optimize
import astropy.io.fits

matplotlib.rcParams.update({'font.size': 18})
matplotlib.rcParams.update({'figure.figsize': [12,8]} )

"""
Explanation: Outline
Glossary
6. Deconvolution in Imaging  
Previous: 6.3 Residuals and Image Quality  
Next: 6.x Further Reading and References
End of explanation
"""


noise_image = "../data/fits/noise_image.fits"
with astropy.io.fits.open(noise_image) as hdu:
    data = hdu[0].data[0,0,...]

fig, (image, hist) = plt.subplots(1, 2, figsize=(18,6))
histogram, bins = np.histogram(data.flatten(), bins=401)

dmin = data.min()
dmax = data.max()
x = np.linspace(dmin, dmax, 401)

im = image.imshow(data)

mean = data.mean()
sigma = data.std()
peak = histogram.max()

gauss = lambda x, amp, mean, sigma: amp*np.exp( -(x-mean)**2/(2*sigma**2))

fitdata = gauss(x, peak, mean, sigma)

plt.plot(x, fitdata)
plt.plot(x, histogram, "o")
plt.yscale('log')
plt.ylim(1)

"""
Explanation: 6.5 Source Finding
In radio astronomy, source finding is the process through which the attributes of radio sources -- such as flux density and mophorlogy -- are measured from data. In this section we will only cover source finding in the image plane.
Source finding techniques usually involve four steps, i) charecterizing the noise (or background estimation), ii) thresholding the data based on knowledge of the noise, iii) fidinding regions in the thresholded image which neighbouring pixels are "similar" (this is that same as blob detection in image processing), and iv) parameterizing these 'blobs' through a function (usually a 2D Gaussian). The source attributes are then estimated from the parameterization of the blobs.
6.5.1 Noise Charecterization
As mentioned before, the radio data we process with source finders is noisy. To charecterise this noise we need to make a few assumptions about its nature, namely we assume that the niose results from some stochastic process and that it can be described by a normal distribution
$$ G(x \, | \, \mu,\sigma^2) = \frac{1}{\sigma \sqrt{2\pi}}\text{exp}\left( \frac{-(x-\mu)^2}{2\sigma^2}\right) $$  
where, $\mu$ is the mean (or expected value) of the varaible $x$, and $\sigma^2$ is the varience of the distribution; $\sigma$ is the standard deviation. Hence, the noise can be parameterized through the mean and the standard deviation. Let us illustrate this with an example. Bellow is a noise image from a MeerKAT simulation, along with a histogram of of the pixels (in log space).
End of explanation
"""


noise_image = "../data/fits/star_model_image.fits"
with astropy.io.fits.open(noise_image) as hdu:
    data = hdu[0].data[0,0,...]

fig, (image, hist) = plt.subplots(1, 2, figsize=(18,6))
histogram, bins = np.histogram(data.flatten(), bins=101)


dmin = data.min()
dmax = data.max()
x = np.linspace(dmin, dmax, 101)

im = image.imshow(data)

mean = data.mean()
sigma_std = data.std()

peak = histogram.max()

gauss = lambda x, amp, mean, sigma: amp*np.exp( -(x-mean)**2/(2*sigma**2))

fitdata_std = gauss(x, peak, mean, sigma_std)

plt.plot(x, fitdata_std, label="STD DEV")

plt.plot(x, histogram, "o", label="Data")
plt.legend(loc=1)

plt.yscale('log')
plt.ylim(1)

"""
Explanation: Now, in reality the noise has to measured in the presence of astrophysical emission. Futhermore, radio iamges are also contaminated by various instrumental effects which can manifest as spurious emission in the image domain. All these factors make it difficult to charercterize the noise in a synthesized image. Since the noise generally dominates the images, the mean and standard deviation of the entire image are still fairly good approximations of the noise. Let us now insert a few sources (image and flux distribution shown below) in the noise image from earlier and then try to estimate noise.
End of explanation
"""


mean = data.mean()
sigma_std = data.std()
sigma_neg = data[data<0].std() * 2
mad_mean = lambda a: np.mean( abs(a - np.mean(a) ))
sigma_mad_median = np.median( abs(data - np.median(data) ))

mad_mean = lambda a: np.mean( abs(a - np.mean(a) ))
sigma_mad_mean = mad_mean(data)

peak = histogram.max()

gauss = lambda x, amp, mean, sigma: amp*np.exp( -(x-mean)**2/(2*sigma**2))

fitdata_std = gauss(x, peak, mean, sigma_std)
fitdata_mad_median = gauss(x, peak, mean, sigma_mad_median)
fitdata_mad_mean = gauss(x, peak, mean, sigma_mad_mean)
fitdata_neg = gauss(x, peak, mean, sigma_neg)

plt.plot(x, fitdata_std, label="STD DEV")
plt.plot(x, fitdata_mad_median, label="MAD Median")
plt.plot(x, fitdata_mad_mean, label="MAD Mean")
plt.plot(x, fitdata_neg, label="Negative STD DEV")
plt.plot(x, histogram, "o", label="Data")
plt.legend(loc=1)

plt.yscale('log')
plt.ylim(1)

"""
Explanation: The pixel statistics of the image are no longer Gaussian as apparent from the long trail of the flux distribution. Constructing a Gaussian model from the mean and standard deviation results in a poor fit (blue line in the figure on the right). A better method to estimate the variance is to measure the dispersion of the data points about the mean (or median), this is the mean/median absolute deviation (MAD) technique. We will refer to the to median absolute deviation as the MAD Median, and the mean absolute deviation as the MAD Mean. A synthesis imaging specfic method to estimate the varience of the noise is to only consider the negative pixels. This works under the assumption that all the astrophysical emission (at least in Stokes I) has a positive flux density. The Figure below shows noise estimates from methods mentioned above.
End of explanation
"""


def gauss2D(x, y, amp, mean_x, mean_y, sigma_x, sigma_y):
    """ Generate a 2D Gaussian image"""
    gx = -(x - mean_x)**2/(2*sigma_x**2)
    gy = -(y - mean_y)**2/(2*sigma_y**2)
    
    return amp * np.exp( gx + gy)

def err(p, xx, yy, data):
    """2D Gaussian error function"""
    return gauss2D(xx.flatten(), yy.flatten(), *p) - data.flatten()

def fit_gaussian(data, psf_pix):
    """Fit a gaussian to a 2D data set"""
    
    width = data.shape[0]
    mean_x, mean_y = width/2, width/2
    amp = data.max()
    sigma_x, sigma_y = psf_pix, psf_pix
    params0 = amp, mean_x, mean_y, sigma_x,sigma_y
        
    npix_x, npix_y = data.shape
    x = np.linspace(0, npix_x, npix_x)
    y = np.linspace(0, npix_y, npix_y)
    xx, yy = np.meshgrid(x, y)
        
        
    params, pcov, infoDict, errmsg, sucess = optimize.leastsq(err, 
                        params0, args=(xx.flatten(), yy.flatten(),
                        data.flatten()), full_output=1)
    
        
    perr = abs(np.diagonal(pcov))**0.5
    model = gauss2D(xx, yy, *params)
    
    return params, perr, model

def source_finder(data, peak, boundary, width, psf_pix):
    """A simple source finding tool"""
    
    # first we make an estimate of the noise. Lets use the MAD mean
    sigma_noise = mad_mean(data)

    # Use noise estimate to set peak and boundary thresholds
    peak_sigma = sigma_noise*peak
    boundary_sigma = sigma_noise*boundary
    
    # Pad the image to avoid hitting the edge of the image
    pad = width*2
    residual = np.pad(data, pad_width=((pad, pad), (pad, pad)), mode="constant")
    model = np.zeros(residual.shape)
    
    # Create slice to remove the padding later on
    imslice = [slice(pad, -pad), slice(pad,-pad)]
    
    catalog = [] 
    
    # We will need to convert the fitted sigma values to a width
    FWHM = 2*np.sqrt(2*np.log(2))
    
    while True:
        
        # Check if the brightest pixel is at least as bright as the sigma_peak
        # Otherwise stop.
        max_pix = residual.max()
        if max_pix<peak_sigma:
            break
        
        xpix, ypix = np.where(residual==max_pix)
        xpix = xpix[0] # Get first element
        ypix = ypix[0] # Get first element
        
        # Make slice that selects box of size width centred around bright brightest pixel
        subim_slice = [ slice(xpix-width/2, xpix+width/2),
                    slice(ypix-width/2, ypix+width/2) ]
        
        # apply slice to get subimage
        subimage = residual[subim_slice]
        
        
        # blank out pixels below the boundary threshold
        mask = subimage > boundary_sigma
        
        # Fit gaussian to submimage
        params, perr, _model = fit_gaussian(subimage*mask, psf_pix)
        
        amp, mean_x, mean_y, sigma_x,sigma_y = params
        amp_err, mean_x_err, mean_y_err, sigma_x_err, sigma_y_err = perr
        
        # Remember to reposition the source in original image
        pos_x = xpix + (width/2 - mean_x) - pad
        pos_y = ypix + (width/2 - mean_y) - pad
        
        # Convert sigma values to FWHM lengths
        size_x = FWHM*sigma_x
        size_y = FWHM*sigma_y
        
        # Add modelled source to model image
        model[subim_slice] = _model
        
        # create new source
        source = (
            amp,
            pos_x,
            pos_y,
            size_x,
            size_y
            )
                  
        # add source to catalogue
        catalog.append(source)
        
        # update residual image
        residual[subim_slice] -= _model       
        
    return catalog, model[imslice], residual[imslice], sigma_noise


"""
Explanation: The MAD and negtive value standard deviation methods produce a better solution to the noise distribution in the presence of sources.
6.5.2 Blob Detection and Charercterization
Once the noise has been estimated, the next step is to find and charecterize sources in the image. Generically in image processing this is known as blob detection. In a simple case during synthesis imaging we define a blob as a group contiguous pixels whose spatial intensity profile can be modelled by a 2D Gaussian function. Of course, more advanced functions could be used. Generally, we would like to group together near by pixels, such as spatially 'close' sky model components from deconvolution, into a single complex source. Our interferometric array has finite spatial resolution, so we can further constrain our blobs not to be significantly smaller than the image resolution. We ndefine two further constraints of a blob, the peak and boundary thresholds. The peak threshold, defined as
$$ 
    \sigma_\text{peak} = n * \sigma,
$$
is the minimum intensity the maximum pixel in a blob must have relative to the imge noise. That is, all blobs with peak pixel lower than $\sigma_\text{peak}$ will be excluded from being considered sources. And the bouandary threshold
$$
    \sigma_\text{boundary} = m * \sigma,
$$
defines the boundary of a blob, $m$ and $n$ are natural numbers with $m$ < $n$. 
6.5.2.1 A simple source finder
We are now in a position to write a simple source finder. To do so we implement the following steps: 

Estimate the image noise and set peak and boundary threshold values.
Blank out all pixel values below the boundary value.
Find Peaks in image.
For each peak, fit a 2D Gaussian and subtract the Gaussian fit from the image.
Repeat until the image has no pixels above the detection threshold.
End of explanation
"""


test_image = "../data/fits/star_model_image.fits"
with astropy.io.fits.open(test_image) as hdu:
    data = hdu[0].data[0,0,...]
    
catalog, model, residual, sigma_noise = source_finder(data, 5, 2, 50, 10)

print "Peak_Flux     Pix_x     Pix_y     Size_x     Size_y"
for source in catalog:
    print "  %.4f     %.1f     %.1f     %.2f     %.2f"%source

fig, (img, mod, res) = plt.subplots(1, 3, figsize=(24,12))
vmin, vmax = sigma_noise, data.max()

im = img.imshow(data, vmin=vmin, vmax=vmax)
img.set_title("Data")

mod.imshow(model, vmin=vmin, vmax=vmax)
mod.set_title("Model")

res.imshow(residual, vmin=vmin, vmax=vmax)
res.set_title("Residual")

cbar_ax = fig.add_axes([0.92, 0.25, 0.02, 0.5])
fig.colorbar(im, cax=cbar_ax, format="%.2g")

"""
Explanation: Using this source finder we can produce a sky model which contains all 17 sources in our test image from earlier in the section.
End of explanation
"""


