using DynamicHMCModels

ProjDir = rel_path_d("..", "scripts", "10")
cd(ProjDir)

d = CSV.read(rel_path("..", "data", "chimpanzees.csv"), delim=';');
df = convert(DataFrame, d);
df[:pulled_left] = convert(Array{Int64}, df[:pulled_left])
df[:prosoc_left] = convert(Array{Int64}, df[:prosoc_left])
df[:condition] = convert(Array{Int64}, df[:condition])
df[:actor] = convert(Array{Int64}, df[:actor])
first(df, 5)

struct m_10_04d_model{TY <: AbstractVector, TX <: AbstractMatrix,
  TA <: AbstractVector}
    "Observations."
    y::TY
    "Covariates"
    X::TX
    "Actors"
    A::TA
    "Number of observations"
    N::Int
    "Number of unique actors"
    N_actors::Int
end

function (problem::m_10_04d_model)(θ)
    @unpack y, X, A, N, N_actors = problem   # extract the data
    @unpack β, α = θ  # works on the named tuple too
    ll = 0.0
    ll += sum(logpdf.(Normal(0, 10), β)) # bp & bpC
    ll += sum(logpdf.(Normal(0, 10), α)) # alpha[1:7]
    ll += sum(
      [loglikelihood(Binomial(1, logistic(α[A[i]] + dot(X[i, :], β))), [y[i]]) for i in 1:N]
    )
    ll
end

N = size(df, 1)
N_actors = length(unique(df[:actor]))
X = hcat(ones(Int64, N), df[:prosoc_left] .* df[:condition]);
A = df[:actor]
y = df[:pulled_left]
p = m_10_04d_model(y, X, A, N, N_actors);
θ = (β = [1.0, 0.0], α = [-1.0, 10.0, -1.0, -1.0, -1.0, 0.0, 2.0])
p(θ)

-305.21943396408915

problem_transformation(p::m_10_04d_model) =
    as( (β = as(Array, size(p.X, 2)), α = as(Array, p.N_actors), ) )

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

LogDensityRejectErrors{InvalidLogDensityException,LogDensityProblems.ForwardDiffLogDensity{TransformedLogDensity{TransformVariables.TransformNamedTuple{(:β, :α),Tuple{TransformVariables.ArrayTransform{TransformVariables.Identity,1},TransformVariables.ArrayTransform{TransformVariables.Identity,1}}},Main.ex-m10.04d.m_10_04d_model{Array{Int64,1},Array{Int64,2},Array{Int64,1}}},ForwardDiff.GradientConfig{ForwardDiff.Tag{getfield(LogDensityProblems, Symbol("##1#2")){TransformedLogDensity{TransformVariables.TransformNamedTuple{(:β, :α),Tuple{TransformVariables.ArrayTransform{TransformVariables.Identity,1},TransformVariables.ArrayTransform{TransformVariables.Identity,1}}},Main.ex-m10.04d.m_10_04d_model{Array{Int64,1},Array{Int64,2},Array{Int64,1}}}},Float64},Float64,9,Array{ForwardDiff.Dual{ForwardDiff.Tag{getfield(LogDensityProblems, Symbol("##1#2")){TransformedLogDensity{TransformVariables.TransformNamedTuple{(:β, :α),Tuple{TransformVariables.ArrayTransform{TransformVariables.Identity,1},TransformVariables.ArrayTransform{TransformVariables.Identity,1}}},Main.ex-m10.04d.m_10_04d_model{Array{Int64,1},Array{Int64,2},Array{Int64,1}}}},Float64},Float64,9},1}}}}(ForwardDiff AD wrapper for TransformedLogDensity of dimension 9, w/ chunk size 9)

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

