try:
  import gi
  gi.require_version('NumCosmo', '1.0')
  gi.require_version('NumCosmoMath', '1.0')
except:
  pass

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

from gi.repository import GObject
from gi.repository import NumCosmo as Nc
from gi.repository import NumCosmoMath as Ncm


"""
Explanation: License
HIPrimTutorial
Mon Jun 01 09:15:00 2020\
 Copyright  2020\
 Sandro Dias Pinto Vitenti &#118;&#105;&#116;&#101;&#110;&#116;&#105;&#64;&#117;&#101;&#108;&#46;&#98;&#114;
---
HIPrimTutorial\
 Copyright (C) 2020 Sandro Dias Pinto Vitenti &#118;&#105;&#116;&#101;&#110;&#116;&#105;&#64;&#117;&#101;&#108;&#46;&#98;&#114;
numcosmo is free software: you can redistribute it and/or modify it
 under the terms of the GNU General Public License as published by the
 Free Software Foundation, either version 3 of the License, or
 (at your option) any later version.
numcosmo is distributed in the hope that it will be useful, but
 WITHOUT ANY WARRANTY; without even the implied warranty of
 MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.
 See the GNU General Public License for more details.
You should have received a copy of the GNU General Public License along
 with this program.  If not, see http://www.gnu.org/licenses/.

HIPrim example
In this tutorial we define a primordial power spectrum for the cosmological perturbations and compute the $C_\ell$'s associated to it. Here you are going to learn:

How to create a new model and registry it in the GObject type system.
How to use this model together with a $\Lambda$CDM cosmological model to compute the perturbations.

The primordial power spectrum:
\begin{equation}
\Delta_k = A_s \left(\frac{k}{k_}\right)^{n_s-1}\left[1+a^2\cos\left(b \frac{k}{k_}+c\right)\right],
\end{equation}
where $k_*$ is the pivotal scale and $(A_s,\, n_s,\, a,\, b,\, c)$ are free parameters.
Loading NumCosmo
The first step is to load both NumCosmo and NumCosmoMath libraries. Here we also load the GObject class that we will use to registry the new class in our GObject type system as well as other modules needed in this example.
End of explanation
"""


mb = Ncm.ModelBuilder.new (Nc.HIPrim, "NcHIPrimExample", "A example primordial model")

"""
Explanation: Creating the model
Since python does not give us access to the guts of the GObject system we need to use the Ncm.ModelBuilder class to implement a child of Ncm.Model.
End of explanation
"""


mb.add_sparam ("A_s", "As", 0.0, 1.0, 0.1, 0.0, 1.0e-9, Ncm.ParamType.FREE)
mb.add_sparam ("n_s", "ns", 0.5, 1.5, 0.1, 0.0, 0.96,   Ncm.ParamType.FREE)

mb.add_sparam ("a", "a", 0.0,   1.0, 0.01, 0.0,   0.5, Ncm.ParamType.FREE)
mb.add_sparam ("b", "b", 0.0, 1.0e4, 0.10, 0.0, 100.0, Ncm.ParamType.FREE)
mb.add_sparam ("c", "c", 0.0,   6.0, 0.10, 0.0,   0.0, Ncm.ParamType.FREE)

"""
Explanation: The line above creates a new Ncm.ModelBuilder that will create a child of Nc.HIPrim called Nc.HIPrimExample. 
The next step is to add parameters to the model, all parameters that will possibly be used in a statistical analysis, see the expression above. The documentation of the Ncm.ModelBuilder class can be found here.
End of explanation
"""


GNcHIPrimExample = mb.create ()
GObject.new (GNcHIPrimExample)
NcHIPrimExampleBase = GNcHIPrimExample.pytype
GObject.type_register (NcHIPrimExampleBase);

