# 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 _check_bilby_attributes(result, attribute):
    """Check the bilby attributes are the same format across different versions.
    Currently this checks that an empty list is returned if the attribute is
    None or if the length is 0. This can be deprecated after a few bilby releases.

    Parameters
    ----------
    result: bilby.core.result.Result
        bilby result file to check
    attribute: str
        attribute you wish to check
    """
    _attr = getattr(result, attribute)
    if (_attr is None) or (not len(_attr)):
        _attr = []
    return _attr


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 + "_abs"] = 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)

    _fixed_keys = _check_bilby_attributes(
        bilby_object, "fixed_parameter_keys"
    )
    _constraint_keys = _check_bilby_attributes(
        bilby_object, "constraint_parameter_keys"
    )
    _search_keys = _check_bilby_attributes(
        bilby_object, "search_parameter_keys"
    )
    if all(i for i in (_constraint_keys, _search_keys, _fixed_keys)):
        for key in (_constraint_keys + _search_keys + _fixed_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
