"""
Classes and methods needed to do hypersurface interpolation over arbitrary parameters.
"""

__all__ = ['HypersurfaceInterpolator', 'run_interpolated_fit', 'prepare_interpolated_fit',
            'assemble_interpolated_fits', 'load_interpolated_hypersurfaces', 'pipeline_cfg_from_states',
            'serialize_pipeline_cfg', 'get_incomplete_job_idx']

__author__ = 'T. Stuttard, A. Trettin'

__license__ = '''Copyright (c) 2014-2017, The IceCube Collaboration

 Licensed under the Apache License, Version 2.0 (the "License");
 you may not use this file except in compliance with the License.
 You may obtain a copy of the License at

   http://www.apache.org/licenses/LICENSE-2.0

 Unless required by applicable law or agreed to in writing, software
 distributed under the License is distributed on an "AS IS" BASIS,
 WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 See the License for the specific language governing permissions and
 limitations under the License.'''

import os
import collections
import copy

import numpy as np
from scipy import interpolate
from .hypersurface import Hypersurface, HypersurfaceParam
from pisa import FTYPE, ureg
from pisa.utils import matrix
from pisa.utils.jsons import from_json, to_json
from pisa.utils.fileio import from_file, to_file
from pisa.core.pipeline import Pipeline
from pisa.core.binning import MultiDimBinning, is_binning
from pisa.core.map import Map
from pisa.core.param import Param, ParamSet
from pisa.utils.resources import find_resource
from pisa.utils.fileio import mkdir
from pisa.utils.log import logging, set_verbosity
from pisa.utils.comparisons import ALLCLOSE_KW
from uncertainties import ufloat, correlated_values
from uncertainties import unumpy as unp

