#!/usr/bin/python3.13 -s
#
# Copyright (C) 2019 Francesco Pannarale
#
# 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.


"""
Workflow generator to run PyGRB offline post-processing.
"""

# =============================================================================
# Preamble
# =============================================================================
import sys
import socket
import logging
import argparse
import os

import pycbc
import pycbc.version
import pycbc.workflow as _workflow
from pycbc.results import layout
from pycbc.results.pygrb_postprocessing_utils import extract_ifos
from pycbc.results import save_fig_with_metadata

__author__ = "Francesco Pannarale  <francesco.pannarale@ligo.org>"
__version__ = pycbc.version.git_verbose_msg
__date__ = pycbc.version.date
__program__ = "pycbc_pygrb_pp_workflow"


# =============================================================================
# Main script
# =============================================================================
# Use the standard workflow command-line parsing routines.
parser = argparse.ArgumentParser(description=__doc__[1:])
pycbc.add_common_pycbc_options(parser)
parser.add_argument("-t", "--trig-files", action="store",
                    required=True, nargs="+",
                    help="The locations of the trigger files "
                         "with assumed order: [ALL_TIMES, ONSOURCE,"
                         " OFFSOURCE, OFFTRIAL_1, ..., OFFTRIAL_N]")
parser.add_argument("-i", "--inj-files", action="store",
                    default=None, nargs="+",
                    help="Location(s) of input injection results file(s)")
parser.add_argument("-b", "--bank-file", action="store",
                    default=None,
                    help="The location of the full template bank file")
parser.add_argument("--segment-dir", action="store",
                    default=None,
                    help="The location of the segment files")
_workflow.add_workflow_command_line_group(parser)
_workflow.add_workflow_settings_cli(parser, include_subdax_opts=True)
args = parser.parse_args()

pycbc.init_logging(args.verbose, format="%(asctime)s: %(levelname)s: %(message)s")

# Store starting run directory
start_rundir = os.getcwd()

# Create the workflow object
logging.info("Generating %s workflow", args.workflow_name)
wflow = _workflow.Workflow(args, name=args.workflow_name)

logging.info("Post-processing output will be generated in %s", args.output_dir)
if not os.path.exists(args.output_dir):
    _workflow.makedir(args.output_dir)
os.chdir(args.output_dir)

# Setup results directory
rdir = layout.SectionNumber('webpage', ['offsource_triggers_vs_time',
                                        'signal_consistency',
                                        'injections',
                                        'loudest_offsource_events',
                                        'exclusion_distances',
                                        'open_box',
                                        'workflow'])
_workflow.makedir(rdir.base)
_workflow.makedir(rdir['workflow'])

# Input trigger files
# Expected structure: [ALL_TIMES, ONSOURCE, OFFSOURCE, OFFTRIAL_1, ..., OFFTRIAL_N]
all_times_file = os.path.join(start_rundir, args.trig_files[0])
onsource_file = os.path.join(start_rundir, args.trig_files[1])
offsource_file = os.path.join(start_rundir, args.trig_files[2])
offtrial_files = [os.path.join(start_rundir, trig_file)
                  for trig_file in args.trig_files[3:]]

logging.info("Using the following trigger files:")
logging.info("All times: %s", all_times_file)
logging.info("Onsource: %s", onsource_file)
logging.info("Offsource: %s", offsource_file)
logging.info("Offtrials: %s", offtrial_files)

bank_file = os.path.join(start_rundir, args.bank_file)

# Input injection files and injection set names
inj_files = [start_rundir+'/'+el for el in args.inj_files]

# IFOs actually used: determined by data availability
ifos = extract_ifos(offsource_file)
wflow.ifos = ifos

plotting_nodes = []
html_nodes = []

seg_filenames = ['bufferSeg.txt', 'offSourceSeg.txt', 'onSourceSeg.txt']
seg_files = [os.path.join(args.segment_dir, f) for f in seg_filenames]

# Logfile of this workflow
wf_log_file = _workflow.File(wflow.ifos, 'workflow-log', wflow.analysis_time,
                             extension='.txt', directory=rdir['workflow'])
