%matplotlib inline
import matplotlib.pyplot as plt
import numpy as np

from halomod import TracerHaloModel
import halomod
import hmf
import emcee
import scipy
from multiprocessing import Pool
import corner

print("halomod version: ", halomod.__version__)
print("hmf version:", hmf.__version__)

"""
Explanation: Parameter Constraints with halomod
In this tutorial, parameters of a minimal galaxy HOD are fitted against power spectrum created using halomod with some error bars
For auto-power spectrum, all parameters and related calculation can be done using TracerHaloModel. Parameters can be easily updated through update, which makes it very easy to combine with data fitting tools. Here we use emcee as an example.
Let's import those (and a few other things we'll need):
End of explanation
"""


hm_fid = TracerHaloModel(hod_model='Zehavi05',hod_params=dict(
{"M_min": 11.6222, "M_1": 12.851, "alpha": 1.049}
))
hm = TracerHaloModel(hod_model='Zehavi05',hod_params=dict(
{"M_min": 11.6222, "M_1": 12.851, "alpha": 1.049}
))

"""
Explanation: Setting up TracerHaloModel
As a simplistic example, we'll be using a Zehavi05 model which is simply:
$$
\big\langle N_{\rm cen}(m) \big\rangle = 1,\; {\rm for}\; m \ge M_{\rm min}
$$
$$
\big\langle N_{\rm cen}(m) \big\rangle = 0,\; {\rm for}\; m < M_{\rm min}
$$
$$
\big\langle N_{\rm sat}(m) \big\rangle = \big(\frac{m}{M_1}\big)^\alpha
$$
which can be specified by hod_model and hod_params arguments:
End of explanation
"""


hm.update(hod_params=dict(
{"M_min": 11.6222, "M_1": 12.851, "alpha": 1.149}))

"""
Explanation: And we've set up our HOD. You can easily update the parameter values:
End of explanation
"""


plt.plot(hm_fid.k_hm, hm_fid.power_auto_tracer,label='fid')
plt.plot(hm.k_hm, hm.power_auto_tracer,label='update')
plt.xscale('log')
plt.yscale('log')

plt.xlabel("k [$Mpc^{-1} h$]")
plt.ylabel(r"$\rm P(k) \ [{\rm Mpc^3}h^{-3}]$")
plt.legend(loc="upper right");

"""
Explanation: Compare it with the fiducial power spectrum to check the update:
End of explanation
"""


pkerr = np.linspace(0.01,0.5,num=80)*hm_fid.power_auto_tracer

plt.errorbar(hm_fid.k_hm, hm_fid.power_auto_tracer,yerr= pkerr)

plt.xscale('log')
plt.yscale('log')

plt.xlabel("k [$Mpc^{-1} h$]")
plt.ylabel(r"$\rm P(k) \ [{\rm Mpc^3}h^{-3}]$");

"""
Explanation: Just like that, we're all set.
MCMC-Fitting
For a naive first-look, let's create some fake error bar:
End of explanation
"""


def autopower(theta,k):
    hm.update(hod_params=theta)
    pk=scipy.interpolate.interp1d(hm.k_hm,hm.power_auto_tracer)
    return pk(k)
def lnprior_flat(params):
    if params[0]>9.00 and params[0]<13.00 and params[1]<16.00 and params[1]>9.00 \
    and params[2]<4 and params[2]>0:
        return 0.0
    else:
        return -np.inf
def log_likelihood_flat(params, x,y,yerr):
    theta=dict({"M_min": params[0], "M_1": params[1], "alpha": params[2]})
    yfit= autopower(theta,x)
    lp = lnprior_flat(params)
    return lp - 0.5 * np.sum(((y-yfit)/yerr)**2)

"""
Explanation: And construct a simple Gaussian likelihood function with flat prior:
End of explanation
"""


pars = [11.6222, 12.851, 1.049]
print(log_likelihood_flat(pars,hm_fid.k_hm,hm_fid.power_auto_tracer,pkerr))