MCMC, adapting ϵ (75 steps)
0.018 s/step ...done
MCMC, adapting ϵ (25 steps)
0.019 s/step ...done
MCMC, adapting ϵ (50 steps)
0.014 s/step ...done
MCMC, adapting ϵ (100 steps)
0.0083 s/step ...done
MCMC, adapting ϵ (200 steps)
step 168 (of 200), 0.006 s/step
0.0058 s/step ...done
MCMC, adapting ϵ (400 steps)
step 184 (of 400), 0.0054 s/step
0.0046 s/step ...done
MCMC, adapting ϵ (50 steps)
0.0053 s/step ...done
MCMC (1000 steps)
step 217 (of 1000), 0.0046 s/step
step 440 (of 1000), 0.0046 s/step
step 656 (of 1000), 0.0046 s/step
step 865 (of 1000), 0.0046 s/step
0.0046 s/step ...done
(NUTS_Transition{Array{Float64,1},Float64}[NUTS_Transition{Array{Float64,1},Float64}([-2.874882811863496, 0.6214739516066847, 2.34082572493419, 9.123485524956058, 2.1327891296788453, 1.888195585865117, 2.7456550231756864, 3.6758876957435755, 5.298660470021998], -298.6905090897032, 4, DynamicHMC.DoubledTurn, 0.9557381412017036, 15), NUTS_Transition{Array{Float64,1},Float64}([-2.455582904138159, 0.09053681543025482, 1.829439651008072, 11.897647927063028, 1.6737512624812945, 2.0936407926274283, 1.8974085309485202, 3.3273398214636747, 3.761846887459364], -304.48484676817316, 3, DynamicHMC.DoubledTurn, 0.9975884386214116, 7), NUTS_Transition{Array{Float64,1},Float64}([2.2568428408022125, 0.023426680448468767, -2.888016780123113, 7.783505293567333, -3.0180308431884546, -2.520824705290234, -2.5033310202907213, -1.8035967350376747, 0.5438482852032088], -301.02483722227026, 4, DynamicHMC.DoubledTurn, 0.9955209645880136, 15), NUTS_Transition{Array{Float64,1},Float64}([0.16225685412424912, -0.05167969942053073, -0.6628302647831538, 7.5830738498951, -0.8890923454556815, -0.6562823375595088, -0.6613291012595293, 0.711173439775603, 2.8809078882615697], -300.9012514140964, 3, DynamicHMC.DoubledTurn, 0.9406106593891579, 7), NUTS_Transition{Array{Float64,1},Float64}([-1.490500653295514, -0.003240126764518847, 1.0807506694854405, 10.880406186431948, 0.7625677107098986, 1.5102780909648252, 0.62491975058959, 1.8824779065869022, 4.283775148881657], -304.5776163837921, 3, DynamicHMC.AdjacentTurn, 0.9124727061021657, 15), NUTS_Transition{Array{Float64,1},Float64}([-1.3507112706959497, 0.23807691359280514, 0.7544131476732987, 8.760683025243992, 0.4981848917470037, 1.1594505018988914, 0.8196886778188258, 1.5878884792786032, 4.146844210343427], -303.2125985053955, 3, DynamicHMC.DoubledTurn, 1.0, 7), NUTS_Transition{Array{Float64,1},Float64}([0.784036877039296, 0.5020369168147679, -1.2404190334273273, 6.5828781976962345, -1.6474166639508117, -2.2423129454532145, -1.5283071113573898, -0.3235644799589122, 0.6368916860280265], -302.40119989253253, 4, DynamicHMC.DoubledTurn, 0.8563850617058293, 15), NUTS_Transition{Array{Float64,1},Float64}([5.51351518559235, 0.3852562967114241, -5.966321399235768, 2.121203333698526, -6.051621127871399, -5.501437974497077, -5.705714901602562, -5.198531173070727, -2.8639309010702503], -300.3992198639924, 4, DynamicHMC.DoubledTurn, 0.9795783223165161, 15), NUTS_Transition{Array{Float64,1},Float64}([6.045158497120189, 0.731160246881031, -6.060660386263717, 7.37729569197586, -6.749900995334494, -6.914693018786695, -6.35437797042012, -5.15001972507434, -4.225201538526421], -302.20731534125633, 4, DynamicHMC.DoubledTurn, 0.9918233073991679, 15), NUTS_Transition{Array{Float64,1},Float64}([6.791970036292155, 0.4534573886458032, -7.247710290339399, 1.3934067749163956, -7.574113968013276, -7.533954049053628, -7.573409316484046, -6.062612984616853, -4.880097462656821], -300.2386748271331, 3, DynamicHMC.DoubledTurn, 0.9801883795734315, 7)  …  NUTS_Transition{Array{Float64,1},Float64}([-3.861533525896566, 0.29736200627113835, 2.822262453259548, 10.445244911419893, 3.3454012944302383, 3.1429695965019464, 3.3919071860792003, 4.075206764167646, 5.746342402326152], -299.3217086399044, 4, DynamicHMC.DoubledTurn, 0.9939263418944215, 15), NUTS_Transition{Array{Float64,1},Float64}([-3.054239338740829, 0.6195194269199636, 2.4526608812363704, 11.755175510961033, 2.693290833966343, 2.535577332388025, 2.695010013693976, 3.650387635704386, 5.219039950143055], -298.99371321489565, 3, DynamicHMC.DoubledTurn, 0.944866025215164, 7), NUTS_Transition{Array{Float64,1},Float64}([5.248303503603661, 0.3171775091455515, -5.435706151257682, 11.175104443825573, -6.176800980808534, -6.073643777123454, -5.615023978352112, -4.661188692672812, -3.2696575002077815], -298.29976959950045, 4, DynamicHMC.AdjacentTurn, 0.9898456564243621, 23), NUTS_Transition{Array{Float64,1},Float64}([2.6118545179084016, 0.18146892522759447, -2.7263242598501956, 9.641871614647643, -3.259441151077631, -3.320340976187703, -3.2282957905886054, -1.9107401016857135, -0.4603858188158778], -298.6754725600163, 4, DynamicHMC.DoubledTurn, 0.5275034251169985, 15), NUTS_Transition{Array{Float64,1},Float64}([-1.9698492513110681, 0.6513465243239012, 1.3469506222792131, 6.793749579911243, 1.152762986445535, 1.1311113581156136, 1.752013866020127, 2.6364578314635283, 3.0607142371036957], -299.42477329901186, 4, DynamicHMC.DoubledTurn, 0.9625343429220289, 15), NUTS_Transition{Array{Float64,1},Float64}([-2.3643494398541347, 0.6337840955683325, 1.840568110477678, 6.210386756350045, 1.5206539714690515, 1.587627491211736, 2.125208485507525, 3.009585829568913, 3.6916835107682426], -298.9941951155513, 3, DynamicHMC.DoubledTurn, 0.9419206418971463, 7), NUTS_Transition{Array{Float64,1},Float64}([2.5561860686070075, 0.579144093509701, -3.242267743736925, 16.83037501221078, -3.7071320470757065, -3.053607300816943, -3.1231400216020386, -2.107027705745131, -0.13378340470260497], -302.15445421947777, 3, DynamicHMC.DoubledTurn, 0.3200614682193005, 7), NUTS_Transition{Array{Float64,1},Float64}([1.9823228679677432, 1.0454454239768927, -2.284691543181772, 17.661800375220693, -3.22312638991393, -3.197692778628491, -2.5669623898878173, -1.6315468132925521, 0.7932899905953908], -302.10798791566214, 3, DynamicHMC.DoubledTurn, 0.926446940704161, 7), NUTS_Transition{Array{Float64,1},Float64}([-2.2829245576710955, 0.6955077870678723, 1.9007473810350672, 16.373661286268646, 1.1462843732382504, 1.4526171708977325, 1.6433698839898667, 2.451717867014099, 5.286294156755354], -304.8103727122961, 3, DynamicHMC.AdjacentTurn, 0.992345538933759, 15), NUTS_Transition{Array{Float64,1},Float64}([-1.2119361219495506, 0.8429988598601122, 1.1714090861268793, 18.688371174186397, 0.26933925357454463, 0.09084550772772104, 0.3970181308211943, 1.1784520145423119, 4.045525027524578], -305.9574438510582, 3, DynamicHMC.AdjacentTurn, 0.9263099155916021, 15)], NUTS sampler in 9 dimensions
  stepsize (ϵ) ≈ 0.198
  maximum depth = 10
  Gaussian kinetic energy, √diag(M⁻¹): [3.402084002860694, 0.4540242132760563, 3.4170271910830716, 6.946458286449433, 3.4251448791547463, 3.4270290016097413, 3.418389293031463, 3.412061225222789, 3.430718857963224]
)

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

