# Jaesub Hong, Nov, 2011
#
# append lines or rows to a fits table
# the data structure of the table is manually handled.
#
# an example usage ---------------------------------------
#
# import modfits as mf
# import struct as st

# fits=mf.modfits("this.fits")	# open the fits file
# fits.rewind(50)		# start from the last 50th lines
# for i in range(100):		# add 100 lines
# 	fits.append("gz%5d" % i + st.pack('>fffff',0.0,0.1,1.1e2,2.2e3,3.3e4) )
# fits.close(1)			# close the file
#
# ---------------------------------------------------------


import re
import os

class modfits:
	def __init__(self, filename):
		self.filename=filename
		self.fsize=os.path.getsize(filename)
		self.fits=open(filename,"r+b")
		self.pos=0
		self.hsize=80
		self.bsize=2880

		self.naxis1 = 1
		self.naxis2 = 1
		self.pos1=0
		self.pos2=0

		self.addbyte=0	# added bytes
		self.addline=0	# added lines
		self.resi=0	# see if added bytes are a multiple of each line
		self.fill=0	# left over to fill the regular buffer size

		leftdata = 0
		while self.pos < self.fsize:
			foundHEnd = False
			while foundHEnd == False:
		 		buffer=self.fits.read(self.hsize)
#				print buffer
#				print buffer[0:6]
				if buffer[0:6] == 'NAXIS1':
					catch = re.search(r'NAXIS1\s*=\s*([0-9]+)',buffer)
					if catch != None: 
						self.naxis1 = int(catch.group(1))
						self.pos1 = self.pos
				if buffer[0:6] == 'NAXIS2':
					catch = re.search(r'NAXIS2\s*=\s*([0-9]+)',buffer)
					if catch != None: 
						self.naxis2 = int(catch.group(1))
						self.pos2 = self.pos
				if buffer[0:6] == 'END   ':
					foundHEnd = True
				self.pos +=self.hsize
#				break
#			break
			lefthdr=self.bsize-(self.pos % self.bsize)
			datasize=self.naxis1*self.naxis2
			leftdata=self.bsize-(datasize % self.bsize)

			self.pos = self.pos+lefthdr+datasize+leftdata
			self.fits.seek(self.pos)
			
#			print self.naxis1, self.naxis2, self.pos1, self.pos2, self.pos, leftdata, self.fsize

		self.pos -= leftdata
		self.fits.seek(self.pos1)
		buffer=self.fits.read(self.hsize)
#		print buffer
		self.fits.seek(self.pos2)
		buffer=self.fits.read(self.hsize)
#		print buffer
		self.fits.seek(self.pos)

	def rewind(self,lines=1):
		# assumes we are at the begining of each line
		self.naxis2 -= lines
		self.pos -= lines*self.naxis1
		self.fits.seek(self.pos)

	def append(self,buffer):
		blen=len(buffer)
		self.addbyte += blen 
		self.fits.write(buffer)

		left = self.addbyte % self.naxis1

		if left != 0: 
			self.resi = 1
			left = self.naxis1 - left
		else: 
			self.resi=0


		self.addline = self.addbyte / self.naxis1 + self.resi
		fsize = (self.addline + self.naxis2) * self.naxis1 
		self.fill = self.bsize - (fsize % self.bsize)+left

	def close(self, verbose=0):
	
		if self.resi == 1:
			print "warning: a possibly wrong buffer size. Expect a multiple of ", self.naxis1

		if verbose>0 :
			print "%d lines are added" % self.addline

		if self.fill !=0: self.fits.write(' '* self.fill)
		self.fits.seek(self.pos2)
		self.fits.write('NAXIS2  = %20d / length of dimension 2' % (self.addline +self.naxis2))
		self.fits.close()