################################################################################
#  R-Code 
#     Description:
#
#     Procedure is an R-implementation of Baum-Welch procedure. The 
#     Viterbi algorithm is described in great details in [Rabiner, 1989] 
#
#     Input: a file of normalized and transformed  probes intensities, 
#	     ordered by their position on the chromosome.
#     
#     Output: the sequence of most probable states for each observation, 
#             i.e. probe location.
#
#
#   Usage:
#
#    All parameters are hard-coded  in the body of the script below. 
#    Arguments:
#     Input file is a two column  file:
#           First column  is the absolute position of probe on the genome.
#           Second column is the normalized and transformed intensity for 
#	    that probe.
#
#     In order to to change the number of states in HMM model, change 
#     <NumStates>, where valid values are in [2:5] range.
#
#     In order to change the number of components in mixture, change n.mix=2 
#     to a desired value. 
#
#    Values:
#     
#    The results are stored  in <out.file.name>  in the following format:
#
#          Signal
#          States
#          Observed intensity
#          Most probable state 
#          Average intensity, corresponding to this state 
#          Reference 
################################################################################

#------------------------------------------------------------------------------
#     User adjustable parameters
#------------------------------------------------------------------------------
data.path           ="data"
res.path            ="res"
out.file.name       ="out.txt"
inp.file.name       ="data.txt"

max.number.iteration=20
NumStates           =3
n.mix               =2

initial.values.means    <-c(-2,0,0.6,1,2)
initial.values.variance <-c(1.3,0.2,0.2,0.5,0.5)

################################################################################
#  <probab> calculates the probability density function for the mixture of several 
#  normal distributions. It calculate the probability of observations given that 
#  the observations are sampled from some distribution that is a mixture of normal 
#  distributions in a known proportions with known mean and variance. 
#
#  Input :
#       o:         vector of observations
#       mean.vec:  vector of distribution means
#       sigma.vec: vector of distribution variances
#       frac.vec:  vector of distribution mixture proportion
#
#  Output :
#       prob.f:    vector of probabilities of observation
#-------------------------------------------------------------------------------

probab<-function(o,mean.vec,sigma.vec,frac.vec){
     n.mix.f=length(mean.vec)
     x=dnorm(o,mean.vec[1],sigma.vec[1])
     x[x<0.0000001] =0.0000001
     prob.f=frac.vec[1]*x
     for(i in 2:n.mix.f){
           x=dnorm(o,mean.vec[i],sigma.vec[i])
           x[x<0.0000001] =0.0000001
           prob.f=prob.f+frac.vec[i]*x
     }
     return(prob.f)
}
################################################################################


#-------------------------------------------------------------------------------
#  Initialization of transition matrices:
#
#     We use two transition matrices for "FAR" and "NEAR" probes distances.
#     By chip design, we have most frequent probe distances of 25 bases, i.e.
#     probes that belong to the same gene and the distances about 1000 b that
#     correspond to the distance between neighboring genes.
#     We assume that transition (change of state) 10 times more likely between
#     genes than within gene.
#      
#-------------------------------------------------------------------------------

transition.probab.near =0.99999
transition.probab.far  =0.9999

a <-array(dim=c(NumStates,NumStates))
a1<-array(dim=c(NumStates,NumStates))

for (i in 1:NumStates){
     for (j in 1:NumStates){
          if(i==j){ 
                a[i,j]=transition.probab.near
                a1[i,j]=transition.probab.far
          }
          else {
                a[i,j]=(1-transition.probab.near )/(NumStates-1)
               a1[i,j]=(1-transition.probab.far)/(NumStates-1)
          }
     }
}

#-------------------------------------------------------------------------------
#  Loading the Data 
#     data is two column file: 
#     first column  is the absolute position of probe on genome 
#     second column is the normalized and transformed intensity
#-------------------------------------------------------------------------------

setwd(data.path)
data                    <-read.table(inp.file.name, header = F,sep = "\t")
probe.pos               <-data[,1]
signal                  <-data[,2]
num.of.observations =   length(signal)

#-------------------------------------------------------------------------------
#  calculating distances between the probes

for (t in (num.of.observations - 1):1){
     distance[t+1]<-probe.pos[t+1]-probe.pos[t]
}
distance[1]=0

#-------------------------------------------------------------------------------
#  Defining variables
#     We model signal probability distribution as mixture of 
#     two normal distributions in proportion defined by  <mix.fraction >,
#     so that for each state we have two set of means and variance parameters:
#         mix.fraction  is the matrix ( NumStates x n.mix )
#         means         is the matrix ( NumStates x n.mix )
#         temp.means    is the matrix ( NumStates x n.mix )
#         sigmas        is the matrix ( NumStates x n.mix )
#
#         forward, backward variables are defined by formulas (18-25) in [Rabiner, 1989]
#         gamma.k  is defined by formula (26-29)  in [Rabiner, 1989]
#         xi       defined by formula (36-37)  in [Rabiner, 1989]
#-------------------------------------------------------------------------------

