<reponame>jereliu/cabernet
"""Class definition for agent to manage data and run models."""
import collections

import numpy as np

import model

import inference.vi as vi
import inference.predictor as predictor
import inference.cdf as cdf

import util.dtype as dtype_util

DEFAULT_GP_LOG_LS_RESID = np.log(0.2).astype(np.float32)
DEFAULT_CDF_LOG_LS_RESID = np.log(0.3).astype(np.float32)

DEFAULT_N_POST_SAMPLE = 500

DEFAULT_CALIB_PERCENTILES_TRAIN = np.linspace(.0001, .9999, num=15).astype(dtype_util.NP_DTYPE)
DEFAULT_CALIB_PERCENTILES_PRED = np.linspace(.0001, .9999, num=100).astype(dtype_util.NP_DTYPE)
DEFAULT_OUTPUT_PERCENTILES = np.array([10, 25, 50, 75, 90]).astype(dtype_util.NP_DTYPE)

BNE_SUMMARY_NAMES = ("mean", "median", "var", "quantiles", "mean_cdf")

BNESummary = collections.namedtuple("BNESummary",
                                    field_names=BNE_SUMMARY_NAMES)
BNESummary.__new__.__defaults__ = (None,) * len(BNESummary._fields)


# TODO(jereliu): model diagnosis for skewness and kurtosis.

