# equilibrium probs
equil=function(P){
        e=eigen(t(P))$vectors[,1]
        e/sum(e)
}

# simulate hidden sequence
hssim=function(n,lambda)
{
        r=dim(lambda)[1]
        states=1:r
        s=vector("numeric",n)
        pi.lam=equil(lambda)
        s[1]=sample(states,1,FALSE,pi.lam)
        for (t in 2:n) {
                s[t]=sample(states,1,FALSE,prob=lambda[s[t-1],])
        }
        s
}

# simulate observed sequence
obssim=function(s,P)
{
        n=length(s)
        r=dim(P)[2]
        q=dim(P)[1]
        states=1:q
        obs=vector("numeric",n)
        for (t in 1:n) {
                obs[t]=sample(states,1,FALSE,prob=P[,s[t]])
        }
        obs
}

lambda=rbind(c(0.999,0.0005,0.0005),c(0.001,0.9980,0.001),c(0.005,0.005,0.99))

s=hssim(10000,lambda)

P=cbind(c(0.6,0.1,0.1,0.2),c(0.25,0.3,0.2,0.25),c(0.1,0.6,0.2,0.1))

obs=obssim(s,P)

# optional - converting to/from another alphabet...
letters=c("a","c","g","t")
numbers=1:4
convert=function(x,frm,to)
{
	to[match(x,frm)]
}
obslets=convert(obs,numbers,letters)

# estimate emmision probs from observed and hidden sequence
Pest=function(s,obs,r,q)
{
est=matrix(0,nrow=q,ncol=r)
for (i in 1:r) {
    est[,i]=table(obs[s==i])
    est[,i]=est[,i]/sum(est[,i])
    }
    est
}

phat=Pest(s,obs,3,4)

# estimate lambda from hidden sequence
lambdaest=function(s,r)
{
	n=length(s)
	est=matrix(0,ncol=r,nrow=r)
	for (t in 2:n) {
	    est[s[t-1],s[t]]=est[s[t-1],s[t]]+1
	    }
	for (i in 1:r) {
	    est[i,]=est[i,]/sum(est[i,])
	}
	est
}

lamhat=lambdaest(s,3)

# marginal likelihood
mll=function(obs,lambda,P)
{
        r=dim(lambda)[1]
        q=dim(P)[1]
        n=length(obs)
        f=equil(lambda)
        ll=0
        for (i in 1:n) {
            fnew=P[obs[i],]*(f%*%lambda)
            xi=sum(fnew)
            ll=ll+log(xi)
            f=fnew/xi
        }
        ll
}

print(mll(obs,lambda,P))
print(mll(obs,lamhat,phat))

# local decoding algorithm
local=function(obs,lambda,P)
{
	r=dim(lambda)[1]
	q=dim(P)[1]
	n=length(obs)
	s=vector("numeric",n)
	f=matrix(0,nrow=r,ncol=n)
	b=matrix(0,nrow=r,ncol=n)
	# forwards
	f0=equil(lambda)
	f[,1]=P[obs[1],]*(f0%*%lambda)
	for (i in 2:n) {
	    f[,i]=P[obs[i],]*(f[,i-1]%*%lambda)
	    f[,i]=f[,i]/sum(f[,i])
	}	
	# backwards
	b[,n]=rep(1,r)
	s[n]=which.max(f[,n]*b[,n])
	for (i in (n-1):1) {
	    b[,i]=(P[obs[i+1],]*b[,i+1])%*%lambda
	    b[,i]=b[,i]/sum(b[,i])
	    s[i]=which.max(f[,i]*b[,i])
	}
	s
}

locest=local(obs,lambda,P)

# global decoding algorithm


# global decoding algorithm
global=function(obs,lambda,P)
{
	r=dim(lambda)[1]
	q=dim(P)[1]
	n=length(obs)
	s=vector("numeric",n)
	f=matrix(0,nrow=r,ncol=n)
	# forwards
	f0=equil(lambda)
	f[,1]=P[obs[1],]*(f0%*%lambda)
	for (i in 2:n) {
	    for (k in 1:r) {
		f[k,i]=P[obs[i],k]*max(f[,i-1]*lambda[,k])
	    }
	    f[,i]=f[,i]/sum(f[,i])
	}	
	# backwards
	s[n]=which.max(f[,n])
	for (i in (n-1):1) {
	    s[i]=which.max(lambda[,s[i+1]]*f[,i])
	}
	s
}

globest=global(obs,lambda,P)

op=par(mfrow=c(3,1))
plot(ts(s),main="Truth")
plot(ts(locest),main="Local decoding")
plot(ts(globest),main="Global decoding")
par(op)

# eof