#!/usr/bin/env python3
import argparse
import pandas    as pnd
import glob
import tqdm
import os
import utils_noroot as utnr

from log_store import log_store

log=log_store.add_logger(f'rx_scripts:print_sample_path')
#--------------------------------
class data:
    sample_path = None
#--------------------------------
def get_args():
    parser = argparse.ArgumentParser(description='Used to merge and print cutflow information related to selection of sample')
    parser.add_argument('-s', '--sample_path' , type=str, help='Path to directory containing JSON files and ROOT files', required=True)
    args = parser.parse_args()
    
    data.sample_path = args.sample_path
#--------------------------------
def get_paths():
    if not os.path.isdir(data.sample_path):
        log.error(f'Cannot find: {data.sample_path}')
        raise

    cfl_wc = f'{data.sample_path}/*.json'
    l_path = glob.glob(cfl_wc)
    nfile  = len(l_path)
    if nfile == 0:
        log.error(f'No JSON file found in: {cfl_wc}')
        raise
    else:
        log.info(f'Found {nfile} JSON files')

    return l_path
#--------------------------------
def get_yield_df(path):
    d_data = utnr.load_json(path)
    d_total= d_data['Total']
    d_pased= d_data['Pased']
    df     = pnd.DataFrame({'Total' : d_total, 'Pased' : d_pased})

    return df
#--------------------------------
def add_df(l_df):
    df_tot = l_df[0]
    l_df   = l_df[1:]

    log.info(f'Merging JSON files')
    for df in tqdm.tqdm(l_df, ascii=' -'):
        df_tot = df_tot.add(df)

    return df_tot
#--------------------------------
def add_efficiency(df):
    df['Efficiency']  = df.Pased / df.Total
    df['Cummulative'] = df.Efficiency.cumprod()

    return df
#--------------------------------
def run():
    l_path = get_paths()
    l_df   = [ get_yield_df(path) for path in l_path ]
    df     = add_df(l_df)
    df     = add_efficiency(df)

    print(df)
#--------------------------------
def main():
    get_args()
    run()
#--------------------------------
if __name__ == '__main__':
    main()

