using StatisticalRethinking
using Turing

Turing.setadbackend(:reverse_diff)

d = CSV.read(joinpath(dirname(Base.pathof(StatisticalRethinking)), "..", "data",
    "reedfrogs.csv"), delim=';')
size(d) # Should be 48x5

# Set number of tanks
d[:tank] = 1:size(d,1)

@model m12_1(density, tank, surv) = begin

    # Number of unique tanks in the data set
    N_tank = length(tank)

    # Set an TArray for the priors/param
    a_tank = Vector{Real}(undef, N_tank)

    # For each tank [1,..,48] set prior N(0,5)
    a_tank ~ [Normal(0,5)]

    logitp = [a_tank[tank[i]] for i = 1:N_tank]
    surv ~ VecBinomialLogit(density, logitp)
end

posterior = sample(m12_1(Vector{Int64}(d[:density]), Vector{Int64}(d[:tank]),
    Vector{Int64}(d[:surv])), Turing.NUTS(4000, 1000, 0.8))
describe(posterior)
#                 Mean           SD         Naive SE       MCSE         ESS
# a_tank[33]    3.7925121923  1.069473492 0.01690986065 0.0181583984 3468.84559
#  a_tank[2]    5.6619233751  2.661975386 0.04208952648 0.0984382340  731.27452
# a_tank[44]   -0.4076755205  0.349650208 0.00552845521 0.0051973560 4000.00000
# a_tank[40]    2.4882548674  0.642093720 0.01015239313 0.0087774011 4000.00000
# a_tank[41]   -2.1368737936  0.588205639 0.00930034776 0.0091150511 4000.00000
# a_tank[42]   -0.6773019465  0.370413780 0.00585675610 0.0052435688 4000.00000
# a_tank[11]    0.9259845291  0.729448255 0.01153358961 0.0077243130 4000.00000
# a_tank[16]    2.5669573506  1.256913239 0.01987354329 0.0240073491 2741.08301
# a_tank[46]   -0.6654705830  0.370650246 0.00586049496 0.0058669658 3991.18148
#  a_tank[8]    2.5385573941  1.151976429 0.01821434663 0.0232427286 2456.47849
# a_tank[14]    0.0006786356  0.680697609 0.01076277421 0.0069893413 4000.00000
# a_tank[47]    2.1388988496  0.532096615 0.00841318620 0.0078518462 4000.00000
# a_tank[43]   -0.5364419223  0.354074917 0.00559841600 0.0058512095 3661.83713
# a_tank[24]    1.7486397297  0.597014797 0.00943963278 0.0071524317 4000.00000
# a_tank[34]    2.9828401383  0.793129631 0.01254048056 0.0113862941 4000.00000
#  a_tank[9]   -0.4362986938  0.655848632 0.01036987739 0.0065852640 4000.00000
# a_tank[28]   -0.5903300678  0.430889431 0.00681296010 0.0061191595 4000.00000
# a_tank[31]   -0.7790724381  0.438852449 0.00693886648 0.0062867441 4000.00000
# a_tank[29]    0.0892476063  0.403319027 0.00637703374 0.0065342700 3809.80961
# a_tank[18]    2.5984354254  0.789739347 0.01248687547 0.0116099713 4000.00000
#  a_tank[1]    2.5169564861  1.202252233 0.01900927690 0.0261521039 2113.38159
# a_tank[38]    6.5730591701  2.577397287 0.04075222931 0.0960321899  720.32596
# a_tank[37]    2.1321270081  0.542445637 0.00857681860 0.0067553411 4000.00000
# a_tank[30]    1.4436385748  0.510944886 0.00807874799 0.0061875938 4000.00000
# a_tank[35]    2.9837870684  0.798955606 0.01263259732 0.0111689405 4000.00000
# a_tank[20]    6.4209696786  2.614898794 0.04134518019 0.0934495614  782.98823
# a_tank[15]    2.4852422308  1.182168671 0.01869172789 0.0235303705 2524.06980
# a_tank[45]    0.5431043960  0.369638488 0.00584449767 0.0046883639 4000.00000
# a_tank[26]    0.0840349555  0.395331366 0.00625073774 0.0067906853 3389.18550
# a_tank[21]    2.6516314895  0.851902518 0.01346976151 0.0118211120 4000.00000
# a_tank[48]   -0.0571820248  0.337815684 0.00534133495 0.0051459323 4000.00000
# a_tank[10]    2.5310386381  1.157968533 0.01830909012 0.0251590089 2118.39263
#  a_tank[3]    0.9295620334  0.735176403 0.01162415958 0.0127307257 3334.85377
#  a_tank[7]    5.7185532309  2.796270760 0.04421292278 0.1047192286  713.02634
#  a_tank[6]    2.5050631208  1.183397045 0.01871115019 0.0243016893 2371.30836
# a_tank[25]   -1.2007745016  0.470201538 0.00743453910 0.0071826339 4000.00000
# a_tank[32]   -0.4234910116  0.414945560 0.00656086537 0.0060452516 4000.00000
#  a_tank[5]    2.5286171889  1.207469316 0.01909176622 0.0250131147 2330.32587
# a_tank[13]    0.9170279823  0.723356937 0.01143727741 0.0091140235 4000.00000
# a_tank[19]    2.1126922998  0.642028977 0.01015136946 0.0086040212 4000.00000
# a_tank[12]    0.4476884510  0.682975434 0.01079878979 0.0070550762 4000.00000
#  a_tank[4]    5.5617118510  2.607791677 0.04123280681 0.1041546582  626.88574
# a_tank[17]    3.4763818329  1.091795631 0.01726280467 0.0199830235 2985.10977
# a_tank[27]   -1.7411510659  0.578372766 0.00914487638 0.0085613085 4000.00000
# a_tank[22]    2.6038964868  0.779764455 0.01232915859 0.0107305461 4000.00000
# a_tank[23]    2.6217848846  0.808680221 0.01278635699 0.0105079615 4000.00000
# a_tank[36]    2.1205525264  0.557014163 0.00880716721 0.0077309477 4000.00000
# a_tank[39]    2.9653158635  0.795061058 0.01257101910 0.0110571371 4000.00000