logfile = logging.FileHandler(filename=wf_log_file.storage_path, mode='w')
logfile.setLevel(logging.INFO)
formatter = logging.Formatter('%(asctime)s: %(levelname)s: %(message)s')
logfile.setFormatter(formatter)
logging.getLogger('').addHandler(logfile)
logging.info("Created log file %s", wf_log_file.storage_path)

# TODO: Pick up inifile, segments plot, GRB time and location, and
#       report IFO responses
# Read the configuration file
# typecast str from command line to File instances
# cp = configuration.WorkflowConfigParser(opts.pp_config_file)

out_dir = rdir.base
_workflow.makedir(out_dir)
files = _workflow.FileList([])

#
# Create information table about the GRB trigger and plot the search grid
#
info_table_node, grb_info_table = _workflow.make_info_table(wflow, out_dir)
html_nodes.append(info_table_node)
files.append(grb_info_table)
#
plot_node, skygrid_plot = _workflow.make_pygrb_plot(wflow,
                                                    'pygrb_plot_skygrid',
                                                    out_dir,
                                                    trig_file=offsource_file)
plotting_nodes.append(plot_node)
files.append(skygrid_plot)
summary_layout = [(grb_info_table[0],), (skygrid_plot[0],)]
layout.two_column_layout(out_dir, summary_layout)


#
# Plot SNR timeseries
#
out_dir = rdir['offsource_triggers_vs_time']
_workflow.makedir(out_dir)

# Coherent/Reweighted/Single IFO/Null SNR vs time
out_dirs_dict = {'coherent': 'offsource_triggers_vs_time/coh_snr_timeseries',
                 'reweighted': 'offsource_triggers_vs_time/reweighted_snr_timeseries',
                 'single': 'offsource_triggers_vs_time/single_ifo_snr_timeseries',
                 'null': 'offsource_triggers_vs_time/null_snr_timeseries'}

# Grab the name of the prefereed injection set for these plots
tuning_inj_set = wflow.cp.get('workflow-pygrb_pp_workflow', 'tuning-inj-set')
logging.info("The tuning injections set is {0}".format(tuning_inj_set))

# Determine which injections result corresponds to the tuining set
tuning_inj_file = None
if inj_files is not None:
    tuning_inj_file = list(filter(lambda x: tuning_inj_set.lower() in x.lower(), inj_files))
    tuning_inj_file = tuning_inj_file[0] if len(tuning_inj_file) == 1 else None
logging.info("The tuning injections results file is {0}".format(tuning_inj_file))

# Loop over timeseries request by the user
timeseries = wflow.cp.get_subsections('pygrb_plot_snr_timeseries')
for snr_type in timeseries:
    out_dir = rdir[out_dirs_dict[snr_type]]
    _workflow.makedir(out_dir)
    files = _workflow.FileList([])
    # Only single SNR timeseries requires looping over IFOs
    ifos_to_loop = ifos if snr_type == 'single' else [None]
    for ifo in ifos_to_loop:
        timeseries_plots = _workflow.FileList([])
        # Plots without and with injections
        for inj_file in set([None, tuning_inj_file]):
            # If the inj_file is used for the plot, include the tuning
            # injection set tag in assembling the plot name
            tags = [snr_type] if inj_file is None else [snr_type, tuning_inj_set]
            plot_node, output_files = \
                _workflow.make_pygrb_plot(wflow, 'pygrb_plot_snr_timeseries',
                                          out_dir, trig_file=offsource_file,
                                          inj_file=inj_file, ifo=ifo,
                                          tags=tags)
            plotting_nodes.append(plot_node)
            # We want a File, not a 1-element list with a File
            # pycbc_pygrb_plot_snr_timeseries produces only one plot: take [0]
            timeseries_plots.append(output_files[0])
        files.append(timeseries_plots)
    layout.two_column_layout(out_dir, files)


