# Script to convert 4-axis 3D FITS files (e.g. VLA) into nice sensible 3-axis
# files.

import numpy
import pyfits

filename = 'VCC2070VLA.fits'
outfile  = 'VCC2070VLA_reshaped.fits'

# Get the shape of the file
fitsfile=pyfits.open(filename)# mode='update')
		
image = fitsfile[0].data
header =fitsfile[0].header
		
naxes = fitsfile[0].header['NAXIS']


if (naxes < 3):
	print 'Your image has less than 3 dimensions. You do not need this program.'
	exit()
if (naxes > 4):
	print 'Your image has more than 4 dimensions. That is silly.'
if (naxes ==4):

	fitsfile=pyfits.open(filename)# mode='update')
	image = fitsfile[0].data
	header =fitsfile[0].header
		

	print 'Your image has 4 dimensions. This is a bit daft, but it could be worse.'
	print 'This script requires data to be in the format [z,y,x] = [RA,Dec,vel]'
	print 'The program will now assume the format is [1,z,y,x] and convert.'
	print 'If your data has a different shape, you will have to convert it yourself.'

	print 'Initial image shape :',image.shape

	
	z = image.shape[1]	# No. channels					
	y = image.shape[2]	# No. x pixels
	x = image.shape[3]	# No. y pixels
	
	newimage = numpy.reshape(image,[z,y,x])
	
	print 'Converted image shape :',newimage.shape
	
	
	pyfits.core.writeto(outfile,newimage,header, clobber=True)
	
	#fitsfile.flush()
	
	# Overwrite the existing file, just to make sure
	#pyfits.core.writeto(fullname,image,header, clobber=True)