using StatisticalRethinking
using CmdStan, StanMCMCChain
gr(size=(500,800))

Plots.GRBackend()

ProjDir = rel_path("..", "chapters", "04")
cd(ProjDir)

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

df2 = filter(row -> row[:age] >= 18, df)
mean_weight = mean(df2[:weight])
df2[:weight] = convert(Vector{Float64}, df2[:weight]) .- mean_weight

352-element Array{Float64,1}:
   2.8351209801136577
  -8.50467901988634
 -13.125647519886343
   8.051428980113656
  -3.7136135198863442
  18.00210348011366
  -6.747010019886339
  10.489485980113656
 -10.120600519886345
   9.497253480113656
   ⋮
   2.8918199801136595
  -5.584680519886341
  -3.9404095198863445
  -4.167205519886345
   2.0413349801136604
 -10.744289519886344
   7.172594480113659
   9.07201098011366
   7.54113798011366

weightsmodel = "
data {
 int < lower = 1 > N; // Sample size
 vector[N] height; // Predictor
 vector[N] weight; // Outcome
}

parameters {
 real alpha; // Intercept
 real beta; // Slope (regression coefficients)
 real < lower = 0 > sigma; // Error SD
}

model {
 height ~ normal(alpha + weight * beta , sigma);
}

generated quantities {
}
"

"\ndata {\n int < lower = 1 > N; // Sample size\n vector[N] height; // Predictor\n vector[N] weight; // Outcome\n}\n\nparameters {\n real alpha; // Intercept\n real beta; // Slope (regression coefficients)\n real < lower = 0 > sigma; // Error SD\n}\n\nmodel {\n height ~ normal(alpha + weight * beta , sigma);\n}\n\ngenerated quantities {\n}\n"

stanmodel = Stanmodel(name="weights", monitors = ["alpha", "beta", "sigma"],model=weightsmodel,
  output_format=:mcmcchain)

heightsdata = [
  Dict("N" => length(df2[:height]), "height" => df2[:height], "weight" => df2[:weight])
]

=====> /home/travis/build/StanJulia/StatisticalRethinking.jl/docs/build/04

1-element Array{Dict{String,Any},1}:
 Dict("height"=>Union{Missing, Float64}[151.765, 139.7, 136.525, 156.845, 145.415, 163.83, 149.225, 168.91, 147.955, 165.1  …  156.21, 160.655, 146.05, 156.21, 152.4, 162.56, 142.875, 162.56, 156.21, 158.75],"weight"=>[2.83512, -8.50468, -13.1256, 8.05143, -3.71361, 18.0021, -6.74701, 10.4895, -10.1206, 9.49725  …  -0.963712, 2.89182, -5.58468, -3.94041, -4.16721, 2.04133, -10.7443, 7.17259, 9.07201, 7.54114],"N"=>352)

rc, chn, cnames = stan(stanmodel, heightsdata, ProjDir, diagnostics=false,
  CmdStanDir=CMDSTAN_HOME)

describe(chn)