#
# Signal consistency plots
#
out_dir = rdir['signal_consistency']
_workflow.makedir(out_dir)
# Chisq veto vs SNR plots
out_dir = rdir['signal_consistency/chi_squared_tests']
_workflow.makedir(out_dir)
files = _workflow.FileList([])
# Loop over vetoes request by the user
vetoes = wflow.cp.get_subsections('pygrb_plot_chisq_veto')
for veto in vetoes:
    # Loop over IFOs only in the case of single detector chi-squares
    ifo_loop = [ifos[0]] if veto == 'network' else ifos
    for ifo in ifo_loop:
        # Plot with and without injections
        for inj_file in set([tuning_inj_file, None]):
            # If the inj_file is used for the plot, include the tuning
            # injection set tag in assembling the plot name
            tags = [veto] if inj_file is None else [veto, tuning_inj_set]
            ifo_arg = None if veto == 'network' else ifo
            plot_node, output_files = \
                _workflow.make_pygrb_plot(wflow, 'pygrb_plot_chisq_veto',
                                          out_dir, trig_file=offsource_file,
                                          inj_file=inj_file, ifo=ifo_arg,
                                          tags=tags)
            plotting_nodes.append(plot_node)
            # We want a File, not a 1-element list with a File
            # pycbc_pygrb_plot_chisq_veto produces only one plot: take [0]
            files.append(output_files[0])
chisq_layout = list(layout.grouper(files, 2))
layout.two_column_layout(out_dir, chisq_layout)

# Single detector chi-square plots: zoomed in and zoomed out
out_dir = rdir['signal_consistency/individual_detector_snrs']
_workflow.makedir(out_dir)
files = _workflow.FileList([])
# Single IFO SNR vs Coherent SNR plots: zoomed in and zoomed out
# Requires looping over IFOs
if wflow.cp.has_section('pygrb_plot_coh_ifosnr'):
    for ifo in ifos:
        # Plots with and without injections
        for inj_file in set([tuning_inj_file, None]):
            sngl_snr_plots = _workflow.FileList([])
            for zoom_tag in [['zoomin'], []]:
                # If the inj_file is used for the plot, include the tuning
                # injection set tag in assembling the plot name
                tags = zoom_tag if inj_file is None else zoom_tag + [tuning_inj_set]
                # Single IFO SNR vs Coherent SNR
                plot_node, output_files = \
                    _workflow.make_pygrb_plot(wflow, 'pygrb_plot_coh_ifosnr',
                                              out_dir,
                                              trig_file=offsource_file,
                                              inj_file=inj_file, ifo=ifo,
                                              tags=tags)
                plotting_nodes.append(plot_node)
                # We want a File, not a 1-element list with a File
                # pycbc_pygrb_plot_cohifo_snr produces only one plot: take [0]
                sngl_snr_plots.append(output_files[0])
            files.append(sngl_snr_plots)
    layout.two_column_layout(out_dir, files)
else:
    msg = 'No pygrb_plot_coh_ifosnr section found in the configuration file. '
    msg += 'No coherent vs single detector SNR plots will be generated.'
    logging.info(msg)

# Null SNR/Overwhitened null stat vs Coherent SNR plots
null_snr_out_dir = rdir['signal_consistency/null_snrs']
_workflow.makedir(null_snr_out_dir)
null_snr_files = _workflow.FileList([])
# Coincident SNR vs Coherent SNR plots
coinc_out_dir = rdir['signal_consistency/coincident_snr']
_workflow.makedir(coinc_out_dir)
coinc_files = _workflow.FileList([])
# Loop over null statistics requested by the user (including coincident SNR)
nstats = wflow.cp.get_subsections('pygrb_plot_null_stats')
for nstat in nstats:
    # Plots with and without injections, zoomed in and zoomed out
    for inj_file in set([tuning_inj_file, None]):
        if nstat == 'coincident':
            out_dir = coinc_out_dir
            files = coinc_files
        else:
            out_dir = null_snr_out_dir
            files = null_snr_files
        null_stats_plots = _workflow.FileList([])
        # for zoom_tag in [['zoomin'], []]: # IDEA
        for zoom_tag in ['zoomin', None]:
            # If the inj_file is used for the plot, include the tuning
            # injection set tag in assembling the plot name
            tags = [nstat]
            tags = tags + [zoom_tag] if zoom_tag is not None else tags
            tags = tags + [tuning_inj_set] if inj_file is not None else tags
            plot_node, output_files = \
                _workflow.make_pygrb_plot(wflow, 'pygrb_plot_null_stats',
                                          out_dir, trig_file=offsource_file,
                                          inj_file=inj_file, tags=tags)
            plotting_nodes.append(plot_node)
            # We want a File, not a 1-element list with a File
            # pycbc_pygrb_plot_null_stats produces only one plot: take [0]
            null_stats_plots.append(output_files[0])
        files.append(null_stats_plots)