"""
Explanation: The arguments are the following: (for the C documentation see here)

the (latex) symbol describing the parameter. 
the string representing the parameter internally, this same string can be used to access the parameter and change its values.
parameter lower-bound.
parameter upper-bound.
(guess) parameter scale, used by statistical algorithms (e.g., best-fit finders) as a first guess step size. As a rule of thumb it should be at the same order of magnitude as the $1\sigma$ error of the current analysis.
absolute tolerance, used when estimating the error on the parameter (useful for parameter that are expected to be exactly zero).
parameter default value.
Whether this parameter should be treated as a free parameter or fixed (Ncm.ParamType.FREE or Ncm.ParamType.FIXED) by default.

The code below register the new class in both python and GObject system following the steps:
1. creates the new GType, defining the object in the GObject system.
2. creates a new empty instance of this GType.
3. using this instance we get the python type.
4. Finally we register this python type back in the GObject system.
End of explanation
"""


class NcHIPrimExample (NcHIPrimExampleBase):
  def do_lnSA_powspec_lnk (self, lnk):
    lnk0 = self.get_lnk_pivot ()
    lnka = lnk - lnk0
    ka = math.exp (lnka)
    As = self.props.As
    ns = self.props.ns
    a  = self.props.a
    b  = self.props.b
    c  = self.props.c
    a2 = a * a

    return (ns - 1.0) * lnka + math.log (As) + math.log1p (a2 * math.cos (b * ka + c)**2) 

GObject.type_register(NcHIPrimExample);

"""
Explanation: Now we can finally implement in python a class representing the model in hand. We implement a single virtual method of the Nc.HIPrim abstract class lnSA_powspec_lnk. All virtual methods, e.g,  method_name, of the abstract classes can be implemented in python by creating a method called do_method_name in the derived class in python.
End of explanation
"""


__name__ = "NcContext"

Ncm.cfg_init ()
Ncm.cfg_set_log_handler (lambda msg: sys.stdout.write (msg) and sys.stdout.flush ())

"""
Explanation: Some notes about the definition above:
- The base class Nc.HIPrim has already a parameter describing pivotal $k_$, see here. Thus, our first step is to get the value of $k_$. 
- In python the values of the parameters can be accessed using the props property, e.g., a parameter created with name As can be accessed via self.props.As.
- The virtual method above should implement the function $\ln[\Delta(\ln k)]$ instead of $\Delta_k$.
- Since the class created is a python class, we need to registry it in the GObject type system, which is what we do in the last line of the cell above.
Initializing the objects
Now that we have our new primordial model we can create the set of objects necessary to compute the $C_\ell$'s. First we need to initialize the NumCosmo library:
End of explanation
"""


prim = NcHIPrimExample ()
print ("# Model set to its default values: ", prim.props.As,prim.props.ns, prim.props.a, prim.props.b, prim.props.c)

"""
Explanation: Then, we create an instance of our primordial model object, and print the current (default) values of its parameter.
End of explanation
"""


cbe_prec = Nc.CBEPrecision.new ()
cbe_prec.props.k_per_decade_primordial = 50.0

cbe = Nc.CBE.prec_new (cbe_prec)

lmax = 2500

Bcbe = Nc.HIPertBoltzmannCBE.full_new (cbe)
Bcbe.set_TT_lmax (lmax)
Bcbe.set_target_Cls (Nc.DataCMBDataType.TT)
Bcbe.set_lensed_Cls (True)


"""
Explanation: Next, we need a CLASS back-end object, but first we create a CLASS precision object. This step is necessary since we have a more complicated primordial power spectrum and for this reason we need a finer simpling of it when computing the other cosmological observables.
End of explanation
"""


cosmo = Nc.HICosmo.new_from_name (Nc.HICosmo, "NcHICosmoDEXcdm")
cosmo.omega_x2omega_k ()
cosmo.param_set_by_name ("Omegak", 0.0)

reion = Nc.HIReionCamb.new ()

cosmo.add_submodel (reion)
cosmo.add_submodel (prim)

"""
Explanation: Above we set the number of knots per decade when sampling the primordial power spectrum, see here for the documentation of the complete set of parameters. The, we create a class back-end object and the Boltzmann object based on it. Finally, we set the lmax, require the computation of the TT, $C_\ell$'s and the lens correction.
The next objects are necessary to close our cosmological model:
End of explanation
"""


mset = Ncm.MSet.new_array ([cosmo])
mset.pretty_log ()

"""
Explanation: Below we create a model set that will contain all models we will work with. Then, we just print all their properties.
End of explanation
"""


