from __future__ import division, print_function
import numpy as np
import scipy.integrate as integrate
import matplotlib.pyplot as plt
plt.ion()

par_mn = 10.0
par_sig = 2.0

#First, an explicit integral... Lets start at 1 milli-arsec, equivalent 
#to our prior being that there is no chance the distance is more than 1 kpc.
par = np.linspace(1,20,1000)

#This is P(data | par), i.e. the likelihood
#Ignore normalisation constants. We'll add them at the end as is usual in
#Bayesian statistics.
par_distr = np.exp(-(par-par_mn)**2/2.0/par_sig**2)
par_distr /= integrate.trapz(par_distr, par)
par_prior = 1.0/par**4.0
par_distr_with_prior = par_distr * par_prior

#Now normalise
par_distr_with_prior /= integrate.trapz(par_distr_with_prior, par)
mean_par = np.trapz(par_distr_with_prior * par, par)
print("Expected value of parallax: {0:5.2f}".format(mean_par))

#To check that this makes sense, lets plot. *** Always plot ***
plt.clf()
plt.plot(par, par_distr, label='P(D | plx)')
plt.plot(par, par_distr_with_prior, label='P(plx | D)')
plt.xlabel(r'Parallax (mas), $\pi$')
plt.ylabel(r'f($\pi$)')
plt.legend()

#How about a median?
cum_distr_with_prior = np.append(0,integrate.cumtrapz(par_distr_with_prior, par))
median_par = np.interp(0.5, cum_distr_with_prior, par)
print("Median value of parallax: {0:5.2f}".format(median_par))

#How about error bars?
lower_par = np.interp(0.16, cum_distr_with_prior, par)
upper_par = np.interp(0.84, cum_distr_with_prior, par)

print("Parallax with error (mas): {0:4.2f} +{1:4.2f}/-{2:4.2f}".\
    format(median_par, upper_par-median_par, median_par - lower_par))
    
print("Distance with error (pc): {0:4.1f} +{1:4.1f}/-{2:4.1f}".\
    format(1e3/median_par, 1e3/lower_par-1e3/median_par, 1e3/median_par-1e3/upper_par))