means            <-array(dim=c(NumStates,n.mix))
temp.means       <-array(dim=c(NumStates,n.mix))
sigmas           <-array(dim=c(NumStates,n.mix))
mix.fraction     <-array(dim=c(NumStates,n.mix))
gamma.k          <-array(dim=c(n.mix,NumStates,num.of.observations))
pi               <-array(0,dim=c(NumStates))
gamma            <-array(dim=c(NumStates,num.of.observations))
xi               <-array(dim=c(NumStates,NumStates,num.of.observations))
sumgamma         <-array(dim=c(NumStates))
forward          <-array(dim=c(NumStates,num.of.observations))
backward         <-array(dim=c(NumStates,num.of.observations))
old.likel         =0

#---------------------- intialization  ---------------------------

for (i in 1:NumStates){
     pi[i]=0
     for (k in 1:n.mix){
          c.k[i,k]     =k/((n.mix+1)*n.mix/2)
          means[i,k]   =initial.values.means[i]+0.5*(k-1)
          sigmas[i,k]  =initial.values.variance[i]+(k-1)
     }
}
pi[1]=1

################################################################################
#   Baum Welsch Procedure
#   Iterative refining of distribution parameters 
################################################################################
flag<-FALSE
for (iteration in 1:max.number.iteration){
     probvec<-probab(signal[1:num.of.observations],means[1,],sigmas[1,],c.k[1,])
     for (i in 2:NumStates) {
          probvec<-cbind(probvec,probab(signal[1:num.of.observations],
                         means[i,],sigmas[i,],c.k[i,]))
     }

#-------------------------------------------------------------------------------
#  forward variable calculation, defined by formulas (18-22) in [Rabiner, 1989] 
#-------------------------------------------------------------------------------
forward[,1] =pi*probvec[1,]
likel.for   =0
s           =1 /sum(forward[,1])
forward[,1] =forward[,1]*s
likel.for   =likel.for+log(s)

for (t in 1:(num.of.observations-1)){
    if( distance[t]<50) {  forward[,t+1]=(forward[,t]%*%a)*probvec[t+1,]}
    else {                 forward[,t+1]=(forward[,t]%*%a1)*probvec[t+1,]}
    s             =1 /sum(forward[,t+1])
    forward[,t+1] =forward[,t+1]*s
    likel.for     =likel.for+log(s)
}

#-------------------------------------------------------------------------------
#  backward variable calculation, defined by formulas (23-25) in [Rabiner, 1989] 
#-------------------------------------------------------------------------------
backward[,num.of.observations]=1
likel.back   =0
s            = 1 /sum(backward[,1])
backward[,1] = backward[,1]*s
likel.back   =likel.back+log(s)

for (t in (num.of.observations - 1):1){
    if( distance[t]<50) { backward[,t]=a %*% (backward[,t+1]*probvec[t+1,])}
    else {                backward[,t]=a %*% (backward[,t+1]*probvec[t+1,])}
    s            = 1 /sum(backward[,t])
    backward[,t] = backward[,t]*s
    likel.back   = likel.back+log(s)
}

#-------------------------------------------------------------------------------
#  P(O|lambda), defined by formulas (27) in [Rabiner, 1989] 
#-------------------------------------------------------------------------------
for (t in 1:(num.of.observations -1)){
    if( distance[t]<50){
        POlambda = c((forward[,t]%*%a)%*%(backward[,t+1]*probvec[t+1,]))
    }
    else{
        POlambda = c((forward[,t]%*%a1)%*%(backward[,t+1]*probvec[t+1,]))
    }

#-------------------------------------------------------------------------------
#  computation of xi[i,j,t], defined by formulas (36-37) in [Rabiner, 1989] 
#-------------------------------------------------------------------------------
    if( distance[t]<50) tempa<-forward[,t]*a/POlambda
    else                tempa<-forward[,t]*a1/POlambda
    tempb<-probvec[t+1,]*backward[,t+1]
    for (j in 1:NumStates) xi[,j,t] = tempa[,j]*tempb[j]
    for.back=sum(forward[,t]*backward[,t])
        for (i in 1:NumStates) {
            gamma[i,t] =sum( xi[i,,t])
            temp.j     =forward[i,t]*backward[i,t]/for.back
            temps      =0
            for (k in 1:n.mix) {
                temps=temps+c.k[i,k]*max(0.00000001,dnorm(signal[t],
                                 mean=means[i,k], sd=sigmas[i,k]) )
            }
            if(temps==0) temp.j=0
            else         temp.j=temp.j/temps
            for (k in 1:n.mix){
                gamma.k[k,i,t]=temp.j*c.k[i,k]*max(0.00000001,dnorm(signal[t],
                                 mean=means[i,k], sd=sigmas[i,k]))
            }
       }
}

#-------------------------------------------------------------------------------
#  sum gamma, defined by formulas (38) in [Rabiner, 1989] 
#-------------------------------------------------------------------------------

for (i in 1:NumStates){
    sumgamma[i]=sum(gamma[i,1:(num.of.observations -1)])
    for (k in 1:n.mix){
        if(sumgamma[i]>0.001){
            c.k[i,k]=sum(gamma.k[k,i,1:(num.of.observations -1)])/sumgamma[i]
        }
        else{
            c.k[i,k]=1/n.mix
        }
    }
}
pi=gamma[,1]

#-------------------------------------------------------------------------------
#  calculation of expectated values of  transition probabilities, defined by 
#  formulas (40) in [Rabiner, 1989]
#-------------------------------------------------------------------------------
if(flag){
        for (i in 1:NumStates){
            for (j in 1:NumStates)a[i,j]=sum(xi[i,j,],na.rm=TRUE)/sumgamma[i]
        }
}

#-------------------------------------------------------------------------------
# calculation of expectated values  means and variance, defined by formulas (40)
# in [Rabiner, 1989]
#-------------------------------------------------------------------------------
for (i in 1:NumStates){
    for (k in 1:n.mix){
        temp.sum=sum(gamma.k[k,i,1:(num.of.observations -1)])
        if(temp.sum>0.000000001){
            temp.means[i,k]= sum(signal[1:(num.of.observations -1)]
                            *gamma.k[k,i,1:(num.of.observations -1)])/ temp.sum
            sigmas[i,k]= sqrt(sum((signal[1:(num.of.observations -1) ]-
                    means[i,k])^2*gamma.k[k,i,1:(num.of.observations -1)])/temp.sum)
        }
        else{
            temp.means[i,k]= 0;
            sigmas[i,k]=1;
        }
    }
}
means<-temp.means
}

