#!/usr/bin/env python3

# created by Victoria Niu on 04/16/2022
# This python script prepares samples of (mc, L) with event category for the test suite

import numpy as np
import sqlite3

from gstlal import rio
from optparse import OptionParser


def parse_command_line():
  parser = OptionParser()
  parser.add_option('-v', '--verbose', action='store_true', help='verbose')
  parser.add_option('--category', help='the category name, e.g., BNS, NSBH, BBH, TERRESTRIAL')
  parser.add_option('--count', help='the number of datapoints we take from the database')
  parser.add_option('--lr-thresh', help='likelihood threshold')
  parser.add_option('--output', metavar='filename', help='output random sample name')

  options, filenames = parser.parse_args()
  process_params = dict(options.__dict__)
    
  return options, process_params, filenames

options, process_params, filenames = parse_command_line()

def pick_events(filenames, category, num):
  """
  This function prepares the randomly-chosen sample for a given event category.
    
  inputs
  ------
  category: string, the event category
  """
    
  mc = []
  lr = []
  for fn in filenames:
    print(fn)
#    D = rio.load(fn);
#    for inj in D.found_injections():
#      if inj['coinc_event']['likelihood'] > float(options.lr_thresh):
#        mc.append(inj['coinc_inspiral']['mchirp'])
#        lr.append(inj['coinc_event']['likelihood'])

    connection = sqlite3.connect(fn)
    ts_id, = connection.cursor().execute("SELECT DISTINCT time_slide_id FROM time_slide where time_slide_id NOT IN (SELECT DISTINCT time_slide_id FROM time_slide where offset !=0);").fetchone()
    for mc_point,lr_point, in connection.cursor().execute("SELECT mchirp, likelihood FROM coinc_inspiral JOIN coinc_event on coinc_event.coinc_event_id == coinc_inspiral.coinc_event_id WHERE coinc_event.time_slide_id == :tsid;", {"tsid": int(ts_id)}):
      if lr_point > float(options.lr_thresh):
        mc.append(mc_point)
        lr.append(lr_point)

  pick = np.random.randint(len(mc), size=num)
  mc = np.array(mc)
  lr = np.array(lr)
  return mc[pick], lr[pick]


mc, lr = pick_events(filenames, options.category, int(options.count))
#names = np.array([options.category for _ in range(len(mc))], dtype=str)
mc_lr_pick = np.vstack((mc, lr)).T

np.savetxt(options.output, mc_lr_pick)
