# Program to transform a lat-long-vel cube into x-y-z.
# Crude - user will have to determine the coordinate system themselves

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

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

infile = 'LABsmall.fits'


outfile=str(infile.split('.fits')[0])+str('_XYZ_clipped')+str('.fits')


FitsFile = pyfits.open(infile)
image = FitsFile[0].data

sizez = image.shape[0]
sizey = image.shape[1]
sizex = image.shape[2]


# Create an empty array that will be the transformed image
# Size is arbitrary
newimage = numpy.zeros((500,500,500))

Rsun = 8.5
Vsun = 220.0
Vpnt = 220.0



for xp in range(0,sizex):
	print xp
	l = -0.5*float(xp) + 180.0
	l = math.radians(l)
	
	cl = cos(l)
	sl = sin(l)
	
	for yp in range(0,sizey):
		b = 0.5*float(yp) - 90.0
		b = math.radians(b)
		
		cb = cos(b)
		sb = sin(b)
		
		for zp in range(0,sizez):
			vel = 1.03*float(zp) - 150.48
			
			flux = image[zp,yp,xp]
			
			
			try:
				R = Vpnt / ( vel/(Rsun*sl) + (Vsun/Rsun) )
				
				d1 = Rsun*cl - sqrt( Rsun*Rsun*cl*cl - Rsun*Rsun + R*R )
				d2 = Rsun*cl + sqrt( Rsun*Rsun*cl*cl - Rsun*Rsun + R*R )
				
				# Reject all points within the solar radius (distance ambiguity) and at low
				# velocities (towards/away from galactic center, all velocity tangential)
				if (d1>-50) and (d1<50) and (d2>-50) and (d2<50) and (R > Rsun) and (abs(vel) > 10.0):
		
					x1 = d1*cl*cb
					y1 = d1*sl*cb
					z1 = d1*sb
					
					x2 = d2*cl*cb
					y2 = d2*sl*cb
					z2 = d2*sb
					
					# Alternative to write to a text file
					#output.write(str(x1)+' '+str(y1)+' '+str(z1)+' '+str(flux)+'\n')
					#if abs(d2-d1)>0.1:
					#	output.write(str(x2)+' '+str(y2)+' '+str(z2)+' '+str(flux)+'\n')
					
					# Convert kpc to pixel values.
					x1 = int(x1*5.0) + 250
					y1 = int(y1*5.0) + 250
					z1 = int(z1*5.0) + 250
					
					x2 = int(x2*5.0) + 250
					y2 = int(y2*5.0) + 250
					z2 = int(z2*5.0) + 250
					
					newimage[z1,y1,x1] = flux
					
					if abs(d2-d1)>0.1:
						newimage[z2,y2,x2] = flux
					
			except:
				pass
				
FitsFile2 = pyfits.writeto(outfile,newimage,header=None)		
			
FitsFile2.close()	
FitsFile.close()
	

#output.close()