# The observations #
####################
n <- 10
sig <- 3 # unknown sd of the noise
# matrix of unknown parameters
W <- matrix(c(1,2,3,0,0,0,0,1,2,3),ncol=5,byrow=TRUE) 
# hidden variables
X <- matrix(rnorm(n*2),ncol=2) 
# the observations
Y <- X %*% W + matrix(rnorm(n*5,sd=sig),nrow=n) 

# 
var.0 <- 2 # current var of noise
W.0 <- matrix(c(1,1,1,1,0,0,2,2,2,2),ncol=5,byrow=TRUE) 
# current matrix of parameters

# An iteration of the EM 
V.inv <- solve(t(W.0)%*%W.0 + var.0*diag(5)) 
# inverse of the variance of an observation vector
E <- Y%*%V.inv%*%t(W.0) 
# expectation of the hidden given the observations
v1 <- diag(2) - W.0%*%V.inv%*%t(W.0)
EY.sum <- t(E)%*%Y
EE.sum <- n*v1 + t(E)%*%E
W.1 <- solve(EE.sum) %*% EY.sum # new matrix of parameters
var.1 <-  mean(Y^2 - (E%*%W.1)*Y) #CORRECTED
# The general formulation:


# A function that implements an iteration of the EM

factor.EM <- function(W.0,var.0,Y)
{
    p <- dim(W.0)[1]
    q <- dim(W.0)[2]
    V.inv <- solve(t(W.0)%*%W.0 + var.0*diag(q)) 
    E <- Y%*%V.inv%*%t(W.0) 
    v1 <- diag(p) - W.0%*%V.inv%*%t(W.0)
    EY.sum <- t(E)%*%Y
    EE.sum <- n*v1 + t(E)%*%E
    W.1 <- solve(EE.sum) %*% EY.sum
    var.1 <- mean(Y^2 - (E%*%W.1)*Y) 
    return(list(W=W.1,v=var.1))
}
    
# check
factor.EM(W.0,var.0,Y)
W.1
var.1

run.EM <- function(W.0,var.0,Y,max.iter=10^3,max.eps=1/10^3)
{
    iter <- 0
    EPS <- 10^3
    while(iter<=max.iter & EPS > max.eps)
    {
        out <- factor.EM(W.0,var.0,Y)
        W.1 <- out$W
        var.1 <- out$v
        iter <- iter +1 
        EPS <- max(c(max(abs(W.1-W.0)),abs(var.1-var.0)))
        W.0 <- W.1
        var.0 <- var.1
     }
     return(list(W=W.1,v=var.1,iter=iter))
}

# check
run.EM(W.0,var.0,Y)


# Estimating the parameters #
#############################
n <- 10^4
sig <- 0.1
W <- matrix(c(1,2,3,0,0,0,0,1,2,3),ncol=5,byrow=TRUE) 
X <- matrix(rnorm(n*2),ncol=2)
Y <- X %*% W + matrix(rnorm(n*5,sd=sig),nrow=n) 

var.0 <- 0.02
W.0 <- W + matrix(rnorm(10,sd=1),2,5) 
run.EM(W.0,var.0,Y)