#!/usr/bin/env python

# Copyright (C) 2011 Ian W. Harry
#
# 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.

# =============================================================================
# Preamble
# =============================================================================

from __future__ import division

import os, sys, time, copy
from argparse import ArgumentParser
from pylal import MultiInspiralUtils,coh_PTF_pyutils,git_version
from glue.ligolw import table,lsctables,utils,ligolw,ilwd
from glue import segments,lal

__author__  = "Ian Harry <ian.harry@astro.cf.ac.uk>"
__version__ = "git id %s" % git_version.id
__date__    = git_version.date

# set up timer
start = int(time.time()*10**6)
elapsed_time = lambda: int(time.time()*10**6-start)

# =============================================================================
# Parse command line
# =============================================================================

def parse_command_line():

  usage = """usage: %prog [options] 

coh_PTF_trig_combiner is designed to coalesce the triggers made from each element of the split template bank used by the search. The required arguments are

--ifo-tag
--user-tag
--input_files
"""

  parser = ArgumentParser(usage)

  parser.add_argument("-v", "--version", action="version", version=__version__)
  parser.add_argument("-i", "--ifo-tag", action="store", type=str,
                      default=None,
                      help="The ifo tag, H1L1 or H1L1V1 for instance")

  parser.add_argument("-u", "--user-tag", action="store", type=str,
                      default='COH_PTF_INSPIRAL_FIRST',
                      help="The user tag, COH_PTF_INSPIRAL_FIRST for instance")

  parser.add_argument("-S", "--slide-tag", action="store", type=str,
                      default='',
                      help="The slide tag, used to differentiate long slides")

  parser.add_argument("-n", "--grb-name", action="store", type=str,
                      default=None, help="GRB name, such as 090802 "+
                                        "(will be appended to output user tag")

  parser.add_argument("-T", "--num-trials", action="store", type=int,
                      default=6,
                      help="The number of off source trials, default: %default")

  parser.add_argument("-p", "--trig-start-time", action="store", type=int,
                      required=True,
                      help="The start time of the analysis segment")

  parser.add_argument("-a", "--segment-dir", action="store", type=str,
                      help="directory holding buffer, on and off source "
                           "segment files")

  parser.add_argument("-o", "--output-dir", action="store", type=str,
                      default=os.getcwd(), help="output directory, "
                                                "default: %default")

  parser.add_argument("-t", "--slide-cache", action="store",default=None,
                      help="Read slide triggers from this cache")

  parser.add_argument("-b", "--verbose", action="store_true", default=False,
                      help="verbose output with microsecond timer, "
                           "default: %default")

  parser.add_argument("-f", "--input-files", nargs="*", action="store",
                      required=True, metavar="TRIGGER FILE",
                      help="read in listed trigger files")

  parser.add_argument("-r", "--short-slides", action="store_true",
                      help="Did analysis use short time slides?")

  args = parser.parse_args()

  if not args.ifo_tag:
    parser.error("must provide --ifo-tag")

  if not args.user_tag:
    parser.error("must provide --user-tag")

  if not args.input_files and not args.slide_cache:
    parser.error("must provide --cache or --slide-cache")

  if not args.segment_dir:
    parser.error("must provide --segment-dir")

  return args

# =============================================================================
# Main function
# =============================================================================

