#!/usr/bin/env python

# Copyright (C) 2013 Ian W. Harry, Alex 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.
"""
Program for running multi-detector workflow analysis through coincidence and
then generate post-processing and plots.
"""
import pycbc
import pycbc.version
__author__  = "Alex Nitz <alex.nitz@ligo.org>"
__version__ = pycbc.version.git_verbose_msg
__date__    = pycbc.version.date
__program__ = "pycbc_offline"

import sys
import socket
import pycbc.events, pycbc.workflow as wf
import os, argparse, ConfigParser, logging
from glue import segments
import numpy, lal, datetime
from pycbc.results import create_versioning_page, static_table, layout
from pycbc.results.versioning import save_fig_with_metadata
from pycbc.results.metadata import html_escape

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

def symlink_result(f, rdir_path):
    symlink_path(f, rdir[rdir_path])

# Log to the screen until we know where the output file is
logging.basicConfig(format='%(asctime)s:%(levelname)s : %(message)s',
    level=logging.INFO)

parser = argparse.ArgumentParser(description=__doc__[1:])
parser.add_argument('--version', action='version', version=__version__)
parser.add_argument('--workflow-name', default='my_unamed_run')
parser.add_argument("-d", "--output-dir", default=None,
                    help="Path to output directory.")
wf.add_workflow_command_line_group(parser)
args = parser.parse_args()

wf.makedir(args.output_dir)

container = wf.Workflow(args, args.workflow_name)
workflow = wf.Workflow(args, args.workflow_name + '-main')
finalize_workflow = wf.Workflow(args, args.workflow_name + '-finalization')

os.chdir(args.output_dir)

rdir = layout.SectionNumber('results', ['analysis_time',
                                 'detector_sensitivity',
                                 'single_triggers',
                                 'coincident_triggers',
                                 'injections',
                                 'search_sensitivity',
                                 'open_box_result',
                                 'workflow',
                                 ])

wf.makedir(rdir.base)
wf.makedir(rdir['workflow'])

wf_log_file = wf.File(workflow.ifos, 'workflow-log', workflow.analysis_time,
                      extension='.txt',
                      directory=rdir['workflow'])

logging.basicConfig(format='%(asctime)s:%(levelname)s : %(message)s',
                    filename=wf_log_file.storage_path,
                    level=logging.INFO,
                    filemode='w')

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)

# put start / end time at top of summary page
time = workflow.analysis_time
s, e = int(time[0]), int(time[1])
s_utc = str(datetime.datetime(*lal.GPSToUTC(s)[0:6]))
e_utc = str(datetime.datetime(*lal.GPSToUTC(e)[0:6]))
time_str = '<center><p><b>GPS Interval [%s,%s). ' %(s,e)
time_str += 'UTC Interval %s - %s. ' %(s_utc, e_utc)
time_str += 'Interval duration = %.3f days.</b></p></center>'\
                                                         %(float(e-s)/86400.0,)
time_file = wf.File(workflow.ifos, 'time', workflow.analysis_time,
                                           extension='.html',
                                           directory=rdir.base)
kwds = { 'title' : 'Search Workflow Duration (Wall Clock Time)',
        'caption' : "Wall clock start and end times for this invocation of "
                    "the workflow. The command line button shows the "
                    "arguments used to invoke the workflow creation script.",
        'cmd' :' '.join(sys.argv), }
save_fig_with_metadata(time_str, time_file.storage_path, **kwds)

# Get segments and find where the data is
wf.makedir(rdir['analysis_time/segment_data'])
science_veto_name = 'segments-science-veto'
primary_veto_name = 'segments-final-veto-group'
secondary_vetoes_name = 'segments-veto-groups'
nonscience_veto_names = [primary_veto_name, secondary_vetoes_name]

science_seg_file, sci_segs, sci_seg_name = wf.get_science_segments(workflow,
                                            rdir['analysis_time/segment_data'])

runtime_names=[science_veto_name]
in_workflow_names = nonscience_veto_names
veto_cat_files = wf.get_files_for_vetoes(workflow,
                                         rdir['analysis_time/segment_data'],
                                         runtime_names=runtime_names,
                                         in_workflow_names=in_workflow_names)

