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

import sys
import math
import matplotlib as mpl
import matplotlib.pyplot as plt
import numpy as np

from IPython.display import HTML

from gi.repository import GObject
from gi.repository import NumCosmo as Nc
from gi.repository import NumCosmoMath as Ncm
%matplotlib inline

__name__ = "NcContext"

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

"""
Explanation: License
FisherMatrixExample
Tue Jan 19 13:48:00 2021\
 Copyright  2021\
 Sandro Dias Pinto Vitenti &#118;&#105;&#116;&#101;&#110;&#116;&#105;&#64;&#117;&#101;&#108;&#46;&#98;&#114;
---
FisherMatrixExample\
 Copyright (C) 2021 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/.
End of explanation
"""


mb = Ncm.ModelBuilder.new (Ncm.Model, "NcModelSimple", "A simple python example model")

#              Symbol     Name     Lim inf Lim sup  Scale   Abstol  Default  FREE or FIXED
mb.add_sparam (r'\alpha', 'alpha', -5.0,   5.0,     0.1,    0.0,    2.0,     Ncm.ParamType.FREE)
mb.add_sparam (r'\beta',  'beta',   0.0,   2.0,     0.1,    0.0,    1.0,     Ncm.ParamType.FREE)

if not obj:
    GNcModelSimple = mb.create ()
    obj = GObject.new (GNcModelSimple)
    ModelSimple = GNcModelSimple.pytype
    GObject.type_register (ModelSimple)
pass

class ModelSimpleF (ModelSimple):
    some_property = GObject.Property (type = str)

    def __init__ (self):
        ModelSimple.__init__ (self)
    
    def f_x (self, x):
        return np.exp (self.props.alpha + self.props.beta * x)

GObject.type_register (ModelSimpleF)

pass

"""
Explanation: Model definition
We define two models below, the abstract class NcModelSimple and its implementation ModelSimpleF that implements a simple function f_x that computes a straight line using two parameters alpha and beta. Most of the code is the boilerplate necessary to connect C and python and do not need to be modified.
End of explanation
"""


class DGSimple (Ncm.DataGaussCov):
    # We need one vector property to save the independent variables x
    xv = GObject.Property (type = Ncm.Vector, flags = GObject.PARAM_READWRITE)

    def __init__ (self, len = 600):
        Ncm.DataGaussCov.__init__ (self, n_points = len)

        if self.np > 0:
            self.xv = Ncm.Vector.new (self.np)
        else:
            self.xv = None
    
        self.cov_init = False
        if self.np > 0:
            self.cov.set_identity ()
            self.xv.set_zero ()

    def do_get_length (self):
        return self.np
    def do_get_dof (self):
        return self.np
    def do_prepare (self, mset):
        return
    def do_mean_func (self, mset, vp):
        mid = mset.get_id_by_ns ("NcModelSimple")
        mF = mset.peek (mid)

        for i in range (self.np):
            x = self.xv.get (i)
            vp.set (i, mF.f_x (x))

        return

    def create_random_cov (self, slm, rng):
        ya = [slm.f_x (x) for x in self.xv.dup_array ()]
        yv = Ncm.Vector.new_array (ya)

        self.cov.fill_rand_cov2 (yv, 0.1, 0.5, 15.0, rng)

GObject.type_register (DGSimple)
pass

"""
Explanation: Data model definition
Here we create a likelihood based on a multivariate Gaussian distribution. The most important method to be implemented is do_mean_func, it must compute the theoretical mean $f(x_i)$ for the Gaussian distribution $$-2\ln L = \sum_{ij}\left[(f(x_i) - y_i)C^{-1}_{ij}(f(x_j) - y_j)\right] + \dots$$
End of explanation
"""


mF = ModelSimpleF ()
mF.props.alpha = 0.9
mF.props.beta  = 0.2

mset = Ncm.MSet.empty_new ()
mset.set (mF)
mset.param_set_all_ftype (Ncm.ParamType.FREE)
mset.prepare_fparam_map ()

rng = Ncm.RNG.seeded_new (None, 123)

dgs = DGSimple (len = 15)
dgs.xv.set_array (np.linspace (0.0, 1.0, dgs.get_size ()))
dgs.create_random_cov (mF, rng)
dgs.resample (mset, rng)

dgs.xv.log_vals ("x: ", "%.3f", True)
dgs.y.log_vals ("y: ", "%.3f", True)
dgs.cov.log_vals ("cov:", "%.3f")

fig = plt.figure (dpi = 120)

x_a = np.linspace (0.0, 1.0, 200)
y_a = mF.f_x (x_a)
cov_d = [dgs.cov.get (i, i) for i in range (15)]

plt.errorbar (dgs.xv.dup_array (), dgs.y.dup_array (), yerr = cov_d, fmt = 'o')
plt.plot (x_a, y_a, '-')

pass

"""
Explanation: Creating data
Now we have a data model but no covariance $C_{ij}$ nor data $y_i$. Let's simulate some. First we create a fiducial model. Then, a Likelihood with $15$ points, an uniformly spaced dependent variable vector and a random covariance matrix. We do not need actual data $y_i$ to compute the Fisher matrix, but we create some using resample so we can compare the Fisher and the observed Fisher matrix.
End of explanation
"""


dset = Ncm.Dataset.new ()
dset.append_data (dgs)
lh = Ncm.Likelihood.new (dset)

fit = Ncm.Fit.new (Ncm.FitType.NLOPT, "ln-neldermead", lh, mset, 
                   Ncm.FitGradType.NUMDIFF_FORWARD)


"""
Explanation: Creating the fit object
End of explanation
"""


mF.props.alpha = 0.5
mF.props.beta  = 0.1


fit.fisher ()
fit.log_covar ()

"""
Explanation: Fisher matrix
Finally we can compute the Fisher matrix at the point $\alpha = 0.5$ and $\beta = 0.1$
End of explanation
"""


fit.obs_fisher ()
fit.log_covar ()

"""
Explanation: Observed Fisher matrix
Now we can compare it with the observed Fisher matrix. They will differ since the observed fisher matrix only approximate the expected real Fisher matrix at the best fit.
End of explanation
"""


fit.run (Ncm.FitRunMsgs.SIMPLE)

fit.fisher ()
fit.log_covar ()

fit.obs_fisher ()
fit.log_covar ()

"""
Explanation: To make them match we first compute the bestfit and then both matrices. Now, the observed Fisher matrix gives a reasonably good approximation.
End of explanation
"""


