# Source the environment
source /global/common/software/desi/users/adematti/cosmodesi_environment.sh main
# You may want to have it in the jupyter-kernel for plots
${COSMODESIMODULES}/install_jupyter_kernel.sh main  # this to be done once

from desilike.theories.galaxy_clustering import ShapeFitPowerSpectrumTemplate,\
                                                FOLPSTracerPowerSpectrumMultipoles
from desilike.observables.galaxy_clustering import TracerPowerSpectrumMultipolesObservable
from desilike.likelihoods import ObservablesGaussianLikelihood

template = ShapeFitPowerSpectrumTemplate(z=0.8)
theory = FOLPSTracerPowerSpectrumMultipoles(template=template)
theory.init.params['b1'].update(prior=dict(dist='norm', loc=1.5, scale=0.5))  # update theory parameters
observable = TracerPowerSpectrumMultipolesObservable(klim={0: [0.02, 0.2, 0.005], 2: [0.02, 0.2, 0.005]},
                                                     data='path_to_data',
                                                     covariance='path_to_covariance',
                                                     wmatrix='path_to_window_matrix',
                                                     theory=theory)
# Here we could give a list of several observables + joint covariance matrix
likelihood = ObservablesGaussianLikelihood(observables=[observable])
for param in likelihood.all_params.select(basename=['alpha*', 'sn*']):
	param.update(derived='.auto')  # analytic marginalization for linear parameters
# We can also sum up independent likelihoods (likelihood1 + likelihood2)
# Call likelihood
likelihood(qpar=0.99, qper=1.01)  # log-posterior

from desilike.theories.galaxy_clustering import ShapeFitPowerSpectrumTemplate,\
                                                FOLPSTracerPowerSpectrumMultipoles
from desilike.observables.galaxy_clustering import TracerPowerSpectrumMultipolesObservable
from desilike.likelihoods import ObservablesGaussianLikelihood

template = ShapeFitPowerSpectrumTemplate(z=0.8)
theory = FOLPSTracerPowerSpectrumMultipoles(template=template)
theory.init.params['b1'].update(prior=dict(dist='norm', loc=1.5, scale=0.5))  # update theory parameters
observable = TracerPowerSpectrumMultipolesObservable(klim={0: [0.02, 0.2, 0.005], 2: [0.02, 0.2, 0.005]},
                                                     data='path_to_data',
                                                     covariance='path_to_covariance',
                                                     wmatrix='path_to_window_matrix',
                                                     theory=theory)
# Here we could give a list of several observables + joint covariance matrix
likelihood = ObservablesGaussianLikelihood(observables=[observable])
for param in likelihood.all_params.select(basename=['alpha*', 'sn*']):
	param.update(derived='.auto')  # analytic marginalization for linear parameters
# We can also sum up independent likelihoods (likelihood1 + likelihood2)
# Call likelihood
likelihood(qpar=0.99, qper=1.01)  # log-posterior

ObservablesGaussianLikelihood(observables=[power_spectrum, bispectrum, gk, kk, ...],
							  covariance=big_total_covariance)

param.update(namespace='kk')

from desilike.emulators import Emulator, TaylorEmulatorEngine, EmulatedCalculator

# theory.pt = perturbation theory part, only depends on template parameters (qpar, qper, df, dm)
emulator = Emulator(theory.pt, engine=TaylorEmulatorEngine(method='finite', order=4))
emulator.set_samples()
emulator.fit()
pt = emulator.to_calculator()
# emulator.save(emulator_fn); pt = EmulatedCalculator.load(emulator_fn)

# Replace perturbation theory part by new one!
theory.init.update(pt=pt)

from desilike.samplers import EmceeSampler

sampler = EmceeSampler(likelihood, chains=8, nwalkers=40, seed=42, save_fn='chain_*.npy')
chains = sampler.run(min_iterations=2000, check={'max_eigen_gr': 0.03})



### Re-load chains, concatenate them
from desilike.samples import Chain
chain = Chain.concatenate([Chain.load('chain_{:d}.npy'.format(i)).remove_burnin(0.5)[::10]
                           for i in range(8)])
# To turn into GetDist samples
chain.to_getdist()

from desilike.profilers import MinuitProfiler

profiler = MinuitProfiler(likelihood, seed=42, save_fn='profiles.npy')
profiler.maximize(niterations=10)  # optimization
profiler.profile(params=['qpar'])  # 1D profile
profiler.contour(params=['qpar'])  # 2D contours

from desilike.theories import Cosmoprimo
from desilike.likelihoods.cmb import TTTEEEHighlPlanckNPIPECamspecLikelihood
from desilike.likelihoods.supernovae import Union3SNLikelihood

cosmo = Cosmoprimo(fiducial='DESI', engine='capse')  # emulator by Marco Bonici
likelihood = TTTEEEHighlPlanckNPIPECamspecLikelihood(cosmo=cosmo)
likelihood += Union3SNLikelihood(cosmo=cosmo)  # in JAX

sampler = NUTSSampler(likelihood, seed=42)
samples = sampler.run()