import astropy.io.fits as pyfits
import astropy.visualization as viz
import matplotlib.pyplot as plt
import numpy as np
%matplotlib inline
plt.rcParams['figure.figsize'] = (10.0, 10.0)

"""
Explanation: Modeling the X-ray image data
In this notebook, we'll take a closer look at the X-ray image data products, and build a simple generative model for the observed data.
End of explanation
"""


imfits = pyfits.open('a1835_xmm/P0098010101M2U009IMAGE_3000.FTZ')
im = imfits[0].data
bkfits = pyfits.open('a1835_xmm/P0098010101M2X000BKGMAP3000.FTZ')
bk = bkfits[0].data
exfits = pyfits.open('a1835_xmm/P0098010101M2U009EXPMAP3000.FTZ')
ex = exfits[0].data

"""
Explanation: A closer look at the data products
If you haven't done so, download the data as in the FirstLook notebook. Here we read in the three provided files.
End of explanation
"""


plt.imshow(viz.scale_image(im, scale='log', max_cut=40), cmap='gray', origin='lower');

"""
Explanation: im is the observed data as an image (counts in each pixel), after standard processing. This displays the image on a log scale, which allows us to simultaneously see both the cluster and the much fainter background and other sources in the field.
End of explanation
"""


plt.imshow(ex, cmap='gray', origin='lower');

"""
Explanation: ex is the exposure map (units of seconds). This is actually the product of the exposure time used to make the image, and a relative sensitivity map accounting for the vignetting of the telescope, dithering, and bad pixels whose data have been excised. Displaying the exposure map on a linear scale makes the vignetting pattern and other features clear.
End of explanation
"""


plt.imshow(bk, cmap='gray', origin='lower');

"""
Explanation: bk is a particle background map. This is not data at all, but a model for the expected counts/pixel in this specific observation due to the quiescent particle background. The map comes out of a blackbox in the processing pipeline -- even though there are surely uncertainties in it, we have no quantitative description of them to work with. Note that the exposure map above does not apply to the particle backround; some particles are vignetted by the telescope optics, but not to the same degree as X-rays. The resulting spatial pattern and the total exposure time are accounted for in bk.
End of explanation
"""


mask = np.loadtxt('a1835_xmm/M2ptsrc.txt')
for reg in mask:
    # this is inefficient but effective
    for i in np.round(reg[1]+np.arange(-np.ceil(reg[2]),np.ceil(reg[2]))):
        for j in np.round(reg[0]+np.arange(-np.ceil(reg[2]),np.ceil(reg[2]))):
            if (i-reg[1])**2 + (j-reg[0])**2 <= reg[2]**2:
                ex[np.int(i-1), np.int(j-1)] = 0.0

"""
Explanation: There are non-cluster sources in this field. To simplify the model-building exercise, we will crudely mask them out for the moment. A convenient way to do this is by setting the exposure map to zero in these locations, since we will not consider such pixels as part of the image data later on. Below, we read in a text file encoding a list of circular regions in the image, and set the pixels within each of those regions in ex to zero.
End of explanation
"""


plt.imshow(ex, cmap='gray', origin='lower');

"""
Explanation: As a sanity check, let's have a look at the modified exposure map. Compare the location of the "holes" to the science image above.
End of explanation
"""


#todo

"""
Explanation: Building a model
Our data are counts, i.e. the number of times a physical pixel in the camera was activated while pointing at the area of sky corresponding to a pixel in our image. We can think of different sky pixels as having different effective exposure times, as encoded by the exposure map. Counts can be produced by

X-rays from our source of interest (the galaxy cluster)
X-rays from other detected sources (i.e. the other sources we've masked out)
X-rays from unresolved background sources (the cosmic X-ray background)
Diffuse X-rays from the Galactic halo and the local bubble (the local X-ray foreground)
Soft protons from the solar wind, cosmic rays, and other undesirables (the particle background)

Of these, the particle background represents a flux of particles that either do not traverse the telescope optics at all, or follow a different optical path than X-rays. In contrast, the X-ray background is vignetted in the same way as X-rays from a source of interest. We will lump these sources (2-4) together, so that our model is composed of a galaxy cluster, the X-ray background, and the particle background.
Since our data are counts in each pixel, our model needs to predict the counts in each pixel. However, physical models will not predict count distributions, but rather intensity (counts per second per pixel per unit effective area of the telescope). The spatial variation of the effective area relative to the aimpoint is accounted for in the exposure map, and we can leave the overall area to one side when fitting (although we would need it to turn our results into physically interesting conclusions).
Cluster model
We will use a common parametric model for the surface brightness of galaxy clusters: the azimuthally symmetric beta model,
$S(r) = S_0 \left[1.0 + \left(\frac{r}{r_c}\right)^2\right]^{-3\beta + 1/2}$,
where $r$ is projected distance from the cluster center. The parameters of this model are

$x_0$, the $x$ coordinate of the cluster center
$y_0$, the $y$ coordinate of the cluster center
$S_0$, the normalization
$r_c$, a radial scale (called the "core radius")
$\beta$, which determines the slope of the profile

Let's plot it:
End of explanation
"""