# Rethinking
#             mean   sd  5.5% 94.5% n_eff Rhat
# a_tank[1]   2.49 1.16  0.85  4.53  1079    1
# a_tank[2]   5.69 2.75  2.22 10.89  1055    1
# a_tank[3]   0.89 0.75 -0.23  2.16  1891    1
# a_tank[4]   5.71 2.70  2.21 10.85   684    1
# a_tank[5]   2.52 1.14  0.92  4.42  1640    1
# a_tank[6]   2.49 1.13  0.94  4.52  1164    1
# a_tank[7]   5.74 2.71  2.25 10.86   777    1
# a_tank[8]   2.52 1.19  0.95  4.42  1000    1
# a_tank[9]  -0.46 0.69 -1.62  0.55  2673    1
# a_tank[10]  2.53 1.19  0.93  4.59  1430    1
# a_tank[11]  0.93 0.72 -0.17  2.11  1387    1
# a_tank[12]  0.47 0.74 -0.63  1.70  1346    1
# a_tank[13]  0.91 0.76 -0.25  2.30  1559    1
# a_tank[14]  0.00 0.66 -1.04  1.06  2085    1
# a_tank[15]  2.50 1.19  0.95  4.40  1317    1
# a_tank[16]  2.50 1.14  0.98  4.31  1412    1
# a_tank[17]  3.49 1.12  1.94  5.49   945    1
# a_tank[18]  2.59 0.75  1.50  3.81  1561    1
# a_tank[19]  2.11 0.64  1.15  3.15  1712    1
# a_tank[20]  6.40 2.57  3.11 11.04   996    1
# a_tank[21]  2.59 0.74  1.54  3.93  1233    1
# a_tank[22]  2.63 0.79  1.49  4.01  1184    1
# a_tank[23]  2.64 0.83  1.45  4.13  1379    1
# a_tank[24]  1.74 0.59  0.85  2.72  1736    1
# a_tank[25] -1.19 0.45 -1.90 -0.50  2145    1
# a_tank[26]  0.09 0.41 -0.53  0.78  2167    1
# a_tank[27] -1.75 0.56 -2.65 -0.88  1666    1
# a_tank[28] -0.58 0.43 -1.25  0.08  1567    1
# a_tank[29]  0.08 0.39 -0.54  0.71  3053    1
# a_tank[30]  1.43 0.49  0.66  2.24  2754    1
# a_tank[31] -0.79 0.44 -1.50 -0.12  1299    1
# a_tank[32] -0.42 0.41 -1.12  0.23  1661    1
# a_tank[33]  3.84 1.08  2.31  5.70   808    1
# a_tank[34]  3.00 0.85  1.83  4.36  1038    1
# a_tank[35]  2.96 0.82  1.82  4.25  1578    1
# a_tank[36]  2.14 0.55  1.31  3.08  1734    1
# a_tank[37]  2.12 0.56  1.31  3.04  1131    1
# a_tank[38]  6.72 2.62  3.45 11.44   706    1
# a_tank[39]  2.95 0.73  1.85  4.08  1509    1
# a_tank[40]  2.48 0.65  1.53  3.61  1731    1
# a_tank[41] -2.15 0.57 -3.11 -1.29  1231    1
# a_tank[42] -0.67 0.35 -1.22 -0.14  1444    1
# a_tank[43] -0.54 0.35 -1.12  0.03  1776    1
# a_tank[44] -0.43 0.34 -1.00  0.10  1735    1
# a_tank[45]  0.54 0.36 -0.04  1.14  1376    1
# a_tank[46] -0.67 0.34 -1.25 -0.15  1619    1
# a_tank[47]  2.14 0.55  1.31  3.04  1916    1
# a_tank[48] -0.06 0.35 -0.61  0.50  1932    1
