################################################################################
# Script to extract emission regions from a Chandra observation of Jupiter
# code taken from a Jupyter notebook, can be run instead of the notebook.
################################################################################

# imports
from astropy.io import fits
import numpy as np
import scipy.ndimage as ndimage

# load the raw image
imgfile = "JupiterACIS2011Oct2.fits"
raw_img = fits.open(imgfile)[0].data

# build the planet mask - this was done by determining the center, radius
# and axial tilt (in the image plane) by eye
center = [236, 550]
radius = 107

axis_angle = 115 * np.pi/180
polar_ratio = 66.854 / 71.492
planet_mask = np.zeros_like(raw_img)

for ii in range(raw_img.shape[0]):
    for jj in range(raw_img.shape[1]):
        dx = ii-center[0]
        dy = jj-center[1]
        # get distance along equatorial and polar axis
        d_equat = dx*np.cos(axis_angle) + dy*np.sin(axis_angle)
        d_polar = -dx*np.sin(axis_angle) + dy*np.cos(axis_angle)
        if (np.sqrt((d_equat)**2 + (d_polar/polar_ratio)**2) < radius):
            planet_mask[ii,jj] = 1


# this mask gives us the planetary emission region - now get the polar regions

# first convolve the input image with a gaussian - this get's rid of any pixels
# with 0 counts
convimg = ndimage.gaussian_filter(raw_img, 4, mode="wrap")

# aurorae have to be inside the planet - apply the mask
convimg *= planet_mask

# now do a simple thresholding
aurora_mask = convimg > 0.6

# get some surface brightness information:
print("Rough surface brightnesses (cts per pixel):")

print("Full image:", np.mean(raw_img))
print("Background only:", np.mean(raw_img[planet_mask != True]))
print("")
print("Full Planet:", np.mean(raw_img[planet_mask == True]))
print("Auroral Region:", np.mean(raw_img[aurora_mask == True]))
print("Planet without aurorae:", np.mean(raw_img[planet_mask * np.logical_not(aurora_mask) == True]))


# we then build two images: An ellipse with a flat intensity, and the auroral 
# regions, blurred and taken from the image
# since I'm not taking the background counts, I'll also make these images
# sub-images of the input

# sub image borders:
x0=100; x1=350
y0=400; y1=700

# intensity of the flat disk
flat_intens = np.mean(raw_img[(planet_mask * np.logical_not(aurora_mask)) == True])

# actual sub images:
disk_img = (planet_mask * flat_intens)[x0:x1, y0:y1]

aurora_img = (ndimage.gaussian_filter((raw_img-flat_intens), 2)*aurora_mask)[x0:x1, y0:y1]

# now build two image extensions
# for that, I need the WCS keys from the original image
orig_header = fits.open(imgfile)[0].header

# now the extensions
# Disk extension
disk_hdu = fits.ImageHDU(data=disk_img)
disk_hdu.name = "DISK"

# copy the header, but adjust the crpix's!
disk_hdu.header["RADESYS"] = orig_header["RADESYS"]
disk_hdu.header["EQUINOX"] = orig_header["EQUINOX"]
disk_hdu.header["CTYPE1"] = orig_header["CTYPE1"]
disk_hdu.header["CTYPE2"] = orig_header["CTYPE2"]
disk_hdu.header["CRVAL1"] = orig_header["CRVAL1"]
disk_hdu.header["CRVAL2"] = orig_header["CRVAL2"]
disk_hdu.header["CUNIT1"] = orig_header["CUNIT1"]
disk_hdu.header["CUNIT2"] = orig_header["CUNIT2"]
disk_hdu.header["CRPIX1"] = orig_header["CRPIX1"] - y0
disk_hdu.header["CRPIX2"] = orig_header["CRPIX2"] - x0
disk_hdu.header["CDELT1"] = orig_header["CD1_1"]
disk_hdu.header["CDELT2"] = orig_header["CD2_2"]

# Aurora extension
aurora_hdu = fits.ImageHDU(data=aurora_img)
aurora_hdu.name = "AURORAE"

# copy the header, but adjust the crpix's!
aurora_hdu.header["RADESYS"] = orig_header["RADESYS"]
aurora_hdu.header["EQUINOX"] = orig_header["EQUINOX"]
aurora_hdu.header["CTYPE1"] = orig_header["CTYPE1"]
aurora_hdu.header["CTYPE2"] = orig_header["CTYPE2"]
aurora_hdu.header["CRVAL1"] = orig_header["CRVAL1"]
aurora_hdu.header["CRVAL2"] = orig_header["CRVAL2"]
aurora_hdu.header["CUNIT1"] = orig_header["CUNIT1"]
aurora_hdu.header["CUNIT2"] = orig_header["CUNIT2"]
aurora_hdu.header["CRPIX1"] = orig_header["CRPIX1"] - y0
aurora_hdu.header["CRPIX2"] = orig_header["CRPIX2"] - x0
aurora_hdu.header["CDELT1"] = orig_header["CD1_1"]
aurora_hdu.header["CDELT2"] = orig_header["CD2_2"]

# create the FITS file
hdulist = fits.HDUList(fits.PrimaryHDU())
hdulist.append(disk_hdu)
hdulist.append(aurora_hdu)

hdulist.writeto("emission_regions.img", overwrite=True)