# -*- coding: utf-8 -*-
import codecs,pymatrices,argparse,os,math

class Tuplesproc:
	def __init__(self, tupFile,outDir,dictFile,matrName,desc,tokens,rf,cf,rel,rtf,ctf):
		self.tupFile = tupFile
		self.outDir = outDir
		self.matrName = matrName
		self.desc = desc
		self.tokens = tokens
		self.col_file = cf
		self.row_file = rf
		self.rel_file = rel
		self.row_tf = rtf
		self.col_tf = ctf
		self.col_filter()
		self.row_filter()
		self.rel_filter()
		self.build_dict(dictFile)
		self.preprocess_tuple_file()
		self.build_col_map()
		self.rows =0
		self.nz =0
	
	def col_filter(self):
		self.col_good = set()
		if self.col_file:
			cFile = codecs.open(self.col_file,"r","utf8")
			for line in cFile:
				self.col_good.add(line.strip())
			cFile.close()

	def row_filter(self):
		self.row_good = set()
		if self.row_file:
			rFile = codecs.open(self.row_file,"r","utf8")
			for line in rFile:
				self.row_good.add(line.strip())
			rFile.close()

	def rel_filter(self):
		self.rel_good = set()
		if self.rel_file:
			rFile = codecs.open(self.rel_file,"r","utf8")
			for line in rFile:
				self.rel_good.add(line.strip())
			rFile.close()

	def build_dict(self,dictionary):
		self.dict_tf = {}
		if dictionary:
			d_file = codecs.open(dictionary,'r','utf8')
			d_file.readline()
			for line in d_file:
				l = line.strip().split(';')
				if l.__len__() ==3:
					self.dict_tf[l[0]+':'+l[1]]=l[2]
			d_file.close()

	def check_tuple(self,row,rel,col):
		if self.col_file:
			if col not in self.col_good:
				return False
		if self.rel_file:
			if rel not in self.rel_good:
				return False
		if self.row_file:
			if row not in self.row_good:
				return False
		if self.dict_tf.has_key(col):
			if self.col_tf > int(self.dict_tf[col]):
				return False
		if self.dict_tf.has_key(row):
			if self.row_tf > int(self.dict_tf[row]):
				return False
		return True

	def check_row(self,row):
		if not self.col_map.has_key(row[1]+';'+row[2]):
			return False
		if row[0] not in self.row_set:
			return False
		return True

	def build_col_map(self):
		self.col_map = {}
		for i in range(0,self.colList.__len__()):
			self.col_map[self.colList[i]] = i
	
	def preprocess_tuple_file(self):
		col_set = set()
		self.row_set = set()
		f = codecs.open(self.tupFile,'r','utf8')
		for line in f:
			l = line.strip().split(';')
			if l.__len__()==4:
				if self.check_tuple(l[0],l[1],l[2]):
					col_set.add(l[1]+';'+l[2])
					self.row_set.add(l[0])
		f.close()
		self.colList = sorted(list(col_set))
		self.cols = self.colList.__len__()

	def write_info(self):
		out = codecs.open(self.outDir + "/" + self.matrName + "/" + "info.txt","w","utf8")
		out.write(codecs.decode(self.desc,'utf8') + "\n" + self.tokens)
		out.close()
	def write_clabel(self):
		out = codecs.open(self.outDir + "/" + self.matrName + "/" + self.matrName +".clabel","w","utf8")
		for el in self.colList:
			t=el.split(';')
			out.write(t[0]+':'+t[1]+'\n')
		out.close()

	def write_boundaries(self):
		out = codecs.open(self.outDir + "/" + self.matrName + "/boundary.txt","w","utf8")
		cur = self.colList[0].split(';')[0]
		for i in range(1,self.colList.__len__()):
			if self.colList[i].split(';')[0]!=cur:
				out.write(cur+" := " + str(i)+'\n')
				cur = self.colList[i].split(';')[0]
		out.write(cur+" := " + str(self.colList.__len__()))
		out.close()

	def process(self):
		if not os.path.exists(self.outDir+"/"+self.matrName):
			os.makedirs(self.outDir+"/"+self.matrName)
		self.write_info()
		self.write_clabel()
		self.write_boundaries()
		self.process_tuples()
		self.process_rows()

	def process_rows(self):
		self.crs_mat_tmp = open(self.outDir + "/" + self.matrName +  "/matrix_crs_tmp.mat","r")
		self.crs_mat = open(self.outDir + "/" + self.matrName + "/matrix_crs.mat","w")
		self.crs_mat.write(str(int(self.rows))+' '+str(int(self.cols))+' '+str(int(self.nz))+'\n')
		self.ccs_mat_ptr = pymatrices.GetNewCCSMatrix('',self.rows,self.cols)
		counter=0
		for line in self.crs_mat_tmp:
			self.crs_mat.write(line.strip()+'\n')
			self.process_row(line.strip(),counter)
			counter += 1
		self.crs_mat.close()
		self.crs_mat_tmp.close()
		os.remove(self.outDir + "/" + self.matrName + "/matrix_crs_tmp.mat")
		self.col_features = codecs.open(self.outDir + "/" + self.matrName + "/column_features.csv","w","utf8")
		self.col_features.write("label;nz;tf;TF;entropy\n")
		for i in range(0,self.colList.__len__()):
			vec = self.ccs_mat_ptr.GetCompressedColumn(i)
			self.write_column(vec,self.colList[i])
		self.col_features.close()
		pymatrices.WriteToSingleFile(self.ccs_mat_ptr,self.outDir + "/" + self.matrName + "/matrix_ccs.mat")
		
	def process_row(self,row,n):
		t = row.split(' ')
		vals = zip(t[0::2],t[1::2])
		for i in vals:
			self.ccs_mat_ptr.SetCell(int(i[1]),n,int(i[0]))

	def process_tuples(self):
		self.row_features = codecs.open(self.outDir + "/" + self.matrName + "/row_features.csv","w","utf8")
		self.row_features.write('label;nz;tf;TF;entropy\n')
		self.crs_mat_tmp = open(self.outDir + "/" + self.matrName + "/matrix_crs_tmp.mat","w")
		self.rlabel = codecs.open(self.outDir + "/" + self.matrName + "/" + self.matrName + ".rlabel","w","utf8")
		tup_file = codecs.open(self.tupFile,"r","utf8")
		cur = ''
		tmp_vec = vectorWrap()
		for line in tup_file:
			tup = line.strip().split(';')
			if tup.__len__()!=4:
				print 'Wrong line: ' + line
				continue
			if self.check_row(tup):
				if tup[0]!=cur:
					if cur != '':
						self.write_row(cur,tmp_vec)
					cur=tup[0]
					tmp_vec = vectorWrap()
				tmp_vec.addValue(int(self.col_map[tup[1]+';'+tup[2]]),int(tup[3]))
		if cur !='':
			self.write_row(cur,tmp_vec)
		tup_file.close()
		self.row_features.close()
		self.crs_mat_tmp.close()
		self.rlabel.close()
	
	def write_column(self,vector,column):
		vec_vals = vector.GetValues()
		tf = str(int(pymatrices.CountVectorSum(vector)))
		TF = tf
		col_spl = column.split(';')
		if self.dict_tf.has_key(col_spl[1]):
			TF = self.dict_tf[col_spl[1]]
		self.col_features.write(col_spl[0]+':'+col_spl[1]+';'+str(vec_vals.__len__())+';'+tf+';'+TF +';'+'%.5f' % pymatrices.CountVectorEntropy(vector)+'\n')
		

	def write_row(self,row,vector):
		vector.count()
		self.rlabel.write(row + '\n')
		tf = str(int(vector.tf))
		TF = tf
		if self.dict_tf.has_key(row):
			TF = self.dict_tf[row]
		self.row_features.write(row+';'+str(vector.nz)+';'+tf+';'+TF +';'+'%.5f' % vector.entropy+'\n')
		
		for k in vector.data.keys():
			self.crs_mat_tmp.write(str(k)+" " + str(vector.data[k]) + " ")	
		self.crs_mat_tmp.write("\n")
		self.nz += vector.nz
		self.rows +=1

