'''
Created on Mar 16, 2011

@author: boban
'''
import codecs
import pymatrices
from scipy.optimize import leastsq
import numpy 

def make_axpb():
	
	def fun(p, x):
		a, b = p
		return  a*x+ b
	
	def err(p, x, y):
		err = y - fun(p, x)
		return err
	
	return fun, err, numpy.array([0.0001, 0.1])

def make_poly3():
	def fun(p, x):
		a,b,c, d = p
		return a*x+b*x**2 + c*x**3 + d
	
	def err(p, x, y):
		err = y - fun(p, x)
		return err
	
	return fun, err, numpy.array([0.1,0.1,1.1,1.1])

def reg(input, output):
	fun, err, p0 = make_axpb()
	#fun, err, p0 = make_poly3()
	params = leastsq(err, p0 , args=(input, output) )[0]
	return fun, params

def load_pair_list(filename):
	return [x.strip().split(';') for x in codecs.open(filename, 'r', 'utf8').readlines()]

def test_reg():
	input = numpy.array([0.3, 0.5, 0.7, 0.2, 0.5])
	print "IN: ", input
	output = input * 1.43 + 3.25
	print "OUT:", output
	fun, params = reg(input, output)
	print "\nFound params:", params
	print "FUN:", fun(params, input)

def write_msr_list(outfilename, matrix, pairs, prefix = 'subst:'):
	"""
	matrix - square, similarity matrix (rows = cols)
	"""
	with codecs.open(outfilename, 'w', 'utf8') as out:
		crs = matrix.GetCRS()
		max_rows = matrix.NumberOfRows()
		if max_rows != matrix.NumberOfColumns():
			raise RuntimeError, "rows != columns"
		for a, b in pairs:
			ap = prefix + a
			bp = prefix + b
			a_ind = matrix.GetRowIndex(ap.encode('utf8'))
			b_ind = matrix.GetRowIndex(bp.encode('utf8'))

			if a_ind > max_rows:
				print ap,'not in matrix'
				continue
			if b_ind > max_rows:
				print bp,'not in matrix'
				continue
			val = crs.GetCell(a_ind, b_ind)
			out.write('%s;%s;%f\n' % (a,b,val) )

def test_sim2vals():
	mdir = '/mnt/data/macierze/'
	mname = 'kgr3_v2_lincos_filtered'
	inlist_filename = './wn_lin.txt'
	outfilename = 'msr_list.txt'
	pairs = load_3list(inlist_filename).keys()
	print "Loaded %d pairs" % len(pairs)
	matrix = pymatrices.LoadSuperMatrix(mdir, mname, False, True)
	write_msr_list(outfilename, matrix, pairs)

def load_3list(filename):
	l = {} # pair -> val
	for line in codecs.open(filename, 'r', 'utf8'):
		line = line.strip()
		tokens = line.split(';')
		if len(tokens) != 3:
			print "skipping line", line
		else:
			tokens = [t.strip().lower() for t in tokens]
			l[ (tokens[0], tokens[1])] = float(tokens[2])
	return l


def aligng_lists(msr_list_filenamet, sim_list_filename):
	msr = load_3list(msr_list_filenamet)
	wn = load_3list(sim_list_filename)
	inters = list(sorted(set(msr.keys()).intersection( set(wn.keys() ))))
	print "Intersection of the two list contains %d items" % len(inters)
	msr_array = numpy.zeros( len(inters) )
	wn_array = numpy.zeros( len(inters) )

	for i,p in enumerate(inters):
		msr_array[i] = msr[p]
		wn_array[i] = wn[p]
	
	return msr_array, wn_array


def final_reg():
	msrf = './msr_list.txt'
	wnf = './wn_lin.txt'
	msr,wn = aligng_lists(msrf, wnf)
	print msr[:10]
	print wn[:10]
	fun, params = reg(msr, wn)
	print "\nFound params:", params

	test_msr = './test_msr.txt'
	test = load_3list(test_msr)
	print 'A; B; MSR; mapped WN'
	for p in test:
		print '%s; %s; %.4f; %.4f' % (p[0], p[1], test[p], fun(params, numpy.array(test[p])))




def main():
	#test_reg()
	#test_sim2vals()

	final_reg()

if __name__ == '__main__':
	main()
