# Program to convert galactic PPV data and convert to PPP. For more
# information on the basic conversion process, see the LABCreatePPP.py
# script, which is designed for smaller files and easier to understand.

# Special version that operates on a series of FITS files, one per
# velocity channel. This makes it much easier to work with very large
# data sets like HI4PI (35 GB), though it needs a more complex script.

# NOTE : This version requires some manual edits ! User should actually
# check through the code before running it as various sections may need
# to (un)commented !

# The simplest approach is to go through every voxel in the PPP cube and
# extract the data from the appropriate file. But for large cubes, this
# can mean opening large FITS files hundreds of millions of times, which
# would mean the program would take months to run. 
# A better approach is  to precompute the velocity at each voxel, then 
# create a list of the FITS files and their associated points that need to
# be extracted. This means that we open a file, extract all its needed
# points, and  move on to the next file. So the maximum number of file-
# opening events will never exceed the number of files (944, instead of 125 
# million in the case of HI4PI).

# This method means we still have to create a very large list, in order to
# match each voxel in the PPP cube with its original file. This can be slow
# to convert and sort, so the program may need to be run in stages.
# First create the full-size empty PPP cube and write the file. Abort the
# code at this point. Then  have the script open this file and iterate over
# a series of different  chunks to fill in the blanks section by section 
# (will have to manually specify the x/y/z ranges to search).

# For a 500^3 voxel cube, can split into 8 subcubes to speed things up and
# avoid blowing the memory out of the water.


import numpy
import math
from math import *
import pyfits
import os


# Change to directory where script is located
abspath = os.path.abspath(__file__)
dname = os.path.dirname(abspath)
os.chdir(dname)


# If doing things in one step, we may want to remove the existing file before
# we begin. If running in stages, leave this commented out.
#try:
#	os.remove('PPPCube.fits')
#except:
#	pass


# Output cube parameters
cubesize = 500
# Pixel size in kpc
ps = 0.15


# If we're starting for the first time, create the empty PPP cube.
#pppcube = numpy.zeros((cubesize,cubesize,cubesize))
#outputcube = pyfits.writeto('VelocityCube.fits',pppcube,header=None)
#stop


# Pixel values run from 0->cubesize, so physically, 0->cubesize*ps
# The two points of interest to us are :
# - The position of the Sun, which we set to be the centre. Its coordinates
#   relative to the origin (in kpc) are therefore x,y,z = (cubesize/2.0)*ps).
# - The galactic centre. We will define this to at the centre in y and z but
#   offset by the solar radius in x.


# Physical parameters
Vsol = 220.0
Vpnt = 220.0
Rsol = 8.5


# Let the Sun be at the centre of the cube. Its position relative to the origin,
# in kpc, is given by :
Rsko = (cubesize/2.0)*ps

# We can then get the coordinates of the galactic centre in kpc relative to the
# origin
gx = Rsko + Rsol
gy = Rsko
gz = Rsko


# If doing things in one step, create the volume. If running in stages, don't
# do this here as the existing file will be opened later instead.
#pppcube = numpy.zeros((cubesize,cubesize,cubesize))


# Create a list to hold the coordinates. Format :
# px, py, pz, x, y, z
# Where px/y/z/ are in the PPP cube and the others are in the original cube.
coordarray = []
# Later we'll use this array to sequentially open each FITS file and extract the
# relavent pixels.


# Open a reference channel file to get velocity information
chanzero = pyfits.open('Channel_000.fits')
header = chanzero[0].header
velzero = -((float(header['CRPIX3']) * float(header['CDELT3'])) / 1000.0) + (float(header['CDELT3'])/1000.0)
deltvel = float(header['CDELT3'])/1000.0
chanzero.close()
# Now we have the exact velocity of the first channel and the velocity resolution, we can easily calculate
# the velocity of any other channel :
# vel = velzero + (nchan+1)*deltvel


# If we're manually running over an existing cube, first we need to open it and get the array
pppfile = pyfits.open('VelocityCube.fits', mode='update')
pppcube = pppfile[0].data
# We also need to specify the pixels we want to iterate over. Do that in the for loops below.
# Otherwise, if doing things voxel by voxel, just set the loops to be range(cubesize).