class HypersurfaceInterpolator(object):
    """Factory for interpolated hypersurfaces.

    After being initialized with a set of hypersurface fits produced at different
    parameters, it uses interpolation to produce a Hypersurface object
    at a given point in parameter space using scipy's `RegularGridInterpolator`.

    The interpolation is piecewise-linear between points. All points must lie on a
    rectilinear ND grid.

    Parameters
    ----------
    interpolation_param_spec : dict
        Specification of interpolation parameter grid of the form::
            interpolation_param_spec = {
                'param1': {"values": [val1_1, val1_2, ...], "scales_log": True/False}
                'param2': {"values": [val2_1, val2_2, ...], "scales_log": True/False}
                ...
                'paramN': {"values": [valN_1, valN_2, ...], "scales_log": True/False}
            }
        where values are given as :obj:`Quantity`.
    hs_fits : list of dict
        list of dicts with hypersurfacesthat were fit at the points of the parameter mesh
        defined by interpolation_param_spec
    ignore_nan : bool
        Ignore empty bins in hypersurfaces. The intercept in those bins is set to 1 and
        all slopes are set to 0.

    Notes
    -----
    Be sure to give a support that covers the entire relevant parameter range and a
    good distance beyond! To prevent minimization failure from NaNs, extrapolation
    is used if hypersurfaces outside the support are requested but needless to say
    these numbers are unreliable.

    See Also
    --------
    scipy.interpolate.RegularGridInterpolator :
        class used for interpolation
    """

    def __init__(self, interpolation_param_spec, hs_fits, ignore_nan=True):
        self.ndim = len(interpolation_param_spec.keys())
        # key ordering is important to guarantee that dimensions stay consistent
        msg = "interpolation params must be specified as a dict with ordered keys"
        assert isinstance(interpolation_param_spec, collections.OrderedDict), msg
        for k, v in interpolation_param_spec.items():
            assert set(v.keys()) == {"values", "scales_log"}
            assert isinstance(v["values"], collections.Sequence)
        self.interp_param_spec = interpolation_param_spec
        reference_hs = hs_fits[0]["hs_fit"]
        # we are going to produce the hypersurface from a state that is the same
        # as the reference, only the coefficients and covariance matrices are
        # injected from the interpolation.
        self._reference_state = copy.deepcopy(reference_hs.serializable_state)
        # for cleanliness we wipe numbers from the original state
        self._reference_state["intercept_sigma"] = np.nan
        self._reference_state["fit_maps_norm"] = None
        self._reference_state["fit_maps_raw"] = None
        self._reference_state["fit_chi2"] = np.nan
        for param in self._reference_state['params'].values():
            param['fit_coeffts_sigma'] = np.full_like(
                param['fit_coeffts_sigma'], np.nan)
        # Instead of holding numbers, these coefficients and covariance matrices are
        # interpolator objects the produce them at the requested point.
        # The shape of fit_coeffts is [binning ..., fit coeffts]
        self.coeff_shape = reference_hs.fit_coeffts.shape
        self.coefficients = None
        # The shape of fit_cov_mat is [binning ..., fit coeffts, fit coeffts]
        self.covars_shape = reference_hs.fit_cov_mat.shape
        self.covars = None

        # We now need to massage the fit coefficients into the correct shape
        # for interpolation.
        # The dimensions of the interpolation parameters come first, the dimensions
        # of the hypersurface coefficients comes last.
        self.interp_shape = tuple(len(v["values"]) for v in self.interp_param_spec.values())
        # dimension is [interp_shape, binning..., fit coeffts]
        self._coeff_z = np.zeros(self.interp_shape + self.coeff_shape)
        # dimension is [interp_shape, binning..., fit coeffts, fit coeffts]
        self._covar_z = np.zeros(self.interp_shape + self.covars_shape)
        # Here we use the same indexing as below in `fit_hypersurfaces`
        for i, idx in enumerate(np.ndindex(self.interp_shape)):
            # As an additional safety measure, we check that the parameters are what
            # we expect to find at this index.
            expected_params = dict(
                (n, self.interp_param_spec[n]["values"][idx[j]])
                for j, n in enumerate(self.interp_param_spec.keys())
            )
            param_values = hs_fits[i]["param_values"]
            msg = ("The stored values where hypersurfaces were fit do not match those"
                   "in the interpolation grid.")
            assert np.all([expected_params[n].m == param_values[n].m
                           for n in self.interp_param_spec.keys()]), msg
            self._coeff_z[idx] = hs_fits[i]["hs_fit"].fit_coeffts
            self._covar_z[idx] = hs_fits[i]["hs_fit"].fit_cov_mat

        grid_coords = list(
            np.array([val.m for val in val_list["values"]])
            for val_list in self.interp_param_spec.values()
        )
        self.param_bounds = [(np.min(grid_vals), np.max(grid_vals))
                             for grid_vals in grid_coords]
        # If a parameter scales as log, we give the log of the parameter to the
        # interpolator. We must not forget to do this again when we call the
        # interpolator later!
        for i, param_name in enumerate(self.interpolation_param_names):
            if self.interp_param_spec[param_name]["scales_log"]:
                grid_coords[i] = np.log10(grid_coords[i])
        self.coefficients = interpolate.RegularGridInterpolator(
            grid_coords,
            self._coeff_z,
            # We disable extrapolation, but clip parameter values inside the valid
            # range.
            bounds_error=True, fill_value=None
        )
        self.covars = interpolate.RegularGridInterpolator(
            grid_coords,
            self._covar_z,
            bounds_error=True, fill_value=None
        )
        # In order not to spam warnings, we only want to warn about non positive
        # semi definite covariance matrices once for each bin. We store the bin
        # indeces for which the warning has already been issued.
        self.covar_bins_warning_issued = []
        self.ignore_nan = ignore_nan

    @property
    def interpolation_param_names(self):
        return list(self.interp_param_spec.keys())

    @property
    def param_names(self):
        return list(self._reference_state["params"].keys())

    @property
    def binning(self):
        binning = self._reference_state["binning"]
        if not is_binning(binning) :
            binning = MultiDimBinning(**binning)
        return binning

    @property
    def num_interp_params(self):
        return len(self.interp_param_spec.keys())

    def get_hypersurface(self, **param_kw):
        """
        Get a Hypersurface object with interpolated coefficients.

        Parameters
        ----------
        **param_kw
            Parameters are given as keyword arguments, where the names
            of the arguments must match the names of the parameters over
            which the hypersurfaces are interpolated. The values
            are given as :obj:`Quantity` objects with units.
        """
        assert set(param_kw.keys()) == set(self.interp_param_spec.keys()), "invalid parameters"
        # getting param magnitudes in the same units as the parameter specification
        x = np.array([
            param_kw[p].m_as(self.interp_param_spec[p]["values"][0].u)
            # we have checked that this is an OrderedDict so that the order of x is not
            # ambiguous here
            for p in self.interp_param_spec.keys()
        ])
        assert len(x) == len(self.param_bounds)
        for i, bounds in enumerate(self.param_bounds):
            x[i] = np.clip(x[i], *bounds)
        # if a parameter scales as log, we have to take the log here again
        for i, param_name in enumerate(self.interpolation_param_names):
            if self.interp_param_spec[param_name]["scales_log"]:
                # We must be strict with raising errors here, because otherwise
                # the Hypersurface will suddenly have NaNs everywhere! This shouldn't
                # happen because we clip values into the valid parameter range.
                if x[i] <= 0:
                    raise RuntimeError("A log-scaling parameter cannot become zero "
                                       "or negative!")
                x[i] = np.log10(x[i])

        state = copy.deepcopy(self._reference_state)
        # fit covariance matrices are stored directly in the state while fit coeffts
        # must be assigned with the setter method...
        # need squeeze here because the RegularGridInterpolator always puts another
        # dimension around the output
        state["fit_cov_mat"] = np.squeeze(self.covars(x))
        assert state["fit_cov_mat"].shape == self.covars_shape
        for idx in np.ndindex(state['fit_cov_mat'].shape):
            if self.ignore_nan: continue
            assert np.isfinite(state['fit_cov_mat'][idx]), ("invalid cov matrix "
                f"element encountered at {param_kw} in loc {idx}")
        # check covariance matrices for symmetry, positive semi-definiteness
        for bin_idx in np.ndindex(state['fit_cov_mat'].shape[:-2]):
            m = state['fit_cov_mat'][bin_idx]
            if self.ignore_nan and np.any(~np.isfinite(m)):
                state['fit_cov_mat'][bin_idx] = np.identity(m.shape[0])
                m = state['fit_cov_mat'][bin_idx]
            assert np.allclose(
                m, m.T, rtol=ALLCLOSE_KW['rtol']*10.), f'cov matrix not symmetric in bin {bin_idx}'
            if not matrix.is_psd(m):
                state['fit_cov_mat'][bin_idx] = matrix.fronebius_nearest_psd(m)
                if not bin_idx in self.covar_bins_warning_issued:
                    logging.warn(
                        f'Invalid covariance matrix fixed in bin: {bin_idx}')
                    self.covar_bins_warning_issued.append(bin_idx)
        hypersurface = Hypersurface.from_state(state)
        coeffts = np.squeeze(self.coefficients(x))  # calls interpolator
        assert coeffts.shape == self.coeff_shape
        # check that coefficients exist and if not replace with default values
        for idx in np.ndindex(self.coeff_shape):
            if self.ignore_nan and ~np.isfinite(coeffts[idx]):
                coeffts[idx] = 1 if idx[-1] == 0 else 0  # set intercept to 1, slopes 0
            assert np.isfinite(coeffts[idx]), ("invalid coeff encountered at "
                f"{param_kw} in loc {idx}")
        # the setter method defined in the Hypersurface class takes care of
        # putting the coefficients in the right place in their respective parameters
        hypersurface.fit_coeffts = coeffts
        return hypersurface

    def _make_slices(self, *xi):
        """Make slices of hypersurfaces for plotting.

        In some covariance matrices, the spline fits are corrected to make
        the matrix positive semi-definite. The slices produced by this function
        include all of those effects.

        Parameters
        ----------
        xi : list of ndarray
            Points at which the hypersurfaces are to be evaluated. The length of the
            list must equal the number of parameters, each ndarray in the list must have
            the same shape (slice_shape).

        Returns
        -------
        coeff_slices : numpy.ndarray
            slices in fit coefficients. Size: (binning..., number of coeffs) + slice_shape
        covar_slices : numpy.ndarray
            slices in covariance matrix elements.
            Size: (binning..., number of coeffs, number of coeffs) + slice_shape
        """
        slice_shape = xi[0].shape
        for x in xi:
            assert x.shape == slice_shape
        assert len(xi) == self.num_interp_params
        coeff_slices = np.zeros(self.coeff_shape + slice_shape)
        covar_slices = np.zeros(self.covars_shape + slice_shape)
        for idx in np.ndindex(slice_shape):
            pars = collections.OrderedDict()
            for i, name in enumerate(self.interpolation_param_names):
                pars[name] = xi[i][idx]
            hs = self.get_hypersurface(**pars)
            slice_idx = (Ellipsis,) + idx
            coeff_slices[slice_idx] = hs.fit_coeffts
            covar_slices[slice_idx] = hs.fit_cov_mat
        return coeff_slices, covar_slices

    def plot_fits_in_bin(self, bin_idx, ax=None, n_steps=20, **param_kw):
        """
        Plot the coefficients as well as covariance matrix elements as a function
        of the interpolation parameters.

        Parameters
        ----------
            bin_idx : tuple
                index of the bin for which to plot the fits
            ax : 2D array of axes, optional
                axes into which to place the plots. If None (default),
                appropriate axes will be generated. Must have at least
                size (n_coeff, n_coeff + 1).
            n_steps : int, optional
                number of steps to plot between minimum and maximum
            **param_kw :
                Parameters to be fixed when producing slices. If the interpolation
                is in N-D, then (N-2) parameters need to be fixed to produce 2D plots
                of the remaining 2 parameters and (N-1) need to be fixed to produce a
                1D slice.
        """
        plot_dim = self.ndim - len(param_kw.keys())
        assert plot_dim in [1, 2], "plotting only supported in 1D or 2D"
        import matplotlib.pyplot as plt
        n_coeff = self.coeff_shape[-1]
        hs_param_names = list(self._reference_state['params'].keys())
        hs_param_labels = ["intercept"] + [f"{p} p{i}" for p in hs_param_names
                                           for i in range(self._reference_state['params'][p]['num_fit_coeffts'])]

        fig = None
        if ax is None:
            fig, ax = plt.subplots(nrows=n_coeff, ncols=n_coeff+1,
                                   squeeze=False, sharex=True,
                                   figsize=(3+(5*(n_coeff+1)), 2+(3*n_coeff)) )
        # remember whether the plots need log scale or not, by default not
        x_is_log = False
        y_is_log = False

        # names of the variables we are plotting
        plot_names = set(self.interpolation_param_names) - set(param_kw.keys())
        if plot_dim == 1:
            x_name = list(plot_names)[0]
        else:
            x_name, y_name = list(plot_names)

        # in both 1D and 2D cases, we always plot at least an x-variable
        x_unit = self.interp_param_spec[x_name]["values"][0].u
        # we need the magnitudes here so that units are unambiguous when we make
        # the linspace/geomspace for plotting
        x_mags = [v.m_as(x_unit) for v in self.interp_param_spec[x_name]["values"]]
        if self.interp_param_spec[x_name]["scales_log"]:
            x_plot = np.geomspace(np.min(x_mags), np.max(x_mags), n_steps)
            x_is_log = True
        else:
            x_plot = np.linspace(np.min(x_mags), np.max(x_mags), n_steps)
        # we put the unit back later
        if plot_dim == 1:
            # To make slices, we need to set any variables we do not plot over to the
            # value given in param_kw.
            slice_args = []
            # We need to make sure that we give the values in the correct order!
            for n in self.interpolation_param_names:
                if n == x_name:
                    slice_args.append(x_plot * x_unit)
                elif n in param_kw.keys():
                    # again, insure that the same unit is used that went into the
                    # interpolation
                    param_unit = self.interp_param_spec[n]["values"][0].u
                    slice_args.append(
                        np.full(x_plot.shape, param_kw[n].m_as(param_unit)) * param_unit
                    )
                else:
                    raise ValueError("parameter neither specified nor plotted")
            coeff_slices, covar_slices = self._make_slices(*slice_args)
        else:
            # if we are in 2D, we need to do the same procedure again for the y-variable
            y_unit = self.interp_param_spec[y_name]["values"][0].u
            y_mags = [v.m_as(y_unit) for v in self.interp_param_spec[y_name]["values"]]
            if self.interp_param_spec[y_name]["scales_log"]:
                # we add one step to the size in y so that transposition is unambiguous
                y_plot = np.geomspace(np.min(y_mags), np.max(y_mags), n_steps + 1)
                y_is_log = True
            else:
                y_plot = np.linspace(np.min(y_mags), np.max(y_mags), n_steps + 1)

            x_mesh, y_mesh = np.meshgrid(x_plot, y_plot)
            slice_args = []
            for n in self.interpolation_param_names:
                if n == x_name:
                    slice_args.append(x_mesh * x_unit)
                elif n == y_name:
                    slice_args.append(y_mesh * y_unit)
                elif n in param_kw.keys():
                    # again, insure that the same unit is used that went into the
                    # interpolation
                    param_unit = self.interp_param_spec[n]["values"][0].u
                    slice_args.append(
                        np.full(x_mesh.shape, param_kw[n].m_as(param_unit)) * param_unit
                    )
                else:
                    raise ValueError("parameter neither specified nor plotted")
            coeff_slices, covar_slices = self._make_slices(*slice_args)

        # first column plots fit coefficients
        for i in range(n_coeff):
            z_slice = coeff_slices[bin_idx][i]
            if plot_dim == 1:
                ax[i, 0].plot(x_plot, z_slice, label='interpolation')
                # Plotting the original input points only works if the interpolation
                # is in 1D. If we are plotting a 1D slice from a 2D interpolation, this
                # does not work.
                # The number of fit points is the first dimension in self._coeff_z
                if plot_dim == self.ndim:
                    slice_idx = (Ellipsis,) + bin_idx + (i,)
                    ax[i, 0].scatter(x_mags, self._coeff_z[slice_idx],
                                     color='k', marker='x', label='fit points')
                ax[i, 0].set_ylabel(hs_param_labels[i])
            else:
                pc = ax[i, 0].pcolormesh(x_mesh, y_mesh, z_slice)
                cbar = plt.colorbar(pc, ax=ax[i, 0])
                cbar.ax.ticklabel_format(style='sci', scilimits=(0, 0))
                ax[i, 0].set_ylabel(y_name)
                ax[i, 0].set_xlabel(x_name)

            # later column plots the elements of the covariance matrix
            for j in range(0, n_coeff):
                z_slice = covar_slices[bin_idx][i, j]
                if plot_dim == 1:
                    ax[i, j+1].plot(x_plot, z_slice, label='interpolation')
                    # Same problem as above, only in 1D case can this be shown
                    # the number of points is the first dim in self._covar_z
                    if plot_dim == self.ndim:
                        coeff_idx = (Ellipsis,) + bin_idx + (i, j)
                        ax[i, j+1].scatter(x_mags, self._covar_z[coeff_idx],
                                           color='k', marker='x', label='fit points')
                else:
                    pc = ax[i, j+1].pcolormesh(x_mesh, y_mesh, z_slice)
                    cbar = plt.colorbar(pc, ax=ax[i, j+1])
                    cbar.ax.ticklabel_format(style='sci', scilimits=(0, 0))
                    ax[i, j+1].set_ylabel(y_name)
                    ax[i, j+1].set_xlabel(x_name)

        if plot_dim == 1:
            # in the 1D case, labels can be placed on the x and y axes
            for j in range(n_coeff+1):
                ax[-1, j].set_xlabel(x_name)
            ax[0, 0].set_title('coefficient')
            for j in range(n_coeff):
                ax[0, j+1].set_title(f'cov. {hs_param_labels[j]}')
        else:
            # in the 2D case, we need separate annotations
            rows = hs_param_labels
            cols = ["coefficient"] + [f"cov. {hl}" for hl in hs_param_labels]
            pad = 20
            for a, col in zip(ax[0], cols):
                a.annotate(col, xy=(0.5, 1), xytext=(0, pad),
                           xycoords='axes fraction', textcoords='offset points',
                           size='x-large', ha='center', va='baseline')

            for a, row in zip(ax[:, 0], rows):
                a.annotate(row, xy=(0, 0.5), xytext=(-a.yaxis.labelpad - pad, 0),
                           xycoords=a.yaxis.label, textcoords='offset points',
                           size='x-large', ha='right', va='center')
        for i, j in np.ndindex((n_coeff, n_coeff+1)):
            if x_is_log: ax[i, j].set_xscale("log")
            if y_is_log: ax[i, j].set_yscale("log")
            ax[i, j].grid()
            if plot_dim == 1:
                ax[i, j].legend()
            # ax[i, j].relim()
            # ax[i, j].autoscale_view()
            if not x_is_log:
                ax[i, j].ticklabel_format(style='sci', scilimits=(0, 0), axis="x")
            if not y_is_log:
                ax[i, j].ticklabel_format(style='sci', scilimits=(0, 0), axis="y")

        if fig is not None :
            fig.tight_layout()
            if plot_dim == 2:
                fig.subplots_adjust(left=0.15, top=0.95)
            return fig
        else :
            return

