#! /usr/bin/env python
"""
util_MassGrid.py <fnames>
    - prints out <m1> <m2> <lnL>  <neff> \sigma/L for each individual ILE *mass point*, to stdout
    - fnames can be an arbitrary mix of .sqlite or .xml files
    - currently *hardcodes* intrinsic (mass) parameter list

Used to generate ascii dumps summarizing runs
"""
import sys
import os
import RIFT.misc.xmlutils as xmlutils
from optparse import OptionParser
from igwn_ligolw import lsctables, table, utils

import numpy as np
import RIFT.misc.weight_simulations as weight_simulations

# optp= OptionParser()
# opts, args = optp.parse_args()

data_at_intrinsic = {}
for fname in sys.argv[1:]:
    fExtension = fname.split(".")[-1]
    # if .sqlite extension
    if fExtension == "sqlite":
        samples = xmlutils.db_to_samples(fname, lsctables.SnglInspiralTable, ['mass1', 'mass2', 'snr', 'tau0', 'event_duration'])
        for row in samples:  # print each individual row, don't reweight yet
            if not (data_at_intrinsic.has_key((row.mass1,row.mass2))):
                data_at_intrinsic[(row.mass1,row.mass2)] = [[row.snr, row.tau0, row.event_duration]]
            else:
                data_at_intrinsic[(row.mass1,row.mass2)].append([row.snr, row.tau0, row.event_duration])
    # otherwise xml: could have just used ligolw_print
    elif fExtension == "xml" or fExtension == "gz":
        samples = table.get_table(utils.load_filename(fname), lsctables.SnglInspiralTable.tableName)
        for row in samples:  # print each individual row, don't reweight yet
            if not (data_at_intrinsic.has_key((row.mass1,row.mass2))):
                data_at_intrinsic[(row.mass1,row.mass2)] = [[row.snr, row.tau0, row.event_duration]]
            else:
                data_at_intrinsic[(row.mass1,row.mass2)].append([row.snr, row.tau0, row.event_duration])


# Loop over each key and print out
for intrinsic in data_at_intrinsic.keys():
    lnL, neff, sigmaOverL =   np.transpose(data_at_intrinsic[intrinsic])
    lnLmax = np.max(lnL)
    sigma = sigmaOverL*np.exp(lnL-lnLmax)  # remove overall Lmax factor, which factors out from the weights constructed from \sigma
    wts = weight_simulations.AverageSimulationWeights(None, None,sigma)   
    lnLmeanMinusLmax = np.log(np.sum(np.exp(lnL - lnLmax)*wts))
    sigmaNetOverL = (np.sqrt(1./np.sum(1./sigma/sigma)))/np.exp(lnLmeanMinusLmax)
    print(intrinsic[0], intrinsic[1], lnLmeanMinusLmax+lnLmax, np.sum(neff), sigmaNetOverL)


