samples_size = 1000
mean = 2
sigma = 1
data = np.random.normal(mean, sigma, samples_size)

f,ax = plt.subplots()
sns.distplot(data)

def create_new_sample(theta):
    mu,sigma = theta
    if sigma<=0:
        sigma=10
    return np.random.normal(mu, sigma, samples_size)


"""
Explanation: Example of using two distance + thresholds in one ABC sampling run
End of explanation
"""


def dist_measure(x, y):
    return [np.abs(np.mean(x, axis=0) - np.mean(y, axis=0)),
            np.abs(np.var(x, axis=0) - np.var(y, axis=0))]


distances = [dist_measure(data, create_new_sample((mean, sigma))) for _ in range(1000)]

dist_labels = ["mean", "var"]

f, ax = plt.subplots()
sns.distplot(np.asarray(distances)[:, 0], axlabel="distances", label=dist_labels[0])
sns.distplot(np.asarray(distances)[:, 1], axlabel="distances", label=dist_labels[1])
plt.title("Variablility of distance from simulations")
plt.legend()

import abcpmc

prior = abcpmc.GaussianPrior(mu=[2.5, 1.5], sigma=np.eye(2) * 0.5)

alpha = 85
T = 100
eps_start = [1.0, 1.0]


sampler = abcpmc.Sampler(N=1000, Y=data, postfn=create_new_sample, dist=dist_measure, threads=7)

def launch():
    eps = abcpmc.ConstEps(T, eps_start)

    pools = []
    for pool in sampler.sample(prior, eps):
        eps_str = ", ".join(["{0:>.4f}".format(e) for e in pool.eps])
        print("T: {0}, eps: [{1}], ratio: {2:>.4f}".format(pool.t, eps_str, pool.ratio))

        for i, (mean, std) in enumerate(zip(*abcpmc.weighted_avg_and_std(pool.thetas, pool.ws, axis=0))):
            print(u"    theta[{0}]: {1:>.4f} \u00B1 {2:>.4f}".format(i, mean,std))

        eps.eps = np.percentile(pool.dists, alpha, axis=0) # reduce eps value
        pools.append(pool)
        if pool.ratio <0.2:
            break
    sampler.close()
    return pools


import time
t0 = time.time()
pools = launch()
print("took", (time.time() - t0))

T = len(pools)
fig, ax = plt.subplots()
lables =["mu", "sigma"]
for i in range(2):
    moments = np.array([abcpmc.weighted_avg_and_std(pool.thetas[:,i], pool.ws, axis=0) for pool in pools])
    ax.errorbar(range(T), moments[:, 0], moments[:, 1], label=lables[i])
ax.hlines(np.mean(data), 0, T, linestyle="dotted", linewidth=0.7)
ax.hlines(np.std(data), 0, T, linestyle="dotted", linewidth=0.7)
ax.legend()
_ = ax.set_xlim([-.5, T])

distances = np.array([pool.dists for pool in pools])
fig, ax = plt.subplots()
ax.errorbar(np.arange(len(distances)), np.mean(distances, axis=1)[:, 0], np.std(distances, axis=1)[:, 0], label=dist_labels[0])
ax.errorbar(np.arange(len(distances)), np.mean(distances, axis=1)[:, 1], np.std(distances, axis=1)[:, 1], label=dist_labels[1])
ax.legend()
ax.set_title("Distances")

fig,ax = plt.subplots()
eps_values = np.array([pool.eps for pool in pools])
ax.plot(eps_values[:, 0], label=dist_labels[0])
ax.plot(eps_values[:, 1], label=dist_labels[1])
ax.set_xlabel("Iteration")
ax.set_ylabel(r"$\epsilon$", fontsize=15)
ax.legend(loc="best")
ax.set_title("Thresholds")

jg = sns.jointplot(pools[-1].thetas[:, 0], 
              pools[-1].thetas[:, 1], 
              kind="kde", 
             )
jg.ax_joint.set_xlabel(lables[0])
jg.ax_joint.set_ylabel(lables[1])

"""
Explanation: Distance measure compares mean + variance
End of explanation
"""


