#!/usr/bin/env python
import sys
import argparse
import collections

def load(filename):
	""" 
	Load an im3shape output file.
	Column names are taken from the header:

	cat = py3shape.output.load(filename)
	print cat.radius
	print cat.stats_signal_to_noise
	e = cat.e1**2 + cat.e2**2

	"""
	return I3Cat(filename)




class I3Cat(collections.Mapping):
	
	def __init__(self, filename, load=True, load_header=False, col_names=None, header=None, forcetxt=False):

		if isinstance(filename, basestring):
			if load:
				import numpy as np
				if (".txt" in filename) or forcetxt:
					if col_names is None:
						self._load_header(filename)
					else:
						self.col_names = col_names
					self.data = np.atleast_2d(np.loadtxt(filename)).T
				elif ".fits" in filename:
					self._load_fromfits(filename)
			elif load_header:
				self._load_header(filename)
		else:
			import numpy as np
			self.data=np.array(filename)
			self.header_lines = header
			self.col_names = col_names

	def _load_fromfits(self,filename):
			import fitsify
			self.data, self.col_names, self.header_lines = fitsify.unfitsify_catalog(filename)

	def _load_header(self,filename):
		self.header_lines = []
		for i,line in enumerate(open(filename)):
			line = line.strip()
			if i==0:
				if not line.startswith('#'):
					raise ValueError("File does not have a # header line naming columns")
				self.col_names = line.lstrip('#').split()
			else:
				if not line.startswith('#'): break
				self.header_lines.append(line)

	def __len__(self):
		return self.data.shape[1]

	def index(self, name):
		return self.col_names.index(name)

	def __getattr__(self, name):
		return self.data[self.index(name)]

	def __getitem__(self, name):
		return self.data[self.index(name)]

	def __iter__(self):
		return iter(self.col_names)

	def cut_rows(self, mask):
		"The mask should be either boolean, being True where to accept the cuts, or a list of indices"
		self.data = self.data[:,mask]

	def remove_cols(self, col_names, ignore_missing=True):
		indices = []
		for col_name in col_names:
			try:
				index = self.col_names.index(col_name)
			except ValueError:
				if ignore_missing:
					continue
				else:
					raise
			indices.append(index)
		for col_name in col_names:
			self.col_names.remove(col_name)
		self.data = np.delete(self.data, indices, axis=0)

	def merge(self, other, header_lines=None):
		"Merge with another catalog, keeping only cols in both"
		import numpy as np
		col_names = [name for name in self.col_names if name in other.col_names]
		merged_data = [
			np.hstack((getattr(self,name), getattr(other,name)))
			for name in col_names
		]
		if header_lines is None:
			header_lines = self.header_lines 
		merged = I3Cat(merged_data, col_names = col_names, header=header_lines)
		return merged

	def sort(self, col_name='identifier'):
		col = getattr(self, col_name)
		import numpy as np
		indices = np.argsort(col)
		self.data = self.data[:, indices]

	def save(self, filename):
		outfile = open(filename, 'w')
		col_names = '    '.join(self.col_names)
		outfile.write('#%s\n' % col_names)
		import numpy as np
		np.savetxt(outfile, self.data.T, delimiter="\t", fmt="%s")
		outfile.close()




description = "Get columns or indices from an im3shape output file, reading col names from the header"
parser = argparse.ArgumentParser(description=description, add_help=True)
parser.add_argument('filename', type=str, help='Filename to read')
parser.add_argument('column', type=str, help='Name of column to get')
parser.add_argument('-i', '--index', action='store_true', help='Do not print out the column, just its index')
parser.add_argument('-o', '--one', action='store_true', help='With -i, output a one-based index instead of zero based')

if __name__=="__main__":
	args = parser.parse_args()
	if args.index:
		i=I3Cat(args.filename, load=False).index(args.column)
		if args.one:
			i+=1
		print i
	else:
		cat=I3Cat(args.filename)
		col = getattr(cat, args.column)
		for x in col:
			print x