################################################################################
#                    END OF BAUM WELSCH ALGORITHM
################################################################################


################################################################################
#                    VITERBI  ALGORITHM
################################################################################

#-------------------------------------------------------------------------------
#  Variables definition 
#    main.means is the mean intensity of state. Since we use the mixture of
#               normal distributions, we need to determine which of the components 
#               is the most frequent being the main one.
#    ap,ap1     is a logarithms of transition matrices
#    delta,psi  are the auxiliary variables, defined by formulas 30 -32 
#               from [Rabiner, 1989]
#    states     is the reconstructed sequence of most probable states
#    intens     is the auxiliary variable for storing intensities corresponding to 
#               reconstructed states

main.means  <-array(dim= c(NumStates))
ap          <-array(dim= c(NumStates,NumStates))
ap1         <-array(dim= c(NumStates,NumStates))
delta       <-array(dim= c(NumStates,2)) 
states      <-array(dim= c(num.of.observations))
intens      <-array(dim= c(num.of.observations))
psi         <-array(0,dim=c(NumStates,num.of.observations))
prob        <-array(0,dim=c(NumStates,num.of.observations))

ap         <-log(a)
ap1        <-log(a1)
for (i in 1:NumStates){
     main.comp     =order(c.k[i,], decreasing=T)[1]
     main.means[i] =means[i,main.comp]
}

#-------------------------------------------------------------------------------
#         forward computation : initialization   formula (32)
#-------------------------------------------------------------------------------
psi[,1]  =1
delta[,1]<-pi+log(probvec[1,])
#-------------------------------------------------------------------------------
#         forward computation : recursion        formula (33)
#-------------------------------------------------------------------------------
for (t in 2:num.of.observations){
    for (j in 1:NumStates){
        if(distance[t]<50) {m = delta[1,1] + ap[1,j] }
        else               {m = delta[1,1] + ap1[1,j]}
        temppsi = 1
        for(i in 1:NumStates){
            if(distance[t]<50) {x = delta[i, 1] + ap[i,j] }
            else               {x = delta[i, 1] + ap1[i,j]}
            if (x > m){
                m       = x
                temppsi = i
            }
        }
        delta[j, 2] = m +log(probvec[t,j])
        psi[j,t]    = temppsi
    }
    delta[,1]<-delta[,2]
}

#-------------------------------------------------------------------------------
#         forward computation : termination       formula (34)
#---------------------------------------------------   -------------------------
s=order(delta[, 2], decreasing=T)[1]
states[num.of.observations] = s
intens[num.of.observations] = main.means[s]

#-------------------------------------------------------------------------------
#         backtracking
#-------------------------------------------------------------------------------

for (t in (num.of.observations - 1):1){
    s         = psi[s,t]
    states[t] = s
    intens[t] = main.means[s]
    prob[,t]  <-forward[,t]*backward[,t]/sum(forward[,t]*backward[,t])
}

################################################################################
#                       END OF VITERBI  ALGORITHM
################################################################################


#-------------------------------------------------------------------------------
#                       saving  result to the file

write.table(cbind( signal[1:num.of.observations], 
                   states[1:num.of.observations], 
                   intens[1:num.of.observations],
                   out.file.name,
                   sep="\t",quote=FALSE, row.names = FALSE,col.names = TRUE)}

################################################################################
#                                     END 
################################################################################