def pipeline_cfg_from_states(state_dict):
    """Recover a pipeline cfg containing PISA objects from a raw state.

    When a pipeline configuration is stored to JSON, the PISA objects turn into
    their serialized states. This function looks through the dictionary returned by
    `from_json` and recovers the PISA objects such as `ParamSet` and `MultiDimBinning`.

    It should really become part of PISA file I/O functionality to read and write
    PISA objects inside dictionaries/lists into a JSON and be able to recover
    them...
    """

    # TODO: Make this a core functionality of PISA

    # This is just a mess... some objects have a `from_state` method, some take the
    # unpacked state dict as input, some take the state...

    pipeline_cfg = collections.OrderedDict()
    for stage_key in state_dict.keys():
        # need to check all of this manually... no automatic way to do it :(
        if stage_key == "pipeline":
            pipeline_cfg[stage_key] = copy.deepcopy(state_dict[stage_key])
            pipeline_cfg[stage_key]["output_key"] = tuple(
                pipeline_cfg[stage_key]["output_key"])
            binning_state = pipeline_cfg[stage_key]["output_binning"]
            pipeline_cfg[stage_key]["output_binning"] = MultiDimBinning(**binning_state)
            continue
        # undo what we did in `serialize_pipeline_cfg` by splitting the keys into tuples
        tuple_key = tuple(stage_key.split("__"))
        pipeline_cfg[tuple_key] = copy.deepcopy(state_dict[stage_key])
        for k in ["calc_mode", "apply_mode", "node_mode"]:
            if k in pipeline_cfg[tuple_key]:
                if isinstance(pipeline_cfg[tuple_key][k], collections.Mapping):
                    pipeline_cfg[tuple_key][k] = MultiDimBinning(
                        **pipeline_cfg[tuple_key][k])
        if "params" in pipeline_cfg[tuple_key].keys():
            pipeline_cfg[tuple_key]["params"] = ParamSet(
                pipeline_cfg[tuple_key]["params"])
    # if any stage takes any other arguments that we didn't think of here, they
    # won't work
    return pipeline_cfg

def serialize_pipeline_cfg(pipeline_cfg):
    """Turn a pipeline configuration into something we can store to JSON.

    It doesn't work by default because tuples are not allowed as keys when storing to
    JSON. All we do is to turn the tuples into strings divided by a double underscore.
    """
    serializable_state = collections.OrderedDict()
    serializable_state["pipeline"] = pipeline_cfg["pipeline"]
    for k in pipeline_cfg.keys():
        if k == "pipeline": continue
        flat_key = "__".join(k)
        serializable_state[flat_key] = pipeline_cfg[k]
    # this isn't _really_ a serializable state, the objects are still PISA objects...
    # bit it will convert correctly when thrown into `to_json`
    return serializable_state

def assemble_interpolated_fits(fit_directory, output_file, drop_fit_maps=False):
    """After all of the fits on the cluster are done, assemble the results to one JSON.

    The JSON produced by this function is what `load_interpolated_hypersurfaces`
    expects.
    """
    assert os.path.isdir(fit_directory), "fit directory does not exist"
    metadata = from_json(os.path.join(fit_directory, "metadata.json"))

    combined_data = collections.OrderedDict()
    combined_data["interpolation_param_spec"] = metadata["interpolation_param_spec"]

    # Loop over grid points
    hs_fits = []
    grid_shape = tuple(metadata["grid_shape"])
    for job_idx, grid_idx in enumerate(np.ndindex(grid_shape)):

        # Load grid point data
        gridpoint_json = os.path.join(fit_directory, f"gridpoint_{job_idx:06d}.json.bz2")
        logging.info(f"Reading {gridpoint_json}")
        gridpoint_data = from_json(gridpoint_json)

        # Check the loaded data
        assert job_idx == gridpoint_data["job_idx"]
        assert np.all(grid_idx == gridpoint_data["grid_idx"])
        # TODO: Offer to run incomplete fits locally
        assert gridpoint_data["fit_successful"], f"job no. {job_idx} not finished"

        # Drop fit maps if requested (can significantly reduce file size)
        if drop_fit_maps :
            for key, hs_state in gridpoint_data["hs_fit"].items() :
                hs_state["fit_maps_raw"] = None
                hs_state["fit_maps_norm"] = None

        # Add grid point data to output file
        hs_fits.append(collections.OrderedDict(
            param_values=gridpoint_data["param_values"],
            hs_fit=gridpoint_data["hs_fit"]
        ))

    # Write the output file
    combined_data["hs_fits"] = hs_fits
    to_file(combined_data, output_file)

def get_incomplete_job_idx(fit_directory):
    """Get job indices of fits that are not flagged as successful."""

    assert os.path.isdir(fit_directory), "fit directory does not exist"
    metadata = from_json(os.path.join(fit_directory, "metadata.json"))
    grid_shape = tuple(metadata["grid_shape"])
    failed_idx = []
    for job_idx, grid_idx in enumerate(np.ndindex(grid_shape)):
        try:
            gridpoint_json = os.path.join(fit_directory,
                                          f"gridpoint_{job_idx:06d}.json.bz2")
            logging.info(f"Reading {gridpoint_json}")
            gridpoint_data = from_json(gridpoint_json)
        except:
            break
        if not gridpoint_data["fit_successful"]:
            failed_idx.append(job_idx)
        job_idx += 1
    return failed_idx

