# -*- coding: utf-8 -*-
'''
Utilities for dealing with k-best files.

Created on 2009-11-16

@author: Bartosz Broda
'''

import progressbar
import codecs

def generate_kbest_locally(k, words, matrix, mdir, cmp_desc):
	import pycomparator, pymatrices
	pMatrix = pymatrices.LoadSuperMatrix(mdir,matrix)
	cmpfun = pycomparator.FunctionFactory.CreateFunction(cmp_desc, pMatrix)
	kbest = {}

	widgets = [progressbar.Percentage(), ' ', progressbar.Bar(), ' ', progressbar.ETA()]
	pbar = progressbar.ProgressBar(widgets=widgets, maxval=len(words)-1).start()
	
	for i,w in enumerate(words):
		kbest[w] = pycomparator.generate_kbest(cmpfun, w, k)
		pbar.update(i)
	pbar.finish()
		
	return kbest

def generate_kbest_on_cluster(k, words, matrix, cmp_fun, adress):
	# TODO skopiowanie kodu z clarin/scr/supermatrixws/supermatrixws/similarity/kbest
	pass


def parse_kbest_txt_file(filename, k, skip_lines = 2):
	"""
	Example file with two lines header:	
	Words from matrix: /home/boban/workdir/macierze/supermatrix2/noun/polaczone/kipi_knigi_rzepa_substAdjP_AmIsubj_Coord_NGen_pq
	Similarity function: rwf_lin(rank_size=10000, uniform_rank = false, partial_order=true, div_factor=1, sports_order=true)
	subst:osoba
		subst:człowiek;0.35057
	 	subst:kobieta;0.268102
	 	subst:mężczyzna;0.245539
	 	subst:podmiot;0.244383
	 	subst:obywatel;0.213787
	"""
	head_word = None
	kbest = {} # word -> [ (word, sim)]
	curr_list = [] # [ (word, sim)]
	for line in codecs.open(filename, 'r', 'utf8'):
		if skip_lines > 0:
			skip_lines -= 1
			continue
		if not line.startswith('\t'):
			line = line.strip()
			if head_word:
				sorted_kbest = list( reversed( sorted(curr_list) ) )[:k]
				kbest[head_word] = [(x,y) for (y,x) in sorted_kbest]
				curr_list = []
			head_word = line
		else:
			line = line.strip()
			word, sim = line.split(';')
			curr_list.append( (float(sim), word) )


	if head_word:
		sorted_kbest = list( reversed( sorted(curr_list) ) )[:k]
		kbest[head_word] = [(x,y) for (y,x) in sorted_kbest]
		curr_list = []
			
	return kbest

def write_kbest(filename, kbest):
	with codecs.open(filename, 'w', 'utf8') as out:
		for head in sorted(kbest):
			out.write( '%s\n' % head )
			for p in kbest[head]:
				out.write( '\t%s\n' % (';'.join([str(x) for x in p])) )


