#! /usr/bin/env python

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

import os
import pesummary
from pesummary.core.cli.inputs import _Input
from pesummary.gw.file.read import read as GWRead
from pesummary.gw.classification import PEPredicates, PAstro
from pesummary.utils.utils import make_dir, logger
from pesummary.utils.exceptions import InputError
from pesummary.core.cli.parser import ArgumentParser as _ArgumentParser

__author__ = ["Charlie Hoy <charlie.hoy@ligo.org>"]
__doc__ = """This executable is used to generate a txt file containing the
source classification probailities"""


class ArgumentParser(_ArgumentParser):
    def _pesummary_options(self):
        options = super(ArgumentParser, self)._pesummary_options()
        options.update(
            {
                "--prior": {
                    "choices": ["population", "default", "both"],
                    "default": "both",
                    "help": (
                        "Prior to use when calculating source classification "
                        "probabilities"
                    )
                },
                "--plot": {
                    "choices": ["bar", "mass_1_mass_2"],
                    "default": "bar",
                    "help": "Name of the plot you wish to make",
                },
            }
        )
        return options


def generate_probabilities(result_files, prior="both", seed=123456789):
    """Generate the classification probabilities

    Parameters
    ----------
    result_files: list
        list of result files
    prior: str
        prior you wish to reweight your samples too
    """
    classifications = []
    if prior == "both":
        _func = "dual_classification"
        _kwargs = {}
    else:
        _func = "classification"
        _kwargs = {"population": True if prior == "population" else False}

    for num, i in enumerate(result_files):
        mydict = {}
        if not _Input.is_pesummary_metafile(i):
            mydict = getattr(
                PEPredicates, "{}_from_file".format(_func)
            )(i, seed=seed, **_kwargs)
            em_bright = getattr(
                PAstro, "{}_from_file".format(_func)
            )(i, seed=seed, **_kwargs)
        else:
            f = GWRead(i)
            label = f.labels[0]
            mydict = getattr(
                 PEPredicates(f.samples_dict[label]), _func
            )(seed=seed, **_kwargs)
            em_bright = getattr(
                PAstro(f.samples_dict[label]), _func
            )(seed=seed, **_kwargs)
        if prior == "both":
            mydict["default"].update(em_bright["default"])
            mydict["population"].update(em_bright["population"])
        else:
            mydict.update(em_bright)
        classifications.append(mydict)
    return classifications


def save_classifications(savedir, classifications, labels):
    """Read and return a list of parameters and samples stored in the result
    files

    Parameters
    ----------
    result_files: list
        list of result files
    classifications: dict
        dictionary of classification probabilities
    """
    import os
    import json

    base_path = os.path.join(savedir, "{}_{}_prior_pe_classification.json")
    for num, i in enumerate(classifications):
        for prior in i.keys():
            with open(base_path.format(labels[num], prior), "w") as f:
                json.dump(i[prior], f)


def make_plots(
    result_files, webdir=None, labels=None, prior=None, plot_type="bar",
    probs=None
):
    """Save the plots generated by PEPredicates

    Parameters
    ----------
    result_files: list
        list of result files
    webdir: str
        path to save the files
    labels: list
        lisy of strings to identify each result file
    prior: str
        Either 'default' or 'population'. If 'population' the samples are reweighted
        to a population prior
    plot_type: str
        The plot type that you wish to make
    probs: dict
        Dictionary of classification probabilities
    """
    if webdir is None:
        webdir = "./"

    for num, i in enumerate(result_files):
        if labels is None:
            label = num
        else:
            label = labels[num]
        f = GWRead(i)
        if not isinstance(f, pesummary.gw.file.formats.pesummary.PESummary):
            f.generate_all_posterior_samples()
        if plot_type == "bar":
            from pesummary.gw.plots.plot import _classification_plot
            if prior == "both":
                probs_func = lambda probs, prior: probs[prior]
            else:
                probs_func = lambda probs, prior: probs
            if prior == "default" or prior == "both":
                fig = _classification_plot(probs_func(probs[num], "default"))
                fig.savefig(
                    os.path.join(
                        webdir,
                        "{}_default_pepredicates_bar.png".format(label)
                    )
                )
            if prior == "population" or prior == "both":
                fig = _classification_plot(probs_func(probs[num], "population"))
                fig.savefig(
                    os.path.join(
                        webdir,
                        "{}_population_pepredicates_bar.png".format(label)
                    )
                )
        elif plot_type == "mass_1_mass_2":
            if prior == "default" or prior == "both":
                fig = PEPredicates.plot(
                    f.samples, f.parameters, population_prior=False
                )
                fig.savefig(
                    os.path.join(
                        webdir, "{}_default_pepredicates.png".format(label)
                    )
                )
            if prior == "population" or prior == "both":
                fig = PEPredicates.plot(f.samples, f.parameters)
                fig.savefig(
                    os.path.join(
                        webdir, "{}_population_pepredicates.png".format(label)
                    )
                )


def main(args=None):
    """Top level interface for `summarypublication`
    """
    parser = ArgumentParser(description=__doc__)
    parser.add_known_options_to_parser(
        ["--webdir", "--samples", "--labels", "--prior", "--plot", "--seed"]
    )
    opts, _ = parser.parse_known_args(args=args)
    if opts.webdir:
        make_dir(opts.webdir)
    else:
        logger.warning(
            "No webdir given so plots will not be generated and "
            "classifications will be shown in stdout rather than saved to file"
        )
    classifications = generate_probabilities(
        opts.samples, prior=opts.prior, seed=opts.seed
    )
    if opts.labels is None:
        opts.labels = []
        for i in opts.samples:
            f = GWRead(i)
            if hasattr(f, "labels"):
                opts.labels.append(f.labels[0])
            else:
                raise InputError("Please provide a label for each result file")
    if opts.webdir:
        if opts.prior != "both":
            _classifications = [{opts.prior: c} for c in classifications]
        else:
            _classifications = classifications
        save_classifications(opts.webdir, _classifications, opts.labels)
    else:
        print(classifications)
        return
    if opts.plot == "bar":
        probs = classifications
    else:
        probs = None
    make_plots(
        opts.samples, webdir=opts.webdir, labels=opts.labels, prior=opts.prior,
        plot_type=opts.plot, probs=probs
    )


if __name__ == "__main__":
    main()
