using DynamicHMCModels

ProjDir = rel_path_d("..", "scripts", "12")

df = CSV.read(rel_path( "..", "data",  "Kline.csv"), delim=';');
size(df) # Should be 10x5

(10, 5)

df[:logpop] = map((x) -> log(x), df[:population]);
df[:society] = 1:10;

first(df[[:total_tools, :logpop, :society]], 5)

struct m_12_06d_model{TY <: AbstractVector, TX <: AbstractMatrix,
  TS <: AbstractVector}
    "Observations (total_tools)."
    y::TY
    "Covariates (logpop)"
    X::TX
    "Society"
    S::TS
    "Number of observations (10)"
    N::Int
    "Number of societies (also 10)"
    N_societies::Int
end

function (problem::m_12_06d_model)(θ)
    @unpack y, X, S, N, N_societies = problem   # extract the data
    @unpack β, α, σ = θ  # β : a, bp, α : a_society
    ll = 0.0
    ll += logpdf(Cauchy(0, 1), σ)
    ll += sum(logpdf.(Normal(0, σ), α)) # α[1:10]
    ll += logpdf.(Normal(0, 10), β[1]) # a
    ll += logpdf.(Normal(0, 1), β[2]) # a
    ll += sum(
      [loglikelihood(Poisson(exp(α[S[i]] + dot(X[i, :], β))), [y[i]]) for i in 1:N]
    )
    ll
end

N = size(df, 1)
N_societies = length(unique(df[:society]))
X = hcat(ones(Int64, N), df[:logpop]);
S = df[:society]
y = df[:total_tools]
p = m_12_06d_model(y, X, S, N, N_societies);
θ = (β = [1.0, 0.25], α = rand(Normal(0, 1), N_societies), σ = 0.2)
p(θ)

-372.9970164178759

problem_transformation(p::m_12_06d_model) =
    as( (β = as(Array, size(p.X, 2)), α = as(Array, p.N_societies), σ = asℝ₊) )

P = TransformedLogDensity(problem_transformation(p), p)
∇P = LogDensityRejectErrors(ADgradient(:ForwardDiff, P));
#∇P = ADgradient(:ForwardDiff, P);

chain, NUTS_tuned = NUTS_init_tune_mcmc(∇P, 1000);