layout.two_column_layout(null_snr_out_dir, null_snr_files)
layout.two_column_layout(coinc_out_dir, coinc_files)

# layout.group_layout(rdir['coincident_triggers'],
#                     closed_box_ifars + all_snrifar + [bank_plot[0][0]])

#
# Found/missed injections plots and tables
#
out_dir = rdir['injections']
_workflow.makedir(out_dir)
# Loop over injection plots requested by the user
inj_sets = wflow.cp.get_subsections('injections')
inj_plots = wflow.cp.get_subsections('pygrb_plot_injs_results')
# The command above also picks up the injection set names so we remove them
# from the set of requested injection plot types
inj_plots = [inj_plot for inj_plot in inj_plots if inj_plot not in inj_sets]
for inj_set in inj_sets:
    inj_file = list(filter(lambda x: inj_set.lower() in x.lower(), inj_files))
    inj_file = inj_file[0] if len(inj_file) == 1 else None
    out_dir = rdir['injections/'+inj_set]
    _workflow.makedir(out_dir)
    files = _workflow.FileList([])
    # Generate plots: loop over inj_plots and found-missed/missed-found
    for inj_plot in inj_plots:
        y_qty, x_qty = inj_plot.split('_')
        ifos_to_loop = [None]
        if y_qty == 'effsitedist':
            ifos_to_loop = ifos
        for ifo in ifos_to_loop:
            for fm_or_mf in ['missed-on-top', 'found-on-top']:
                tags = [y_qty, x_qty, fm_or_mf, inj_set]
                plot_node, output_files = \
                    _workflow.make_pygrb_plot(wflow, 'pygrb_plot_injs_results',
                                              out_dir,
                                              trig_file=offsource_file,
                                              ifo=ifo, inj_file=inj_file,
                                              tags=tags)
                plotting_nodes.append(plot_node)
                # We want a File, not a 1-element list with a File
                # pycbc_pygrb_plot_injs_results produces only one plot, so
                # take [0]
                files.append(output_files[0])
    # Generate quiet-found and missed-found html tables
    # TODO:  pass loudest-onsource-trigger/loudest-offsource-triggers
    # arguments if inj_set=None
    inj_layout = list(layout.grouper(files, 2))
    layout.two_column_layout(out_dir, inj_layout)

    # TODO: Follow up of loudest N quiet-found injections:


#
# FAP distributions
#
out_dir = rdir['loudest_offsource_events']
_workflow.makedir(out_dir)
files = []
# Loop over statistics requested by the user
stats = wflow.cp.get_subsections('pygrb_plot_stats_distribution')
for stat in stats:
    plot_node, output_file = \
        _workflow.make_pygrb_plot(wflow, 'pygrb_plot_stats_distribution',
                                  out_dir, trig_file=offsource_file,
                                  inj_file=inj_file, tags=[stat])
    plotting_nodes.append(plot_node)
    # We want a File, not a 1-element list with a File
    # pycbc_pygrb_plot_stats_distribution produces only one plot: take [0]
    files.append(output_file[0])

# TODO: Include minifollow-ups materiale


#
# Exclusion distance and efficiency plots based on offtrials
#
out_dir = rdir['exclusion_distances']
_workflow.makedir(out_dir)

