# Takes particle data from gf and grids it. User specifies the distance, 
# telescope resolution, survey sensitivity, and maximum size of the data 
# cube (will truncate anything outside this volume).
# Returns a position-velocity data cube in S/N units (no actual noise is added yet)

import numpy
from numpy import *
import pyfits

# Maximum size of cube in pixels
# x,y are space, z is velocity
Nx = 200
Ny = 200
Nz = 200

# Particle mass in linear solar units. Ideally, use all particles, otherwise
# just multiply this accordingly.
Mp = 3.915E5

# Distance of simulated particles in Mpc
D = 17.0

# Telescope spatial resolution, arcminutes
Ps = 3.5

# Telescope velocity resolution, km/s
V = 5.0

# Observation sensitivity rms, mJy
rms = 0.5

# Input filename
filename = 'GASposvel.dat'

# Specify which frame to grid
f = 150

# Name of FITS file to create
fitsname = 'FitsFile.fits'

# END OF USER INPUTS


rms = rms / 1000.0

# Open file and skip to the correct frame
infile = open(filename,'r')
nparts = int(infile.readline())
nframs = int(infile.readline())

count = 0

for i in range(0,f):
	for j in range(0,nparts):
		line = infile.readline()
		count = count +1
		
# Declare blank array to hold particles x,y,v values
particles = numpy.zeros((3,nparts))

# Now	read in those values to the array
for i in range(0,nparts):
	line = infile.readline()
	x,y,z,vx,vy,vz = line.split()
	particles[0,i] = float(x)
	particles[1,i] = float(z)
	particles[2,i] = float(vy)

infile.close()

# Find the minimum, median and maxmimum positions of the particles
# Using mean rather than median - median tends to cutoff quite a lot
xmin = numpy.amin(particles[0,:])
xmax = numpy.amax(particles[0,:])
xmed = numpy.mean(particles[0,:])

ymin = numpy.amin(particles[1,:])
ymax = numpy.amax(particles[1,:])
ymed = numpy.mean(particles[1,:])	

vmin = numpy.amin(particles[2,:])
vmax = numpy.amax(particles[2,:])
vmed = numpy.mean(particles[2,:])
	
# Determine the pixel resolution in kpc
Sp = ((Ps/60.0)/360.0) * 2* pi * D * 1000.0

# Check the size of the grid needed to hold the entire simulation
# If larger than specifiec limit, restrict it
rangexpix = int((xmax - xmin) / Sp)
if rangexpix > Nx:
	rangexpix = Nx
	
rangeypix = int((ymax - ymin) / Sp)
if rangeypix > Ny:
	rangeypix = Ny
	
rangevpix = int((vmax - vmin) / V)
if rangevpix > Nz:
	rangevpix = Nz

# Transform minimum coordinate so that the box will be centered on the median
# particle position
xmin = xmed - Sp*(float(rangexpix)/2.0)
ymin = ymed - Sp*(float(rangeypix)/2.0)
vmin = vmed - V*(float(rangevpix)/2.0)
	
# Create a blank array for the FITS images
image = numpy.zeros((rangevpix,rangeypix,rangexpix))


# Go through all the particles and check where they should be, adding them
# to the image array in the appropriate location
for i in range(0,nparts):
	print i
	x = particles[0,i]
	y = particles[1,i]
	v = particles[2,i]
	
	xcell = int((x - xmin)/Sp)
	ycell = int((y - ymin)/Sp)
	vcell = int((v - vmin)/V)
	
	if (xcell > 0) and (xcell < rangexpix):
		if (ycell > 0) and (ycell < rangeypix):
			if (vcell > 0) and (vcell < rangevpix):
				image[vcell,ycell,xcell] = image[vcell,ycell,xcell] + 1.0

# Now convert that image to S/N
image = (image*Mp) / (2.36E5*D*D*rms*V)
image = image / (rms/1000.0)				

# Finally, save the image to a FITS file
FitsFile = pyfits.writeto(fitsname,image,header=None)