sci_ok_seg_file, sci_ok_segs, sci_ok_seg_name = wf.get_analyzable_segments(\
                                            workflow, sci_segs, veto_cat_files,
                                            rdir['analysis_time/segment_data'])

datafind_files, analyzable_file, analyzable_segs, analyzable_name = \
                                           wf.setup_datafind_workflow(workflow,
                                                     sci_ok_segs, "datafind",
                                                     seg_file=science_seg_file)

cum_veto_files, veto_names, ind_cats = wf.get_cumulative_veto_group_files(\
                                            workflow, 'segments-veto-groups',
                                            veto_cat_files,
                                            rdir['analysis_time/segment_data'],
                                            execute_now=False)

final_veto_file, final_veto_name, ind_cats = \
                                   wf.get_cumulative_veto_group_files(workflow,
                                            'segments-final-veto-group',
                                            veto_cat_files,
                                            rdir['analysis_time/segment_data'],
                                            execute_now=False)

# setup gating files if provided
gate_files = wf.setup_gating_workflow(workflow, output_dir="gating")

# add the gate files to the jobs that use them
for job in ['calculate_psd', 'inspiral', 'single_template',
            'plot_singles_timefreq']:
  for gate_file in gate_files:
      ifo_gate = gate_file.cache_entry.url
      try:
          workflow.cp.set('{}-{}'.format(job, gate_file.ifo.lower()),
                          'gating-file', ifo_gate)
      except ConfigParser.SectionError:
          workflow.cp.add_section('{}-{}'.format(job,
                                  gate_file.ifo.lower()))
          workflow.cp.set('{}-{}'.format(job, gate_file.ifo.lower()),
                          'gating-file', ifo_gate)

# Precalculated PSDs
precalc_psd_files = wf.setup_psd_workflow(workflow, analyzable_segs,
                                            datafind_files, "psdfiles")

# Template bank stuff
hdfbank = wf.setup_tmpltbank_workflow(workflow, analyzable_segs,
                                      datafind_files, output_dir="bank",
                                      psd_files=precalc_psd_files,
                                      return_format='hdf')

splitbank_files_fd = wf.setup_splittable_workflow(workflow, hdfbank,
                                                  out_dir="bank",
                                                  tags=['full_data'])
splitbank_files_inj = wf.setup_splittable_workflow(workflow, hdfbank,
                                                   out_dir="bank",
                                                   tags=['injections'])

bank_plot = [(wf.make_template_plot(workflow, hdfbank[0],
              rdir['coincident_triggers']),)]

# setup the injection files
inj_files, inj_tags = wf.setup_injection_workflow(workflow,
                                                  output_dir="inj_files")

######################## Setup the FULL DATA run ##############################
tag = output_dir = "full_data"
ctags = [tag, 'full']

# setup the matchedfilter jobs
ind_insps = insps = wf.setup_matchedfltr_workflow(workflow, analyzable_segs,
                                   datafind_files, splitbank_files_fd,
                                   output_dir, tags = [tag])

insps = wf.merge_single_detector_hdf_files(workflow, hdfbank[0],
                                           insps, output_dir, tags=[tag])

# setup sngl trigger distribution fitting jobs
# 'statfiles' is list of files used in calculating coinc statistic
statfiles = []
statfiles += wf.setup_trigger_fitting(workflow, insps, hdfbank,
                                      final_veto_file, final_veto_name)

# Calculate the inspiral psds and make plots
psd_files = []
trig_generated_name = 'TRIGGERS_GENERATED'
trig_generated_segs = {}
data_analysed_name = 'DATA_ANALYSED'
data_analysed_segs = {}
insp_files_seg_dict = segments.segmentlistdict()

