# 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)

# This version does files in multiple subdirectories

import numpy
from numpy 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)

# 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 = 8.5E9/19986.0

# Distance of simulated particles in Mpc
D = 17.0

# Telescope spatial resolution, arcminutes
Ps = 3.5

# Telescope velocity resolution, km/s
V = 10.0

# Observation sensitivity rms, mJy
rms = 0.6

# Input filename
filename = 'GASposvel2.dat'

# Specify how many frames to grid
f = 199

# No. particles
nparts = 19987

# No. frames
nframs = 200

# Base name of FITS file to create
fitsname = 'FitsFile'

# END OF USER INPUTS


# Get rms in Jy
rms = rms / 1000.0

# Make a list of the subdirectories
dirlist = next(os.walk('.'))[1]


# Go through each directory. If the required file is found, make a subdirectory to hold
# the images, and store the fits files there.

for dirc in dirlist:
	print dirc
	os.chdir(dname+'/'+dirc)
	
	if str(os.path.isfile(filename)) == 'True':
		try:
			os.mkdir('FitsImages')
		except:
			pass	

		# Open file and read in the basic parameters
		infile = open(filename,'r')
		nparts = int(infile.readline())
		nframs = int(infile.readline())

		#count = 0

		# Skip to a specific frame
		#for i in range(0,f):
		#	for j in range(0,nparts):
		#		line = infile.readline()
		#		count = count +1


		for frame in range(0,f):
			print frame		
			# 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):
				#print i
				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)


			# 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.median(particles[0,:])

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

			vmin = numpy.amin(particles[2,:])
			vmax = numpy.amax(particles[2,:])
			vmed = numpy.median(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
	
			# Force range to be equivalent to 1 Mpc
			rangexpix = 60
	
			#rangeypix = int((ymax - ymin) / Sp)
			#if rangeypix > Ny:
			#	rangeypix = Ny
	
			rangeypix = 60
	
			#rangevpix = int((vmax - vmin) / V)
			#if rangevpix > Nz:
			#	rangevpix = Nz
	
			# Force velocity range equivalent to 1000 km/s
			rangevpix = 100

			# 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):
				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)				

			# Finally, save the image to a FITS file
			os.chdir(dname+'/'+dirc+'/FitsImages/')
			fitsfilename = fitsname+'_'+str(frame).zfill(3)+'.fits'
			FitsFile = pyfits.writeto(fitsfilename,image,header=None, clobber=True)

		infile.close()