MCMC, adapting ϵ (75 steps)
0.0026 s/step ...done
MCMC, adapting ϵ (25 steps)
0.0027 s/step ...done
MCMC, adapting ϵ (50 steps)
0.0034 s/step ...done
MCMC, adapting ϵ (100 steps)
0.0014 s/step ...done
MCMC, adapting ϵ (200 steps)
0.00081 s/step ...done
MCMC, adapting ϵ (400 steps)
0.00055 s/step ...done
MCMC, adapting ϵ (50 steps)
0.00069 s/step ...done
MCMC (1000 steps)
0.00042 s/step ...done
(NUTS_Transition{Array{Float64,1},Float64}[NUTS_Transition{Array{Float64,1},Float64}([1.6659063217846828, 0.20547983964637204, 0.033239179099069656, -0.03653906514092166, -0.05196698908115623, 0.19088301320675644, 0.006148216447083141, -0.1342597360464994, 0.12387898235993758, -0.25745579831012444, 0.09008373285696991, -0.09997949245727245, -1.6196046622162208], -42.15524149399916, 4, DynamicHMC.DoubledTurn, 1.0, 15), NUTS_Transition{Array{Float64,1},Float64}([1.0698500695054172, 0.27499627274347693, -0.20789292481460475, 0.0984071774089927, 0.057107875588127105, 0.35466563114812105, -0.029642605520146748, -0.5449981591287166, -0.039917216922391995, -0.4573469466752317, 0.27579373768626353, -0.4061565057130561, -1.055594208938716], -42.49114008600646, 4, DynamicHMC.DoubledTurn, 0.9550078799405995, 15), NUTS_Transition{Array{Float64,1},Float64}([1.0866179871609887, 0.24862127392970132, -0.23565853734952755, -0.06183538123869904, 0.15801133506309462, 0.30743635698287014, -0.007325892036212034, 0.08213680778654427, 0.34903122326169195, -0.031752390644454356, 0.520871970087058, 0.19042293663975354, -1.3979432867248232], -45.50434776625452, 4, DynamicHMC.DoubledTurn, 0.9899263310496017, 15), NUTS_Transition{Array{Float64,1},Float64}([1.12459894786011, 0.27649508880077855, -0.26467586616932565, -0.1241535249136635, -0.13848040114080765, 0.15993196968990922, 0.23585846000110314, -0.28320832766055065, 0.029393524641839495, -0.3352853074390302, -0.04032193441532315, -0.3187576284619706, -1.611707264687893], -50.26096046150815, 4, DynamicHMC.DoubledTurn, 0.9843472057553645, 15), NUTS_Transition{Array{Float64,1},Float64}([1.5016485694926036, 0.24846668880749706, -0.4995454897409589, -0.1642305009049892, -0.33767161029178083, 0.04078305508738424, 0.1990423086021051, -0.29935451628282783, -0.09489971063671818, -0.3225533875104829, -0.15308781298672233, -0.30796930298082986, -1.9183853213472106], -53.28430996233725, 3, DynamicHMC.DoubledTurn, 0.6703403927564862, 7), NUTS_Transition{Array{Float64,1},Float64}([0.6550896056450812, 0.3082462182225012, -0.18113907311544591, 0.041846609978144926, 0.16168155993623645, 0.4247890456139156, 0.175320386341817, -0.35861776406579915, 0.18602044976176357, -0.26446536229418777, 0.050037857903707364, -0.23776270456342286, -0.7494928541080879], -54.460468491709996, 3, DynamicHMC.DoubledTurn, 1.0, 7), NUTS_Transition{Array{Float64,1},Float64}([0.9412589743809462, 0.2746001603524749, -0.17694273478104552, 0.04118234986565986, 0.122333904025743, 0.4710015762943421, 0.1972755595757228, -0.4813913501386997, 0.2874948900363817, -0.27488866998914646, 0.07607014518072938, -0.14733393848656, -0.6515454949628215], -42.29101879679805, 4, DynamicHMC.DoubledTurn, 0.9702483587419126, 15), NUTS_Transition{Array{Float64,1},Float64}([1.3060232504537899, 0.2412044102841871, 0.04426043396609602, 0.16019180329102228, 0.12603641543345764, 0.5120584885494569, 0.027717912347812075, -0.4925128258791912, 0.3054162057890045, -0.11870675362402999, 0.25752402408642566, 0.06327462970993226, -1.3836263147859489], -44.32653906228223, 4, DynamicHMC.DoubledTurn, 0.8812909777233838, 15), NUTS_Transition{Array{Float64,1},Float64}([0.3433842619465165, 0.32259088348259884, 0.027180913275021305, 0.24751681371055606, 0.18723916490935738, 0.6339162050089239, 0.2791358543116875, -0.42212017369804816, 0.18822098561814754, -0.20394685441529367, 0.2492015158748102, -0.17811683703027273, -1.248165927672158], -43.47861112141645, 4, DynamicHMC.DoubledTurn, 0.8681566506780692, 15), NUTS_Transition{Array{Float64,1},Float64}([-0.09187615730039934, 0.3947483205020631, 0.10781707359872769, 0.17832662730647178, 0.19562475234742957, 0.7140374744268541, 0.0998054498113936, -0.4217741067893409, 0.29207680610409775, -0.2681833149373481, 0.2963853623697165, -0.39366346149338693, -1.0385197254771574], -43.94754844595677, 4, DynamicHMC.DoubledTurn, 0.9305172527326503, 15)  …  NUTS_Transition{Array{Float64,1},Float64}([0.7806791025779878, 0.3171765829829895, -0.1427960227723332, 0.2692321859493546, -0.2811185616985394, 0.1768589109265467, -0.002098781427978311, -0.7213796288347022, -0.11061914398547144, -0.29409325784861334, 0.09829404404437159, -0.3908837901525087, -0.8463402768005482], -43.449182798161324, 4, DynamicHMC.DoubledTurn, 0.9932872336871893, 15), NUTS_Transition{Array{Float64,1},Float64}([0.43232425767894084, 0.3402106097596222, 0.0517115626709098, 0.2856286196146034, -0.19271335974201487, 0.3129106594829696, 0.03679364323936983, -0.6579700267988844, -0.032480559090427526, -0.15389796933255606, 0.08829124456553006, -0.4194976227670935, -0.6570366048291864], -42.76581727579139, 4, DynamicHMC.DoubledTurn, 1.0, 15), NUTS_Transition{Array{Float64,1},Float64}([0.7361513454716555, 0.2943783013460889, -0.16867023371606565, 0.03263018613134995, 0.18273229640331057, 0.20552298491961438, 0.039940051952336346, -0.30924153502346635, 0.2426693178700715, -0.18840276152006014, 0.0687132063364661, 0.00954637047276706, -2.0669216809951108], -47.7310062175095, 3, DynamicHMC.AdjacentTurn, 0.9878168715619594, 15), NUTS_Transition{Array{Float64,1},Float64}([1.432041478194573, 0.21976606568162152, 0.17625041273697267, 0.02863181351535702, -0.16479024194068198, 0.5555570683689169, 0.09653001098223836, -0.06674394571371683, 0.0974529494258842, -0.06744438569737385, 0.5124071991846443, -0.050075053780351125, -0.7609031327515956], -45.362019255261536, 4, DynamicHMC.DoubledTurn, 0.904617755987145, 15), NUTS_Transition{Array{Float64,1},Float64}([0.9127423996655671, 0.28551344851171145, -0.37424640882345944, 0.02719323902903284, -0.2091149289352628, 0.07063965673480574, 0.10424360303553973, -0.4537537021126407, 0.22459527643701252, -0.09949350648495012, 0.1748622599487036, -0.21020608955582795, -1.6552899798844518], -46.02668189135178, 4, DynamicHMC.DoubledTurn, 0.945034865118797, 15), NUTS_Transition{Array{Float64,1},Float64}([2.6710247408608283, 0.07648273766469361, -0.400493445883669, -0.0402570450347318, -0.2553621020744352, 0.6693571673309716, -0.055510606436739315, -0.5707395439864997, 0.24755420944367876, 0.11515659866602618, 0.3558130509554679, 0.500413954505436, -0.39658492401355766], -49.36008312773766, 4, DynamicHMC.DoubledTurn, 0.8253786126212656, 15), NUTS_Transition{Array{Float64,1},Float64}([-0.04250260825909888, 0.3670685299243436, -0.057591839558898454, 0.3658093773108945, 0.6096593578816878, 0.27554700908504226, 0.18234997078379805, -0.05833951812303252, 0.25999588758247405, -0.17317807792876416, 0.5208524299006505, -0.16599215987577043, -1.2553095812537574], -52.010659295912006, 4, DynamicHMC.DoubledTurn, 0.9974991850313892, 15), NUTS_Transition{Array{Float64,1},Float64}([2.0799239895588033, 0.1782260968859048, -0.30542950062311597, -0.2093963888137574, -0.7168263571761891, 0.3312055098211429, -0.08634440919780198, -0.6227375832476112, -0.0355761969186952, -0.15647378572351694, 0.04583333250124953, -0.10705086115788312, -1.3456369543630624], -47.211367958992234, 4, DynamicHMC.DoubledTurn, 0.9975694302618962, 15), NUTS_Transition{Array{Float64,1},Float64}([0.22133663014687768, 0.3516593692347825, -0.16676461597041609, 0.4488569347587557, 0.4214502942871093, 0.061241665373066044, 0.05224698853011892, -0.3855734153426507, 0.14993135703366745, -0.5137586487882556, 0.18713547578128145, -0.21755271087977307, -0.73681503060344], -52.72811149264219, 4, DynamicHMC.DoubledTurn, 0.9305905084791871, 15), NUTS_Transition{Array{Float64,1},Float64}([0.13039050865375795, 0.37074640776280393, -0.02289766883560957, 0.3834832014862093, 0.2837410393811572, 0.23635899887369516, -0.10726095289978116, -0.36213521261656306, 0.11467443394591886, -0.5029342107736857, 0.2802343472703936, -0.30499684667137716, -0.5769554781802271], -50.72379400838408, 4, DynamicHMC.DoubledTurn, 0.9451002669597475, 15)], NUTS sampler in 13 dimensions
  stepsize (ϵ) ≈ 0.232
  maximum depth = 10
  Gaussian kinetic energy, √diag(M⁻¹): [0.6947788475520215, 0.07818368451289717, 0.23624468429731635, 0.2035309099237096, 0.17239490256897558, 0.19173419119563512, 0.18055418404807888, 0.18649951347375304, 0.18493174758228506, 0.16272035384413086, 0.1515353968004349, 0.27234985151194924, 0.41638260114702064]
)

