import sys
import pyfits
import numpy

from astrometry.util.fits import pyfits_writeto

def is_sdss_idr(hdu):
	hdr = hdu.header
	return ((hdr.get('SIMPLE', True) == False)
			and ('SDSS' in hdr)
			and ('UNSIGNED' in hdr))
			#and hdr.get('SDSS', None) == pyfits.UNDEFINED
			#and hdr.get('UNSIGNED', None) == pyfits.UNDEFINED)

def is_sdss_idr_file(infile):
	p = pyfits.open(infile)
	rtn = is_sdss_idr(p[0])
	p.close()
	return rtn

# Takes a pyfits HDU object (which is unchanged) and returns a new
# pyfits object containing the fixed file.
def fix_sdss_idr(hdu):
	hdr = hdu.header.copy()

	if hdr.get('SIMPLE', True):
		print 'SIMPLE = T: not an SDSS idR file.'
		return hdu
	print 'Setting SIMPLE = True'
	hdr.remove('SIMPLE')
	hdr.set('SIMPLE', True, 'FITS compliant (via fix-sdss-idr.py)')

	if 'SDSS' in hdr:
		print 'Setting SDSS = True'
		hdr.remove('SDSS')
		hdr.set('SDSS', True, 'SDSS (via fix-sdss-idr.py)')
	else:
		print 'No SDSS header card: not an SDSS idR file.'
		return hdu

	if 'UNSIGNED' in hdr:
		print 'Setting UNSIGNED = True'
		hdr.remove('UNSIGNED')
		hdr.set('UNSIGNED', True, 'SDSS unsigned ints')
	else:
		print 'No UNSIGNED header card: not an SDSS idR file.'
		return hdu

	#hdr._updateHDUtype()
	#newhdu = hdr._hdutype(data=pyfits.DELAYED, header=hdr)
	newhdu = pyfits.PrimaryHDU(data=pyfits.DELAYED, header=hdr)
	## HACK - encapsulation violation
	newhdu._file = hdu._file
	#newhdu._ffile = hdu._ffile
	newhdu._datLoc = hdu._datLoc

	newhdu.data = newhdu.data.astype(numpy.int32)
	newhdu.data[newhdu.data < 0] += 2**16
	print 'data type:', newhdu.data.dtype
	print 'data range:', newhdu.data.min(), 'to', newhdu.data.max()
	return newhdu

def fix_sdss_idr_file(infile, outfile):
	print 'Reading', infile
	newhdu = fix_sdss_idr(pyfits.open(infile)[0])
	print 'Writing', outfile
	pyfits_wireto(newhdu, outfile)


if __name__ == '__main__':
	if len(sys.argv) != 3:
		print 'Usage: %s <input-fits-file> <output-fits-file>' % sys.argv[0]
		sys.exit(-1)
	infile = sys.argv[1]
	outfile = sys.argv[2]

	fix_sdss_idr_file(infile, outfile)
	sys.exit(0)