def run_interpolated_fit(fit_directory, job_idx, skip_successful=False):
    """Run the hypersurface fit for a grid point.

    If `skip_successful` is true, do not run if the `fit_successful` flag is already
    True.
    """

    #TODO a lot of this is copied from fit_hypersurfaces in hypersurface.py, would be safer to make more OAOO
    #TODO Copy the param value storage stuff from fit_hypersurfaces across in the meantime

    assert os.path.isdir(fit_directory), "fit directory does not exist"

    gridpoint_json = os.path.join(fit_directory, f"gridpoint_{job_idx:06d}.json.bz2")
    gridpoint_data = from_json(gridpoint_json)

    if skip_successful and gridpoint_data["fit_successful"]:
        logging.info(f"Fit at job index {job_idx} already successful, skipping...")
        return

    metadata = from_json(os.path.join(fit_directory, "metadata.json"))

    interpolation_param_spec = metadata["interpolation_param_spec"]

    # this is a pipeline configuration in the form of an OrderedDict
    nominal_dataset = metadata["nominal_dataset"]
    # Why can we still not load PISA objects from JSON that are inside a dict?! Grrr...
    nominal_dataset["pipeline_cfg"] = pipeline_cfg_from_states(
        nominal_dataset["pipeline_cfg"]
    )
    # this is a list of pipeline configurations
    sys_datasets = metadata["sys_datasets"]
    for sys_dataset in sys_datasets:
        sys_dataset["pipeline_cfg"] = pipeline_cfg_from_states(
            sys_dataset["pipeline_cfg"]
        )
    # this is a dict of param_name : value pairs
    param_values = gridpoint_data["param_values"]
    # we do a redundant check to make sure the parameter values at this grid point are
    # correct
    interpolation_param_names = metadata["interpolation_param_names"]
    grid_shape = tuple(metadata["grid_shape"])
    # the grid point index of this job
    grid_idx = list(np.ndindex(grid_shape))[job_idx]
    for i, n in enumerate(interpolation_param_names):
        ms = "Inconsistent parameter values at grid point!"
        assert interpolation_param_spec[n]["values"][grid_idx[i]] == param_values[n], ms

    # now we need to adjust the values of the parameter in all pipelines for this point
    logging.info(f"updating pipelines with parameter values: {param_values}")
    for dataset in [nominal_dataset] + sys_datasets:
        for stage_cfg in dataset["pipeline_cfg"].values():
            if "params" not in stage_cfg.keys(): continue
            for param in interpolation_param_names:
                if param in stage_cfg["params"].names:
                    stage_cfg["params"][param].value = param_values[param]

    # these are the parameters of the hypersurface, NOT the ones we interpolate them
    # over!
    hypersurface_params = []
    for param_state in metadata["hypersurface_params"]:
        hypersurface_params.append(HypersurfaceParam.from_state(param_state))

    def find_hist_stage(pipeline):
        """Locate the index of the hist stage in a pipeline."""
        hist_idx_found = False
        for i, s in enumerate(pipeline.stages):
            if s.__class__.__name__ == "hist":
                hist_idx = i
                hist_idx_found = True
                break
        if not hist_idx_found:
            raise RuntimeError("Could not find histogram stage in pipeline, aborting.")
        return hist_idx

    # We create Pipeline objects, get their outputs and then forget about the Pipeline
    # object on purpose! The memory requirement to hold all systematic sets at the same
    # time is just too large, especially on the cluster. The way we do it below we
    # only need enough memory for one dataset at a time.

    for dataset in [nominal_dataset] + sys_datasets:
        pipeline = Pipeline(dataset["pipeline_cfg"])
        dataset["mapset"] = pipeline.get_outputs()
        # get the un-weighted event counts as well so that we can exclude bins
        # with too little statistics
        # First, find out which stage is the hist stage
        hist_idx = find_hist_stage(pipeline)
        pipeline.stages[hist_idx].unweighted = True
        dataset["mapset_unweighted"] = pipeline.get_outputs()
    del pipeline

    # Merge maps according to the combine regex, if one was provided
    combine_regex = metadata["combine_regex"]
    if combine_regex is not None:
        for dataset in [nominal_dataset] + sys_datasets:
            dataset["mapset"] = dataset["mapset"].combine_re(combine_regex)
            dataset["mapset_unweighted"] = dataset["mapset_unweighted"].combine_re(combine_regex)

    minimum_mc = metadata["minimum_mc"]
    # Remove bins (i.e. set their count to zero) that have too few MC events
    for dataset in sys_datasets + [nominal_dataset]:
        for map_name in dataset["mapset"].names:
            insuff_mc = dataset["mapset_unweighted"][map_name].nominal_values < minimum_mc
            # Setting the hist to zero sets both nominal value and std_dev to zero
            dataset["mapset"][map_name].hist[insuff_mc] = 0.

    hypersurface_fit_kw = metadata["hypersurface_fit_kw"]
    hypersurfaces = collections.OrderedDict()
    log = metadata["log"]  # flag determining whether hs fit is run in log-space or not
    for map_name in nominal_dataset["mapset"].names:
        nominal_map = nominal_dataset["mapset"][map_name]
        nominal_param_values = nominal_dataset["sys_params"]

        sys_maps = [sys_dataset["mapset"][map_name] for sys_dataset in sys_datasets]
        sys_param_values = [sys_dataset["sys_params"] for sys_dataset in sys_datasets]

        hypersurface = Hypersurface(
            # Yes, this MUST be a deepcopy! Otherwise weird memory overwrites happen
            # and all the numbers get jumbled across the hypersurfaces of different maps
            params=copy.deepcopy(hypersurface_params),
            initial_intercept=0. if log else 1.,  # Initial value for intercept
            log=log
        )

        hypersurface.fit(
            nominal_map=nominal_map,
            nominal_param_values=nominal_param_values,
            sys_maps=sys_maps,
            sys_param_values=sys_param_values,
            norm=True,
            # Is the space or loading time really a problem?
            # keep_maps=False,  # it would take a lot more space otherwise
            **hypersurface_fit_kw
        )

        logging.debug("\nFitted hypersurface report:\n%s" % hypersurface)
        hypersurfaces[map_name] = hypersurface

    gridpoint_data["hs_fit"] = hypersurfaces
    gridpoint_data["fit_successful"] = True

    to_json(gridpoint_data, gridpoint_json)

def prepare_interpolated_fit(
    nominal_dataset, sys_datasets, params, fit_directory, interpolation_param_spec,
    combine_regex=None, log=False, minimum_mc=0, **hypersurface_fit_kw
):
    '''
    Writes steering files for fitting hypersurfaces on a grid of arbitrary parameters.
    The fits can then be run on a cluster with `run_interpolated_fit`.

    Parameters
    ----------
    nominal_dataset : dict
        Definition of the nominal dataset. Specifies the pipleline with which the maps
        can be created, and the values of all systematic parameters used to produced the
        dataset.
        Format must be:
            nominal_dataset = {
                "pipeline_cfg" = <pipeline cfg file (either cfg file path or dict)>),
                "sys_params" = { param_0_name : param_0_value_in_dataset, ..., param_N_name : param_N_value_in_dataset }
            }
        Sys params must correspond to the provided HypersurfaceParam instances provided
        in the `params` arg.

    sys_datasets : list of dicts
        List of dicts, where each dict defines one of the systematics datasets to be
        fitted. The format of each dict is the same as explained for `nominal_dataset`

    params : list of HypersurfaceParams
        List of HypersurfaceParams instances that define the hypersurface. Note that
        this defined ALL hypersurfaces fitted in this function, e.g. only supports a
        single parameterisation for all maps (this is almost always what you want).

    output_directory : str
        Directory in which the fits will be run. Steering files for the fits to be run
        will be stored here.

    combine_regex : list of str, or None
        List of string regex expressions that will be used for merging maps. Used to
        combine similar species. Must be something that can be passed to the
        `MapSet.combine_re` function (see that functions docs for more details). Choose
        `None` is do not want to perform this merging.

    interpolation_param_spec : collections.OrderedDict
        Specification of parameter grid that hypersurfaces should be interpolated over.
        The dict should have the following form::
            interpolation_param_spec = {
                'param1': {"values": [val1_1, val1_2, ...], "scales_log": True/False}
                'param2': {"values": [val2_1, val2_2, ...], "scales_log": True/False}
                ...
                'paramN': {"values": [valN_1, valN_2, ...], "scales_log": True/False}
            }
        The hypersurfaces will be fit on an N-dimensional rectilinear grid over
        parameters 1 to N. The flag `scales_log` indicates that the interpolation over
        that parameter should happen in log-space.

    minimum_mc : int, optional
        Minimum number of un-weighted MC events required in each bin.

    hypersurface_fit_kw : kwargs
        kwargs will be passed on to the calls to `Hypersurface.fit`
    '''

    # Take (deep) copies of lists/dicts to avoid modifying the originals
    # Useful for cases where this function is called in a loop (e.g. leave-one-out tests)
    nominal_dataset = copy.deepcopy(nominal_dataset)
    sys_datasets = copy.deepcopy(sys_datasets)
    params = copy.deepcopy(params)

    # Check types
    assert isinstance(sys_datasets, collections.Sequence)
    assert isinstance(params, collections.Sequence)
    assert isinstance(fit_directory, str)
    # there must not be any ambiguity between fitting the hypersurfaces and
    # interpolating them later
    msg = "interpolation params must be specified as a dict with ordered keys"
    assert isinstance(interpolation_param_spec, collections.OrderedDict), msg
    for k, v in interpolation_param_spec.items():
        assert set(v.keys()) == {"values", "scales_log"}
        assert isinstance(v["values"], collections.Sequence)
        # We need to extract the magnitudes from the Quantities to avoid a
        # UnitStrippedWarning. For some reason, doing `np.min(v["values"])` messes up
        # the data structure inside the values in a way that can cause a crash when we
        # try to serialize the values later. Lesson: Stripping units inadvertently can
        # have strange, unforeseen consequences.
        mags = [x.m for x in v["values"]]
        if v["scales_log"] and np.min(mags) <= 0:
            raise ValueError("A log-scaling parameter cannot be equal to or less "
                "than zero!")

    # Check output format and path
    assert os.path.isdir(fit_directory), "fit directory does not exist"

    # Check formatting of datasets is as expected
    all_datasets = [nominal_dataset] + sys_datasets
    for dataset in all_datasets:
        assert isinstance(dataset, collections.Mapping)
        assert "pipeline_cfg" in dataset
        assert isinstance(dataset["pipeline_cfg"], (str, collections.Mapping))
        assert "sys_params" in dataset
        assert isinstance(dataset["sys_params"], collections.Mapping)

        dataset["pipeline_cfg"] = serialize_pipeline_cfg(dataset["pipeline_cfg"])

    # Check params
    assert len(params) >= 1
    for p in params:
        assert isinstance(p, HypersurfaceParam)

    # Report inputs
    msg = "Hypersurface fit details :\n"
    msg += f"  Num params            : {len(params)}\n"
    msg += f"  Num fit coefficients  : {sum([p.num_fit_coeffts for p in params])}\n"
    msg += f"  Num datasets          : 1 nominal + {len(sys_datasets)} systematics\n"
    msg += f"  Nominal values        : {nominal_dataset['sys_params']}\n"
    msg += "Hypersurface fits are prepared on the following grid:\n"
    msg += str(interpolation_param_spec)
    logging.info(msg)

    # because we require this to be an OrderedDict, there is no ambiguity in the
    # construction of the mesh here
    param_names = list(interpolation_param_spec.keys())
    grid_shape = tuple(len(v["values"]) for v in interpolation_param_spec.values())

    # We store all information needed to run a fit in metadata
    metadata = collections.OrderedDict(
        interpolation_param_spec=interpolation_param_spec,
        interpolation_param_names=param_names,  # convenience
        grid_shape=grid_shape,  # convenience
        nominal_dataset=nominal_dataset,
        sys_datasets=sys_datasets,
        hypersurface_params=params,
        combine_regex=combine_regex,
        log=log,
        minimum_mc=minimum_mc,
        hypersurface_fit_kw=hypersurface_fit_kw
    )

    to_json(metadata, os.path.join(fit_directory, "metadata.json"))

    # we write on JSON file for each grid point
    for job_idx, grid_idx in enumerate(np.ndindex(grid_shape)):
        # Although this is technically redundant, we store the parameter values
        # explicitly for each grid point.
        param_values = {}
        for i, n in enumerate(param_names):
            param_values[n] = interpolation_param_spec[n]["values"][grid_idx[i]]

        gridpoint_data = {
            "param_values": param_values,
            "hs_fit": None,
            "job_idx": job_idx,
            "grid_idx": grid_idx,
            "fit_successful": False
        }
        to_json(gridpoint_data, os.path.join(fit_directory,
            f"gridpoint_{job_idx:06d}.json.bz2"))

    logging.info(f"Grid fit preparation complete! Total number of jobs: {job_idx+1}")
    return job_idx+1  # zero-indexing

