library(rstan)
rstan_options(auto_write=T)
options(mc.cores = parallel::detectCores())

n <- 100 #number of samples
N <- 3 
M <- 3
H <- 2 #true inner dim


##true parameters
TA <- matrix(c(0.5,1.5, 0,2, 1,2), nrow=N, ncol=H, byrow=T)
TB <- matrix(c(1,0, 2,0, 3,7), nrow=H, ncol=M, byrow=T)

## sample is taken from poisson dist
x <- array(0, dim=c(N,M,n))
for(i in 1:n){
  for(j in 1:N){
    for(k in 1:M){
      x[j,k,i] <- rpois(1, (TA%*%TB)[j,k])
    }
  }
}

## learner inner dim
learning_H <- H+1

#hyperparam for prior
alpha <- 2
beta <- 2

nmf_data <- list(n=n, N=N, M=M,H=learning_H, x=x, alpha=alpha, beta=beta)

nmf_fit <- stan(file="nmf_poisson.stan", data=nmf_data, chains=4, iter=2000)

pred_AB <- matrix(0, nrow=N, ncol=M)

nmf_param <- extract(nmf_fit, permuted=T)
parameter_num <- dim(nmf_param$A)[1]

for(i in 1:parameter_num){
  pred_AB <- pred_AB + nmf_param$A[i,,]%*%nmf_param$B[i,,]/parameter_num
}