posterior = TransformVariables.transform.(Ref(problem_transformation(p)), get_position.(chain));
posterior[1:5]

5-element Array{NamedTuple{(:β, :α, :σ),Tuple{Array{Float64,1},Array{Float64,1},Float64}},1}:
 (β = [1.6659063217846828, 0.20547983964637204], α = [0.033239179099069656, -0.03653906514092166, -0.05196698908115623, 0.19088301320675644, 0.006148216447083141, -0.1342597360464994, 0.12387898235993758, -0.25745579831012444, 0.09008373285696991, -0.09997949245727245], σ = 0.19797695138374946)
 (β = [1.0698500695054172, 0.27499627274347693], α = [-0.20789292481460475, 0.0984071774089927, 0.057107875588127105, 0.35466563114812105, -0.029642605520146748, -0.5449981591287166, -0.039917216922391995, -0.4573469466752317, 0.27579373768626353, -0.4061565057130561], σ = 0.3479855897119586)
 (β = [1.0866179871609887, 0.24862127392970132], α = [-0.23565853734952755, -0.06183538123869904, 0.15801133506309462, 0.30743635698287014, -0.007325892036212034, 0.08213680778654427, 0.34903122326169195, -0.031752390644454356, 0.520871970087058, 0.19042293663975354], σ = 0.24710466510986218)
 (β = [1.12459894786011, 0.27649508880077855], α = [-0.26467586616932565, -0.1241535249136635, -0.13848040114080765, 0.15993196968990922, 0.23585846000110314, -0.28320832766055065, 0.029393524641839495, -0.3352853074390302, -0.04032193441532315, -0.3187576284619706], σ = 0.1995466441558585)
 (β = [1.5016485694926036, 0.24846668880749706], α = [-0.4995454897409589, -0.1642305009049892, -0.33767161029178083, 0.04078305508738424, 0.1990423086021051, -0.29935451628282783, -0.09489971063671818, -0.3225533875104829, -0.15308781298672233, -0.30796930298082986], σ = 0.14684387648125144)