# First precalculate which voxels to extract. No FITS files are opened here since the
# calculation can be done purely analytically, given the known coordinate system. Of
# course, no data from the FITS files is extracted - we'll do that afterwards.
# This is the part in which the user has to manually alter the pixel ranges to search,
# if running in stages. Set the values in the for loops. Alternatively just use
# range(cubesize) if it can be done in one step.
# All this part does is set the array values so that later we can sort them to extract
# the necessary values.
for x in range(0,251):
	print x
	for y in range(0,251):
		for z in range(0,251):
			# Get local coordinates in kpc (measured from origin)
			xl = x*ps
			yl = y*ps
			zl = z*ps

			# Distance from GC in kpc
			R = sqrt( (xl-gx)**2.0 + (yl-gy)**2.0 + (zl-gz)**2.0 )
           
			# Ensure R is resolved
			if R > ps:

				# Calculate latitude and longitude. Convention in HI4PI file is
				# that longitude runs from +180 (x=0) to -180 (x=720).
				# All cordinates must be relative to centre of galaxy.
				# atan2 function uses combination of x and y to determine angle
				# (ordinary atan has to be carefully set according to quadrant)
				# Longitude is measured relative to Sun, which is at centre of
				# cube - NOT the centre of the galaxy gx (y and z same for Sun
				# and gc)
				# Longitude runs from +180 to -180 so no need for additional
				# checks
				glong = degrees(atan2(yl-gy,xl-Rsko))  				  

				# Latitude runs from -90 (y=0) to +90 (y=360).    
				glat = -degrees(acos( (zl-gz)/R) ) + 90.0

				# Inside solar circle we have distance ambiguity so avoid that
				# region
				if R > Rsol:
					vel = ( (Vpnt/R) - (Vsol/Rsol) )*Rsol*sin(radians(glong))
                
					# We now have all the galactic coordinates and just need
					# to convert them to image coordinates.

					# For HI4PI data :
					# l = -0.0833333*float(xp) + 180.083333
					# Longitude is x axis.
					# xmin = 0 = 180; xmax = 720 = -180

					# b = -90.0 + (180.0/2160.0)*(float(yp)+1.0)
					# Latitude is y axis.
					# ymin = 0 = -90; ymax = 360 = +90

					# vel = -((float(head['CRPIX3']) * float(head['CDELT3'])) / 1000.0) + (float(head['CDELT3'])/1000.0)
					# But we already determined the velocity of channel 0 and deltav earlier so :
					# vel = velzero + nchan*deltvel

					xp = int((glong - 180.0 - (180.0/2160.0)) / (-180.0/2160.0))
					yp = int(((glat + 90.0)/(180.0/2160.0))-1.0)
					zp = int((vel - velzero)/deltvel) + 1

					# Should be possible to do a better interpolation of z
					# than simple integer rounding !

					# It's still possible to output other cubes if needed, provided we set the array and file
					# correctly elsewhere.
					#pppcube[z,y,x] = vel

					# If the calculated pixels in the original files are sensible, and correspond to coordinates
					# which are good (i.e. avoiding low velocities and longitudes), add them to the array.
					if (xp > 0) and (xp < 4322):
						if (yp > 0) and (yp < 2162):
							if (zp > 0) and (zp < 944):
								if abs(vel) > 10.0:
									if ((glong > 10.0) and (glong < 160.0)) or ((glong < -10.0) and (glong > -160.0)):										
										coordarray.append([x,y,z,xp,yp,zp])


# Now sort the array by the channel needed from the extracted files
print 'Converting and sorting array :'
coordarray = numpy.asarray(coordarray)
sortedarray = coordarray[coordarray[:,5].argsort()]


# Get the first channel to initialise everything
oldchan = sortedarray[0,5]
chanfile = 'Channel_'+str(oldchan).zfill(3)+'.fits'
fitsfile = pyfits.open(chanfile)
image = fitsfile[0].data

# Go through part of the sorted array
print 'Extracting data :'
for i in range(len(sortedarray)):
	newchan = sortedarray[i,5]
	
	# If the new channel number has changed, close the old file and open the new one
	if newchan <> oldchan:
		print 'Working on channel ',newchan
		fitsfile.close()
		chanfile = 'Channel_'+str(newchan).zfill(3)+'.fits'
		fitsfile = pyfits.open(chanfile)
		image = fitsfile[0].data
		oldchan = newchan
	
	xp = sortedarray[i,3]
	yp = sortedarray[i,4]
	
	# Files are usually cubes so use z = 0 to get the first slice. This is assuming the files were converted
	# using miriad, which tends to output 2 channel cubes instead of individual slices.
	flux = image[0,yp,xp]
	
	cx = sortedarray[i,0] 
	cy = sortedarray[i,1]
	cz = sortedarray[i,2]
	
	pppcube[cz,cy,cx] = flux 



# Update existing cube
pppfile.flush()
pppfile.close()