#!/bin/env python
import argparse, logging, pycbc.version, ligo.segments, numpy
from dqsegdb.apicalls import dqsegdbQueryTimes as query
#from pycbc.workflow.segment import cat_to_veto_def_cat as convert_cat

def sane(seg_list):
    """ Convert list of len two lists containing strs to segment list """
    segs = ligo.segments.segmentlist([])
    for seg in seg_list:
        segs.append(ligo.segments.segment(int(seg[0]), int(seg[1])))
    return segs
    
def parse_veto_definer(veto_def_filename):
    """ Parse a veto definer file from the filename and return a dictionary 
    indexed by ifo and veto definer category level.
    
    Parameters
    ----------
    veto_def_filename: str
        The path to the veto definer file
    
    Returns:
        parsed_definition: dict   
            Returns a dictionary first indexed by ifo, then category level, and
            finally a list of veto definitions.
    """
    from ligo.lw import table, utils as ligolw_utils
    from pycbc.io.ligolw import LIGOLWContentHandler as h

    indoc = ligolw_utils.load_filename(veto_def_filename, False, contenthandler=h)
    veto_table = table.Table.get_table(indoc, 'veto_definer')

    ifo = veto_table.getColumnByName('ifo')
    name = veto_table.getColumnByName('name')
    version = numpy.array(veto_table.getColumnByName('version'))
    category = numpy.array(veto_table.getColumnByName('category'))
    start = numpy.array(veto_table.getColumnByName('start_time'))
    end = numpy.array(veto_table.getColumnByName('end_time'))
    start_pad = numpy.array(veto_table.getColumnByName('start_pad'))
    end_pad = numpy.array(veto_table.getColumnByName('end_pad'))
    
    data = {}
    for i in range(len(veto_table)):
        if ifo[i] not in data:
            data[ifo[i]] = {}
        
        if category[i] not in data[ifo[i]]:
            data[ifo[i]][category[i]] = []
            
        veto_info = {'name': name[i],
                     'version': version[i],
                     'start': start[i],
                     'end': end[i],
                     'start_pad': start_pad[i],
                     'end_pad': end_pad[i],
                     }
        data[ifo[i]][category[i]].append(veto_info)
    return data

def get_vetoes(veto_def, ifo, server, veto_name, start_default, end_default):
    """ Cycle through the veto name string and collect the vetoes for the
    selected categories. Return the final segment list
    """
    raise ValueError("This code needs updating to work with the new segment "
                     "interface. If it's still used please fix this, "
                     "otherwise we can just remove this code.")
    veto_segments = ligo.segments.segmentlist([])
    for cat in veto_name:
        #cat = convert_cat(cat)
        flags = veto_def[ifo][cat]
        
        for flag in flags:
            start = flag['start'] if flag['start'] >= start_default else start_default
            end = flag['end'] if flag['end'] !=0 else end_default

            raw_segs = sane(query("https", server, ifo,
                                  flag['name'], flag['version'],
                                  'active', start, end)[0]['active'])
                                  
            for rseg in raw_segs:
                s, e = rseg[0] + flag['start_pad'], rseg[1] + flag['end_pad']
                veto_segments.append(ligo.segments.segment(s, e))
    return veto_segments.coalesce()
    

parser = argparse.ArgumentParser()
parser.add_argument('--version', action='version', version=pycbc.version.git_verbose_msg)
parser.add_argument('--verbose', action='store_true')

parser.add_argument('--gps-start-time', type=int, required=True,
                    help="integer gps start time")
parser.add_argument('--gps-end-time', type=int, required=True,
                    help="integer gps end time")
parser.add_argument('--veto-definer', type=str, required=True,
                    help="path to veto definer xml file")
parser.add_argument('--science-veto-levels', type=str,
                    help="Veto levels to apply by removing strain data before analysis ex. '1' for CAT1 veto")
parser.add_argument('--trigger-veto-levels', type=str,
                    help="Veto levels to apply by removing triggers from analyzed times ex. '12H' for CAT 1 and CAT2 vetoes plus hardware injections")
parser.add_argument('--segment-server', type=str,
                    help="segment server string")
parser.add_argument('--science-names', nargs=2, 
                    help="name of the segment flag IFO:NAME:VERSION to use for science")
                      
group = parser.add_argument_group("pycbc_inspiral options that determine padding and minimum time analyzable.")
group.add_argument('--segment-length', type=int)
group.add_argument('--min-analysis-segments', type=int)
group.add_argument('--pad-data', type=int)
group.add_argument('--segment-start-pad', type=int)
group.add_argument('--segment-end-pad', type=int)

args = parser.parse_args()


analysis_start_pad =  args.segment_start_pad + args.pad_data
analysis_end_pad = args.segment_end_pad + args.pad_data
minimum_segment_length = ((args.segment_length - args.segment_start_pad 
                          - args.segment_end_pad) * args.min_analysis_segments 
                          + analysis_start_pad + analysis_end_pad)

pycbc.init_logging(args.verbose)

ifo_segs = []

veto_def = parse_veto_definer(args.veto_definer)

# Read in the science segments for the requested time
for science_name in args.science_names:
    ifo, name, version = science_name.split(':')
    
    logging.info("For IFO: %s, querying science time (%s, %s)" % (ifo, name, version))
    segments = sane(query("https", args.segment_server, ifo, name, version, 
                         'active', args.gps_start_time, args.gps_end_time)[0]['active'])

    #trim segments to the request time
    request = ligo.segments.segment(args.gps_start_time, args.gps_end_time)
    segments = (ligo.segments.segmentlist([request]) & segments)

    # apply cat 1 vetoes here
    logging.info('Found %ss of data' % abs(segments))
    segments = segments.coalesce()

    cat1_segs = get_vetoes(veto_def, ifo, 
                           args.segment_server, 
                           args.science_veto_levels,
                           args.gps_start_time,
                           args.gps_end_time,
                           ).coalesce()

    segments -= cat1_segs
    logging.info('Found %ss after applying CAT1 vetoes' % abs(segments))
    # remove short segments, and account for filter padding
    logging.info('Removing segments shorter than %ss' % minimum_segment_length)
    lsegments = ligo.segments.segmentlist([])
    segments = segments.coalesce()
    for seg in segments:
        if abs(seg) >= minimum_segment_length:
            start = seg[0] + analysis_start_pad
            end = seg[1] - analysis_end_pad
            lsegments.append(ligo.segments.segment(start, end))
    segments = lsegments
    logging.info('Found %ss after applying removing padding / short segments' % abs(segments))
    
    # apply vetoes that remove triggers here  
    segments = segments.coalesce()
    vtrig_segs = get_vetoes(veto_def, ifo, 
                           args.segment_server, 
                           args.trigger_veto_levels,
                           args.gps_start_time,
                           args.gps_end_time,
                           ).coalesce()
    segments -= vtrig_segs
    
    logging.info('Found %ss after applying trigger vetoes' % abs(segments))
    segments.coalesce()
    
    ifo_segs += [segments]
    
coinc_time = abs(ifo_segs[0] & ifo_segs[1])
print("Available Coincident Time from %s-%s" % (args.gps_start_time, args.gps_end_time))
print("%s seconds, %5.5f days" % (coinc_time, coinc_time / 86400.0))