#!/usr/bin/env python

# Copyright (C) 2019 Collin D. Capano, Christopher M. Biwer
#
# 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 that generates a posterior file and plots from one or more
inference samples files.
"""

import argparse
import h5py
import logging
import os
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 shlex
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 pegasus_workflow
from pycbc.workflow import datafind
from pycbc.workflow import plotting
from pycbc import __version__
import pycbc.workflow.inference_followups as inffu


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 which gives the configuration file(s) that were
    used for the event, and a ``samples-files`` option which gives one or
    more samples files from which to extract a posterior. 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.

    To specify multiple configuration or samples files, the files should be
    space (or new-line) separated. Relative or absolute paths may be used.

    Example:

    .. code-block:: ini

       [event-gw150914_09_50_45]
       label = GW150914+09:50:45UTC
       config-files = inference-150914_09h_50m_45UTC.ini
       samples-files = run1/H1L1-INFERENCE_150914_09H_50M_45UTC-5-1.hdf
                       run2/H1L1-INFERENCE_150914_09H_50M_45UTC-5-1.hdf
                       run3/H1L1-INFERENCE_150914_09H_50M_45UTC-5-1.hdf

    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) that were given in the config file.
    config-files : list of lists
        List of the configuration file(s) for each event.
    samples-files : lsit of lists
        List of the samples file(s) for each event.
    """
    # get the events
    events = cp.get_subsections('event')
    labels = []
    config_files = []
    samples_files = []
    for event in events:
        section = '-'.join(['event', event])
        if cp.has_option(section, 'label'):
            label = cp.get(section, 'label')
        else:
            label = event
        cf = shlex.split(cp.get(section, 'config-files'))
        sf = shlex.split(cp.get(section, 'samples-files'))
        labels.append(label)
        config_files.append(list(map(os.path.abspath, cf)))
        samples_files.append(list(map(os.path.abspath, sf)))
    return events, labels, config_files, samples_files


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("--version", action="version", version=__version__,
                    help="Prints version information.")
opts = parser.parse_args()

posterior_file_dir = 'posterior_files'
config_file_dir = 'config_files'
config_file_tmplt = 'inference-{}.ini'

# 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)

# the file we'll store all output files in
filedir = 'output'

# 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")

# get the events
events, labels, infconfig_files, samples_files = \
    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 = {}
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]
    samples_filelist = samples_files[num_event]

    config_fnames = infconfig_files[num_event]

    # write the configuration file to the config files directory
    cp = configuration.WorkflowConfigParser(config_fnames)
    config_file = 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])

    # convert the samples files to workflow File types
    for ii, fname in enumerate(samples_filelist):
        pfile = pegasus_workflow.File(fname)
        pfile.PFN(fname, "local")
        # set the storage path to be the same
        pfile.storage_path = fname
        samples_filelist[ii] = pfile
    samples_filelist = core.FileList(samples_filelist)

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

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

    # files for detector_sensitivity summary subsection
    base = "detector_sensitivity"
    # we'll just use the first file, and assume the rest are the same
    psd_plot = plotting.make_spectrum_plot(
        workflow, [samples_filelist[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)

# 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]))
