#!/usr/bin/env python

# Copyright (C) 2016 Christopher M. Biwer, Alexander Harvey Nitz
#
# This program is free software; you can redistribute it and/or modify it
# under the terms of the GNU General Public License as published by the
# Free Software Foundation; either version 3 of the License, or (at your
# option) any later version.
#
# This program is distributed in the hope that it will be useful, but
# WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU General
# Public License for more details.
#
# You should have received a copy of the GNU General Public License along
# with this program; if not, write to the Free Software Foundation, Inc.,
# 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301, USA.
"""Creates a DAX for a parameter estimation workflow.
"""

import argparse
import h5py
import logging
import os
import shlex
import Pegasus.DAX3 as dax
import pycbc
import pycbc.workflow.minifollowups as mini
import pycbc.workflow.pegasus_workflow as wdax
import socket
import sys
import numpy
from ligo import segments
from pycbc import results
from pycbc.results import layout
from pycbc.types import MultiDetOptionAction
from pycbc.types import MultiDetOptionAppendAction
from pycbc.workflow import configuration
from pycbc.workflow import core
from pycbc.workflow import datafind
from pycbc.workflow import plotting
from pycbc import __version__
import pycbc.workflow.inference_followups as inffu
from pycbc.workflow.jobsetup import PycbcInferenceExecutable


def read_events_from_config(cp):
    """Gets events to load from a config file.

    Each event should have its own section with header ``[event-{{NAME}}]``,
    where ``NAME`` is a unique identifier. The section must have a
    ``config-files`` option that gives the configuration file(s) to use for the
    event. To specify multiple configuration files, the files should be
    space (or new-line) separated. Relative or absolute paths may be used.

    Overrides and deletions for the config file(s) may be provided
    via ``config-overrides`` and ``config-delete``, respectively. These should
    be space (or new line) separated. The overrides should have format
    ``SECTION:OPTION:VALUE`` and the deletions ``SECTION:OPTION``.

    The section may also have a ``label`` option that provides a human-readable
    label for the results page. If no ``label`` option is provided, the
    ``{{NAME}}`` will be used.

    Multiple instances of an inference job may be invoked by adding a ``nruns``
    option, and setting it to an integer larger than 1. In this case, multiple
    inference analyses will be carried out on the event with the same settings,
    only differing by the seed given to ``pycbc_inference``. The output from
    the multiple runs will be combined into a single posterior file, with all
    posterior plots generated from the posterior file. If the ``nruns`` option
    isn't provided, a single inference job will run for each event.

    The order that events are provided in the configuration file will be the
    order they are presented on the results page.

    Example:

    .. code-block:: ini

       [event-gw150914]                                                                
       label = GW150914+09:50:45UTC                                                    
       config-files = sampler.ini                                    
                      gw150914_like.ini                              
                      o1_data.ini                                    
       config-overrides = data:trigger-time:1126259462.43                              
       nruns = 3

    Parameters
    ----------
    cp : pycbc.workflow.configuration.WorkflowConfigParser
        Configuration file giving the events.

    Returns
    -------
    events : list of str
        The names of the event(s) that were given in the config file.
    labels : list of str
        The labels to use for the event(s).
    config_opts : list
        The options for loading the inference config file, as parsed
        argparse.ArgumentParser options.
    nruns : list of int
        The number of inference jobs to create for the event.
    """
    # get the events
    events = cp.get_subsections('event')
    # create a dummy command-line parser for getting the config files and
    # options
    cfparser = argparse.ArgumentParser()
    configuration.add_workflow_command_line_group(cfparser)
    # lists for storing output
    labels = []
    cpopts = []
    nruns = []
    for event in events:
        section = '-'.join(['event', event])
        # get the label
        if cp.has_option(section, 'label'):
            label = cp.get(section, 'label')
        else:
            label = event
        labels.append(label)
        # convert the config-file options to a command line string
        cli = cp.section_to_cli(section, skip_opts=['label', 'nruns'])
        cpopts.append(cfparser.parse_args(shlex.split(cli)))
        # get the number of times to run the event
        if cp.has_option(section, 'nruns'):
            nrun = int(cp.get(section, 'nruns'))
        else:
            nrun = 1
        if nrun < 1:
            raise ValueError('nruns must be >= 1')
        nruns.append(nrun)
    return events, labels, cpopts, nruns