Omegabh2 = 0.02222
Omegach2 = 0.1197
H0       = 67.31
h        = H0 / 100.0
ns       = 0.9655
ln1010As = 3.089
tau      = 0.078

Omegab   = Omegabh2 / h**2
Omegac   = Omegach2 / h**2
As       = math.exp (ln1010As) * 1.0e-10

pi_H0     = mset.param_get_by_full_name ("NcHICosmo:H0")
pi_Omegab = mset.param_get_by_full_name ("NcHICosmo:Omegab")
pi_Omegac = mset.param_get_by_full_name ("NcHICosmo:Omegac")
pi_As     = mset.param_get_by_full_name ("NcHIPrim:As")
pi_ns     = mset.param_get_by_full_name ("NcHIPrim:ns")

reion.set_z_from_tau (cosmo, tau)

prim.props.a = 0.0
prim.props.b = 0.0
prim.props.c = 0.0

mset.param_set_pi ([pi_H0, pi_Omegab, pi_Omegac, pi_As, pi_ns], [H0, Omegab, Omegac, As, ns])
mset.pretty_log ()

"""
Explanation: To change the values of anyone of these parameters execute the command below: (these parameters correspond to Planck 2015 TT + LowP bestfit)
End of explanation
"""


ps   = Nc.PowspecMLCBE.new ()
dist = Nc.Distance.new (2.0)

z_min = 0.0
z_max = 2.0

k_min = 2.65e-2
k_max = 1.0e1

ps.set_kmin (k_min)
ps.set_kmax (k_max)

ps.require_zi (z_min)
ps.require_zf (z_max)

ps.prepare (cosmo)
dist.prepare (cosmo)

print (cosmo.Omega_b0h2 ())
print (cosmo.Omega_c0h2 ())
print (dist.theta100CMB (cosmo))
print (reion.get_tau (cosmo))
print (ps.sigma_tophat_R (cosmo, 1.0e-7, 0.0, 8.0 / h))

prim.props.a = 0.5
prim.props.b = 100.0
prim.props.c = 0.0


"""
Explanation: We can now compute some derivative quantities from this parametrization.
End of explanation
"""


Bcbe.prepare (cosmo)

Cls_vec = Ncm.Vector.new (lmax + 1)

Bcbe.get_TT_Cls (Cls_vec)

ell = np.array (list (range (2, lmax + 1)))

Cls_a = np.array (Cls_vec.dup_array ()[2:])
Dls_a = ell * (ell + 1.0) * Cls_a


"""
Explanation: Above we created a Nc.HICosmoDEXcdm object as the basic cosmological model with its default values and reparametrized $\Omega_{DE0} \to \Omega_{k0}$ where we set $\Omega_{k0} = 0$. Then, we add both the primordial model and the reionization model as sub-models of cosmo.
Now that we have all the necessary components we can "prepare" the Bcbe object using our models:
End of explanation
"""


plt.figure (figsize=(14, 7))

plt.xscale('log')
plt.plot (ell, Dls_a, 'r', label="default parameters")

plt.xlabel (r'$\ell$')
plt.ylabel (r'$D_\ell$')
leg = plt.legend (loc = 'best')


"""
Explanation: Above we prepare Bcbe, create a new vector to hold the $C_\ell$s, and an array representing the values of $\ell$. Then, we define the array Dls defined by $D_\ell = \ell(\ell+1)C_\ell$. Below we plot our result:
End of explanation
"""


prim.props.a = 0.0

Bcbe.prepare (cosmo)
Bcbe.get_TT_Cls (Cls_vec)

ClsPlaw_a = np.array (Cls_vec.dup_array ()[2:])
DlsPlaw_a = ell * (ell + 1.0) * ClsPlaw_a

"""
Explanation: To compare with the simple power-law case we just need to set the parameter $a=0$ (see here).
End of explanation
"""


plt.figure (figsize=(14, 7))

plt.xscale('log')
plt.plot (ell, Dls_a,     'r', label="default parameters")
plt.plot (ell, DlsPlaw_a, 'b', label="power-law")

plt.xlabel (r'$\ell$')
plt.ylabel (r'$D_\ell$')
leg = plt.legend (loc = 'best')


"""
Explanation: And plot both together:
End of explanation
"""