"""
Explanation: Test against the fiducial values:
End of explanation
"""


pars= np.array([11.6222, 12.851, 1.049])

ndim,nwalkers = 3,100
pos = [pars+1e-2*np.random.uniform(0,1,ndim)*pars for i in range(nwalkers)]
with Pool() as pool:
    sampler = emcee.EnsembleSampler(nwalkers, ndim, log_likelihood_flat, \
            args=(hm_fid.k_hm,hm_fid.power_auto_tracer,pkerr),pool=pool)
    #start = time.time()
    sampler.run_mcmc(pos, 1000, progress=True)
    #end = time.time()
    #multi_time = end - start
    #print("Multiprocessing took {0:.1f} seconds".format(multi_time))

"""
Explanation: Just for a taste of it, let's run the mcmc with 100 walkers and 1000 steps:
End of explanation
"""


samples = sampler.get_chain().reshape((-1,3))

mmin,m1,alpha = map(lambda v: (v[1],v[2]-v[1],v[1]-v[0]),\
                    zip(*np.percentile((samples),[16,50,84],axis=0)))

"""
Explanation: Getting the chains and extracting posterior:
End of explanation
"""


print(mmin)
print(m1)
print(alpha)

"""
Explanation: The 1-d confidence interval:
End of explanation
"""


fig = corner.corner(samples);

"""
Explanation: And the posterior distribution as a corner plot:
End of explanation
"""


def autopower(theta,k):
    hm.update(hod_params=theta)
    pk=scipy.interpolate.interp1d(hm.k_hm,hm.power_auto_tracer)
    return pk(k), np.array([hm.mean_tracer_den,hm.bias_effective_tracer])
def lnprior_flat(params):
    if params[0]>9.00 and params[0]<13.00 and params[1]<16.00 and params[1]>9.00 \
    and params[2]<4 and params[2]>0:
        return 0.0
    else:
        return -np.inf
def log_likelihood_flat(params, x,y,yerr):
    theta=dict({"M_min": params[0], "M_1": params[1], "alpha": params[2]})
    yfit, arr= autopower(theta,x)
    lp = lnprior_flat(params)
    return lp - 0.5 * np.sum(((y-yfit)/yerr)**2),arr

"""
Explanation: Constraining derived parameters
Notice that halomod provides a large pool of tools to calculate various quantities of potential interest, making it easier to perform posterior analysis on some parameters derived from your choice of halo model. For example, let's rewrite the likelihood function to add a blob containing galaxy number density(mean_tracer_den) and linear galaxy bias(bias_effective_tracer):
End of explanation
"""


pars = [11.6222, 12.851, 1.049]
print(log_likelihood_flat(pars,hm_fid.k_hm,hm_fid.power_auto_tracer,pkerr))

"""
Explanation: Now the likelihood function also returns the blob:
End of explanation
"""


pars= np.array([11.6222, 12.851, 1.049])

ndim,nwalkers = 3,100
pos = [pars+1e-2*np.random.uniform(0,1,ndim)*pars for i in range(nwalkers)]
with Pool() as pool:
    sampler = emcee.EnsembleSampler(nwalkers, ndim, log_likelihood_flat, \
            args=(hm_fid.k_hm,hm_fid.power_auto_tracer,pkerr),pool=pool)
    #start = time.time()
    sampler.run_mcmc(pos, 1000, progress=True)
    #end = time.time()
    #multi_time = end - start
    #print("Multiprocessing took {0:.1f} seconds".format(multi_time))

"""
Explanation: Rerun the same fitting:
End of explanation
"""


blob = sampler.get_blobs().reshape((-1,2))

ng,b0 = map(lambda v: (v[1],v[2]-v[1],v[1]-v[0]),\
                    zip(*np.percentile((blob),[16,50,84],axis=0)))

"""
Explanation: Extracting the blob:
End of explanation
"""


print(ng)
print(b0)

"""
Explanation: The 1$\sigma$ confidence interval is:
End of explanation
"""


