# -*- coding: utf-8 -*-
import codecs, sys, random

import pairs, kbest

class Evaluator(object):
	
	def __init__(self, wnpairs, k = 20, flex = None, minsim = 0.000001):
		self.wnparis = wnpairs
		self.flex = flex
		self.not_evaluated = []
		self.k = k
		self.kbest = {} # word -> [(word, value), ...], should be initialized in prepare_kbest
		self.wordnet_words = set([])
		self.minsim = minsim
		
		self._prepare_wn_pairs()
		self.prepare_kbest()
		
	def _prepare_wn_pairs(self):
		flex_pairs = {}
		if self.flex:
			for w in self.wnparis:
				w_flex = '%s:%s' % (self.flex,w)
				flex_pairs[w_flex] = ['%s:%s' % (self.flex,x) for x in self.wnparis[w]]
		else:
			flex_pairs = self.wnparis
		
		self.wnparis = flex_pairs
		
		for w in self.wnparis:
			self.wordnet_words.add(w)
			for kw in self.wnparis[w]:
				self.wordnet_words.add(kw)
			
	def prepare_kbest(self):
		raise NotImplementedError, "Prepare kbest in Evaluator - this method should be overridden" 
	
	def evaluate(self, only_in_wordnet = False):
		"""
		Params:
		- only_in_wordnet - will remove all words from kbest list that are not in wordnet pairs
		"""
		p = 0
		r = 0
		all_count = 0 
		for w in self.kbest:
			if only_in_wordnet:
				in_matrix = set([x[0] for x in self.kbest[w] 
								if x[0] in self.wordnet_words
								and x[1] > self.minsim
								] [:self.k])
			else:
				in_matrix = set([x[0] for x in self.kbest[w] if x[1] > self.minsim][:self.k])
			
			if not w in self.wnparis or len(in_matrix) == 0:
				self.not_evaluated.append(w)
				continue
			
			in_wordnet = set(self.wnparis[w])
			intersection = in_matrix.intersection(in_wordnet)
			
			p += float(len(intersection))/float(len(in_matrix))
			r += float(len(intersection))/float(len(in_wordnet))
			all_count += 1
			
		all_count = float(all_count)
		precision = 0.0
		recall = 0.0
		if p > 0.0:
			precision = 100*p / all_count
		if r > 0.0:
			recall = 100*r / all_count
		return (precision, recall)
	
class ListEvaluator(Evaluator):
	
	def __init__(self, kbest_list, wnpairs, k = 20, minsim = 0.000001, flex = None):
		super(ListEvaluator, self).__init__(wnpairs, k, flex, minsim)
		self.kbest = kbest_list
			
	def prepare_kbest(self):
		pass

class CutEvaluator(Evaluator):

	def __init__(self, matrixdir, matrixname, wnpairs, cmpdesc, k = 20, minsim = 0.000001, flex = "subst", sample = 2400):
		self.matrixdir = matrixdir
		self.matrixname = matrixname
		self.cmpdesc = cmpdesc
		self.sample = sample
		super(CutEvaluator, self).__init__(wnpairs, k, flex, minsim)

	def prepare_matrix_words(self):
		self.matrix = {}
		import os.path
		labels = os.path.join(self.matrixdir, os.path.join(self.matrixname,self.matrixname + ".rlabel"))
		for line in open(labels,"r"):
			line = line.strip()
			self.matrix[line] = 1
			self.matrix[line[6:]] = 1

	def prepare_kbest(self):

		self.prepare_matrix_words()
		
		tmplist = []
		for w in self.wordnet_words:
			tmplist.append(w)
		tmplist.sort()
		random.seed(0)
		tmplist.shuffle()		
		
		kbest_words = set([])
		passed = 0
		for w in tmplist:
			if w not in self.matrix:
				continue
			kbest_words.add(w)
			passsed = passed + 1
			if passed>self.sample:
				break
		
		self.kbest = kbest.generate_kbest_locally(self.k, kbest_words, self.matrixname, self.matrixdir, self.cmpdesc)	
		kbest.write_kbest(self.cmpdesc+str(self.k)+"best.txt",self.kbest)

	def evaluate(self, only_in_wordnet = False):
		"""
		Params:
		- only_in_wordnet - will remove all words from kbest list that are not in wordnet pairs
		"""
		p = 0
		r = 0
		all_count = 0 
		for w in self.kbest:
			if only_in_wordnet:
				in_matrix = set([x[0] for x in self.kbest[w] 
								if x[0] in self.wordnet_words
								and x[1] > self.minsim
								] [:self.k])
			else:
				in_matrix = set([x[0] for x in self.kbest[w] if x[1] > self.minsim][:self.k])
			
			if not w in self.wnparis or len(in_matrix) == 0:
				self.not_evaluated.append(w)
				continue
			
			in_wordnet = set(self.wnparis[w])
			intersection = in_matrix.intersection(in_wordnet)
			
			p += float(len(intersection))/float(len(in_matrix))
			r += float(len(intersection))/float(len(in_wordnet))
			all_count += 1
			if (all_count > self.sample):
				break
			
		all_count = float(all_count)
		precision = 0.0
		recall = 0.0
		if p > 0.0:
			precision = 100*p / all_count
		if r > 0.0:
			recall = 100*r / all_count
		return (precision, recall)
	

def load_words(filename, skip_lines = 0):
	words = codecs.open(filename, 'r', 'utf8').readlines()[skip_lines:]
	return sorted([w.strip() for w in words])

def main():
	if len(sys.argv) != 5:
		print "Usage: python wnp_eval.py pairs_fileordir kbest_file k true|false"
		sys.exit(1)
		
	pairs_filename = sys.argv[1]
	kbest_filename = sys.argv[2]
	k = int(sys.argv[3])
	
	only_wordnet = False
	if sys.argv[4].lower().strip() == 'true':
		only_wordnet = True
	
	print "Loading kbest file", kbest_filename 
	kbest_lists = kbest.parse_kbest_txt_file(kbest_filename, k)
	print "Opening", pairs_filename
	loaded_pairs = pairs.load(pairs_filename)
	print "Loaded %d pairs" % (len(loaded_pairs))
	print "Loaded kbest lists for %s head words" % len(kbest_lists) 
	ev = ListEvaluator(kbest_lists, loaded_pairs, k = k, flex = 'subst')
	print "Evaluation underway..."
	print "Precision: %.2f%%\nRecall: %.2f%%" % ev.evaluate(only_wordnet) 
	print "Skipped %d words..." % len(ev.not_evaluated)
		

if __name__ == '__main__':
	main()