for ifo, files in zip(*ind_insps.categorize_by_attr('ifo')):
    trig_generated_segs[ifo] = segments.segmentlist([f.segment for f in files])
    data_analysed_segs[ifo] = \
        segments.segmentlist([f.metadata['data_seg'] for f in files])

    # Remove duplicates from splitbank
    trig_generated_segs[ifo] = \
        segments.segmentlist(set(trig_generated_segs[ifo]))
    data_analysed_segs[ifo] = \
        segments.segmentlist(set(data_analysed_segs[ifo]))

    insp_files_seg_dict[ifo + ":" + trig_generated_name] = \
                                                       trig_generated_segs[ifo]
    insp_files_seg_dict[ifo + ":" + data_analysed_name] = \
                                                        data_analysed_segs[ifo]

    if datafind_files:
        frame_files = datafind_files.find_output_with_ifo(ifo)
    else:
        frame_files = None
    psd_files += [wf.setup_psd_calculate(workflow, frame_files, ifo,
              data_analysed_segs[ifo], data_analysed_name, 'psds')]

insp_files_seg_file = wf.SegFile.from_segment_list_dict('INSP_SEGMENTS',
                 insp_files_seg_dict, valid_segment=workflow.analysis_time,
                 extension='xml', directory=rdir['analysis_time/segment_data'])

s = wf.make_spectrum_plot(workflow, psd_files, rdir['detector_sensitivity'],
                          precalc_psd_files=precalc_psd_files)
r = wf.make_range_plot(workflow, psd_files, rdir['detector_sensitivity'],
                       require='summ')
r2 = wf.make_range_plot(workflow, psd_files, rdir['detector_sensitivity'],
                        exclude='summ')

det_summ = [(s, r[0] if len(r) != 0 else None)]
layout.two_column_layout(rdir['detector_sensitivity'],
                     det_summ + list(layout.grouper(r2, 2)))

############################## Setup the injection runs #######################

inj_coincs = wf.FileList()
files_for_combined_injfind = []
for inj_file, tag in zip(inj_files, inj_tags):
    ctags = [tag, 'inj']
    output_dir = '%s_coinc' % tag

    if workflow.cp.has_option_tags('workflow-injections',
                                   'compute-optimal-snr', tags=[tag]):
        optimal_snr_file = wf.compute_inj_optimal_snr(
                workflow, inj_file, psd_files, 'inj_files', tags=[tag])
        file_for_injfind = optimal_snr_file
    else:
        file_for_injfind = inj_file

    if workflow.cp.has_option_tags('workflow-injections', 'inj-cut', tags=[tag]):
        file_for_vetoes = wf.cut_distant_injections(
                workflow, file_for_injfind, 'inj_files', tags=[tag])
    else:
        file_for_vetoes = inj_file

    if workflow.cp.has_option_tags('workflow-injections', 'strip-injections',
                                   tags=[tag]):
        small_inj_file = wf.veto_injections(workflow, file_for_vetoes,
                             insp_files_seg_file, trig_generated_name,
                             "inj_files", tags=[tag])
    else:
        small_inj_file = file_for_vetoes

    files_for_combined_injfind.append(file_for_injfind)

    # setup the matchedfilter jobs
    insps = wf.setup_matchedfltr_workflow(workflow, analyzable_segs,
                                     datafind_files, splitbank_files_inj,
                                     output_dir, injection_file=small_inj_file,
                                     tags = [tag])

    insps = wf.merge_single_detector_hdf_files(workflow, hdfbank[0],
                                               insps, output_dir, tags=[tag])


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

# Create a page to contain a dashboard link
dashboard_file = wf.File(workflow.ifos, 'DASHBOARD', workflow.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
wf.makedir(rdir['workflow/dax'])
wf.makedir(rdir['workflow/input_map'])
wf.makedir(rdir['workflow/output_map'])
wf.makedir(rdir['workflow/planning'])

wf.make_results_web_page(finalize_workflow, os.path.join(os.getcwd(),
rdir.base))

container += workflow
container += finalize_workflow

import Pegasus.DAX3 as dax
dep = dax.Dependency(parent=workflow.as_job, child=finalize_workflow.as_job)
container._adag.addDependency(dep)

container.save()

# Protect the open box results folder
# os.chmod(rdir['open_box_result'], 0700)
logging.info("Written dax.")

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