from IPython.core.pylabtools import figsize
import numpy as np
from matplotlib import pyplot as plt
from os import makedirs
from urllib.request import urlretrieve
import pymc3 as pm
import theano.tensor as tt
from pymc3 import model_to_graphviz
from pymc3 import forestplot


## analysis 1 P208
urlretrieve("http://peach.l.chiba-u.ac.jp/course_folder/z15N50.csv", "data/z15N50.csv")
data01 = np.loadtxt("data/z15N50.csv", skiprows=1)
print(data01)

with pm.Model() as model:
    theta = pm.Beta('theta', alpha = 1, beta = 1)
    obs = pm.Bernoulli('obs', theta, observed = data01)

model_to_graphviz(model)

with model:
    trace = pm.sample(10000)


# checking results
figsize(12.5, 4)
pm.plots.traceplot(trace, var_names=["theta"])
pm.plots.plot_posterior(trace, var_names=["theta"])
pm.plots.autocorrplot(trace, var_names=["theta"])


# Analysis 2 p208
resp = [1,0,1,1,1,1,1,0,0,0,1,0,0,1,0]
subj = [0,0,0,0,0,0,0,0,1,1,1,1,1,1,1]

from pymc3 import model_to_graphviz
N_subj=2
with pm.Model() as model2:

    theta = pm.Beta('theta', alpha = 2, beta = 2, shape=N_subj)

    y = pm.Bernoulli('y', theta,  observed=resp)
    
    diff_theta = pm.Deterministic('delta_theta', theta[0] - theta[1])
model_to_graphviz(model2)

# mcmc sampling
with model2:
    trace = pm.sample(10000)

# viz & checking results
pm.plots.traceplot(trace, var_names=["theta"])
pm.plots.plot_posterior(trace, var_names=["theta"])
pm.plots.plot_posterior(trace, var_names=["delta_theta"])
pm.plots.autocorrplot(trace, var_names=["theta"])


### analysis 3 p240
N_subj = len(data.s.unique())
resp = data.y
subj = data.s - 1
with pm.Model() as model3:
    omega = pm.Beta('omega',1,1)
    kappaM2 = pm.Gamma('kappaM2', 0.01, 0.01)
    kappa = pm.Deterministic('kappa', kappaM2 + 2)
    theta = pm.Beta('theta', 
                    alpha = omega*(kappa-2)+1,
                    beta = (1-omega)*(kappa-2)+1, 
                    shape = N_subj)



    y = pm.Bernoulli('y', theta[subj],  observed=resp)
    
model_to_graphviz(model3)
### analysis 4 p253
N_player = len(data.Player.unique())
N_pos = 9
Ns = data.AtBats
Zs = data.Hits
subj = data.PlayerNumber - 1
pos = data.PriPosNumber - 1

with pm.Model() as model4:
    omegaO = pm.Beta('omegaO',1,1)
    kappaM2O = pm.Gamma('kappaM2', 0.01, 0.01)
    kappaO = pm.Deterministic('kappaO', kappaM2O + 2)
    
    omegaPOS = pm.Beta('omegaPOS', 
                       alpha = omegaO*(kappaO-2)+1, 
                       beta = (1-omegaO)*(kappaO-2)+1, 
                       shape = N_pos)
    kappaM2POS = pm.Gamma('kappaM2POS', 0.01, 0.01,
                           shape = N_pos)
    kappaPOS = pm.Deterministic('kappaPOS', kappaM2POS + 2)

   
    theta = pm.Beta('theta', 
                    alpha = omegaPOS[pos]*(kappaPOS[pos]-2)+1, 
                   beta = (1-omegaPOS[pos])*(kappaPOS[pos]-2)+1, 
                    shape = N_player)

    y = pm.Binomial('y', n = Ns, p = theta,  observed = Zs)
    
model_to_graphviz(model4)