#! /usr/bin/env python

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

import os
import numpy as np

from pesummary.utils.samples_dict import MultiAnalysisSamplesDict
from pesummary.utils.utils import logger
from pesummary.core.cli.parser import ArgumentParser as _ArgumentParser
from pesummary.core.cli.inputs import _Input
from pesummary.io import read

__author__ = ["Charlie Hoy <charlie.hoy@ligo.org>"]
__doc__ = """This executable is used to combine posterior samples. This is
different from 'summarycombine' as 'summarycombine' combines N files into a single
metafile containing N analyses while 'summarycombine_posteriors' combines N
posterior samples and creates a single file containing a single analysis"""


class ArgumentParser(_ArgumentParser):
    def _pesummary_options(self):
        options = super(ArgumentParser, self)._pesummary_options()
        options.update(
            {
                "--labels": {
                    "nargs": "+",
                    "required": True,
                    "help": (
                        "Analyses you wish to combine. If only a file "
                        "containing more than one analysis is provided, please "
                        "pass the labels in that file that you wish to "
                        "combine. If multiple single analysis files are "
                        "provided, please pass unique labels to distinguish "
                        "each analysis. If a file containing more than one "
                        "analysis is provided alongside a single analysis "
                        "file, or multiple files containing more than one "
                        "analysis each, only a single analysis can be combined "
                        "from each file"
                    )
                },
                "--weights": {
                    "nargs": "+",
                    "type": float,
                    "help": (
                        "Weights to assign to each analysis. Must be same "
                        "length as labels"
                    )
                },
                "--use_all": {
                    "action": "store_true",
                    "default": False,
                    "help": "Use all posterior samples (do not weight)"
                },
                "--shuffle": {
                    "action": "store_true",
                    "default": False,
                    "help": "Shuffle the combined samples"
                },
                "--file_format": {
                    "type": str,
                    "default": "dat",
                    "help": "Format of output file"
                },
                "--filename": {
                    "type": str,
                    "help": "Name of the output file"
                },
                "--outdir": {
                    "type": str,
                    "default": "./",
                    "help": "Directory to save the file",
                },
                "--add_to_existing": {
                    "action": "store_true",
                    "default": False,
                    "help": (
                        "Add the combined samples to an existing metafile. "
                        "Only used when a PESummary metafile is provided via "
                        "the `--samples` option. If this option is provided, "
                        "the `--file_format` and `--filename` options are "
                        "ignored"
                    )
                }
            }
        )
        return options


class Input(_Input):
    """Class to handle the core command line arguments

    Parameters
    ----------
    opts: argparse.Namespace
        Namespace object containing the command line options

    Attributes
    ----------
    result_files: list
        list of result files passed
    labels: list
        list of labels used to distinguish the result files
    """
    def __init__(self, opts):
        self.opts = opts
        self.seed = self.opts.seed
        self.result_files = self.opts.samples
        self.mcmc_samples = False
        self.add_to_existing = False
        cond = np.sum([self.is_pesummary_metafile(f) for f in self.result_files])
        if cond > 1:
            raise ValueError(
                "Can only combine analyses from a single PESummary metafile"
            )
        elif cond == 1 and len(self.result_files) > 1:
            raise ValueError(
                "Can only combine analyses from a single PESummary metafile "
                "or multiple non-PESummary metafiles"
            )
        self.pesummary = False
        if self.is_pesummary_metafile(self.result_files[0]):
            self.pesummary = True
            self._labels = self.opts.labels
        else:
            self.labels = self.opts.labels


def main(args=None):
    """Top level interface for `summarycombine_posteriors`
    """
    parser = ArgumentParser()
    parser.add_known_options_to_parser(
        [
            "--add_to_existing", "--labels", "--weights", "--samples",
            "--outdir", "--filename", "--seed", "--shuffle", "--file_format",
            "--use_all"
        ]
    )
    opts, unknown = parser.parse_known_args(args=args)
    args = Input(opts)
    if not args.pesummary:
        samples = {
            label: samples for label, samples in
            zip(args.labels, args.result_files)
        }
        mydict = MultiAnalysisSamplesDict.from_files(
            samples, disable_prior=True, disable_injection_conversion=True
        )
    else:
        mydict = read(
            args.result_files[0], disable_prior=True,
            disable_injection_conversion=True
        ).samples_dict
    combined = mydict.combine(
        use_all=opts.use_all, weights=opts.weights, labels=args.labels,
        shuffle=opts.shuffle, logger_level="info"
    )
    if opts.add_to_existing and args.pesummary:
        from .summarymodify import _Input as _ModifyInput
        from pesummary.gw.file.meta_file import _GWMetaFile
        from pathlib import Path

        class PESummaryInput(_ModifyInput):
            def __init__(self, samples):
                self.samples = samples
                self.data = None

        _args = PESummaryInput(args.result_files)
        data = _args.data
        data["{}_combined".format("_".join(args.labels))] = {
            "posterior_samples": combined.to_structured_array(),
            "meta_data": {
                "sampler": {"nsamples": combined.number_of_samples},
                "meta_data": {"combined": ", ".join(args.labels)}
            }
        }
        input_file = Path(args.result_files[0])
        if opts.filename is None:
            logger.warning(
                "No filename specified. Saving samples to '{}' in '{}' by "
                "default.".format(
                    input_file.stem + "_combined" + input_file.suffix,
                    opts.outdir
                )
            )
            opts.filename = input_file.stem + "_combined" + input_file.suffix
        _metafile = os.path.join(opts.outdir, opts.filename)
        logger.info("Saving samples to file '{}'".format(_metafile))
        if os.path.isfile(_metafile):
            raise FileExistsError(
                "Trying to write samples to '{}' but the file exists. Please "
                "specify a different out directory".format(_metafile)
            )
        if input_file.suffix in [".h5", ".hdf5"]:
            _GWMetaFile.save_to_hdf5(
                data, list(data.keys()), None, _metafile, no_convert=True
            )
        else:
            _GWMetaFile.save_to_json(data, _metafile)
        return
    elif opts.add_to_existing:
        logger.warning(
            "Can only use the `--add_to_existing` option when a PESummary "
            "metafile is provided via the `--samples` option. Writing "
            "combined samples to a `dat` file"
        )
    logger.info(
        "Saving samples to file '{}'".format(
            os.path.join(opts.outdir, opts.filename)
        )
    )
    combined.write(
        file_format=opts.file_format, filename=opts.filename, outdir=opts.outdir
    )


if __name__ == "__main__":
    main()
