#!/usr/bin/env python3

# created by Victoria Niu on 04/16/2022
# This program inputs (mc, L) of events with the category that we have known,
# and tests the accuracy of our p_astro_model

import numpy as np
import sys
sys.path.insert(1, '../')

from pastro import pastro
from optparse import OptionParser

def parse_command_line():
  parser = OptionParser()
  parser.add_option('-v', '--verbose', action='store_true', help='verbose')
  parser.add_option('--p-thresh', help='the pass-rate threshold')
  parser.add_option('--model-file', help = 'provide the model h5 file')
  
  options, filenames = parser.parse_args()
  process_params = dict(options.__dict__)
  
  return options, process_params, filenames

options, process_params, filenames = parse_command_line()
model = pastro.pastro_model.from_h5(options.model_file)
model.finalize(model.prior())


counts_raw = {cat:0. for cat in model.categories}
counts = {cat:0. for cat in model.categories}
for cate in model.categories:
  data = np.loadtxt('test_data/'+cate.lower()+'.txt')
  for mc, lr in data:
    pa = model(mc, lr)
    counts_raw = {cat:counts_raw[cat] + pa[cat] for cat in pa}
    if pa[cate]==max(pa.values()) and (max(pa.values())>float(options.p_thresh)):
      counts[cate] += 1
print('raw count: ', counts_raw)
print('counts above probability threshold: ', counts)

