__author__ = "aymgal, johannesulf"

import os
import shutil
import numpy as np
import sys
import copy

from lenstronomy.Sampling.Samplers.base_nested_sampler import NestedSampler

__all__ = ["DyPolyChordSampler"]


class DyPolyChordSampler(NestedSampler):
    """Wrapper for dynamical nested sampling algorithm DyPolyChord by E. Higson, M.
    Hobson, W. Handley, A. Lasenby.

    papers : arXiv:1704.03459, arXiv:1804.06406
    doc : https://dypolychord.readthedocs.io
    """

    def __init__(
        self,
        likelihood_module,
        prior_type="uniform",
        prior_means=None,
        prior_sigmas=None,
        width_scale=1,
        sigma_scale=1,
        output_dir=None,
        output_basename="-",
        resume_dyn_run=False,
        polychord_settings=None,
        remove_output_dir=False,
        use_mpi=False,
    ):  # , num_mpi_procs=1):
        """
        :param likelihood_module: likelihood_module like in likelihood.py (should be callable)
        :param prior_type: 'uniform' of 'gaussian', for converting the unit hypercube to param cube
        :param prior_means: if prior_type is 'gaussian', mean for each param
        :param prior_sigmas: if prior_type is 'gaussian', std dev for each param
        :param width_scale: scale the widths of the parameters space by this factor
        :param sigma_scale: if prior_type is 'gaussian', scale the gaussian sigma by this factor
        :param output_dir: name of the folder that will contain output files
        :param output_basename: prefix for output files
        :param resume_dyn_run: if True, previous resume files will not be deleted so that previous run can be resumed
        :param polychord_settings: settings dictionary to send to pypolychord. Check dypolychord documentation for details.
        :param remove_output_dir: remove the output_dir folder after completion
        :param use_mpi: Use MPI computing if `True`
        """
        self._check_install()
        super(DyPolyChordSampler, self).__init__(
            likelihood_module,
            prior_type,
            prior_means,
            prior_sigmas,
            width_scale,
            sigma_scale,
        )

        # if use_mpi:
        #     mpi_str = 'mpirun -np {}'.format(num_mpi_procs)
        # else:
        #     mpi_str = None
        if polychord_settings == None:
            polychord_settings = {}
        self._use_mpi = use_mpi

        self._output_dir = output_dir
        self._is_master = True

        if self._use_mpi:
            from mpi4py import MPI

            self._comm = MPI.COMM_WORLD

            if self._comm.Get_rank() != 0:
                self._is_master = False
        else:
            self._comm = None

        if not resume_dyn_run:
            if self._is_master:
                if os.path.exists(self._output_dir):
                    shutil.rmtree(self._output_dir, ignore_errors=True)
                os.mkdir(self._output_dir)

        self._output_basename = output_basename
        self._settings = copy.deepcopy(polychord_settings)
        self._settings["file_root"] = self._output_basename
        self._settings["base_dir"] = self._output_dir

        if self._all_installed:
            # create the dyPolyChord callable object
            self._sampler = self._RunPyPolyChord(
                self.log_likelihood, self.prior, self.n_dims
            )
        else:
            self._sampler = None

        self._rm_output = remove_output_dir
        self._has_warned = False

    def run(self, dynamic_goal, kwargs_run):
        """Run the DyPolyChord dynamical nested sampler.

        see https://dypolychord.readthedocs.io for content of kwargs_run

        :param dynamic_goal: 0 for evidence computation, 1 for posterior computation
        :param kwargs_run: kwargs directly passed to dyPolyChord.run_dypolychord
        :return: samples, means, logZ, logZ_err, logL, ns_run
        """
        print("prior type :", self.prior_type)
        print("parameter names :", self.param_names)

        if self._all_installed:
            # TODO : put a default dynamic_goal ?
            # dynamic_goal = 0 for evidence-only, 1 for posterior-only

            self._dyPolyChord.run_dypolychord(
                self._sampler,
                dynamic_goal,
                settings_dict_in=self._settings,
                comm=self._comm,
                **kwargs_run
            )

            if self._is_master:
                ns_run = self._ns_process_run(
                    self._settings["file_root"], self._settings["base_dir"]
                )

        else:
            # in case DyPolyChord or NestCheck was not compiled properly, for unit tests
            ns_run = {
                "theta": np.zeros((1, self.n_dims)),
                "logl": np.zeros(1),
                "output": {
                    "logZ": np.zeros(1),
                    "logZerr": np.zeros(1),
                    "param_means": np.zeros(self.n_dims),
                },
            }
            self._write_equal_weights(ns_run["theta"], ns_run["logl"])

        if self._is_master:
            samples, logL = self._get_equal_weight_samples()
            # logL     = ns_run['logl']
            # samples_w = ns_run['theta']
            logZ = ns_run["output"]["logZ"]
            logZ_err = ns_run["output"]["logZerr"]
            means = ns_run["output"]["param_means"]

            print("The log evidence estimate using the first run is {}".format(logZ))
            print("The estimated mean of the first parameter is {}".format(means[0]))

            if self._rm_output:
                shutil.rmtree(self._output_dir, ignore_errors=True)

            return samples, means, logZ, logZ_err, logL, ns_run
        else:
            sys.exit(0)

    def log_likelihood(self, args):
        """Compute the log-likelihood given list of parameters.

        :param args: parameter values
        :return: log-likelihood (from the likelihood module)
        """
        return super().log_likelihood(args), []

    def _get_equal_weight_samples(self):
        """Inspired by pymultinest's Analyzer, because DyPolyChord has more or less the
        same output conventions as MultiNest."""
        file_name = "{}_equal_weights.txt".format(self._output_basename)
        file_path = os.path.join(self._output_dir, file_name)
        data = np.loadtxt(file_path, ndmin=2)
        logL = -0.5 * data[:, 0]
        samples = data[:, 1:]
        return samples, logL

    def _write_equal_weights(self, samples, logL):
        # write fake output file for unit tests
        file_name = "{}_equal_weights.txt".format(self._output_basename)
        file_path = os.path.join(self._output_dir, file_name)
        data = np.zeros((samples.shape[0], 1 + samples.shape[1]), dtype=float)
        data[:, 0] = -2.0 * logL
        data[:, 1:] = samples
        np.savetxt(file_path, data, fmt="% .14E")

    def _check_install(self):
        try:
            import dyPolyChord
            from dyPolyChord import pypolychord_utils
        except ImportError:
            dyPolyChord = None
            pypolychord_utils = None
            print(
                "Warning : dyPolyChord not properly installed. \
                   You can get it from : https://github.com/ejhigson/dyPolyChord"
            )
            dypolychord_installed = False
        else:
            dypolychord_installed = True
            self._dyPolyChord = dyPolyChord
            self._RunPyPolyChord = pypolychord_utils.RunPyPolyChord

        try:
            from nestcheck import data_processing
        except ImportError:
            print(
                "Warning : nestcheck not properly installed (results might be unexpected). \
                   You can get it from : https://github.com/ejhigson/nestcheck"
            )
            nestcheck_installed = False
        else:
            nestcheck_installed = True
            self._ns_process_run = data_processing.process_polychord_run

        self._all_installed = dypolychord_installed and nestcheck_installed
