# Licensed under an MIT style license -- see LICENSE.md

import os
import numpy as np
from pesummary.core.file.formats.base_read import SingleAnalysisRead
from pesummary.core.plots.latex_labels import latex_labels
from pesummary import conf
from pesummary.utils.utils import logger

__author__ = ["Charlie Hoy <charlie.hoy@ligo.org>"]


def bilby_version(path=None, data=None):
    """Extract the version of bilby used to create a result file

    Parameters
    ----------
    path: str, optional
        path to the bilby result file
    data: iterable, optional
        the contents of the bilby result file, opened with either h5py or
        json libraries
    """
    if path is None and data is None:
        raise ValueError(
            "Pass either the path to the result file, or result file object "
            "opened with either h5py or json"
        )
    elif data is None:
        import json
        import h5py
        try:
            data = h5py.File(path, "r")
        except Exception:
            with open(path, "r") as f:
                data = json.load(f)
    option1 = lambda data: data["version"][()]
    option2 = lambda data: data["version"]
    for opt in [option1, option2]:
        try:
            version = opt(data)
            if isinstance(version, (list, np.ndarray)):
                return version[0]
            return version
        except Exception:
            pass
    raise ValueError("Unable to extract version from bilby result file")


def _load_bilby(path):
    """Wrapper for `bilby.core.result.read_in_result`

    Parameters
    ----------
    path: str
        path to the bilby result file you wish to load
    """
    from bilby.core.result import read_in_result
    from bilby.core.utils import get_version_information
    result_version = bilby_version(path=path)
    if isinstance(result_version, bytes):
        result_version = result_version.decode("utf-8")
    env_version = 'bilby={}'.format(get_version_information())
    if result_version != env_version:
        logger.warning(
            "Result file was written with version {} while your environment "
            "has version {}. This may cause problems when reading the result "
            "file.".format(result_version, env_version)
        )
    return read_in_result(filename=path)


def read_bilby(
    path, disable_prior=False, complex_params=[], latex_dict=latex_labels,
    nsamples_for_prior=None, _bilby_class=None, **kwargs
):
    """Grab the parameters and samples in a bilby file

    Parameters
    ----------
    path: str
        path to the result file you wish to read in
    disable_prior: Bool, optional
        if True, do not collect prior samples from the `bilby` result file.
        Default False
    complex_params: list, optional
        list of parameters stored in the bilby result file which are complex
        and you wish to store the `amplitude` and `angle` as seperate
        posterior distributions
    latex_dict: dict, optional
        list of latex labels for each parameter
    nsamples_for_prior: int, optional
        number of samples to draw from the analytic priors
    """
    if _bilby_class is None:
        _bilby_class = Bilby
    bilby_object = _load_bilby(path)
    posterior = bilby_object.posterior
    _original_keys = posterior.keys()
    for key in _original_keys:
        for param in complex_params:
            if param in key and any(np.iscomplex(posterior[key])):
                posterior[key + "_amp"] = abs(posterior[key])
                posterior[key + "_angle"] = np.angle(posterior[key])
                posterior[key] = np.real(posterior[key])
            elif param in key:
                posterior[key] = np.real(posterior[key])
    # Drop all non numeric bilby data outputs
    posterior = posterior.select_dtypes(include=[float, int])
    parameters = list(posterior.keys())
    samples = posterior.to_numpy().real
    injection = bilby_object.injection_parameters
    if injection is None:
        injection = {i: j for i, j in zip(
            parameters, [float("nan")] * len(parameters))}
    else:
        for i in parameters:
            if i not in injection.keys():
                injection[i] = float("nan")
    _injection = injection.copy()
    for key, item in _injection.items():
        if isinstance(item, (str, np.str_)):
            injection.pop(key)

    if all(i for i in (
           bilby_object.constraint_parameter_keys,
           bilby_object.search_parameter_keys,
           bilby_object.fixed_parameter_keys)):
        for key in (
                bilby_object.constraint_parameter_keys
                + bilby_object.search_parameter_keys
                + bilby_object.fixed_parameter_keys):
            if key not in latex_dict:
                label = bilby_object.get_latex_labels_from_parameter_keys(
                    [key])[0]
                latex_dict[key] = label
    try:
        extra_kwargs = _bilby_class.grab_extra_kwargs(bilby_object)
    except Exception:
        extra_kwargs = {"sampler": {}, "meta_data": {}}
    extra_kwargs["sampler"]["nsamples"] = len(samples)
    extra_kwargs["sampler"]["pe_algorithm"] = "bilby"
    try:
        version = bilby_object.version
        if isinstance(version, (list, np.ndarray)):
            version = version[0]
    except Exception as e:
        version = None

    data = {
        "parameters": parameters,
        "samples": samples.tolist(),
        "injection": injection,
        "version": version,
        "kwargs": extra_kwargs
    }
    if bilby_object.meta_data is not None:
        if "command_line_args" in bilby_object.meta_data.keys():
            data["config"] = {
                "config": bilby_object.meta_data["command_line_args"]
            }
    if not disable_prior:
        logger.debug("Drawing prior samples from bilby result file")
        if nsamples_for_prior is None:
            nsamples_for_prior = len(samples)
        prior_samples = Bilby.grab_priors(
            bilby_object, nsamples=nsamples_for_prior
        )
        data["prior"] = {"samples": prior_samples}
        if len(prior_samples):
            data["prior"]["analytic"] = prior_samples.analytic
    else:
        try:
            _prior = bilby_object.priors
            data["prior"] = {
                "samples": {},
                "analytic": {key: str(item) for key, item in _prior.items()}
            }
        except (AttributeError, KeyError):
            pass
    return data