def load_interpolated_hypersurfaces(input_file):
    '''
    Load a set of interpolated hypersurfaces from a file.

    Analogously to "load_hypersurfaces", this function returns a
    collection with a HypersurfaceInterpolator object for each Map.

    Parameters
    ----------
    input_file : str
        A JSON input file as produced by fit_hypersurfaces if interpolation params
        were given. It has the form::
            {
                interpolation_param_spec = {
                    'param1': {"values": [val1_1, val1_2, ...], "scales_log": True/False}
                    'param2': {"values": [val2_1, val2_2, ...], "scales_log": True/False}
                    ...
                    'paramN': {"values": [valN_1, valN_2, ...], "scales_log": True/False}
                },
                'hs_fits': [
                    <list of dicts where keys are map names such as 'nue_cc' and values
                    are hypersurface states>
                ]
            }

    Returns
    -------
    collections.OrderedDict
        dictionary with a :obj:`HypersurfaceInterpolator` for each map
    '''
    assert isinstance(input_file, str)

    logging.info(f"Loading interpolated hypersurfaces from file: {input_file}")

    # Load the data from the file
    input_data = from_file(input_file)

    # check the file contents
    assert set(['interpolation_param_spec', 'hs_fits']).issubset(
        set(input_data.keys())), 'missing keys'

    # input_data['hs_fits'] is a list of dicts, each dict contains "param_values"
    # and "hs_fit"
    map_names = None
    logging.info("Reading file complete, generating hypersurfaces...")
    for hs_fit_dict in input_data['hs_fits']:
        # this is still not the actual Hypersurface, but a dict with the (linked)
        # maps and the HS fit for the map...
        hs_state_maps = hs_fit_dict["hs_fit"]
        if map_names is None:
            map_names = list(hs_state_maps.keys())
        else:
            assert set(map_names) == set(hs_state_maps.keys()), "inconsistent maps"
        # When data is recovered from JSON, the object states are not automatically
        # converted to the corresponding objects, so we need to do it manually here.
        for map_name in map_names:
            hs_state_maps[map_name] = Hypersurface.from_state(hs_state_maps[map_name])

    logging.info(f"Read hypersurface maps: {map_names}")

    # Now we have a list of dicts where the map names are on the lower level.
    # We need to convert this into a dict of HypersurfaceInterpolator objects.
    output = collections.OrderedDict()
    for m in map_names:
        hs_fits = [{"param_values": fd["param_values"], "hs_fit": fd['hs_fit'][m]} for fd in input_data['hs_fits']]
        output[m] = HypersurfaceInterpolator(input_data['interpolation_param_spec'], hs_fits)

    return output


# -*- coding: utf-8 -*-
"""
Created on Mon Apr  1 20:20:31 2019
基于Keras 的条件随机场分词算法
@author: 李畅
"""

from keras.models import Sequential
from keras.layers import Embedding,Bidirectional,LSTM,BatchNormalization
from keras.preprocessing.sequence import pad_sequences
from keras import optimizers
from keras_contrib.layers import CRF
from keras.callbacks import TensorBoard
import numpy as np
from os import path,rename,remove
import pickle

mcrf_datas='./mcrf.data' # 以二进制保存的mcrf类数据

