import stan
import matplotlib.pyplot as plt
import arviz as az

import nest_asyncio
nest_asyncio.apply()      #   To use PyStan3 on Notebook

N = 10
k = 7
stan_code = """
    data {
        int N;
        int k;
    }
    parameters {
        real<lower=0, upper=1> theta;
    }
    model {
        k ~ binomial(N, theta);
    }
"""
sm = stan.build(stan_code, data = {'N':N, 'k':k})
fit = sm.sample(num_samples=10000)

i_data = az.from_pystan(posterior = fit, posterior_model = sm)

smry = az.summary(i_data)
print('Summary =\n',smry)

az.plot_trace(i_data)
plt.show()

d_fm = fit.to_frame()

az.plot_kde(d_fm['theta'])
plt.xlabel(r'$\theta$', fontsize = 16)
plt.show()

Building: found in cache, done.
Sampling:   0%
Sampling:  25% (11000/44000)
Sampling:  50% (22000/44000)
Sampling:  75% (33000/44000)
Sampling: 100% (44000/44000)
Sampling: 100% (44000/44000), done.
Messages received during sampling:
  Gradient evaluation took 2.2e-05 seconds
  1000 transitions using 10 leapfrog steps per transition would take 0.22 seconds.
  Adjust your expectations accordingly!
  Gradient evaluation took 1.3e-05 seconds
  1000 transitions using 10 leapfrog steps per transition would take 0.13 seconds.
  Adjust your expectations accordingly!
  Gradient evaluation took 1.4e-05 seconds
  1000 transitions using 10 leapfrog steps per transition would take 0.14 seconds.
  Adjust your expectations accordingly!
  Gradient evaluation took 1.3e-05 seconds
  1000 transitions using 10 leapfrog steps per transition would take 0.13 seconds.
  Adjust your expectations accordingly!

Summary =
         mean    sd  hdi_3%  hdi_97%  mcse_mean  mcse_sd  ess_bulk  ess_tail  \
theta  0.666  0.13   0.423    0.898      0.001    0.001   14497.0   15777.0   

       r_hat  
theta    1.0