def main(trigFiles, ifoTag, userTag, segdir, outdir, grbTag=None, numTrials=6,\
         timeSlides=None, slideTag=None, trigStart=None, shortSlides=False,
         verbose=False):
  
  if verbose: sys.stdout.write("Getting segments...\n")

  # get ifos
  ifos = sorted(lsctables.instrument_set_from_ifos(ifoTag))

  #
  # construct segments
  #

  # get segments
  segs           = coh_PTF_pyutils.readSegFiles(segdir)
  trialTime      = abs(segs['on'])
  
  # FIXME: Should the offtrials always be adjacent?
  # construct off source trials
  offSourceSegs  = segments.segmentlist()
  offSourceSegs.append(segs['off'])
  offSourceSegs.append(segments.segment(trigStart, trigStart + trialTime))
  for i in xrange(numTrials-1):
    t = i+2
    offSourceSegs.append(segments.segment(tuple(map(trialTime.__add__,\
                                                    offSourceSegs[t-1]))))
 
  if verbose: 
    sys.stdout.write("Constructed %d off source trials with segments at %d:\n"\
                     % (numTrials, elapsed_time()))
    for i in xrange(numTrials):
      sys.stdout.write("%s\n" % str(offSourceSegs[i+1]))

  #
  # load triggers
  #

  # get trigger files 
  if verbose: sys.stdout.write("\nFinding files...\n")

  if timeSlides:
    slideCache = lal.Cache.fromfile(open(timeSlides, 'r'))
    slideCache = slideCache.sieve(ifos=ifoTag, description=userTag)
    if slideTag:
      slideCache = slideCache.sieve(description=slideTag)
    if len(slideCache)<1:
      sys.stderr.write("No files found at %d.\n" % elapsed_time())
      sys.exit(1)
    else:
      sys.stdout.write("%d files found at %d.\n"\
                       % (len(slideCache), elapsed_time()))

  # set columns
  ifoAtt = { 'G1':'g', 'H1':'h1', 'H2':'h2', 'L1':'l', 'V1':'v', 'T1':'t' }
  # set up columns
  validcols = lsctables.MultiInspiralTable.validcolumns.keys()
  cols = ['end_time', 'end_time_ns', 'ifos', 'process_id', 'ra', 'dec', 'snr',\
          'null_statistic', 'null_stat_degen', 'mass1', 'mass2', 'mchirp']
  cols.extend(['time_slide_id'])
  # add snr
  cols.extend(['snr_%s' % ifoAtt[i] for i in ifos\
               if 'snr_%s' % ifoAtt[i] in validcols])
  # add chisq
  cols.extend(['chisq', 'chisq_dof'])
  cols.extend(['chisq_%s' % ifoAtt[i] for i in ifos\
               if 'chisq_%s' % ifoAtt[i] in validcols])
  # add bank chisq
  cols.extend(['bank_chisq', 'bank_chisq_dof'])
  #cols.extend(['bank_chisq_%s' % ifoAtt[i] for i in ifos\
  #             if 'bank_chisq_%s' % ifoAtt[i] in validcols])
  # add auto chisq  
  cols.extend(['cont_chisq', 'cont_chisq_dof'])
  #cols.extend(['cont_chisq_%s' % ifoAtt[i] for i in ifos\
  #             if 'cont_chisq_%s' % ifoAtt[i] in validcols])
  # add sigmasq
  cols.extend(['sigmasq_%s' % ifoAtt[i] for i in ifos\
               if 'sigmasq_%s' % ifoAtt[i] in validcols])
  # add amplitude terms
  cols.extend(['amp_term_%d' % i for i in xrange(1,11)])
  # set columns
  cols = [c for c in cols if c in\
          lsctables.MultiInspiralTable.validcolumns.keys()]
  lsctables.MultiInspiralTable.loadcolumns = cols

  # prepare triggers
  trigs = {}
 
  if verbose: sys.stdout.write("\nLoading all triggers...\n")

  trigs['ALL_TIMES'] = lsctables.New(lsctables.MultiInspiralTable,
          columns=lsctables.MultiInspiralTable.loadcolumns)
  
  # load triggers
  if not timeSlides and not shortSlides:
    trigs['ALL_TIMES'].extend(MultiInspiralUtils.ReadMultiInspiralFromFiles(\
            [e for e in trigFiles]))
    zeroLagTrigs = None

    timeSlideDoc = utils.load_filename(trigFiles[0],
            gz=trigFiles[0].endswith(".gz"),
            contenthandler = lsctables.use_in(ligolw.LIGOLWContentHandler))
    timeSlideTable = table.get_table(timeSlideDoc,
            lsctables.TimeSlideTable.tableName)
    segmentTable = table.get_table(timeSlideDoc,
            lsctables.SegmentTable.tableName)
    timeSlideSegmentMap = table.get_table(timeSlideDoc,
            lsctables.TimeSlideSegmentMapTable.tableName)
  elif shortSlides:
    zeroLagTrigs = {}
    zeroLagTrigs['ALL_TIMES'] = copy.deepcopy(trigs['ALL_TIMES'])
    
    tempTrigs, slideDictList, segmentDict, timeSlideTable, segmentTable, \
            timeSlideSegmentMap = \
            MultiInspiralUtils.ReadMultiInspiralTimeSlidesFromFiles(\
            trigFiles, generate_output_tables=True)
    trigs['ALL_TIMES'].extend(tempTrigs)

    zeroIDs = [slideDictList.index(slide) for slide in slideDictList \
               if all(v == 0 for v in slide.values())]
    zeroLagTrigs['ALL_TIMES'].extend([trig for trig in trigs['ALL_TIMES'] \
                                      if int(trig.time_slide_id) in zeroIDs])
  else:
    trigs['ALL_TIMES'] = lsctables.New(lsctables.MultiInspiralTable,\
                               columns=lsctables.MultiInspiralTable.loadcolumns)
    if trigFiles:
      inpFiles = [e for e in trigFiles]
    else:
      inpFiles = []
    inpFiles.extend([e for e in slideCache])
    if not trigFiles:
      cache = slideCache
    tempTrigs,slideDictList,segmentDict,timeSlideTable,\
        segmentTable,timeSlideSegmentMap = \
            MultiInspiralUtils.ReadMultiInspiralTimeSlidesFromFiles(inpFiles,\
                generate_output_tables=True)
    trigs['ALL_TIMES'].extend(tempTrigs)
    # Also want to separate out the zerolag trigs
    zeroLagTrigs = {}
    zeroLagTrigs['ALL_TIMES'] = lsctables.New(lsctables.MultiInspiralTable,\
                               columns=lsctables.MultiInspiralTable.loadcolumns)
    # Identify the zero lag slide
    for slideID,dict in enumerate(slideDictList):
      for ifo in ifos:
        if dict[ifo] != 0:
          break
      else:
        zeroLagID = slideID
        break
    else:
      # No zero lag trigs here
      zeroLagTrigs = None

    # Add the zero-lag triggers
    if zeroLagTrigs:
      zeroLagTrigs['ALL_TIMES'].extend(trig for trig in trigs['ALL_TIMES'] if \
                                       int(trig.time_slide_id) == zeroLagID)

  # Temporary hack to allow the code to work with tables that dont have
  # the new time slide ID column.
  if 'time_slide_id' not in trigs['ALL_TIMES'].columnnames:
    trigs['ALL_TIMES'].appendColumn('time_slide_id')
  for trig in trigs['ALL_TIMES']:
    if not hasattr(trig, 'time_slide_id'):
      trig.time_slide_id = ilwd.ilwdchar("multi_inspiral:time_slide_id:0")

  if verbose:
    sys.stdout.write("%d total triggers found at %d.\n"
                     % (len(trigs['ALL_TIMES']), elapsed_time()))
    if zeroLagTrigs is not None:
      sys.stdout.write("%d zerolag triggers found at %d.\n"
                       % (len(zeroLagTrigs['ALL_TIMES']), elapsed_time()))

  # 
  # load search summary table
  #

  # get search summary table from old file
  oldxml   = utils.load_filename(trigFiles[0],
                                  gz = trigFiles[0].endswith("gz"), contenthandler = lsctables.use_in(ligolw.LIGOLWContentHandler))
  oldSearchSummTable = table.get_table(oldxml, "search_summary")

  #
  # bin triggers
  #

  if verbose: sys.stdout.write("\nSeparating triggers by end time...\n")

  if shortSlides or timeSlides:
    zeroLagTrigs['ONSOURCE'] = lsctables.New(lsctables.MultiInspiralTable,\
                               columns=lsctables.MultiInspiralTable.loadcolumns)
    zeroLagTrigs['OFFSOURCE'] = copy.deepcopy(zeroLagTrigs['ONSOURCE'])
    for i in range(numTrials):
      zeroLagTrigs['OFFTRIAL_%d' % (i+1)] = \
              copy.deepcopy(zeroLagTrigs['ONSOURCE'])

    # separate into correct bins
    for trig in zeroLagTrigs['ALL_TIMES']:
      if trig.get_end() not in segs['buffer']:
        zeroLagTrigs['OFFSOURCE'].append(trig)
      elif trig.get_end() in segs['on']:
        zeroLagTrigs['ONSOURCE'].append(trig)
      for i in xrange(numTrials):
        if trig.get_end() in offSourceSegs[i+1]:
          zeroLagTrigs['OFFTRIAL_%d' % (i+1)].append(trig)

  trigs['OFFSOURCE'] = lsctables.New(lsctables.MultiInspiralTable,\
          columns=lsctables.MultiInspiralTable.loadcolumns)
  if shortSlides or timeSlides:
    for trig in trigs['ALL_TIMES']:
      slideID = int(trig.time_slide_id)
      for ifo in ifos:
        offset = slideDictList[slideID][ifo]
        if trig.get_end()+offset in segs['buffer']:
          break
        else:
          trigs['OFFSOURCE'].append(trig)
  else:
    trigs['ONSOURCE'] = lsctables.New(lsctables.MultiInspiralTable,\
            columns=lsctables.MultiInspiralTable.loadcolumns)
    for i in xrange(numTrials):
      trigs['OFFTRIAL_%d' % (i+1)] = copy.deepcopy(trigs['ONSOURCE'])

    for trig in trigs['ALL_TIMES']:
      if trig.get_end() not in segs['buffer']:
        trigs['OFFSOURCE'].append(trig)
      elif trig.get_end() in segs['on']:
        trigs['ONSOURCE'].append(trig)
      for i in xrange(numTrials):
        if trig.get_end() in offSourceSegs[i+1]:
          trigs['OFFTRIAL_%d' % (i+1)].append(trig)
  
  if verbose: sys.stdout.write("Done at %d.\n" % (elapsed_time()))
  #
  # write combined triggers to file
  #

  if verbose: sys.stdout.write("\nWriting new xml files...\n")

  # prepare xmldocument 
  xmldoc = ligolw.Document()
  xmldoc.appendChild(ligolw.LIGO_LW())

  # append process params table
  xmldoc = coh_PTF_pyutils.append_process_params(xmldoc, sys.argv, __version__,\
                                                 __date__)
 
  # get search summary table from old file
  xmldoc.childNodes[-1].appendChild(oldSearchSummTable)

  # Add time slide table
  xmldoc.childNodes[-1].appendChild(timeSlideTable)
  # Add segment table
  xmldoc.childNodes[-1].appendChild(segmentTable)
  # And the map between the two
  xmldoc.childNodes[-1].appendChild(timeSlideSegmentMap)

  # construct new xml file names
  start = int(trigFiles[0].split('-')[-2])
  end = int(trigFiles[0].split('-')[-1].split('.')[0]) + start
  if timeSlides:
    tsUserTag = '%s_TIMESLIDES' %(userTag)
    if grbTag:
      tsUserTag = '%s_GRB%s' % (tsUserTag, grbTag)
  if grbTag:
    userTag = '%s_GRB%s' % (userTag, grbTag)

  # write new xml files
  if not (shortSlides or timeSlides):
    for trigTime in trigs.keys():
      xmlName = '%s/%s-%s_%s-%d-%d.xml.gz'\
                % (outdir, ifoTag, userTag, trigTime, start, end-start)
      xmldoc.childNodes[-1].appendChild(trigs[trigTime])
      utils.write_filename(xmldoc, xmlName, gz = xmlName.endswith('gz'))
      xmldoc.childNodes[-1].removeChild(trigs[trigTime]) 
      if verbose: 
        sys.stdout.write("%s written at %d\n" % (xmlName, elapsed_time()))

  elif shortSlides and not timeSlides:
    for trigTime in trigs.keys():
      xmlName = '%s/%s-%s_%s-%d-%d.xml.gz'\
                % (outdir, ifoTag, userTag, trigTime, start, end-start)
      xmldoc.childNodes[-1].appendChild(trigs[trigTime])
      utils.write_filename(xmldoc, xmlName, gz = xmlName.endswith('gz'))
      xmldoc.childNodes[-1].removeChild(trigs[trigTime])
      if verbose:
        sys.stdout.write("%s written at %d\n" % (xmlName, elapsed_time()))
    for trigTime in zeroLagTrigs.keys():
      if trigTime in ['OFFSOURCE', 'ALL_TIMES']:
        xmlName = '%s/%s-%s_ZEROLAG_%s-%d-%d.xml.gz'\
                  % (outdir, ifoTag, userTag, trigTime[:3], start, end-start)
      else:
        xmlName = '%s/%s-%s_%s-%d-%d.xml.gz'\
                  % (outdir, ifoTag, userTag, trigTime, start, end-start)
      xmldoc.childNodes[-1].appendChild(zeroLagTrigs[trigTime])
      utils.write_filename(xmldoc, xmlName, gz = xmlName.endswith('gz'))
      xmldoc.childNodes[-1].removeChild(zeroLagTrigs[trigTime])
      if verbose:
        sys.stdout.write("%s written at %d\n" % (xmlName, elapsed_time()))

  else:
    for trigTime in trigs.keys():
      if slideTag:
        fullTag = '%s-%s_%s' %(ifoTag, tsUserTag, slideTag)
      else:
        fullTag = '%s-%s' %(ifoTag, tsUserTag)
      xmlName = '%s/%s_%s-%d-%d.xml.gz'\
                % (outdir, fullTag, trigTime, start, end-start)
      xmldoc.childNodes[-1].appendChild(trigs[trigTime])
      utils.write_filename(xmldoc, xmlName, gz = xmlName.endswith('gz'))
      xmldoc.childNodes[-1].removeChild(trigs[trigTime])
      if verbose:
        sys.stdout.write("%s written at %d\n" % (xmlName, elapsed_time()))


if __name__=='__main__':

  args = parse_command_line()

  outdir      = os.path.abspath(args.output_dir)
  segdir      = os.path.abspath(args.segment_dir)
  ifoTag      = args.ifo_tag
  userTag     = args.user_tag
  trigFiles   = args.input_files
  GRBname     = args.grb_name
  numTrials   = args.num_trials
  timeSlides  = args.slide_cache
  slideTag    = args.slide_tag
  trigStart   = args.trig_start_time
  shortSlides = args.short_slides
  verbose     = args.verbose

  main(trigFiles, ifoTag, userTag, segdir, outdir, grbTag=GRBname,
       numTrials=numTrials, timeSlides=timeSlides, slideTag=slideTag,
       trigStart=trigStart, shortSlides=shortSlides, verbose=verbose)
  if verbose: sys.stdout.write("Done at %d.\n" % (elapsed_time()))
