from __future__ import absolute_import

from ..utils import check_directory_exists_and_if_not_mkdir
from .base_sampler import NestedSampler
from ..utils import logger


class Pymultinest(NestedSampler):
    """
    bilby wrapper of pymultinest
    (https://github.com/JohannesBuchner/PyMultiNest)

    All positional and keyword arguments (i.e., the args and kwargs) passed to
    `run_sampler` will be propagated to `pymultinest.run`, see documentation
    for that class for further help. Under Keyword Arguments, we list commonly
    used kwargs and the bilby defaults.

    Keyword Arguments
    ------------------
    npoints: int
        The number of live points, note this can also equivalently be given as
        one of [nlive, nlives, n_live_points]
    importance_nested_sampling: bool, (False)
        If true, use importance nested sampling
    sampling_efficiency: float or {'parameter', 'model'}, ('parameter')
        Defines the sampling efficiency
    verbose: Bool
        If true, print information information about the convergence during
    resume: bool
        If true, resume run from checkpoint (if available)

    """
    default_kwargs = dict(importance_nested_sampling=False, resume=True,
                          verbose=True, sampling_efficiency='parameter',
                          n_live_points=500, n_params=None,
                          n_clustering_params=None, wrapped_params=None,
                          multimodal=True, const_efficiency_mode=False,
                          evidence_tolerance=0.5,
                          n_iter_before_update=100, null_log_evidence=-1e90,
                          max_modes=100, mode_tolerance=-1e90,
                          outputfiles_basename=None, seed=-1,
                          context=0, write_output=True, log_zero=-1e100,
                          max_iter=0, init_MPI=False, dump_callback=None)

    def _translate_kwargs(self, kwargs):
        if 'n_live_points' not in kwargs:
            for equiv in self.npoints_equiv_kwargs:
                if equiv in kwargs:
                    kwargs['n_live_points'] = kwargs.pop(equiv)

    def _verify_kwargs_against_default_kwargs(self):
        """
        Test the length of the directory where multinest will write the output.

        This is an issue with MultiNest that we can't solve here.
        https://github.com/JohannesBuchner/PyMultiNest/issues/115
        """
        if not self.kwargs['outputfiles_basename']:
            self.kwargs['outputfiles_basename'] =\
                '{}/pm_{}/'.format(self.outdir, self.label)
        if self.kwargs['outputfiles_basename'].endswith('/') is False:
            self.kwargs['outputfiles_basename'] = '{}/'.format(
                self.kwargs['outputfiles_basename'])
        if len(self.kwargs['outputfiles_basename']) > (100 - 22):
            logger.warning(
                'The length of {} exceeds 78 characters. '
                ' Post-processing will fail because the file names will be cut'
                ' off. Please choose a shorter "outdir" or "label".'
                .format(self.__kwargs['outputfiles_basename']))
        check_directory_exists_and_if_not_mkdir(
            self.kwargs['outputfiles_basename'])
        NestedSampler._verify_kwargs_against_default_kwargs(self)

    def run_sampler(self):
        import pymultinest
        self._verify_kwargs_against_default_kwargs()
        out = pymultinest.solve(
            LogLikelihood=self.log_likelihood, Prior=self.prior_transform,
            n_dims=self.ndim, **self.kwargs)

        self.result.sampler_output = out
        self.result.samples = out['samples']
        self.result.log_evidence = out['logZ']
        self.result.log_evidence_err = out['logZerr']
        self.result.outputfiles_basename = self.kwargs['outputfiles_basename']
        return self.result