class vectorWrap:
	def __init__(self):
		self.data = {}
		self.nz = 0
		self.tf = 0

	def addValue(self,n,value):
		if self.data.has_key(n):
			self.data[n] += value
		else:
			self.data[n] = value
			self.nz +=1
		self.tf += value
	
	def count(self):
		self.entropy = 0
		for i in self.data.values():
			self.entropy += (float(i)/self.tf)*math.log(float(i)/self.tf)
		self.entropy = -self.entropy

def main(inputFile,outputDir,dictionary,matrix_name,desc,tokens,rf,cf,rel,rtf,ctf,ram,srt):
	if not srt:
		res = os.system("sort "+inputFile+" -S " + ram + " -o " + inputFile + '.sorted' + " '-t;' -k 1,1")
		if res != 0:
			print "Sort error"
			return
		inputFile = inputFile+'.sorted'
	tupProc = Tuplesproc(inputFile,outputDir,dictionary,matrix_name,desc,tokens,rf,cf,rel,rtf,ctf)
	tupProc.process()

if __name__ == '__main__':
	parser = argparse.ArgumentParser(description='Build matrix from tuples')
	parser.add_argument('-i', '--input', dest = 'inFile', required=True, 
			help = 'File with tuples')
	parser.add_argument('-o','--out', dest = 'outDir', required=True, 
			help = 'Output directory')
	parser.add_argument('--dict', dest = 'dict',
			help = 'File with (global) term frequency dictionary')
	parser.add_argument('--matrix-name', dest = 'm_name', required=True, 
			help = 'Name of output matrix')
	parser.add_argument('--matrix-description', dest = 'desc', default='No description given', 
			help = '(=No description given) Optional description of a matrix')
	parser.add_argument('--tokens', dest = 'tokens', default ='0',
			help = '(=0) Number of tokens in a corpus')
	parser.add_argument('--row-file', dest = 'row_file',
			help = 'File which contains set of rows to write in SuperMatrix')
	parser.add_argument('--col-file', dest = 'col_file',
			help = 'File which contains set of columns to write in SuperMatrix')
	parser.add_argument('--relation-file', dest = 'rel_file',
			help = 'File which contains set of relation to write in SuperMatrix')
	parser.add_argument('--row-min-tf', dest = 'row_min_tf',type=int,default =0,
			help = '(=0) Minimum value of term frequency in rows features')
	parser.add_argument('--col-min-tf', dest = 'col_min_tf', type=int,default =0, 
			help = '(=0) Minimum value of term frequency in columns features')
	parser.add_argument('--sorted', action = 'store_true', default = False, 
			help = 'Set this flag if tuples are already sorted')
	parser.add_argument('--mem-size', dest = 'ram', default ='10G', 
			help = '(=10G) Size of available memory (used for sorting)')

	args = parser.parse_args()

	main(args.inFile,args.outDir,args.dict,args.m_name,args.desc,args.tokens,args.row_file,args.col_file,args.rel_file,args.row_min_tf,args.col_min_tf,args.ram,args.sorted)