5-element Array{NamedTuple{(:β, :α),Tuple{Array{Float64,1},Array{Float64,1}}},1}:
 (β = [-2.874882811863496, 0.6214739516066847], α = [2.34082572493419, 9.123485524956058, 2.1327891296788453, 1.888195585865117, 2.7456550231756864, 3.6758876957435755, 5.298660470021998])
 (β = [-2.455582904138159, 0.09053681543025482], α = [1.829439651008072, 11.897647927063028, 1.6737512624812945, 2.0936407926274283, 1.8974085309485202, 3.3273398214636747, 3.761846887459364])
 (β = [2.2568428408022125, 0.023426680448468767], α = [-2.888016780123113, 7.783505293567333, -3.0180308431884546, -2.520824705290234, -2.5033310202907213, -1.8035967350376747, 0.5438482852032088])
 (β = [0.16225685412424912, -0.05167969942053073], α = [-0.6628302647831538, 7.5830738498951, -0.8890923454556815, -0.6562823375595088, -0.6613291012595293, 0.711173439775603, 2.8809078882615697])
 (β = [-1.490500653295514, -0.003240126764518847], α = [1.0807506694854405, 10.880406186431948, 0.7625677107098986, 1.5102780909648252, 0.62491975058959, 1.8824779065869022, 4.283775148881657])

