#!/usr/bin/env python
import optparse, sys, math
from collections import defaultdict

optparser = optparse.OptionParser()
optparser.add_option("-b", "--bitext", dest="bitext", default="data/dev-test-train.de-en", help="Parallel corpus (default data/dev-test-train.de-en)")
optparser.add_option("-t", "--threshold", dest="threshold", default=0.5, type="float", help="Threshold for aligning with Dice's coefficient (default=0.5)")
optparser.add_option("-n", "--num_sentences", dest="num_sents", default=sys.maxint, type="int", help="Number of sentences to use for training and alignment")
(opts, _) = optparser.parse_args()

sys.stderr.write("Reading in the data...")
#bitext is a list of [German],[English] pairs
bitext = [[sentence.strip().split() for sentence in pair.split(' ||| ')] for pair in open(opts.bitext)][:opts.num_sents]
e_count = []
fe_pair = defaultdict(list)
#gets a count of each English word and the English translations for every foriegn (German) word
for (n, (f, e)) in enumerate(bitext):
	for e_j in set(e):
		e_count.append(e_j)
		for f_i in set(f):
			fe_pair[f_i].append(e_j)
	if n % 500 == 0: sys.stderr.write(".")

#initialize the probabilities to 1/(length of the number of translations for that word)
def init_em(e_group,f_group,fe):
	sys.stderr.write("\nInitializing EM probabilities...")
	p = defaultdict(float)
	for f in f_group:
		for e in e_group:
			p[(f,e)] = 1.0/len(fe[f])
	return p

#the training
def em(p, bitext, EM_times, e_count, fe_pair):
	sys.stderr.write("\nTraining with IBM Model 1...")
	itt = 0
	while True:
		log = 0
		count = defaultdict(float)
		for (n, (f, e)) in enumerate(bitext):
			norm = 1.0/len(f)
			total_source = defaultdict(float)
			for e_word in set(e):
				for f_word in set(f):
					total_source[e_word] += p[(f_word, e_word)] * norm
				log = math.log(total_source[e_word])
				for f_word in set(f):
					count[(f_word, e_word)] += (p[(f_word, e_word)]*norm) / total_source[e_word]
			if n % 500 == 0: sys.stderr.write('.')
		for f_word in fe_pair.keys():
#			sys.stderr.write("\n"+f_word)
			total = sum(count.values())
			for e_word in e_count:
				p[(f_word, e_word)] = count[(f_word, e_word)] / total
		itt+=1
		if itt >= EM_times:
			return p #converged
						

#initialize p(e_i|g_i)'s
p = init_em(e_count, fe_pair.keys(), fe_pair)
p_final = em(p, bitext, 5, e_count, fe_pair)
del(e_count)
del(fe_pair)

#for (k, (f_i, e_j)) in enumerate(fe_count.keys()):
#  IBM1[(f_i,e_j)] = 2.0 * fe_count[(f_i, e_j)] / (f_count[f_i] + e_count[e_j])
#  if k % 5000 == 0:
#	sys.stderr.write(".")
#sys.stderr.write("\n")

sys.stderr.write("\nCalculating predicted values...")
#output & thresholding
for (f, e) in bitext:
	for (i, f_i) in enumerate(f):
		val = 0
		loc = 0
		for (j, e_j) in enumerate(e):
			if p[(f_i,e_j)] > val:
				val = p[(f_i,e_j)]
				loc = j
		sys.stdout.write("%i-%i " % (i,loc))
	sys.stdout.write("\n")
	if n % 500 == 0: sys.stderr.write('.')
