using StatisticalRethinking, Distributed, JLD, LinearAlgebra
using Mamba

howell1 = CSV.read(joinpath(dirname(Base.pathof(StatisticalRethinking)), "..", "data", "Howell1.csv"), delim=';')
df = convert(DataFrame, howell1);

df2 = filter(row -> row[:age] >= 18, df);

data = Dict(
  :x => convert(Array{Float64,1}, df2[:weight]),
  :y => convert(Array{Float64,1}, df2[:height])
)

Dict{Symbol,Array{Float64,1}} with 2 entries:
  :y => [151.765, 139.7, 136.525, 156.845, 145.415, 163.83, 149.225, 168.91, 14…
  :x => [47.8256, 36.4858, 31.8648, 53.0419, 41.2769, 62.9926, 38.2435, 55.48, …

logfgrad = function(x::DenseVector)
  b0 = x[1]
  b1 = x[2]
  logs2 = x[3]
  r = data[:y] .- b0 .- b1 .* data[:x]
  logf = (-0.5 * length(data[:y]) - 0.001) * logs2 -
           (0.5 * dot(r, r) + 0.001) / exp(logs2) -
           0.5 * b0^2 / 1000 - 0.5 * b1^2 / 1000
  grad = [
    sum(r) / exp(logs2) - b0 / 1000,
    sum(data[:x] .* r) / exp(logs2) - b1 / 1000,
    -0.5 * length(data[:y]) - 0.001 + (0.5 * dot(r, r) + 0.001) / exp(logs2)
  ]
  logf, grad
end

# MCMC Simulation with No-U-Turn Sampling

n = 5000
burnin = 1000
sim = Mamba.Chains(n, 3, start = (burnin + 1), names = ["b0", "b1", "s2"])
theta = NUTSVariate([0.0, 0.0, 0.0], logfgrad)
for i in 1:n
  sample!(theta, adapt = (i <= burnin))
  if i > burnin
    sim[i, :, 1] = [theta[1:2]; exp(theta[3])]
  end
end

describe(sim)#-

Iterations = 1001:5000
Thinning interval = 1
Chains = 1
Samples per chain = 4000

Empirical Posterior Estimates:
       Mean          SD       Naive SE       MCSE         ESS
b0 113.55063076 2.044149868 0.0323208473 0.1141517020  320.67174
b1   0.91226476 0.044723043 0.0007071334 0.0024656788  328.99531
s2  26.02196258 1.954176725 0.0308982470 0.0506641103 1487.73929

Quantiles:
       2.5%        25.0%       50.0%       75.0%       97.5%
b0 109.19893751 112.3235613 113.558488 114.91439736 117.6074127
b1   0.82371558   0.8842933   0.910974   0.93910943   1.0080595
s2  22.50528704  24.6187215  25.903577  27.30734239  30.0766665