posterior_β = mean(first, posterior)

2-element Array{Float64,1}:
 1.144203183321902
 0.42431412961884746

posterior_α = mean(last, posterior)

7-element Array{Float64,1}:
 -1.5920086375579863
 10.324009010610762
 -1.8901209779683208
 -1.8883570624345418
 -1.597935116728296
 -0.6691088734633035
  0.8941455781850367

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

1×9 Array{Float64,2}:
 609.907  1000.0  602.623  444.439  …  605.916  607.51  608.433  618.404

NUTS_statistics(chain)

Hamiltonian Monte Carlo sample of length 1000
  acceptance rate mean: 0.92, min/25%/median/75%/max: 0.0 0.9 0.96 0.99 1.0
  termination: AdjacentDivergent => 1% AdjacentTurn => 18% DoubledTurn => 81%
  depth: 1 => 0% 2 => 1% 3 => 39% 4 => 59% 5 => 0%

rethinking = "
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 -0.74503184 0.26613979 0.0042080396 0.0060183398 1000
a.2 10.77955494 5.32538998 0.0842018089 0.1269148045 1000
a.3 -1.04982353 0.28535997 0.0045119373 0.0049074219 1000
a.4 -1.04898135 0.28129307 0.0044476339 0.0056325117 1000
a.5 -0.74390933 0.26949936 0.0042611590 0.0052178124 1000
a.6  0.21599365 0.26307574 0.0041595927 0.0045153523 1000
a.7  1.81090866 0.39318577 0.0062168129 0.0071483527 1000
 bp  0.83979926 0.26284676 0.0041559722 0.0059795826 1000
bpC -0.12913322 0.29935741 0.0047332562 0.0049519863 1000
";

"\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\na.1 -0.74503184 0.26613979 0.0042080396 0.0060183398 1000\na.2 10.77955494 5.32538998 0.0842018089 0.1269148045 1000\na.3 -1.04982353 0.28535997 0.0045119373 0.0049074219 1000\na.4 -1.04898135 0.28129307 0.0044476339 0.0056325117 1000\na.5 -0.74390933 0.26949936 0.0042611590 0.0052178124 1000\na.6  0.21599365 0.26307574 0.0041595927 0.0045153523 1000\na.7  1.81090866 0.39318577 0.0062168129 0.0071483527 1000\n bp  0.83979926 0.26284676 0.0041559722 0.0059795826 1000\nbpC -0.12913322 0.29935741 0.0047332562 0.0049519863 1000\n"

[posterior_β, posterior_α]

2-element Array{Array{Float64,1},1}:
 [1.144203183321902, 0.42431412961884746]
 [-1.5920086375579863, 10.324009010610762, -1.8901209779683208, -1.8883570624345418, -1.597935116728296, -0.6691088734633035, 0.8941455781850367]