# Program to convert galactic PPV data and convert to PPP.
# Since PPV does not map linearly to PPP, a straightforward conversion
# leaves ugly gaps due to finite resolution. An alternative approach is
# to start with am empty PPP cube. The grid size and coordinate system
# can be arbitray : 
# - Let the centre be the position of the Sun
# - Let the z axis be the depth of the disc
# At any point, we can therefore determine R, l, b.
# The coordinate system also determines the position of the centre of
# the galaxy and therefore the physical resolution.

# This version is calibrated to the coordinate system of the Leiden-
# Argentine-Bonn all-sky HI survey. To use other data, the user will
# have to work outt he coordinate conversion and correct the equations.

# Since the LAB data set is relatively small, this version takes in the
# entire cube. Another version is available for larger data sets that
# need to be processed in smaller chunks.

# More information on converting between PPV and PPP can be found here :
# http://ircamera.as.arizona.edu/astr_250/Lectures/Lec_22sml.htm


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)

try:
    os.remove('PPPcube.fits')
except:
    pass


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

# 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 (units of km/s and kpc)
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


# Open the LAB cube
LabFITS = pyfits.open('lab_small.fits')

image = LabFITS[0].data

pppcube = numpy.zeros((cubesize,cubesize,cubesize))


# Program now checks every pixel in the PPP cube, so this is necessarily slow
# (at least I can't think of an obvious way to simplify this). At each point,
# it checks if the conversion is meaningful, since :
# - R close to zero will give silly results
# - R within solar circle has a distance ambiguity giving two solutions, so
# this can't be sampled
# - Solutions will be garbage at low velocities (when longitude close to 0 or
# 180 so that velocity is entirely across plane of the sky except for random
# motions).


for x in range(cubesize):
    print(x)
    for y in range(cubesize):
        for z in range(cubesize):
            # 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 LAB 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).
					 # This gives the same coordinates as the LAB data, which is very
					 # convenient.
                # 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)
                glong = degrees(atan2(yl-gy,xl-Rsko))
                

                # In LAB file, 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 the LAB data :
                    # l = -0.5*float(xp) + 180.0
                    # Longitude is x axis.
                    # xmin = 0 = 180; xmax = 720 = -180
                    
                    # b = 0.5*float(yp) - 90.0
                    # Latitude is y axis.
                    # ymin = 0 = -90; ymax = 360 = +90
                    
                    # vel = 1.03*float(zp) - 150.48
                    xp = int( (glong-180.0) * -2.0 )
                    yp = int( (glat + 90.0) * 2.0 )
                    zp = int( (vel + 150.48) / 1.03 )
                    # Should be possible to do a better interpolation of z
                    # than simple integer rounding !

                    # For testing purposes it can be helpful to output longitude,
						  # latitude or velocity instead of flux. If so, the next nested
						  # ifs should be commented out.
                    #pppcube[z,y,x] = glong
                
                 
                    # Only do the flux extraction if the coordinates are
                    # within the cube and velocity > 10 km/s
                    if (xp > 0) and (xp < 721):
                        if (yp > 0) and (yp < 360):
                            if (zp > 0) and (zp < 301):
                                if abs(vel) > 10.0:
												if ((glong > 10.0) and (glong < 160.0)) or ((glong < -10.0) and (glong > -160.0)):
													pppcube[z,y,x] = image[zp,yp,xp]
            
                            
# Write cube
fitsfile = pyfits.writeto('PPPCube.fits',pppcube,header=None)