from edward.models import Bernoulli, Beta, Empirical, Uniform

N = 100

def build_fair_dataset(N):
    pheads = tf.constant(0.5)
    c = Bernoulli(probs=pheads, sample_shape = N)
    return sess.run([pheads, c])

def build_unfair_dataset(N):
    pheads = tf.constant(0.05)
    c = Bernoulli(probs=pheads, sample_shape = N)
    return sess.run([pheads, c])

def build_dataset(N):
    pheads = Uniform(low=0.0,high=1.0)
    c = Bernoulli(probs=pheads, sample_shape = N)
    return sess.run([pheads, c])

x = tf.range(-0.2, 1.2, 0.001)

plt.plot(*sess.run([x, Uniform(low=0.0,high=1.0).prob(x)]));
plt.ylim((-0.2,1.2))
plt.title('Uniform distribution')

pheads_true,c_train = build_fair_dataset(N)

pheads_true

c_train

sum(c_train == 0)

sum(c_train == 1)

"""
Explanation: Coin
End of explanation
"""


# declaring priors

pheads_fair = Beta(concentration1=1000.0, concentration0=1000.0) # blue
pheads_unfair = Beta(concentration1=0.1, concentration0=0.1) # green

pheads_unknown = Beta(concentration1 = 1.0, concentration0=1.0)

x = tf.range(0.0,1.0,0.001)
plt.plot(*sess.run([x, pheads_fair.prob(x)]));
plt.plot(*sess.run([x, pheads_unfair.prob(x)]));
plt.plot(*sess.run([x, pheads_unknown.prob(x)]));
plt.axvline(x=pheads_true);

# Forward model
pheads = pheads_unknown
c = Bernoulli(probs=pheads, sample_shape=N)

"""
Explanation: inference
End of explanation
"""


# inference
pheads_cond = ed.complete_conditional(pheads)
pheads_post = ed.copy(pheads_cond, {c: c_train})

sess.run({key:val for key, val in six.iteritems(pheads_post.parameters) if isinstance(val,tf.Tensor)})

# criticism
mean, stddev = sess.run([pheads_post.mean(),pheads_post.stddev()])
print(" exact posterior mean: " + str(mean))
print("exact posterior std" + str(stddev))

x = tf.range(0.0,1.0,0.001)
plt.plot(*sess.run([x,pheads.prob(x)]));
plt.plot(*sess.run([x,pheads_post.prob(x)]));
plt.axvline(x=pheads_true);

"""
Explanation: exact solution
End of explanation
"""


pheads_true,c_train = build_fair_dataset(100)

fig = plt.figure()
ax = plt.axes(xlim=(-0.05,1.05), ylim=(-1.0,11.0))

def go(pheads_prior, sample_shape, c_train, recursion=1):
    # model 
    c = Bernoulli(probs=pheads_prior,
                 sample_shape=sample_shape)
    
    # INFERENCE
    pheads_cond = ed.complete_conditional(pheads_prior)
    pheads_post = ed.copy(pheads_cond,{c:c_train[:sample_shape]})
    
    # CRITICISM
    ax.plot(*sess.run([x,pheads_post.prob(x)]));
    
    print("finished recursion "+str(recursion))
    recursion += 1
    # RECURSION
    if len(c_train[sample_shape:]) >= sample_shape:
        go(pheads_post, sample_shape, c_train[sample_shape:],recursion)
        
pheads_prior = Beta(concentration1=0.1, concentration0=0.1)
ax.plot(*sess.run([x, pheads_prior.prob(x)])); 
plt.axvline(x=pheads_true);
go(pheads_prior,33,c_train)

"""
Explanation: RECURSIVE INFERENCE
End of explanation
"""


# BACKWARD MODEL
T = 10000
q_pheads = Empirical(params=tf.Variable(tf.ones([T])*0.5))

# INFERENCE
proposal_pheads = Beta(concentration1=1.0,
                       concentration0=1.0)
inference = ed.MetropolisHastings(latent_vars = {pheads: q_pheads},
                                 proposal_vars={pheads: proposal_pheads},
                                 data={c:c_train})
inference.run()

# criticism
mean, stddev = sess.run([q_pheads.mean(),q_pheads.stddev()])
print("inferred posterior mean:")
print(mean)
print("inferred posterior std:")
print(stddev)

plt.plot(q_pheads.params.eval());
plt.axhline(y=pheads_true)

def lags(x):
    mean = tf.reduce_mean(x)
    var = tf.cast(tf.size(x) - 1, tf.float32) *tf.reduce_mean(tf.square(x-mean))
    ret = tf.map_fn(lambda k:tf.cond(tf.equal(k,0),
                                     lambda:var, 
                                     lambda:tf.reduce_sum((x[:-k] - mean) * (x[k:]-mean))),
                                     tf.range(0,tf.size(x)),
                                     dtype=tf.float32)
    return ret / var

plt.plot(lags(q_pheads.params).eval());

