#!/usr/bin/python2.7
# -*- coding: utf-8 -*-

import random
import math
import sys
import pickle
from copy import deepcopy

class HMM:
	#Constructor
	def __init__(self, nbrState, observations, start_p = None, trans_p = None, emission_p = None):
		self.nbrState = nbrState
		self.observations = observations
		
		#Init transition probabilities
		if trans_p == None:
			self.trans_p = []
			for i in range(nbrState):
				somme=0.0
				self.trans_p.append([])
				for j in range(nbrState):
					self.trans_p[i].append(1.0/nbrState + (random.random() - 0.5)/(3*nbrState))
					somme=somme+self.trans_p[i][j]
				somme = 1.0-somme
				for j in range(nbrState):
					self.trans_p[i][j]=self.trans_p[i][j]+somme/nbrState
		else:
			self.trans_p = trans_p

		#Init start probabilities
		if start_p == None:
			self.start_p = []
			somme=0.0
			for i in range(nbrState):
				self.start_p.append(1.0/nbrState + (random.random() - 0.5)/(3*nbrState))
				somme=somme+self.start_p[i]
			somme = 1.0 - somme
			for i in range(nbrState):
				self.start_p[i] = self.start_p[i] + somme/nbrState
		else:
			self.start_p = start_p

		#Init emission probabilities
		if emission_p == None:
			self.emission_p = []
			nbrObs = len(observations)
			for i in range(nbrState):
				somme=0.0
				self.emission_p.append([])
				for j in range(nbrObs):
					self.emission_p[i].append(1.0/nbrObs + (random.random() - 0.5)/(3*nbrObs))
					somme=somme+self.emission_p[i][j]
				somme=1.0-somme
				for j in range(nbrObs):
					self.emission_p[i][j]=self.emission_p[i][j]+somme/nbrObs
		else:
			self.emission_p = emission_p
	
	#print method of the HMM
	def impress(self):
		print "Nombre d'état: %d, Nombre d'observations: %d" % (self.nbrState, len(self.observations))
		print "Probabilité de départ"
		for i in range(self.nbrState):
			print "%8d" % (i+1),
		print
		print "   ",
		for i in range(self.nbrState):
			print "%.8s" % ("%f" % self.start_p[i]),
		print
		print
		print "Probabilité de transitions"
		for i in range(self.nbrState):
			print "%8d" % (i+1),
		print
		for i in range(self.nbrState):
			print "%3d" % (i+1),
			for j in range(self.nbrState):
				print "%.8s" % ("%f" % self.trans_p[i][j]),
			print
		print
		print "Probabilité d'émissions (transposée)"
		print "   ",
		for i in range(self.nbrState):
			print "%8d" % (i+1),
		print
		for i in range(len(self.observations)):
			print "%7s" % (self.observations[i]),
			for j in range(self.nbrState):
				print "%.8s" % ("%f" % self.emission_p[j][i]),
			print
	
	#Compute and return the probability of the sequence passed in parameter
	def probability(self, obs_seq):
		N = self.nbrState
		M = len(self.observations)
		T = len(obs_seq)
		#the alpha-pass
		alpha = []
		c=[]
		#init
		c.append(0)
		alpha.append([])
		for i in range(N):
			alpha[0].append(self.start_p[i]*self.emission_p[i][obs_seq[0]])
			c[0]=c[0]+alpha[0][i]
		c[0]=1.0/c[0]
		for i in range(N):
			alpha[0][i]=c[0]*alpha[0][i]
			
		#compute
		for t in range(1,T):
			c.append(0.0)
			alpha.append([])
			for i in range(N):
				alpha[t].append(0.0)
				for j in range(N):
					alpha[t][i]=alpha[t][i]+alpha[t-1][j]*self.trans_p[i][j]
				alpha[t][i]=alpha[t][i]*self.emission_p[i][obs_seq[t]]
				c[t]=c[t]+alpha[t][i]
			c[t]=1.0/c[t]
			for i in range(N):
				alpha[t][i]=c[t]*alpha[t][i]
		
		logProb=0.0
		for i in range(T):
			logProb=logProb+math.log(c[i])
		logProb=-logProb
		return logProb

	#Train the HMM to best fit the given sequence, i.e try to maximize the sequence probability
	def training(self, obs_seq, maxIter=100000):
		N = self.nbrState
		M = len(self.observations)
		T = len(obs_seq)
		print "Ajustement du HMM par rapport à %d observations" % T
		iters = 0
		oldLogProb = float("-inf")
		logProb = 0.0
		while iters < maxIter:
			iters=iters+1
			print "Itération %d, état du HMM:" % iters
			self.impress()
			#the alpha-pass
			alpha = []
			c=[]
			#init
			c.append(0)
			alpha.append([])
			for i in range(N):
				alpha[0].append(self.start_p[i]*self.emission_p[i][obs_seq[0]])
				c[0]=c[0]+alpha[0][i]
			c[0]=1.0/c[0]
			for i in range(N):
				alpha[0][i]=c[0]*alpha[0][i]
			
			#compute
			for t in range(1,T):
				c.append(0.0)
				alpha.append([])
				for i in range(N):
					alpha[t].append(0.0)
					for j in range(N):
						alpha[t][i]=alpha[t][i]+alpha[t-1][j]*self.trans_p[i][j]
					alpha[t][i]=alpha[t][i]*self.emission_p[i][obs_seq[t]]
					c[t]=c[t]+alpha[t][i]
				c[t]=1.0/c[t]
				for i in range(N):
					alpha[t][i]=c[t]*alpha[t][i]

			#Compute log[P(Obs_seq)]
			logProb=0.0
			for i in range(T):
				logProb=logProb+math.log(c[i])
			logProb=-logProb
			print "Log[P(ObsSeq)] = %f\n" % logProb
			if logProb > oldLogProb:
				oldLogProb=logProb
			else:
				break

			#beta-pass
			beta = []
			beta.append([])
			for i in range(N):
				beta[0].append(c[T-1])
			for t in range(T-2,-1,-1):
				beta.insert(0,[])
				for i in range(N):
					beta[0].append(0.0)
					for j in range(N):
						beta[0][i]=beta[0][i]+self.trans_p[i][j]*self.emission_p[j][obs_seq[t+1]]*beta[1][j]
					beta[0][i]=c[t]*beta[0][i]

			#compute gama-ti and gama-tij
			gamaTi = []
			gamaTij = []
			for t in range(T-1):
				denom=0.0
				for i in range(N):
					for j in range(N):
						denom=denom+alpha[t][i]*self.trans_p[i][j]*self.emission_p[j][obs_seq[t+1]]*beta[t+1][j]
				gamaTi.append([])
				gamaTij.append([])
				for i in range(N):
					gamaTi[t].append(0.0)
					gamaTij[t].append([])
					for j in range(N):
						gamaTij[t][i].append((alpha[t][i]*self.trans_p[i][j]*self.emission_p[j][obs_seq[t+1]]*beta[t+1][j])/denom)
						gamaTi[t][i] = gamaTi[t][i] + gamaTij[t][i][j]
			
			#Save start, trans, emission probabilities
			oldStart_p = deepcopy(self.start_p)
			oldTrans_p = deepcopy(self.trans_p)
			oldEmission_p = deepcopy(self.emission_p)


			#Re-estimate start, transition and emission probabilities
			#Start
			for i in range(N):
				self.start_p[i]=gamaTi[0][i]

			#Transition
			for i in range(N):
				for j in range(N):
					numer = 0.0
					denom = 0.0
					for t in range(T-1):
						numer = numer+gamaTij[t][i][j]
						denom = denom + gamaTi[t][i]
					self.trans_p[i][j]=numer/denom

			#Emission
			for i in range(N):
				for j in range(M):
					numer = 0.0
					denom = 0.0
					for t in range(T-1):
						if obs_seq[t] == j:
							numer = numer+gamaTi[t][i]
						denom=denom+gamaTi[t][i]
					self.emission_p[i][j]=numer/denom

		self.start_p = oldStart_p
		self.trans_p = oldTrans_p
		self.emission_p = oldEmission_p

	#Convert a sequence of observation in a sequence of observation index usable by others methods
	def convertObsSeq(self, obs_seq):
		obsNumSeq = []
		for obs in obs_seq:
			obsNumSeq.append(self.observations.index(obs))
		return obsNumSeq