def to_bilby(
    parameters, samples, label=None, analytic_priors=None, cls=None,
    meta_data=None, **kwargs
):
    """Convert a set of samples to a bilby object

    Parameters
    ----------
    parameters: list
        list of parameters
    samples: 2d list
        list of samples. Columns correspond to a given parameter
    label: str, optional
        The label of the analysis. This is used in the filename if a filename
        if not specified
    """
    from bilby.core.result import Result
    from bilby.core.prior import Prior, PriorDict
    from pandas import DataFrame

    if cls is None:
        cls = Result
    if analytic_priors is not None:
        priors = PriorDict._get_from_json_dict(analytic_priors)
        search_parameters = priors.keys()
    else:
        priors = {param: Prior() for param in parameters}
        search_parameters = parameters
    posterior_data_frame = DataFrame(samples, columns=parameters)
    bilby_object = cls(
        search_parameter_keys=search_parameters, samples=samples, priors=priors,
        posterior=posterior_data_frame, label="pesummary_%s" % label,
    )
    return bilby_object


def _write_bilby(
    parameters, samples, outdir="./", label=None, filename=None, overwrite=False,
    extension="json", save=True, analytic_priors=None, cls=None,
    meta_data=None, **kwargs
):
    """Write a set of samples to a bilby file

    Parameters
    ----------
    parameters: list
        list of parameters
    samples: 2d list
        list of samples. Columns correspond to a given parameter
    outdir: str, optional
        directory to write the dat file
    label: str, optional
        The label of the analysis. This is used in the filename if a filename
        if not specified
    filename: str, optional
        The name of the file that you wish to write
    overwrite: Bool, optional
        If True, an existing file of the same name will be overwritten
    extension: str, optional
        file extension for the bilby result file. Default json.
    save: Bool, optional
        if True, save the bilby object to file
    """
    bilby_object = to_bilby(
        parameters, samples, label=None, analytic_priors=None, cls=None,
        meta_data=None, **kwargs
    )
    if save:
        _filename = os.path.join(outdir, filename)
        bilby_object.save_to_file(filename=_filename, extension=extension)
    else:
        return bilby_object


def write_bilby(
    parameters, samples, outdir="./", label=None, filename=None, overwrite=False,
    extension="json", save=True, analytic_priors=None, cls=None,
    meta_data=None, labels=None, **kwargs
):
    """Write a set of samples to a bilby file

    Parameters
    ----------
    parameters: list
        list of parameters
    samples: 2d list
        list of samples. Columns correspond to a given parameter
    outdir: str, optional
        directory to write the dat file
    label: str, optional
        The label of the analysis. This is used in the filename if a filename
        if not specified
    filename: str, optional
        The name of the file that you wish to write
    overwrite: Bool, optional
        If True, an existing file of the same name will be overwritten
    extension: str, optional
        file extension for the bilby result file. Default json.
    save: Bool, optional
        if True, save the bilby object to file
    """
    from pesummary.io.write import _multi_analysis_write

    func = _write_bilby
    if not save:
        func = to_bilby
    return _multi_analysis_write(
        func, parameters, samples, outdir=outdir, label=label,
        filename=filename, overwrite=overwrite, extension=extension,
        save=save, analytic_priors=analytic_priors, cls=cls,
        meta_data=meta_data, file_format="bilby", labels=labels,
        _return=True, **kwargs
    )


def config_from_file(path):
    """Extract the configuration file stored within a bilby result file

    Parameters
    ----------
    path: str
        path to the bilby result file you wish to load
    """
    bilby_object = _load_bilby(path)
    return config_from_object(bilby_object)


def config_from_object(bilby_object):
    """Extract the configuration file stored within a `bilby.core.result.Result`
    object (or alike)

    Parameters
    ----------
    bilby_object: bilby.core.result.Result (or alike)
        a bilby.core.result.Result object (or alike) you wish to extract the
        configuration file from
    """
    config = {}
    if bilby_object.meta_data is not None:
        if "command_line_args" in bilby_object.meta_data.keys():
            config = {
                "config": bilby_object.meta_data["command_line_args"]
            }
    return config