make: `/home/travis/build/StanJulia/StatisticalRethinking.jl/docs/build/04/tmp/weights' is up to date.

Length of data array is not equal to nchains,
all chains will use the first data dictionary.

Calling /home/travis/cmdstan/bin/stansummary to infer across chains.

Inference for Stan model: weights_model
4 chains: each with iter=(1000,1000,1000,1000); warmup=(0,0,0,0); thin=(1,1,1,1); 4000 iterations saved.

Warmup took (0.24, 0.17, 0.17, 0.22) seconds, 0.80 seconds total
Sampling took (0.17, 0.19, 0.21, 0.12) seconds, 0.69 seconds total

                Mean     MCSE  StdDev    5%   50%   95%    N_Eff  N_Eff/s    R_hat
lp__            -747  2.7e-02     1.2  -750  -747  -746  2.0e+03  2.9e+03  1.0e+00
accept_stat__   0.90  1.7e-03    0.11  0.67  0.94   1.0  4.8e+03  6.9e+03  1.0e+00
stepsize__      0.83  3.6e-02   0.051  0.75  0.86  0.89  2.0e+00  2.9e+00  4.9e+13
treedepth__      2.2  1.0e-01    0.72   1.0   2.0   4.0  4.9e+01  7.1e+01  1.0e+00
n_leapfrog__     6.0  6.9e-01     7.4   3.0   3.0    19  1.1e+02  1.6e+02  1.0e+00
divergent__     0.00     -nan    0.00  0.00  0.00  0.00     -nan     -nan     -nan
energy__         749  4.2e-02     1.7   747   749   752  1.6e+03  2.4e+03  1.0e+00
alpha            155  4.3e-03    0.27   154   155   155  4.1e+03  5.9e+03  1.0e+00
beta            0.91  6.7e-04   0.042  0.84  0.91  0.97  3.9e+03  5.6e+03  1.0e+00
sigma            5.1  2.8e-03    0.19   4.8   5.1   5.4  4.7e+03  6.8e+03  1.0e+00

Samples were drawn using hmc with nuts.
For each parameter, N_Eff is a crude measure of effective sample size,
and R_hat is the potential scale reduction factor on split chains (at
convergence, R_hat=1).

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

Empirical Posterior Estimates:
          Mean         SD       Naive SE        MCSE      ESS
alpha 154.6005795 0.274394689 0.0043385610 0.00427689432 1000
 beta   0.9055396 0.041703493 0.0006593901 0.00062124134 1000
sigma   5.1012100 0.194238569 0.0030711814 0.00289619716 1000

Quantiles:
          2.5%        25.0%       50.0%       75.0%      97.5%
alpha 154.08797500 154.4070000 154.6000000 154.792000 155.1360250
 beta   0.82597698   0.8770465   0.9053655   0.934285   0.9868225
sigma   4.74118650   4.9652575   5.0965000   5.232095   5.4926210

clip_38s_example_output = "

Samples were drawn using hmc with nuts.
For each parameter, N_Eff is a crude measure of effective sample size,
and R_hat is the potential scale reduction factor on split chains (at
convergence, R_hat=1).

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

Empirical Posterior Estimates:
          Mean         SD       Naive SE       MCSE     ESS
alpha 113.82267275 1.89871177 0.0300212691 0.053895503 1000
 beta   0.90629952 0.04155225 0.0006569987 0.001184630 1000
sigma   5.10334279 0.19755211 0.0031235731 0.004830464 1000

Quantiles:
          2.5%       25.0%       50.0%       75.0%       97.5%
alpha 110.1927000 112.4910000 113.7905000 115.1322500 117.5689750
 beta   0.8257932   0.8775302   0.9069425   0.9357115   0.9862574
sigma   4.7308260   4.9644050   5.0958800   5.2331875   5.5133417

"

"\n\nSamples were drawn using hmc with nuts.\nFor each parameter, N_Eff is a crude measure of effective sample size,\nand R_hat is the potential scale reduction factor on split chains (at\nconvergence, R_hat=1).\n\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\nalpha 113.82267275 1.89871177 0.0300212691 0.053895503 1000\n beta   0.90629952 0.04155225 0.0006569987 0.001184630 1000\nsigma   5.10334279 0.19755211 0.0031235731 0.004830464 1000\n\nQuantiles:\n          2.5%       25.0%       50.0%       75.0%       97.5%\nalpha 110.1927000 112.4910000 113.7905000 115.1322500 117.5689750\n beta   0.8257932   0.8775302   0.9069425   0.9357115   0.9862574\nsigma   4.7308260   4.9644050   5.0958800   5.2331875   5.5133417\n\n"

plot(chn)

xi = -16.0:0.1:18.0
rws, vars, chns = size(chn[:, 1, :])
alpha_vals = convert(Vector{Float64}, reshape(chn.value[:, 1, :], (rws*chns)))
beta_vals = convert(Vector{Float64}, reshape(chn.value[:, 2, :], (rws*chns)))
yi = mean(alpha_vals) .+ mean(beta_vals)*xi

scatter(df2[:weight], df2[:height], lab="Observations")
plot!(xi, yi, lab="Regression line")