#install.packages("astsa")
library("astsa")
options(repr.plot.width=15, repr.plot.height=8) #ajusta tamaño de graficas

### Ejemplo en R

# Simula una cadena de Markov de timepo discreto con matriz P y estado inicial i.
# Si P es n x n, i debe estar en {1,...,n}
# num iters es la cantidad de pasos de tiempo

mc.sim <- function(P, i=1, num.iters = 50 ) {
    
    # Vector de salida. Se inicializa vacio. 
    states <- numeric(num.iters)
    
    # El primero es el estado inicial 
    states[1] <- i

    for(k in 1:(num.iters-1)) {
       
        # Si estoy en el X_k, las probabilidades de salto son la fila X_k de P.
        p <- P[states[k], ]

        # Sorteo con esas probabilidades el proximo estado
        states[k+1] <-  which(rmultinom(1, 1, p) == 1)
    }
    
    return(states)
}

P <- matrix(c(0.3, 0.7, 0.4, 0.6),byrow=TRUE, nrow=2)

P

states <- mc.sim(P,i=1,num.iters=50)
plot(states, col=4, cex=2, lwd=2)

n=200

P <- matrix(c(0.95, 0.05, 0.1, 0.9),byrow=TRUE, nrow=2)

mu= c(2,-2)
sigma = 1

states <- mc.sim(P,1,n)
observations <- numeric(n)

for(k in 1:n) {
    observations[k] <-  rnorm(1, mean=mu[states[k]], sd=sigma)
}

#grafico las observaciones
tsplot(observations,col=4, lwd=2)

#grafico las lineas con la media de cada estado para ver la cadena escondida
lines(mu[states], col=5, lwd=4)

hist(observations,breaks=20, col=5)

install.packages("HiddenMarkov")
library("HiddenMarkov")

#construyo los parametros
P <- matrix(c(0.95, 0.05, 0.1, 0.9),byrow=TRUE, nrow=2)
delta <- c(1, 0)

#construyo la cadena y le doy la distribucion en cada estado (normal, medias y desvios por estado)
hmm <- dthmm(NULL, P, delta, "norm", list(mean=c(2, -2), sd=c(1,1)))
hmm <- simulate(hmm, nsim=500)

#en hmm$x queda guardada la simulación (variable y). No se guarda la variable escondida.
tsplot(hmm$x, col=4, lwd=2)

hist(hmm$x, breaks=20, col=5)

#Genero las observaciones a mano para poder saber el estado que dio lugar
# a las mismas (HiddenMarkov no lo guarda, ver comentario anterior)

n=500
P <- matrix(c(0.95, 0.05, 0.1, 0.9),byrow=TRUE, nrow=2)

mu= c(2,-2)
sigma = c(1,1)

states <- mc.sim(P,1, n)
observations <- numeric(n)

for(k in 1:n) {
    observations[k] <-  rnorm(1, mean=mu[states[k]], sd=sigma[states[k]])
}

tsplot(observations, lwd=2, col=4)
lines(mu[states], lwd=2, col=5)

#aplico forward-backward: devuelve el log de las probabilidades.
fb<- forwardback(observations, P, delta, "norm", list(mean=mu, sd=sigma))

#calculo la verosimilitud logaritmica de cada estado 
#utilizando la parte forward (logalpha) y la parte backward (logbeta)

ll = fb$logalpha+fb$logbeta

#Construyo un vector de resultados
xhat = numeric(n)

#Pongo el estado más probable (el de mayor verosimilitud) en cada lugar de xhat
for (k in 1:n) {
    
    if (ll[k,1]>ll[k,2]) {
        
        xhat[k] = 1
    
    } else {
    
        xhat[k] = 2
    
    }
    
}

#Ploteo la estimacion (puntos) y el estado real (lineas)
plot(xhat)
lines(states)
    

### Ejemplo. 

#Decodifico
xhat <- Viterbi(hmm)


#Ploteo la estimacion (puntos) y el estado real (lineas)
plot(xhat)
lines(states)

tsplot(observations)
hist(observations,breaks=40)

x_ini<- ifelse(observations>=0, 1, 2)
tsplot(x_ini)

# Esta funcion permite estimar la matriz de transiciones contando los saltos de estado.
trans.matrix <- function(X)
{
    X<-t(as.matrix(X))
    tt <- table( c(X[,-ncol(X)]), c(X[,-1]) )
    tt <- tt / rowSums(tt)
}

#Construyo una estimacion inicial de la matriz
P_ini=trans.matrix(x_ini)
P_ini

#Asigno probabilidad 1 al estado donde inicia la cadena.
delta_ini = as.numeric(x_ini[1] == (1:2))
delta_ini

#Extraigo la media y la desviacion estandar del comportamiento en cada estado asignado
mu1 = mean(observations[x_ini==1])
mu2 = mean(observations[x_ini==2])
sigma1 = sd(observations[x_ini==1])
sigma2 = sd(observations[x_ini==2])

mu_ini = c(mu1,mu2)
sigma_ini = c(sigma1,sigma2)

mu1
mu2
sigma1
sigma2

hmm_ini <- dthmm(observations, P_ini, delta_ini, "norm", list(mean=mu_ini, sd=sigma_ini))

hmm_fit <- BaumWelch(hmm_ini)

summary(hmm_fit)

#Parametros estimados. Notar que ajustan mejor a los valores con los cuales se genero el proceso
hmm_fit$Pi
hmm_fit$delta
hmm_fit$pm

xhat = Viterbi(hmm_fit)

plot(xhat)
lines(states)

medias_est = hmm_fit$pm$mean

tsplot(observations)
lines(medias_est[xhat],col=2,lw=2)

baum_viterbi_normal <- function(hmm_ini,tol=1e-3,maxiter=10) {
    
    paso = Inf #incremento de la verosimilitud
    k = 0 #no. de iteracion
    N = nrow(hmm_ini$P) #no. de estados
    
    hmm_fit = hmm_ini #en hmm_fit guardo la salida. Arranca igual a hmm_ini
    
    while (paso>tol && k<maxiter) {

        lL_ini =  logLik(hmm_fit) #guardo la verosimilitud anterior
        x = Viterbi(hmm_fit) #recupero los estados mas probables
      
        #Reestimo los parametros:
        delta = as.numeric(x[1] == (1:N))
        P = trans.matrix(x)

        #Construyo un vector para medias y varianzas y los lleno con los estimadores filtrando por estado
        mus = vector()
        sigmas = vector()
        for (i in 1:N) {
            mus[i] = mean(hmm_fit$x[x==i])
            sigmas[i] = sd(hmm_fit$x[x==i])
        }
 
        hmm_fit = dthmm(hmm_fit$x, Pi=P, delta=delta, distn="norm", pm= list(mean=mus, sd=sigmas))
        paso = abs(logLik(hmm_fit)-lL_ini) #calculo cuanto mejora la verosimilitud
        k = k+1 #incremento el contador de iteraciones
        
        print(paste("Paso: ",k))
        print(paste("Verosimilud: ",logLik(hmm)))
    }
    return(hmm_fit)
}

hmm_fit2 = baum_viterbi_normal(hmm_ini)

#Notar que esta estimación es levemente distinta a la de Baum-Welch.
summary(hmm_fit2)

xhat2 = Viterbi(hmm_fit2)

plot(xhat2)
lines(states)

medias_est_2 = hmm_fit$pm$mean

tsplot(observations)
lines(medias_est_2[xhat2],col=2,lw=2)