#Start the demo
text = ""
for line in sys.stdin:
		text = text + line
obsSeq = list(text)

if len(sys.argv) == 2:
	f = open(sys.argv[1],"r")
	monHmm = pickle.load(f)
	f.close()
	#En raison de l'absence de 'w' dans la séquence d'entrainement, sa probabilité a été réduite à 0.
	#Pour éviter le bug sur un texte anglais (qui contient des 'w', lui) on 'pique' un peu de chance au 'v' pour la donner au 'w'.
	#Cela ne modifie pas le HMM de facon visible (10^-15) mais lui permet de fonctionner si il rencontre une séquence contenant un 'w'.
	monHmm.emission_p[0][22] = monHmm.emission_p[0][22] + 1e-15
	monHmm.emission_p[0][21] = monHmm.emission_p[0][21] - 1e-15
	monHmm.impress()
else:
	monHmm = HMM(2, map(chr, range(97,123))+[' '])
	#Lancement du training sans limite
	monHmm.training(monHmm.convertObsSeq(obsSeq[0:50000]))
	saveHmm = open("save_hmm.sav","w")
	pickle.dump(monHmm,saveHmm)
	saveHmm.close()

for i in range(0,300000,50000):
	print "Log[P(Obs[%d:%d])] = %f" % (i,i+50000, monHmm.probability(monHmm.convertObsSeq(obsSeq[i:i+50000])))
