using Distributions
using Turing
using Stan

# Load data; loaded data is a list of dict named `ldastandata`
include(Pkg.dir("Turing")*"/stan-models/lda-stan.data.jl")
topicdata = ldastandata[1]

# Load model
include(Pkg.dir("Turing")*"/stan-models/lda.model.jl")
#= NOTE: loaded model is defined as below
@model ldamodel(K, V, M, N, w, doc, beta, alpha) = begin
  theta = Vector{Vector{Real}}(M)
  for m = 1:M
    theta[m] ~ Dirichlet(alpha)
  end

  phi = Vector{Vector{Real}}(K)
  for k = 1:K
    phi[k] ~ Dirichlet(beta)
  end

  phi_dot_theta = [log([dot(map(p -> p[i], phi), theta[m]) for i = 1:V]) for m=1:M]
  for n = 1:N
    Turing.acclogp!(vi, phi_dot_theta[doc[n]][w[n]])
  end
end
=#

setchunksize(100)    # increase AD chunk-size to 100

samples = sample(ldamodel(data=topicdata), NUTS(1000, 0.65))

# Load visualization script for topic models; visualization function is called `vis_topic_res`
include(Pkg.dir("Turing")*"/stan-models/topic_model_vis_helper.jl")

@doc vis_topic_res  # show the usage of the visualization function

vis_topic_res(samples, topicdata["K"], topicdata["V"], 1000)

# Load data; loaded data is a list of dict named `nbstandata`
include(Pkg.dir("Turing")*"/stan-models/MoC-stan.data.jl")
topicdata2 = nbstandata[1]

# Load model
include(Pkg.dir("Turing")*"/stan-models/MoC.model.jl")
#= NOTE: loaded model is defined as below
@model nbmodel(K, V, M, N, z, w, doc, alpha, beta) = begin
  theta ~ Dirichlet(alpha)

  phi = Array{Any}(K)
  for k = 1:K
    phi[k] ~ Dirichlet(beta)
  end

  log_theta = log(theta)
  Turing.acclogp!(vi, sum(log_theta[z[1:M]]))

  log_phi = map(x->log(x), phi)
  for n = 1:N
    Turing.acclogp!(vi, log_phi[z[doc[n]]][w[n]])
  end

  phi
end

=#

samples2 = sample(nbmodel(data=topicdata2), NUTS(1000, 0.65))

vis_topic_res(samples2, topicdata2["K"], topicdata2["V"], 1000)

using ProgressMeter

x,n = 1,10
p = Progress(n)
for iter = 1:10
    x *= 2
    sleep(0.5)
    ProgressMeter.next!(p; showvalues = [(:iter,iter), (:x,x)])
end