posterior_β = mean(posterior[i].β for i in 1:length(posterior))
posterior_α = mean(posterior[i].α for i in 1:length(posterior))
posterior_σ = mean(posterior[i].σ for i in 1:length(posterior))

0.312727459700007

ess = mapslices(effective_sample_size, get_position_matrix(chain); dims = 1)
ess

1×13 Array{Float64,2}:
 783.007  879.438  797.294  867.788  …  662.52  772.524  1000.0  604.106

NUTS_statistics(chain)

Hamiltonian Monte Carlo sample of length 1000
  acceptance rate mean: 0.94, min/25%/median/75%/max: 0.4 0.91 0.97 0.99 1.0
  termination: AdjacentTurn => 6% DoubledTurn => 94%
  depth: 3 => 4% 4 => 95% 5 => 1%

m_12_6_result = "
Iterations = 1:1000
Thinning interval = 1
Chains = 1,2,3,4
Samples per chain = 1000

Empirical Posterior Estimates:
                            Mean                SD               Naive SE             MCSE            ESS
            a          1.076167468  0.7704872560 0.01218247319 0.0210530022 1000.000000
           bp         0.263056273  0.0823415805 0.00130193470 0.0022645077 1000.000000
  a_society.1   -0.191723568  0.2421382537 0.00382854195 0.0060563054 1000.000000
  a_society.2    0.054569029  0.2278506876 0.00360263570 0.0051693148 1000.000000
  a_society.3   -0.035935050  0.1926364647 0.00304584994 0.0039948433 1000.000000
  a_society.4    0.334355037  0.1929971201 0.00305155241 0.0063871707  913.029080
  a_society.5    0.049747513  0.1801287716 0.00284808595 0.0043631095 1000.000000
  a_society.6   -0.311903245  0.2096126337 0.00331426674 0.0053000536 1000.000000
  a_society.7    0.148637507  0.1744680594 0.00275858223 0.0047660246 1000.000000
  a_society.8   -0.164567976  0.1821341074 0.00287979309 0.0034297298 1000.000000
  a_society.9    0.277066965  0.1758237250 0.00278001719 0.0055844175  991.286501
 a_society.10   -0.094149204  0.2846206232 0.00450024719 0.0080735022 1000.000000