x = tf.range(0.0,1.0,0.001)
plt.plot(*sess.run([x,pheads.prob(x)]));
plt.plot(*sess.run([x,pheads_cond.prob(x)],
                  {c:c_train}));
plt.hist(q_pheads.params.eval(),
        bins=100,range=(0.0,1.0),
        normed = True);
plt.axvline(x=pheads_true);

"""
Explanation: approximate inference :MCMC: metroplis hastings:
End of explanation
"""


# BACKWARD MODEL
T = 10000
q_pheads = Empirical(params=tf.Variable(tf.ones([T])*0.5))

# INFERENCE 
inference = ed.Gibbs(latent_vars={pheads: q_pheads},
                    data={c: c_train})
inference.run()

# CRITICISM
mean, stddev = sess.run([q_pheads.mean(),q_pheads.stddev()])
print('Inferred posterior mean:')
print(mean)
print('Inferred posterior std:')
print(stddev)

plt.plot(q_pheads.params.eval());
plt.axhline(y=pheads_true)

x = tf.range(0.0,1.0,0.001)
plt.plot(*sess.run([x,pheads.prob(x)]));
plt.plot(*sess.run([x,pheads_cond.prob(x)],
                  {c:c_train}));
plt.hist(q_pheads.params.eval(),
        bins=100,range=(0.0,1.0),
        normed = True);
plt.axvline(x=pheads_true);

plt.plot(lags(q_pheads.params).eval());

"""
Explanation: MCMC: Gibbs
End of explanation
"""


# BACKWARD MODEL
T = 10000 # number of empirical samples
q_pheads = Empirical(params=tf.Variable(tf.ones([T])*.5))

# INFERENCE
inference = ed.HMC(latent_vars={pheads: q_pheads},
                  data = {c: c_train})
inference.run(step_size=1.0 / N, n_steps=20)

# criticism
mean, stddev = sess.run([q_pheads.mean(), q_pheads.stddev()])
print("Inferred posterior mean:")
print(mean)
print("Inferred posterior stddev")
print(stddev)

plt.plot(q_pheads.params.eval())
plt.axhline(y=pheads_true)

plt.plot(lags(q_pheads.params).eval())

x = tf.range(0.0, 1.0, 0.001)
plt.plot(*sess.run([x, pheads.prob(x)]))
plt.plot(*sess.run([x, pheads_cond.prob(x)],
                  {c:c_train}))
plt.hist(q_pheads.params.eval(),
        bins=100, range=(0.0, 1.0),
        normed=True)
plt.axvline(x=pheads_true)

"""
Explanation: MCMC: Hamiltonian monte carlo
Metropolis hasting and Gibbs will fail in high_dimensional paramter space while HMC will not. 
HMC unitlises the geometry of the important regions of the posterior (where the probability mass is concentrated) to guide the jump proposal.
For Gibbs you need to be able to derive the complete conditional algebraically. Most of the times this is not the case, so Gibbs cannot be used. Hamiltonian monte carlo does not have this restriction. However, it only works for continuous variables. This is because it works with derivatives (i.e. we need smooth functions).
End of explanation
"""


# BACKWARD MODEL
q_pheads_concentration1 = tf.nn.softplus(tf.Variable(51+ tf.random_normal([])))
q_pheads_concentration0 = tf.nn.softplus(tf.Variable(51+ tf.random_normal([])))
q_pheads = Beta(concentration1=q_pheads_concentration1,
               concentration0=q_pheads_concentration0)

x = tf.range(-5.0, 5.0, 0.001)
plt.plot(*sess.run([x,tf.nn.softplus(x)]))

# inference
inference = ed.KLqp(latent_vars={pheads: q_pheads},
                   data={c:c_train})
inference.run(n_samples=20, n_iter=1000)

sess.run({key: val for key, val in six.iteritems(q_pheads.parameters) if isinstance(val,tf.Tensor)})

plt.plot(*sess.run([x,pheads.prob(x)]))
plt.plot(*sess.run([x, pheads_cond.prob(x)],
                  {c:c_train}))
plt.plot(*sess.run([x, q_pheads.prob(x)]))
plt.axvline(x=pheads_true)

# criticism
mean, stddev = sess.run([q_pheads.mean(),q_pheads.stddev()])
print("Inferred posterior mean:")
print(mean)
print("Inferred posterior stddev")
print(stddev)

## A/B/... Testing

"""
Explanation: variational inference (VI)
Good alternative to markov chain monte carlo methods. Class of algorithm that cast posterior inference as optimisation:
Approximate the model paramters, that is, the posterior p(params|data = obs):
a. build a variational backward model, q(params: lambda), over the params with some variational paramters lambda.
b. Match q(params; lambda) to p(params|obs) by optimising over lambda, 
lambda* = argmin Closeness(p(params|obs), q(params;lambda))
c. use q(params;lambda*) instead of p(params|obs)
d. criticise the variational model, revise it (-->a)
End of explanation
"""