# Offtrials and injection sets requested by the user
num_trials = int(wflow.cp.get('trig_combiner', 'num-trials'))
offtrials = ["offtrial_%s" % (i+1) for i in range(num_trials)]
inj_sets = wflow.cp.get_subsections('injections')
out_dir = rdir['exclusion_distances']
_workflow.makedir(out_dir)
for i, offtrial in enumerate(offtrials):
    out_dir = rdir[f'exclusion_distances/{offtrial}']
    _workflow.makedir(out_dir)
    files = _workflow.FileList([])
    for inj_set in inj_sets:
        inj_file = list(filter(lambda x: inj_set.lower() in x.lower(), inj_files))
        inj_file = inj_file[0] if len(inj_file) == 1 else None
        tags = [offtrial, inj_set]
        plot_node, output_files = \
            _workflow.make_pygrb_plot(wflow, 'pygrb_efficiency',
                                      out_dir, trig_file=offsource_file,
                                      onsource_file=offtrial_files[i],
                                      inj_file=inj_file,
                                      bank_file=bank_file,
                                      seg_files=seg_files, tags=tags)
        plotting_nodes.append(plot_node)
        files.extend(output_files)
    eff_layout = list(layout.grouper(files, 2))
    layout.two_column_layout(out_dir, eff_layout, offtrial)

# TODO: add onsource calculation by using the true onsource file

#
# Trigger Followups (TODO)
#

# Make room for throughput histograms (TODO)
base = rdir['workflow/throughput']
_workflow.makedir(base)

# Save global config file
base = rdir['workflow/configuration']
_workflow.makedir(base)
ini_file_path = os.path.join(base, args.workflow_name+'.ini')
with open(ini_file_path, 'w') as ini_fh:
    wflow.cp.write(ini_fh)
ini_file = _workflow.FileList([_workflow.File(wflow.ifos, '',
                                              wflow.analysis_time,
                                              file_url='file://' + ini_file_path)])
layout.single_layout(base, ini_file)

# Create versioning information
_workflow.make_versioning_page(
    wflow,
    wflow.cp,
    rdir['workflow/version'],
)

# Create the final log file
log_file_html = _workflow.File(wflow.ifos, 'WORKFLOW-LOG', wflow.analysis_time,
                               extension='.html', directory=rdir['workflow'])

# Create a page to contain a dashboard link
dashboard_file = _workflow.File(wflow.ifos, 'DASHBOARD', wflow.analysis_time,
                                extension='.html', directory=rdir['workflow'])
dashboard_str = """<center><p style="font-size:20px"><b><a href="PEGASUS_DASHBOARD_URL" target="_blank">Pegasus Dashboard Page</a></b></p></center>"""
kwds = {'title': 'Pegasus Dashboard',
        'caption': "Link to Pegasus Dashboard",
        'cmd': "PYCBC_SUBMIT_DAX_ARGV", }
save_fig_with_metadata(dashboard_str, dashboard_file.storage_path, **kwds)

# Create pages for the submission script to write data
_workflow.makedir(rdir['workflow/dax'])
_workflow.makedir(rdir['workflow/input_map'])
_workflow.makedir(rdir['workflow/output_map'])
_workflow.makedir(rdir['workflow/planning'])

logging.info("Path for make_results_web_page: %s",
             os.path.join(os.getcwd(), rdir.base))
_workflow.make_results_web_page(
    wflow, os.path.join(os.getcwd(), rdir.base), template='red',
    explicit_dependencies=plotting_nodes+html_nodes)

# Protect the open box results folder
out_dir = rdir['open_box']
_workflow.makedir(out_dir)
os.chmod(out_dir, 0o0700)
# TODO: follow up loudest offsource trigger

# Close the log and flush to the html file
logging.shutdown()
with open(wf_log_file.storage_path, "r") as logfile:
    logdata = logfile.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(), args.workflow_name, socket.gethostname(), logdata)
kwds = {'title': 'Workflow Generation Log',
        'caption': "Log of the workflow script %s" % sys.argv[0],
        'cmd': ' '.join(sys.argv), }
save_fig_with_metadata(log_str, log_file_html.storage_path, **kwds)
layout.single_layout(rdir['workflow'], ([dashboard_file, log_file_html]))

# Go back to the run directory and save the workflow
os.chdir(start_rundir)
wflow.save()
logging.info("Dax written.")