sigma_society    0.310352849  0.1374834682 0.00217380450 0.0057325226  575.187461
";

"\nIterations = 1:1000\nThinning interval = 1\nChains = 1,2,3,4\nSamples per chain = 1000\n\nEmpirical Posterior Estimates:\n                            Mean                SD               Naive SE             MCSE            ESS\n            a          1.076167468  0.7704872560 0.01218247319 0.0210530022 1000.000000\n           bp         0.263056273  0.0823415805 0.00130193470 0.0022645077 1000.000000\n  a_society.1   -0.191723568  0.2421382537 0.00382854195 0.0060563054 1000.000000\n  a_society.2    0.054569029  0.2278506876 0.00360263570 0.0051693148 1000.000000\n  a_society.3   -0.035935050  0.1926364647 0.00304584994 0.0039948433 1000.000000\n  a_society.4    0.334355037  0.1929971201 0.00305155241 0.0063871707  913.029080\n  a_society.5    0.049747513  0.1801287716 0.00284808595 0.0043631095 1000.000000\n  a_society.6   -0.311903245  0.2096126337 0.00331426674 0.0053000536 1000.000000\n  a_society.7    0.148637507  0.1744680594 0.00275858223 0.0047660246 1000.000000\n  a_society.8   -0.164567976  0.1821341074 0.00287979309 0.0034297298 1000.000000\n  a_society.9    0.277066965  0.1758237250 0.00278001719 0.0055844175  991.286501\n a_society.10   -0.094149204  0.2846206232 0.00450024719 0.0080735022 1000.000000\nsigma_society    0.310352849  0.1374834682 0.00217380450 0.0057325226  575.187461\n"

[posterior_β, posterior_α, [posterior_σ]]

3-element Array{Array{Float64,1},1}:
 [1.0941196362899368, 0.2620867405557856]
 [-0.1978720976177207, 0.03417371585057167, -0.05023649623231359, 0.33273038499792107, 0.048244074699847685, -0.32028075664581157, 0.1376474346638074, -0.1825132938847455, 0.27500196776994396, -0.09796467099368024]
 [0.312727459700007]