class MCRF:
    def __init__(self,train_file,validate_file,model_file=\
            './crf_model.h5',embedding_dim=128,birnn_units=128,feat_filt_freq=3,\
            train_max_len=64,test_max_len=1024,batch_size=64,epochs=10):
        """
        @param train_file : 训练文件
        @param validate_file : 验证文件
        @param model_file: 模型文件,h5类型文件
        @param embedding_dim: 隐层维度
        @param birnn_units: RNN 单元个数
        @param feat_filt_freq: 将出现频次低于该值的特征过滤掉
        @param train_max_len: 训练集中的最大长度
        @param test_max_len: 测试文件一行最长的长度
        @param batch_size: 训练过程中的批次大小
        @param epochs: 训练的轮次
        """
        self.train_file=train_file
        self.validate_file = validate_file
        self.model_file=model_file
        self.feat_filt_freq=feat_filt_freq
        self.embedding_dim=embedding_dim
        self.birnn_units=birnn_units
        self.test_max_len=test_max_len
        self.batch_size=batch_size
        self.epochs=epochs
        self.tags=['B','M','E','S'] # 字符所有可能的标记
        self.max_len=train_max_len # 最长的一段文字的长度
        self.vocab=[]
        (self.train_x,self.train_y),(self.test_x,self.test_y)=self._preprocess()
        pass

    def _process_data(self, data):
        """
        数据处理
        """
        word_index=dict((w,i) for i,w in enumerate(self.vocab))
        x=[[word_index.get(w[0], 1) for w in s] for s in data] # 对未登录词，index=0
        y_chunk=[[self.tags.index(w[1]) for w in s] for s in data]
        x= pad_sequences(x,self.max_len) # 在 x 的左侧填充0，使得各个sample的长度相同
        y_chunk=pad_sequences(y_chunk,self.max_len,value=-1) # 左侧填充-1
        y_chunk=np.expand_dims(y_chunk,2) # 扩充维度，每个标记为一个list[]
        return x,y_chunk

    def _preprocess(self):
        """
        数据预处理
        """
        feat_freqs={}
        data=[] # 二维矩阵，一行为一段文字
        with open(self.train_file,encoding='utf-8',mode='r') as fr:
            sample_len=0 # 记录当前一段文字的长度
            one_sample=[]
            for line in fr.readlines():
                line=line.strip()
                if not line or len(line) < 3:
                    sample_len=0
                    data.append(one_sample)
                    one_sample=[]
                    continue
                sample_len +=1
                word,tag=line.split()
                one_sample.append([word,tag])
                if (word,tag) not in feat_freqs.keys():
                    feat_freqs.update({(word,tag):0})
                feat_freqs[(word,tag)] +=1
        # 特征值过滤
        self.vocab=[key[0] for key,val in feat_freqs.items() if val > \
                    self.feat_filt_freq]
        train_datas = self._process_data(data)
        del data,feat_freqs
        data=[]
        #读取测试文件
        with open(self.validate_file,encoding='utf-8',mode='r') as fr:
            one_example=[]
            for line in fr.readlines():
                line=line.strip()
                if not line or len(line) < 3:
                    data.append(one_example)
                    one_example=[]
                    continue
                word,tag = line.split()
                one_example.append([word,tag])
        test_datas=self._process_data(data)
        return train_datas,test_datas

    def _create_model(self):
        """
        创建训练使用的模型
        """
        model=Sequential()
        model.add(Embedding(len(self.vocab),self.embedding_dim,mask_zero=True))
        model.add(Bidirectional(LSTM(self.birnn_units//2, return_sequences=True)))
        crf=CRF(len(self.tags),sparse_target=True)
        model.add(crf)
        model.summary()
#        model.compile('adam',loss=crf.loss_function,metrics=[crf.accuracy])
        rms_prop=optimizers.RMSprop(lr=0.01,decay=1e-4)
        model.compile(optimizer=rms_prop,loss=crf.loss_function,metrics=\
                      [crf.accuracy])
        return model

    def train(self,model_exists=False):
        """
        模型训练
        """
        model=self._create_model()
        if not model_exists:
            # 创建tensorboard 回调
            tb_callback=TensorBoard(log_dir='./logs',
                    histogram_freq=0,
                    batch_size=64,
                    write_graph=True,
                    write_grads=True,
                    write_images=True,
                    embeddings_freq=0,
                    embeddings_layer_names=None,
                    embeddings_metadata=None)
            model.fit(self.train_x, self.train_y, batch_size=self.batch_size,epochs=self.epochs,\
                 validation_data=[self.test_x,self.test_y],callbacks=[tb_callback])
            model.save(self.model_file)
        else:
            model.load_weights(self.model_file)
        return model

    def _predict_str_preprocess(self,string):
        """
        对待处理的字符串预处理
        """
        word_index=dict((w,i) for i,w in enumerate(self.vocab))
        x=[word_index.get(w[0], 1) for w in string]
        length = len(x)
        x=pad_sequences([x],self.test_max_len)
        return x,length

    def predict(self,string,model):
        """
        对序列的标记进行预测
        @param string: 待处理的一串字符
        @param model: 模型
        """
        string,length=self._predict_str_preprocess(string)
        raw=model.predict(string)[0][-length:]
        result=[np.argmax(row) for row in raw]
        result_tags=[self.tags[i] for i in result]
        return result,result_tags

    def predict_file(self,testfile):
        """
        对测试文件4词位标记
        """
        if not path.exists(self.model_file):
            model = self.train(model_exists=False)
        else:
            model = self.train(model_exists=True)
        tmpfile=path.join(path.dirname(testfile),'tmp')
        with open(testfile,encoding='utf-8',mode='r') as fr, open(tmpfile,\
                 encoding='utf-8',mode='w') as fw:
            prog=0
            for line in fr.readlines():
                prog +=1
                if prog % 100==0:
                    print('processing line{}'.format(prog))
                line=line.strip()
                if not line:
                    print('',file=fw)
                    continue
                _, tags=self.predict(line,model)
                line=list(line)
                # 对该行文字分词
                ans=[]
                tmpstr=''
                for i in range(len(line)):
                    if tags[i]=='S':
                        tmpstr=''
                        ans.append(line[i])
                    elif tags[i] == 'B':
                        tmpstr=''
                        tmpstr=line[i]
                    elif tags[i] == 'M':
                        tmpstr += line[i]
                    else:
                        tmpstr += line[i]
                        ans.append(tmpstr)
                        tmpstr=''
                print('  '.join(ans),file=fw)
        remove(testfile)
        rename(tmpfile,testfile)
        pass

if not path.exists(mcrf_datas):
    mcrf=MCRF('./pku_training.utf8','./pku_validate.utf8')
else:
    with open(mcrf_datas,mode='rb') as fr:
        mcrf=pickle.load(fr)
test_file='./pku_test.utf8'
mcrf.predict_file(test_file)


"""### FastJet Tools

The [FastJet package](http://fastjet.fr/) provides, among other things, fast
jet clustering utilities. It is written in C++ and includes a Python interface
that is easily installed at compile time by passing the `--enable-pyext` flag
to `configure`. If you use this module for published research, please [cite
FastJet appropriately](http://fastjet.fr/about.html).

The core of EnergyFlow does not rely on FastJet, and hence it is not required
to be installed, but the following utilities are available assuming that
`import fastjet` succeeds in your Python environment (if not, no warnings or
errors will be issued but this module will not be usable).
"""
from __future__ import absolute_import, division, print_function

import numpy as np

from energyflow.utils.generic_utils import import_fastjet

fj = import_fastjet()

__all__ = []

if fj:

    __all__ = ['pjs_from_ptyphims', 'ptyphims_from_pjs', 'cluster', 'softdrop']

    def pjs_from_ptyphims(ptyphims):
        """Converts particles in hadronic coordinates to FastJet PseudoJets.

        **Arguments**

        - **ptyphims** : _2d numpy.ndarray_
            - An array of particles in hadronic coordinates.

        **Returns**

        - _list_ of _fastjet.PseudoJet_
            - A list of PseudoJets corresponding to the particles in the given
            array.
        """

        pjs = []
        for ptyphim in ptyphims:
            pj = fj.PseudoJet()
            pj.reset_PtYPhiM(*ptyphim[:4])
            pjs.append(pj)

        return pjs

    def ptyphims_from_pjs(pjs, mass=True):
        """Extracts hadronic four-vectors from FastJet PseudoJets.

        **Arguments**

        - **pjs** : _list_ of _fastjet.PseudoJet_
            - An iterable of PseudoJets.
        - **mass** : _bool_
            - Whether or not to include the mass in the extracted four-vectors.

        **Returns**

        - _numpy.ndarray_
            - An array of four-vectors corresponding to the given PseudoJets as
            `[pT, y, phi, m]`, where the mass is optional.
        """

        if mass:
            return np.asarray([[pj.pt(), pj.rap(), pj.phi(), pj.m()] for pj in pjs])
        else:
            return np.asarray([[pj.pt(), pj.rap(), pj.phi()] for pj in pjs])

    def cluster(pjs, algorithm='ca', R=fj.JetDefinition.max_allowable_R):
        """Clusters a list of PseudoJets according to a specified jet
        algorithm and jet radius.

        **Arguments**

        - **pjs** : _list_ of _fastjet.PseudoJet_
            - A list of Pseudojets representing particles or other kinematic
            objects that are to be clustered into jets.
        - **algorithm** : {'kt', 'antikt', 'ca', 'cambridge', 'cambridge_aachen'}
            - The jet algorithm to use during the clustering. Note that the
            last three options all refer to the same strategy and are provided
            because they are all used by the FastJet Python package.
        - **R** : _float_
            - The jet radius. The default value corresponds to
            `max_allowable_R` as defined by the FastJet python package.

        **Returns**

        - _list_ of _fastjet.PseudoJet_
            - A list of PseudoJets corresponding to the clustered jets.
        """

        algorithm_l = algorithm.lower()
        if algorithm_l  == 'kt':
            jet_alg = fj.kt_algorithm
        elif algorithm_l == 'antikt':
            jet_alg = fj.antikt_algorithm
        elif algorithm_l in {'ca', 'cambridge', 'cambridge_aachen'}:
            jet_alg = fj.cambridge_algorithm
        else:
            raise ValueError("algorithm '{}' not understood".format(algorithm))

        return fj.JetDefinition(jet_alg, R)(pjs)

    def softdrop(jet, zcut=0.1, beta=0, R=1.0):
        r"""Implements the SoftDrop grooming algorithm on a jet that has been
        found via clustering. Specifically, given a jet, it is recursively
        declustered and the softer branch removed until the SoftDrop condition
        is satisfied:

        $$
        \frac{\min(p_{T,1},p_{T,2})}{p_{T,1}+p_{T,2}} > z_{\rm cut}
        \left(\frac{\Delta R_{12}}{R}\right)^\beta
        $$

        where $1$ and $2$ refer to the two PseudoJets declustered at this stage.
        See the [SoftDrop paper](https://arxiv.org/abs/1402.2657) for a
        complete description of SoftDrop. If you use this function for your
        research, please cite [1402.2657](https://doi.org/10.1007/
        JHEP05(2014)146).

        **Arguments**

        - **jet** : _fastjet.PseudoJet_
            - A FastJet PseudoJet that has been obtained from a suitable
            clustering (typically Cambridge/Aachen).
        - **zcut** : _float_
            - The $z_{\rm cut}$ parameter of SoftDrop. Should be between `0`
            and `1`.
        - **beta** : _int_ or _float_
            - The $\beta$ parameter of SoftDrop.
        - **R** : _float_
            - The jet radius to use for the grooming. Only relevant if `beta!=0`.

        **Returns**

        - _fastjet.PseudoJet_
            - The groomed jet. Note that it will not necessarily have all of
            the same associated structure as the original jet, but it is
            suitable for obtaining kinematic quantities, e.g. [$z_g$](/docs/
            obs/#zg_from_pj).
        """

        parent1, parent2 = fj.PseudoJet(), fj.PseudoJet()
        if not jet.has_parents(parent1, parent2):
            return jet

        pt1, pt2 = parent1.pt(), parent2.pt()
        z = min(pt1, pt2)/(pt1 + pt2)

        if z >= (zcut if beta == 0 else zcut * (parent1.delta_R(parent2)/R)**beta):
            return jet
        else:
            return softdrop(parent1 if pt1 >= pt2 else parent2, zcut=zcut, beta=beta, R=R)


import ccllib as lib
import numpy as np
from warnings import warn
from pyutils import check

# Configuration types
transfer_function_types = {
    'none':             lib.none,
    'emulator':         lib.emulator,
    'fitting_function': lib.fitting_function,
    'eisenstein_hu':    lib.eisenstein_hu,
    'bbks':             lib.bbks,
    'boltzmann':        lib.boltzmann,
    'boltzmann_camb':   lib.boltzmann_camb,
    'camb':             lib.boltzmann_camb,
    'boltzmann_class':  lib.boltzmann_class,
    'class':            lib.boltzmann_class
}

matter_power_spectrum_types = {
    'halo_model':   lib.halo_model,
    'halomodel':    lib.halo_model,
    'halofit':      lib.halofit,
    'linear':       lib.linear
}

mass_function_types = {
    'angulo':   lib.angulo,
    'tinker':   lib.tinker,
    'tinker10': lib.tinker10,
    'watson':   lib.watson
}

# Error types
error_types = {
    lib.CCL_ERROR_MEMORY:       'CCL_ERROR_MEMORY',
    lib.CCL_ERROR_LINSPACE:     'CCL_ERROR_LINSPACE',
    lib.CCL_ERROR_INCONSISTENT: 'CCL_ERROR_INCONSISTENT',
    lib.CCL_ERROR_SPLINE:       'CCL_ERROR_SPLINE',
    lib.CCL_ERROR_SPLINE_EV:    'CCL_ERROR_SPLINE_EV',
    lib.CCL_ERROR_INTEG:        'CCL_ERROR_INTEG',
    lib.CCL_ERROR_ROOT:         'CCL_ERROR_ROOT',
    lib.CCL_ERROR_CLASS:        'CCL_ERROR_CLASS',
    lib.CCL_ERROR_COMPUTECHI:   'CCL_ERROR_COMPUTECHI',
    lib.CCL_ERROR_MF:           'CCL_ERROR_MF',
    lib.CCL_ERROR_HMF_INTERP:   'CCL_ERROR_HMF_INTERP',
    lib.CCL_ERROR_PARAMETERS:   'CCL_ERROR_PARAMETERS',
}

class Parameters(object):
    """The Parameters class contains cosmological parameters.

    """

    def __init__(self, Omega_c=None, Omega_b=None, h=None, A_s=None, n_s=None,
                 Omega_k=0., Omega_n=0., w0=-1., wa=0., sigma8=None,
                 z_mg=None, df_mg=None):
        """
        Creates a set of cosmological parameters.

        Note:
            Although some arguments default to `None`, they will raise a
            ValueError inside this function if not specified, so they are not
            optional.

        Args:
            Omega_c (float): Cold dark matter density fraction.
            Omega_b (float): Baryonic matter density fraction.
            h (float): Hubble constant divided by 100 km/s/Mpc; unitless.
            A_s (float): Power spectrum normalization. Optional if sigma8 is
                         specified.
            n_s (float): Primordial scalar perturbation spectral index.
            Omega_k (float, optional): Curvature density fraction. Defaults to 0.
            Omega_n (float, optional): Massless neutrino density fraction.
                                       Defaults to 0.
            w0 (float, optional): First order term of dark energy equation of
                                  state. Defaults to -1.
            wa (float, optional): Second order term of dark energy equation of
                                  state. Defaults to 0.
            sigma8 (float): Variance of matter density perturbations at 8 Mpc/h
                            scale. Optional if A_s is specified.
            df_mg (:obj: array_like): Perturbations to the GR growth rate as a
                                      function of redshift, Delta f. Used to
                                      implement simple modified growth
                                      scenarios.
            z_mg (:obj: array_like): Array of redshifts corresponding to df_mg.

        """
        # Set current ccl_parameters object to None
        self.parameters = None

         # Set nz_mg (no. of redshift bins for modified growth fns.)
        if z_mg is not None and df_mg is not None:
            # Get growth array size and do sanity check
            z_mg = np.atleast_1d(z_mg)
            df_mg = np.atleast_1d(df_mg)
            assert z_mg.size == df_mg.size
            nz_mg = z_mg.size
        else:
            # If one or both of the MG growth arrays are set to zero, disable
            # all of them
            if z_mg is not None or df_mg is not None:
                raise ValueError("Must specify both z_mg and df_mg.")
            z_mg = None
            df_mg = None
            nz_mg = -1

        # Check to make sure specified amplitude parameter is consistent
        if (A_s is None and sigma8 is None) \
        or (A_s is not None and sigma8 is not None):
            raise ValueError("Must set either A_s or sigma8.")

        # Set norm_pk to either A_s or sigma8
        norm_pk = A_s if A_s is not None else sigma8

        # The C library decides whether A_s or sigma8 was the input parameter
        # based on value, so we need to make sure this is consistent too
        if norm_pk >= 1e-5 and A_s is not None:
            raise ValueError("A_s must be less than 1e-5.")

        if norm_pk < 1e-5 and sigma8 is not None:
            raise ValueError("sigma8 must be greater than 1e-5.")

        # Check if any compulsory parameters are not set
        compul = [Omega_c, Omega_b, Omega_k, Omega_n, w0, wa, h, norm_pk, n_s]
        names = ['Omega_c', 'Omega_b', 'Omega_k', 'Omega_n', 'w0', 'wa',
                 'h', 'norm_pk', 'n_s']
        for nm, item in zip(names, compul):
            if item is None:
                raise ValueError("Necessary parameter '%s' was not set "
                                 "(or set to None)." % nm)

        # Create new instance of ccl_parameters object
        if nz_mg == -1:
            # Create ccl_parameters without modified growth
            self.parameters = lib.parameters_create(
                                    Omega_c, Omega_b, Omega_k, Omega_n,
                                    w0, wa, h, norm_pk, n_s,
                                    -1, None, None)
        else:
            # Create ccl_parameters with modified growth arrays
            self.parameters = lib.parameters_create_vec(
                                    Omega_c, Omega_b, Omega_k, Omega_n,
                                    w0, wa, h, norm_pk, n_s,
                                    z_mg, df_mg)

    def __getitem__(self, key):
        """Access parameter values by name.

        """
        try:
            val = getattr(self.parameters, key)
        except AttributeError:
            raise KeyError("Parameter '%s' not recognized." % key)
        return val

    def __setitem__(self, key, val):
        """Set parameter values by name.

        """
        raise NotImplementedError("Parameters objects are immutable; create a "
                                  "new Parameters() instance instead.")

        try:
            # First check if the key already exists (otherwise the parameter
            # would be silently added to the ccl_parameters class instance)
            getattr(self.parameters, key)
        except AttributeError:
            raise KeyError("Parameter '%s' not recognized." % key)

        # Set value of parameter
        setattr(self.parameters, key, val)
        # TODO: Should update/replace CCL objects appropriately

    def __str__(self):
        """Output the parameters that were set, and their values.

        """
        params = ['Omega_c', 'Omega_b', 'Omega_m', 'Omega_n', 'Omega_k',
                  'w0', 'wa', 'H0', 'h', 'A_s', 'n_s', 'Omega_g', 'T_CMB',
                  'sigma_8', 'Omega_l', 'z_star', 'has_mgrowth']

        vals = ["%15s: %s" % (p, getattr(self.parameters, p)) for p in params]
        string = "Parameters\n----------\n"
        string += "\n".join(vals)
        return string

class Cosmology(object):
    """Wrapper for the ccl_cosmology object.

    Includes cosmological parameters and cached data.

    """

    def __init__(self,
                 params=None, config=None,
                 Omega_c=None, Omega_b=None, h=None, A_s=None, n_s=None,
                 Omega_k=0., Omega_n=0., w0=-1., wa=0., sigma8=None,
                 z_mg=None, df_mg=None,
                 transfer_function='boltzmann_class',
                 matter_power_spectrum='halofit',
                 mass_function='tinker'):
        """Creates a wrapper for ccl_cosmology.

        TODO: enumerate transfer_function and
        matter_power_spectrum options.

        Args:
            params (:obj:`Parameters`): Cosmological parameters object.
            config (:obj:`ccl_configuration`, optional): Configuration for how
            to use CCL. Takes precident over any other passed in configuration.
            Defaults to None.
            transfer_function (:obj:`str`, optional): The transfer function to
            use. Defaults to `boltzmann_class`.
            matter_power_spectrum (:obj:`str`, optional): The matter power
            spectrum to use. Defaults to `halofit`.
            mass_function (:obj:`str`, optional): The mass function to use.
            Defaults to `tinker` (2010).

        """

        # Use either input cosmology parameters or Parameters() object
        if params is None:
            # Create new Parameters object
            params = Parameters(Omega_c=Omega_c, Omega_b=Omega_b, h=h, A_s=A_s,
                                n_s=n_s, Omega_k=Omega_k, Omega_n=Omega_n,
                                w0=w0, wa=wa, sigma8=sigma8, z_mg=z_mg,
                                df_mg=df_mg)
            self.params = params
            params = params.parameters # We only need the ccl_parameters object
        elif isinstance(params, lib.parameters):
            # Raise an error if ccl_parameters given directly
            raise TypeError("Must pass a Parameters() object, not ccl_parameters.")
        elif isinstance(params, Parameters):
            # Parameters object given directly
            self.params = params

            # Warn if any cosmological parameters were specified at the same
            # time as a Parameters() object; they will be ignored
            argtest = [Omega_c==None, Omega_b==None, h==None, A_s==None,
                       n_s==None, Omega_k==0., Omega_n==0., w0==-1., wa==0.,
                       sigma8==None, z_mg==None, df_mg==None]

            if not all(arg == True for arg in argtest):
                warn("Cosmological parameter kwargs are ignored if 'params' is "
                     "not None", UserWarning)
        else:
            raise TypeError("'params' is not a valid Parameters object.")

        # Check that the ccl_configuration-related arguments are valid
        if config is not None:
            # User passed a ccl_configuration object; ignore other arguments
            # and use this

            # Check that input object is of the correct type
            if not isinstance(config, lib.configuration):
                raise TypeError("'config' is not a valid ccl_configuration "
                                "object.")

            # Store ccl_configuration for later access
            self.configuration = config

        else:
            # Construct a new ccl_configuration object from kwargs

            # Check validity of configuration-related arguments
            if transfer_function not in transfer_function_types.keys():
                raise ValueError( "'%s' is not a valid transfer_function type. "
                                  "Available options are: %s" \
                                 % (transfer_function,
                                    transfer_function_types.keys()) )
            if matter_power_spectrum not in matter_power_spectrum_types.keys():
                raise ValueError( "'%s' is not a valid matter_power_spectrum "
                                  "type. Available options are: %s" \
                                 % (matter_power_spectrum,
                                    matter_power_spectrum_types.keys()) )
            if mass_function not in mass_function_types.keys():
                raise ValueError( "'%s' is not a valid mass_function type. "
                                  "Available options are: %s" \
                                 % (mass_function,
                                    mass_function_types.keys()) )

            # Assign values to new ccl_configuration object
            config = lib.configuration()

            config.transfer_function_method = \
                            transfer_function_types[transfer_function]
            config.matter_power_spectrum_method = \
                            matter_power_spectrum_types[matter_power_spectrum]
            config.mass_function_method = \
                            mass_function_types[mass_function]

            # Store ccl_configuration for later access
            self.configuration = config

        # Create new ccl_cosmology instance
        self.cosmo = lib.cosmology_create(self.params.parameters, config)

        # Check status
        if self.cosmo.status != 0:
            raise RuntimeError("(%d): %s" \
                               % (self.cosmo.status, self.cosmo.status_message))

    def __del__(self):
        """Free the ccl_cosmology instance that this Cosmology object is managing.

        """
        lib.cosmology_free(self.cosmo)

    def __str__(self):
        """Output the cosmological parameters that were set, and their values,
        as well as the status of precomputed quantities and the internal CCL
        status.

        """
        # String of cosmo parameters, from self.params (Parameters object)
        param_str = self.params.__str__()

        # String containing precomputation statuses
        precomp_stats = [
            ('has_distances', self.has_distances()),
            ('has_growth',    self.has_growth()),
            ('has_power',     self.has_power()),
            ('has_sigma',     self.has_sigma()),
            ]
        precomp_stat = ["%15s: %s" % stat for stat in precomp_stats]
        precomp_str = "\n".join(precomp_stat)

        # String from internal CCL status
        status_str = self.status()

        # Return composite string
        string = param_str
        string += "\n\nPrecomputed data\n----------------\n"
        string += precomp_str
        string += "\n\nStatus\n------\n"
        string += status_str
        return string

    def __getitem__(self, key):
        """Access cosmological parameter values by name.

        """
        return self.params.__getitem__(key)

    def compute_distances(self):
        """Interfaces with src/compute_background.c: ccl_cosmology_compute_distances().
        Sets up the splines for the distances.

        """
        status = 0
        status = lib.cosmology_compute_distances(self.cosmo, status)
        check(status)

    def compute_growth(self):
        """Interfaces with src/ccl_background.c: ccl_cosmology_compute_growth().
        Sets up the splines for the growth function.

        """
        status = 0
        status = lib.cosmology_compute_growth(self.cosmo, status)
        check(status)

    def compute_power(self):
        """Interfaces with src/ccl_power.c: ccl_cosmology_compute_power().
        Sets up the splines for the power spectrum.

        """
        status = 0
        status = lib.cosmology_compute_power(self.cosmo, status)
        check(status)

    def has_distances(self):
        """Checks if the distances have been precomputed.

        Returns:
            True if precomputed, False otherwise.

        """
        return bool(self.cosmo.computed_distances)

    def has_growth(self):
        """Checks if the growth function has been precomputed.

        Returns:
            True if precomputed, False otherwise.

        """
        return bool(self.cosmo.computed_growth)

    def has_power(self):
        """Checks if the power spectra have been precomputed.

        Returns:
            True if precomputed, False otherwise.

        """
        return bool(self.cosmo.computed_power)

    def has_sigma(self):
        """Checks if sigma8 has been computed.

        Returns:
            True if precomputed, False otherwise.

        """
        return bool(self.cosmo.computed_sigma)

    def status(self):
        """Get error status of the ccl_cosmology object.

        Note: error status is all currently under development.

        Returns:
            :obj:`str` containing the status message.

        """
        # Get status ID string if one exists
        if self.cosmo.status in error_types.keys():
            status = error_types[self.cosmo.status]
        else:
            status = self.cosmo.status

        # Get status message
        msg = self.cosmo.status_message

        # Return status information
        return "status(%s): %s" % (status, msg)


# Copyright Ryan-Rhys Griffiths and Aditya Raymond Thawani 2020
# Author: Ryan-Rhys Griffiths
"""
Script for comparing against human performance on a set of 5 molecules with Tanimoto MOGP.
"""

import argparse

import gpflow
from gpflow.ci_utils import ci_niter
from gpflow.mean_functions import Constant
from gpflow.utilities import print_summary
import numpy as np
from sklearn.metrics import r2_score, mean_squared_error, mean_absolute_error

from data_utils import transform_data, TaskDataLoader, featurise_mols
from kernels import Tanimoto

def main(path, representation):
    """
    :param path: str specifying path to dataset.
    :param representation: str specifying the molecular representation. One of ['fingerprints, 'fragments', 'fragprints']
    """

    task = 'e_iso_pi'  # task always e_iso_pi with human performance comparison
    data_loader = TaskDataLoader(task, path)
    smiles_list, y = data_loader.load_property_data()
    X = featurise_mols(smiles_list, representation)

    # 5 test molecules

    test_smiles = ['BrC1=CC=C(/N=N/C2=CC=CC=C2)C=C1',
                   'O=[N+]([O-])C1=CC=C(/N=N/C2=CC=CC=C2)C=C1',
                   'CC(C=C1)=CC=C1/N=N/C2=CC=C(N(C)C)C=C2',
                   'BrC1=CC([N+]([O-])=O)=CC([N+]([O-])=O)=C1/N=N/C2=CC([H])=C(C=C2[H])N(CC)CC',
                   'ClC%11=CC([N+]([O-])=O)=CC(C#N)=C%11/N=N/C%12=CC([H])=C(C=C%12OC)N(CC)CC']

    # and their indices in the loaded data
    test_smiles_indices = [116, 131, 168, 221, 229]

    X_train = np.delete(X, np.array(test_smiles_indices), axis=0)
    y_train = np.delete(y, np.array(test_smiles_indices))
    X_test = X[[116, 131, 168, 221, 229]]

    # experimental wavelength values in EtOH. Main csv file has 400nm instead of 407nm because measurement was
    # under a different solvent
    y_test = y[[116, 131, 168, 221, 229]]
    y_test[2] = 407.

    y_train = y_train.reshape(-1, 1)
    y_test = y_test.reshape(-1, 1)

    # #  We standardise the outputs but leave the inputs unchanged
    #
    # _, y_train, _, y_test, y_scaler = transform_data(X_train, y_train, X_test, y_test)

    X_train = X_train.astype(np.float64)
    X_test = X_test.astype(np.float64)

    data_loader_z_iso_pi = TaskDataLoader('z_iso_pi', path)
    data_loader_e_iso_n = TaskDataLoader('e_iso_n', path)
    data_loader_z_iso_n = TaskDataLoader('z_iso_n', path)

    smiles_list_z_iso_pi, y_z_iso_pi = data_loader_z_iso_pi.load_property_data()
    smiles_list_e_iso_n, y_e_iso_n = data_loader_e_iso_n.load_property_data()
    smiles_list_z_iso_n, y_z_iso_n = data_loader_z_iso_n.load_property_data()

    y_z_iso_pi = y_z_iso_pi.reshape(-1, 1)
    y_e_iso_n = y_e_iso_n.reshape(-1, 1)
    y_z_iso_n = y_z_iso_n.reshape(-1, 1)

    X_z_iso_pi = featurise_mols(smiles_list_z_iso_pi, representation)
    X_e_iso_n = featurise_mols(smiles_list_e_iso_n, representation)
    X_z_iso_n = featurise_mols(smiles_list_z_iso_n, representation)

    output_dim = 4  # Number of outputs
    rank = 1  # Rank of W
    feature_dim = len(X_train[0, :])

    tanimoto_active_dims = [i for i in range(feature_dim)]  # active dims for Tanimoto base kernel.

    # We define the Gaussian Process Regression Model using the Tanimoto kernel

    m = None

    def objective_closure():
        return -m.log_marginal_likelihood()

    # Augment the input with zeroes, ones, twos, threes to indicate the required output dimension
    X_augmented = np.vstack((np.append(X_train, np.zeros((len(X_train), 1)), axis=1),
                             np.append(X_z_iso_pi, np.ones((len(X_z_iso_pi), 1)), axis=1),
                             np.append(X_e_iso_n, np.ones((len(X_e_iso_n), 1)) * 2, axis=1),
                             np.append(X_z_iso_n, np.ones((len(X_z_iso_n), 1)) * 3, axis=1)))

    X_test = np.append(X_test, np.zeros((len(X_test), 1)), axis=1)
    X_train = np.append(X_train, np.zeros((len(X_train), 1)), axis=1)

    # Augment the Y data with zeroes, ones, twos and threes that specify a likelihood from the list of likelihoods
    Y_augmented = np.vstack((np.hstack((y_train, np.zeros_like(y_train))),
                             np.hstack((y_z_iso_pi, np.ones_like(y_z_iso_pi))),
                             np.hstack((y_e_iso_n, np.ones_like(y_e_iso_n) * 2)),
                             np.hstack((y_z_iso_n, np.ones_like(y_z_iso_n) * 3))))

    y_test = np.hstack((y_test, np.zeros_like(y_test)))

    # Base kernel
    k = Tanimoto(active_dims=tanimoto_active_dims)
    # set_trainable(k.variance, False)

    # Coregion kernel
    coreg = gpflow.kernels.Coregion(output_dim=output_dim, rank=rank, active_dims=[feature_dim])

    # Create product kernel
    kern = k * coreg

    # This likelihood switches between Gaussian noise with different variances for each f_i:
    lik = gpflow.likelihoods.SwitchedLikelihood([gpflow.likelihoods.Gaussian(), gpflow.likelihoods.Gaussian(),
                                                 gpflow.likelihoods.Gaussian(), gpflow.likelihoods.Gaussian()])

    # now build the GP model as normal
    m = gpflow.models.VGP((X_augmented, Y_augmented), mean_function=Constant(np.mean(y_train[:, 0])), kernel=kern,
                          likelihood=lik)

    # fit the covariance function parameters
    maxiter = ci_niter(1000)
    gpflow.optimizers.Scipy().minimize(m.training_loss, m.trainable_variables, options=dict(maxiter=maxiter),
                                       method="L-BFGS-B", )
    print_summary(m)

    # mean and variance GP prediction

    y_pred, y_var = m.predict_f(X_test)

    # Output Standardised RMSE and RMSE on Train Set

    y_pred_train, _ = m.predict_f(X_train)
    train_rmse_stan = np.sqrt(mean_squared_error(y_train, y_pred_train))
    train_rmse = np.sqrt(mean_squared_error(y_train, y_pred_train))
    print("\nStandardised Train RMSE: {:.3f}".format(train_rmse_stan))
    print("Train RMSE: {:.3f}".format(train_rmse))

    r2 = r2_score(y_test[:, 0], y_pred)
    rmse = np.sqrt(mean_squared_error(y_test[:, 0], y_pred))
    mae = mean_absolute_error(y_test[:, 0], y_pred)
    per_molecule = np.diag(abs(y_pred - y_test[:, 0]))

    print("\n Averaged test statistics are")
    print("\nR^2: {:.3f}".format(r2))
    print("RMSE: {:.3f}".format(rmse))
    print("MAE: {:.3f}".format(mae))
    print("\nAbsolute error per molecule is {} ".format(per_molecule))

if __name__ == '__main__':

    parser = argparse.ArgumentParser()

    parser.add_argument('-p', '--path', type=str, default='../dataset/photoswitches.csv',
                        help='Path to the photoswitches.csv file.')
    parser.add_argument('-r', '--representation', type=str, default='fragprints',
                        help='str specifying the molecular representation. '
                             'One of [fingerprints, fragments, fragprints].')

    args = parser.parse_args()

    main(args.path, args.representation)