"""Function to setup a general GP."""
from .general_preprocess import GeneralPrepreprocess
from .general_kernel import general_kernel, smooth_kernel
from catlearn.regression import GaussianProcess


class GeneralGaussianProcess(object):
    """Define a general setup for the Gaussin process.

    This should not be used to try and obtain highly accurate solutions. Though
    it should give a reasonable model.
    """

    def __init__(self, clean_type='eliminate', dimension='single',
                 kernel='general'):
        """Initialize the class.

        Parameters
        ----------
        clean_type : str
            Define method for handling missing data. Currently only elimination
            implemented.
        dimension : str
            The number of parameters to return. Can be 'single', or 'all'.
        kernel : str
            Pass 'smooth' if a smooth but non-linear function is expected.
        """
        self.clean_type = clean_type
        self.dimension = dimension
        self.kernel = kernel

    def train_gaussian_process(self, train_features, train_targets):
        """Generate a general gaussian process model.

        Parameters
        ----------
        train_features : array
            The array of training features.
        train_targets : array
            A list of training target values.

        Returns
        -------
        gp : object
            The trained Gaussian process.
        """
        train_features, train_targets = self._process_train_data(
            train_features, train_targets)

        if 'smooth' in self.kernel:
            kdict = smooth_kernel(train_features, self.dimension)
        else:
            kdict = general_kernel(train_features, self.dimension)

        self.gp = GaussianProcess(
            train_fp=train_features, train_target=train_targets,
            kernel_list=kdict, regularization=1e-2,
            optimize_hyperparameters=True, scale_data=False
            )

        return self.gp

    def gaussian_process_predict(self, test_features):
        """Function to make GP predictions on tests data.

        Parameters
        ----------
        test_features : array
            The array of test features.

        Returns
        -------
        prediction : dict
            The prediction data generated by the Gaussian process.
        """
        test_features = self.cleaner.transform(test_features)

        pred = self.gp.predict(test_fp=test_features)

        return pred

    def _process_train_data(self, train_features, train_targets):
        """Prepare the data.

        Parameters
        ----------
        train_features : array
            The array of training features.
        train_targets : array
            A list of training target values.

        Returns
        -------
        train_features : array
            The array of cleaned and scaled training features.
        train_targets : array
            A list of cleaned training target values.
        """
        self.cleaner = GeneralPrepreprocess(clean_type=self.clean_type)
        train_features, train_targets, _ = self.cleaner.process(
            train_features, train_targets)

        return train_features, train_targets
