"""
Tools for reading im3shape tables using the astropy package
Used for post-processing DES data, but you could use load 
catalogs more generally.
"""
import astropy.io.fits
import astropy.table
import astropy.io.registry
import numpy as np



class I3TableHeader(astropy.io.ascii.basic.CommentedHeaderHeader):
	"""Header class for im3shape catalogs.
	Exactly the same as the standard header for ascii files except that
	metadata in the comments is stored
	"""
	def update_meta(self, lines, meta):
		table_meta = meta['table']
		for line in lines[1:]:
			line=line.strip()
			if not line.startswith('#'):
				break
			if not line:
				continue
			line = line.strip('#').strip()
			if ' ' in line:
				space = line.index(' ')
			elif '\t' in line:
				space = line.index('\t')
			key = line[:space].strip().strip(':')
			value = line[space:].strip()
			table_meta[key] = value


class I3TableReader(astropy.io.ascii.CommentedHeader):
	""" Reader for im3shape table. Same as standard one except
		it uses the header code defined about and tabs instead of spaces
		for the separator.
	"""
	def __init__(self):
		super(I3TableReader, self).__init__()
		self.header = I3TableHeader()
		self.header.data = self.data
		self.data.header = self.header
		self.header.start_line = 0
		self.data.start_line = 0
		self.header.comment = r'\s*#'
		self.header.write_comment = '# '
		self.data.comment = r'\s*#'
		self.data.write_comment = '# '
		self.header.splitter.delimiter = '\t'
		self.data.splitter.delimiter = '\t'



class I3Table(astropy.table.Table):
	""" Final im3shape table class.  Same as a standard table except it supports
		attribute style column access ("table.e1") and does not try to print itself
		out when repr'd interactively.
	"""
	def __getattr__(self, name):
		cols = self.__dict__['columns']
		if name in cols:
			return self[name]
		raise AttributeError("No property or column %s"%name)
	def __repr__(self):
		return "<im3shape catalog %d cols %d rows>" % (len(self.colnames), len(self))

	@classmethod
	def _construct_subclass_reader_wrapper(cls,function):
		def new_function(*args, **kwargs):
			return cls(function(*args,**kwargs))
		if function.__doc__:
			new_function.__doc__ = function.__doc__+"\nWrapped programmatically to return "+cls.__name__
		return new_function

	@classmethod
	def _register_subclass_io(cls,parent_class=astropy.table.Table):
		reader_registrations = []
		#Look at all the existing readers and if they are 
		#registered for the parent class then record the name
		#and function they use
		for (name,parent),reader in astropy.io.registry._readers.items():
			if parent_class==parent:
				reader_registrations.append((name,reader))
		#register all those functions for the new class too,
		#except that we need to wrap the function to return an instance
		#of our class instead
		for (name,reader) in reader_registrations:
			new_reader = cls._construct_subclass_reader_wrapper(reader)
			astropy.io.registry.register_reader(name, cls, new_reader)
		#Now do exactly the same for the writers, except no wrapping needed
		writer_registrations = []
		#Get existing
		for (name,parent),writer in astropy.io.registry._writers.items():
			if parent_class==parent:
				writer_registrations.append((name,writer))
		#register new
		for (name,writer) in writer_registrations:
			astropy.io.registry.register_writer(name, cls, writer)

I3Table._register_subclass_io()
#_write_function=astropy.io.fits.connect.write_table_fits
#astropy.io.registry.register_writer('fits', I3Table, _write_function)

def load_catalogs(filenames, quiet=True):
	""" Load a collection of im3shape output files, preserving the metadata 
	lines and columns"""
	tables = []
	reader = I3TableReader()
	nfilename = len(filenames)
	for i,filename in enumerate(filenames):
		#read this table using our special reader that
		#saves the #comments at the top as metadata
		t = reader.read(filename)


#		if len(np.unique(t['identifier'])) < 2:
#			index = int(filename.split('-')[1].split('.')[0])
#			import warnings; warnings.simplefilter('once'); warnings.warn('fixing ids in GREAT-DES');
#			t['identifier']= index + np.arange(len(t))
		
		if not quiet:
			print '%d/%d   %s \tn_gal=%d' % (i+1,nfilename,filename,len(t))

		#Drop any zero length tables as they can mess up the stacking later if you have string cols
		if len(t)==0:
			continue
		
		#only keep metadata from first table
		if i>0 and len(t)>0: 
			t.meta.clear()
		#add this to our list of tables
		tables.append(t)
	#stack all the tables we have together

	if len(tables)==1:
		table = tables[0]
	else:
		table = astropy.table.vstack(tables)
		del tables
	#The "first" and "count" parameters, if present, 
	#do not make sense if we loaded more than one file
	if len(filenames)>1:
		if 'first' in table.meta: del table.meta['first']
		if 'count' in table.meta: del table.meta['count']
	#Convert to our table format with the new attribute access syntax
	return I3Table(table)