class BNE(object):
    def __init__(self,
                 X, y, base_pred,
                 X_new, base_pred_new,
                 X_calib=None, y_calib=None,
                 X_calib_induce=None, base_pred_calib_induce=None,
                 log_ls_system=DEFAULT_GP_LOG_LS_RESID,
                 log_ls_random=DEFAULT_CDF_LOG_LS_RESID,
                 calib_percentiles_train=DEFAULT_CALIB_PERCENTILES_TRAIN,
                 calib_percentiles_pred=DEFAULT_CALIB_PERCENTILES_PRED,
                 pred_percentiles=DEFAULT_OUTPUT_PERCENTILES):
        """Initializer."""
        # names of summary statistics
        self.summary_names = BNE_SUMMARY_NAMES
        self.posterior_summary = BNESummary()

        # initialize data
        self.X_train = X
        self.y_train = y
        self.base_pred_train = base_pred

        self.X_calib = self.X_train if X_calib is None else X_calib
        self.y_calib = self.y_train if y_calib is None else y_calib
        self.X_calib_induce = self.X_calib if X_calib_induce is None else X_calib_induce
        self.base_pred_calib_induce = (self.base_pred_train if base_pred_calib_induce is None
                                       else base_pred_calib_induce)

        self.X_test = X_new
        self.base_pred_test = base_pred_new

        # initialize hyper parameters
        self.log_ls_system = log_ls_system
        self.log_ls_random = log_ls_random

        self.calib_percentiles_train = calib_percentiles_train
        self.calib_percentiles_pred = calib_percentiles_pred
        self.pred_percentiles = pred_percentiles

        # initialize internal parameters
        # model
        self.system_model = None
        self.random_model = None

        # initialize model estimator, predictor and sessions
        self.system_estimator = None
        self.random_estimator = None

        self.system_predictor = None
        self.random_predictor = None
        self.random_summarizer = None

        self.system_session = None
        self.random_session = None

        # initialize containers for posterior samples at training/predict locations
        self.system_model_sample_train = None
        self.random_model_sample_train = None

        self.system_model_sample_calib = None  # for training random component model
        self.system_model_sample_pred = None
        self.random_model_sample_pred = None

        self.system_model_quantile_calib = None  # for training random component model
        self.system_model_quantile_pred = None
        self.random_model_quantile_pred = None

    def run_model(self, system_model_kwargs, random_model_kwargs):
        """Estimates and generates predictive samples/quantiles for full model."""
        self.run_system_model(**system_model_kwargs)
        self.run_random_model(**random_model_kwargs)

    def run_system_model(self,
                         log_ls_system=None, restart_model=True,
                         n_sample=DEFAULT_N_POST_SAMPLE,
                         **estimate_kwargs):
        """Estimates and generates predictive samples/quantiles for sys model."""
        # define model
        if log_ls_system:
            self.log_ls_system = log_ls_system

        if not self.system_model or restart_model:
            self.system_model = (
                model.HierarchicalGP(X=self.X_train, y=self.y_train,
                                     base_pred=self.base_pred_train,
                                     resid_log_ls=self.log_ls_system)
            )

        # ESTIMATION
        self._estimate_system_model(**estimate_kwargs)

        # PREDICTION
        self.system_predictor = predictor.Predictor(estimator=self.system_estimator)

        # sample posterior, train, calibration and pred location
        self.system_model_sample_train = self._sample_system_model(n_sample)

        self.system_model_sample_calib = (
            self._pred_sample_system_model(X_new=self.X_calib_induce,
                                           base_pred_new=self.base_pred_calib_induce))
        self.system_model_sample_pred = (
            self._pred_sample_system_model(X_new=self.X_test,
                                           base_pred_new=self.base_pred_test))

        # sample quantiles, calibration location (for training random component model)
        self.system_model_quantile_calib = (
            self._pred_quantiles_system_model(sample_dict=self.system_model_sample_calib,
                                              perc_eval=self.calib_percentiles_train))

        # sample quantiles, predictive location (for prediction output)
        self.system_model_quantile_pred = (
            self._pred_quantiles_system_model(sample_dict=self.system_model_sample_pred,
                                              perc_eval=self.calib_percentiles_pred))

    def run_random_model(self,
                         log_ls_random=None, restart_model=True,
                         n_sample=DEFAULT_N_POST_SAMPLE,
                         **estimate_kwargs):
        """Estimates and generates predictive samples/quantiles for rand model."""
        # define model
        if log_ls_random:
            self.log_ls_random = log_ls_random

        if not self.random_model or restart_model:
            self.random_model = (
                model.MonoGP(X=self.X_calib, y=self.y_calib,
                             X_induce=self.X_calib_induce,
                             cdf_sample_induce=self.system_model_quantile_calib,
                             log_ls=self.log_ls_random)
            )

        # ESTIMATION
        self._estimate_random_model(**estimate_kwargs)

        # PREDICTION
        self.random_predictor = predictor.Predictor(estimator=self.random_estimator)

        # sample posterior, train and pred locations
        self.random_model_sample_train = self._sample_random_model(n_sample)
        self.random_model_sample_pred = (
            self._pred_sample_random_model(X_new=self.X_test,
                                           quant_pred_new=self.system_model_quantile_pred))

        # SUMMARIZE
        self.random_summarizer = (
            cdf.CDFMoments(estimator=self.random_estimator,
                           cdf_sample_dict=self.random_model_sample_pred))
        self.posterior_summary = BNESummary(**self._pred_summary_random_model())

    def _estimate_system_model(self,
                               step_size=1e-2,
                               max_steps=50000, model_dir=None,
                               save_step=500, verbose=True,
                               restart_estimate=False,
                               **vi_kwargs):
        """Estimates systematic component model."""
        # estimation
        if not self.system_estimator:
            # in case of first run, configure estimator graph
            self.system_estimator = vi.VIEstimator(model=self.system_model)
            self.system_estimator.config(step_size=step_size, **vi_kwargs)

        elif restart_estimate:
            # in case of re-run and want to restart, re-config
            # estimator graph and erase current session.
            self.system_estimator.config(step_size=step_size, **vi_kwargs)
            self.system_session = None
        else:
            # in case of re-run and want to reuse current setting, do nothing
            pass

        self.system_session = self.system_estimator.run(sess=self.system_session,
                                                        max_steps=max_steps,
                                                        model_dir=model_dir,
                                                        save_step=save_step,
                                                        verbose=verbose)

    def _estimate_random_model(self,
                               step_size=1e-2,
                               max_steps=10000, model_dir=None,
                               save_step=500, verbose=True,
                               restart_estimate=False,
                               **vi_kwargs):
        """Estimates random component model."""
        # estimation
        if not self.random_estimator:
            # in case of first run, configure estimator graph
            self.random_estimator = vi.VIEstimator(model=self.random_model)
            self.random_estimator.config(step_size=step_size, **vi_kwargs)
        elif restart_estimate:
            # in case of re-run and want to restart, re-config
            # estimator graph and erase current session.
            self.random_estimator.config(step_size=step_size, **vi_kwargs)
            self.random_session = None
        else:
            # in case of re-run and want to reuse current setting, do nothing
            pass

        self.random_session = self.random_estimator.run(sess=self.random_session,
                                                        max_steps=max_steps,
                                                        model_dir=model_dir,
                                                        save_step=save_step,
                                                        verbose=verbose)

    def _sample_system_model(self, n_sample):
        """Generates posterior samples for model parameters at training locations."""
        if not self.system_predictor:
            raise ValueError("Predictor for systematic component model empty.")
        if not self.system_session:
            raise ValueError("Session for systematic component model empty.")

        # posterior sampling for training sample
        # in-sample posterior sample, training locations
        self.system_predictor.config(sample_type="post_sample",
                                     rv_dict=self.system_estimator.model.model_param,
                                     n_sample=n_sample)
        sample_dict = self.system_predictor.run(sess=self.system_session)
        return sample_dict["post_sample"]

    def _sample_random_model(self, n_sample):
        """Generates posterior samples for model parameters at training locations."""
        if not self.random_predictor:
            raise ValueError("Predictor for random component model is empty.")
        if not self.random_session:
            raise ValueError("Session for systematic component model is empty.")

        # posterior sampling for training sample
        # in-sample posterior sample, training locations
        self.random_predictor.config(sample_type="post_sample",
                                     rv_dict=self.random_estimator.model.model_param,
                                     n_sample=n_sample)
        sample_dict = self.random_predictor.run(sess=self.random_session)
        return sample_dict["post_sample"]

    def _pred_sample_system_model(self, X_new, base_pred_new):
        """Generates posterior samples at predictive locations."""
        if not self.system_model_sample_train:
            raise ValueError("Train sample for systematic component model empty.")
        if not self.system_session:
            raise ValueError("Session for systematic component model empty.")

        self.system_predictor.config(sample_type="pred_sample",
                                     X_new=X_new,
                                     base_pred_new=base_pred_new,
                                     post_sample_dict=self.system_model_sample_train)

        sample_dict = self.system_predictor.run(sess=self.system_session)

        return sample_dict["pred_sample"]

    def _pred_sample_random_model(self, X_new, quant_pred_new):
        """Generates posterior samples at predictive locations."""
        if not self.system_model_sample_train:
            raise ValueError("Train sample for systematic component model empty.")
        if not self.system_session:
            raise ValueError("Session for systematic component model empty.")

        self.random_predictor.config(sample_type="pred_sample",
                                     X_new=X_new,
                                     quant_dict_new=quant_pred_new,
                                     post_sample_dict=self.random_model_sample_train,
                                     verbose=True)

        sample_dict = self.random_predictor.run(sess=self.random_session)

        return sample_dict["pred_sample"]

    def _pred_quantiles_system_model(self, sample_dict, perc_eval):
        """Generates predictive quantiles at predictive locations."""
        if not self.system_model_sample_pred:
            raise ValueError("Predictive sample for systematic component model empty.")
        if not self.system_session:
            raise ValueError("Session for systematic component model empty.")

        self.system_predictor.config(sample_type="pred_quant",
                                     perc_eval=perc_eval,
                                     sample_dict=sample_dict)

        sample_dict = self.system_predictor.run(sess=self.system_session)
        return sample_dict["pred_quant"]

    def _pred_summary_random_model(self):
        """Computes summary statistics from the estimated CDFs."""
        for summary_name in self.summary_names:
            self.random_summarizer.config(summary_name,
                                          percentiles=self.pred_percentiles / 100.)

        return self.random_summarizer.run()

    @property
    def log_ls_system(self):
        return self.__log_ls_system

    @log_ls_system.setter
    def log_ls_system(self, value):
        self.__log_ls_system = value.astype(dtype_util.NP_DTYPE)

    @property
    def log_ls_random(self):
        return self.__log_ls_random

    @log_ls_random.setter
    def log_ls_random(self, value):
        self.__log_ls_random = value.astype(dtype_util.NP_DTYPE)