def event_slug(label):
    """Slugifies an event label."""
    return label.replace(' ', '_').replace(':', '_').replace('+', '_')


def symlink_path(f, path):
    """ Symlinks a path.
    """
    if f is None:
        return
    try:
        os.symlink(f.storage_path, os.path.join(path, f.name))
    except OSError:
        pass


# command line parser
parser = argparse.ArgumentParser(description=__doc__[1:])
# add option groups
configuration.add_workflow_command_line_group(parser)
# workflow options
core.add_workflow_settings_cli(parser, include_subdax_opts=True)
parser.add_argument("--seed", type=int, default=0,
                    help="Seed to use for inference job(s). If multiple "
                         "events are analyzed, the seed will be incremented "
                         "by one for each event.")
# version option
parser.add_argument("--version", action="version", version=__version__,
                    help="Prints version information.")


# parser command line
opts = parser.parse_args()

# configuration files
config_file_tmplt = 'inference-{}.ini'
config_file_dir = 'config_files'
# the directory we'll store samples files to
samples_file_dir = 'samples_files'
# the directory we'll store posterior files to
posterior_file_dir = 'posterior_files'

# make data output directory
if opts.output_dir is None:
    opts.output_dir = opts.workflow_name + '_output'
core.makedir(opts.output_dir)
core.makedir('{}/{}'.format(opts.output_dir, config_file_dir))
core.makedir('{}/{}'.format(opts.output_dir, posterior_file_dir))

# set the dax file name
dax_file = opts.dax_file
if dax_file is None:
    dax_file = "{}.dax".format(opts.workflow_name)

# log to terminal until we know where the path to log output file
log_format = "%(asctime)s:%(levelname)s : %(message)s"
logging.basicConfig(format=log_format, level=logging.INFO)

# create workflow and sub-workflows
container = core.Workflow(opts, opts.workflow_name)
workflow = core.Workflow(opts, opts.workflow_name + "-main")
finalize_workflow = core.Workflow(opts, opts.workflow_name + "-finalization")

# read the events to analyze
events, labels, cpopts, nruns = read_events_from_config(workflow.cp)

# change working directory to the output
origdir = os.path.abspath(os.curdir)
os.chdir(opts.output_dir)

# figure out what diagnostic jobs there are
diagnostics = inffu.get_diagnostic_plots(workflow)

# sections for output HTML pages
rdir = layout.SectionNumber("results",
                            ["detector_sensitivity", "priors", "posteriors"] +
                            diagnostics +
                            ["config_files", "workflow"])

# make results directories
core.makedir(rdir.base)
core.makedir(rdir["workflow"])
core.makedir(rdir["config_files"])

# create files for workflow log
log_file_txt = core.File(workflow.ifos, "workflow-log", workflow.analysis_time,
                         extension=".txt", directory=rdir["workflow"])
log_file_html = core.File(workflow.ifos, "WORKFLOW-LOG",
                          workflow.analysis_time,
                          extension=".html", directory=rdir["workflow"])

# switch saving log to file
logging.basicConfig(format=log_format, level=logging.INFO,
                    filename=log_file_txt.storage_path, filemode="w")
log_file = logging.FileHandler(filename=log_file_txt.storage_path, mode="w")
log_file.setLevel(logging.INFO)
formatter = logging.Formatter(log_format)
log_file.setFormatter(formatter)
logging.getLogger("").addHandler(log_file)
logging.info("Created log file %s" % log_file_txt.storage_path)