def prior_samples_from_file(path, cls="PriorDict", nsamples=5000, **kwargs):
    """Return a dict of prior samples from a `bilby` prior file

    Parameters
    ----------
    path: str
        path to a `bilby` prior file
    cls: str, optional
        class you wish to read in the prior file
    nsamples: int, optional
        number of samples to draw from a prior file. Default 5000
    """
    from bilby.core import prior

    if isinstance(cls, str):
        cls = getattr(prior, cls)
    _prior = cls(filename=path)
    samples = _prior.sample(size=nsamples)
    return _bilby_prior_dict_to_pesummary_samples_dict(samples, prior=_prior)


def prior_samples_from_bilby_object(bilby_object, nsamples=5000, **kwargs):
    """Return a dict of prior samples from a `bilby.core.result.Result`
    object

    Parameters
    ----------
    bilby_object: bilby.core.result.Result
        a bilby.core.result.Result object you wish to draw prior samples from
    nsamples: int, optional
        number of samples to draw from a prior file. Default 5000
    """
    samples = bilby_object.priors.sample(size=nsamples)
    return _bilby_prior_dict_to_pesummary_samples_dict(
        samples, prior=bilby_object.priors
    )


def _bilby_prior_dict_to_pesummary_samples_dict(samples, prior=None):
    """Return a pesummary.utils.samples_dict.SamplesDict object from a bilby
    priors dict
    """
    from pesummary.utils.samples_dict import SamplesDict

    _samples = SamplesDict(samples)
    if prior is not None:
        analytic = {key: str(item) for key, item in prior.items()}
        setattr(_samples, "analytic", analytic)
    return _samples


class Bilby(SingleAnalysisRead):
    """PESummary wrapper of `bilby` (https://git.ligo.org/lscsoft/bilby). The
    path_to_results_file argument will be passed directly to
    `bilby.core.result.read_in_result`. All functions therefore use `bilby`
    methods and requires `bilby` to be installed.

    Parameters
    ----------
    path_to_results_file: str
        path to the results file that you wish to read in with `bilby`.
    disable_prior: Bool, optional
        if True, do not collect prior samples from the `bilby` result file.
        Default False
    remove_nan_likelihood_samples: Bool, optional
        if True, remove samples which have log_likelihood='nan'. Default True

    Attributes
    ----------
    parameters: list
        list of parameters stored in the result file
    samples: 2d list
        list of samples stored in the result file
    samples_dict: dict
        dictionary of samples stored in the result file keyed by parameters
    input_version: str
        version of the result file passed.
    extra_kwargs: dict
        dictionary of kwargs that were extracted from the result file
    injection_parameters: dict
        dictionary of injection parameters extracted from the result file
    prior: dict
        dictionary of prior samples keyed by parameters. The prior functions
        are evaluated for 5000 samples.
    pe_algorithm: str
        name of the algorithm used to generate the posterior samples

    Methods
    -------
    to_dat:
        save the posterior samples to a .dat file
    to_latex_table:
        convert the posterior samples to a latex table
    generate_latex_macros:
        generate a set of latex macros for the stored posterior samples
    """
    def __init__(self, path_to_results_file, **kwargs):
        super(Bilby, self).__init__(path_to_results_file, **kwargs)
        self.load(self._grab_data_from_bilby_file, **kwargs)

    @staticmethod
    def grab_priors(bilby_object, nsamples=5000):
        """Draw samples from the prior functions stored in the bilby file
        """
        try:
            return prior_samples_from_bilby_object(
                bilby_object, nsamples=nsamples
            )
        except Exception as e:
            logger.info("Failed to draw prior samples because {}".format(e))
            return {}

    @staticmethod
    def grab_extra_kwargs(bilby_object):
        """Grab any additional information stored in the lalinference file
        """
        f = bilby_object
        kwargs = {"sampler": {
            conf.log_evidence: np.round(f.log_evidence, 2),
            conf.log_evidence_error: np.round(f.log_evidence_err, 2),
            conf.log_bayes_factor: np.round(f.log_bayes_factor, 2),
            conf.log_noise_evidence: np.round(f.log_noise_evidence, 2)},
            "meta_data": {}, "other": f.meta_data}
        return kwargs

    @staticmethod
    def _grab_data_from_bilby_file(path, **kwargs):
        """Load the results file using the `bilby` library
        """
        return read_bilby(path, **kwargs)

    def add_marginalized_parameters_from_config_file(self, config_file):
        """Search the configuration file and add the marginalized parameters
        to the list of parameters and samples

        Parameters
        ----------
        config_file: str
            path to the configuration file
        """
        pass