config_files = {}
seed = opts.seed
# loop over the events to be analyzed
for num_event, event in enumerate(events):
    # slugify the event name so it can be used in file names
    event = event_slug(event)
    label = labels[num_event]
    cpopt = cpopts[num_event]
    nrun = nruns[num_event]

    # create a sub workflow for this event
    # we need to go back to the original directory to do this for all the file
    # references to work correctly
    os.chdir(origdir)
    sub_workflow = core.Workflow(opts,
                                 "{}-{}".format(opts.workflow_name, event))
    # load the inference config file
    cp = configuration.WorkflowConfigParser.from_cli(cpopt)

    # now go back to the output
    os.chdir(opts.output_dir)

    # write the configuration file to the config files directory
    config_file = sub_workflow.save_config(config_file_tmplt.format(event),
                                           config_file_dir, cp)[0]

    # create sym links to config file for results page
    base = "config_files/{}".format(event)
    layout.single_layout(rdir[base], [config_file])
    symlink_path(config_file, rdir[base])

    # make node(s) for running sampler
    samples_files = []
    inference_exe = PycbcInferenceExecutable(sub_workflow.cp, "inference",
                                             ifos=sub_workflow.ifos,
                                             out_dir=samples_file_dir)
    for nn in range(nrun):
        tags = opts.tags + [event]
        if nrun > 1:
            tags.append(str(nn))
        node, samples_file = inference_exe.create_node(
            config_file, seed=seed, tags=tags,
            analysis_time=sub_workflow.analysis_time)
        samples_files.append(samples_file)
        # add node to workflow
        sub_workflow += node
        # increment the seed
        seed = seed + 1

    # create the posterior file and plots
    posterior_file, summary_files, _, _ = inffu.make_posterior_workflow(
        sub_workflow, samples_files, config_file, event, rdir,
        posterior_file_dir=posterior_file_dir, tags=opts.tags)

    # create the diagnostic plots
    _ = inffu.make_diagnostic_plots(sub_workflow, diagnostics, samples_files,
                                    event, rdir, tags=opts.tags)

    # files for detector_sensitivity summary subsection
    base = "detector_sensitivity"
    psd_plot = plotting.make_spectrum_plot(
        sub_workflow, [samples_files[0]], rdir[base],
        tags=opts.tags+[event],
        hdf_group="data")

    # build the summary page
    zpad = int(numpy.ceil(numpy.log10(len(samples_files))))
    layout.two_column_layout(rdir.base, summary_files,
                             unique=str(num_event).zfill(zpad),
                             title=label, collapse=True)

    # build the psd page
    layout.single_layout(rdir['detector_sensitivity'], [psd_plot],
                         unique=str(num_event).zfill(zpad),
                         title=label, collapse=True)

    # add the sub workflow to the main workflow
    workflow += sub_workflow


# create versioning HTML pages
results.create_versioning_page(rdir["workflow/version"], container.cp)

# create node for making HTML pages
plotting.make_results_web_page(finalize_workflow,
    os.path.join(os.getcwd(), rdir.base))

# add sub-workflows to workflow
container += workflow
container += finalize_workflow

# make finalize sub-workflow depend on main sub-workflow
dep = dax.Dependency(parent=workflow.as_job, child=finalize_workflow.as_job)
container._adag.addDependency(dep)

# write dax
container.save(filename=dax_file, output_map_path=opts.output_map,
               transformation_catalog_path=opts.transformation_catalog)

# save workflow configuration file
base = rdir["workflow/configuration"]
core.makedir(base)
wf_ini = workflow.save_config("workflow.ini", base, container.cp)
layout.single_layout(base, wf_ini)

# close the log and flush to the html file
logging.shutdown()
with open (log_file_txt.storage_path, "r") as log_file:
    log_data = log_file.read()
log_str = """
<p>Workflow generation script created workflow in output directory: %s</p>
<p>Workflow name is: %s</p>
<p>Workflow generation script run on host: %s</p>
<pre>%s</pre>
""" % (os.getcwd(), opts.workflow_name, socket.gethostname(), log_data)
kwds = {"title" : "Workflow Generation Log",
        "caption" : "Log of the workflow script %s" % sys.argv[0],
        "cmd" : " ".join(sys.argv)}
results.save_fig_with_metadata(log_str, log_file_html.storage_path, **kwds)
layout.single_layout(rdir["workflow"], ([log_file_html]))
