#!/usr/bin/env python3
import numpy as np
import re, os, sys
import glob
import logging
import argparse

from ldc.waveform.source import SourceMaker
import ldc.io.yml as ymlio
import ldc.io.hdf5 as h5io
from ldc.utils.logging import init_logger, close_logger
from ldc.waveform.waveform.hphc import HpHc

if __name__ == "__main__":
    import argparse
    parser = argparse.ArgumentParser()
    parser.add_argument('-c', '--config', type=str, help= "Path to configuration file")
    parser.add_argument('-i', '--input', type=str, help= "Input catalogs", default="")
    parser.add_argument('-o', '--output', type=str, help= "Output catalog")
    parser.add_argument('-n', '--nsource', type=int, help= "Number of source", default=0)
    parser.add_argument('-l', '--log', type=str, default="", help="Log file")
    args = parser.parse_args()

    logger = init_logger(args.log)
    cfg = ymlio.load_config(args.config)
    if not args.input:
        args.input = cfg["catalogs"]
    if not args.nsource:
        args.nsource = cfg["nsource"]

    cfg["redshifted_mass"] = True
    cfg["non_precessing"] = False

        
    if os.path.isdir(args.input):
        logger.info("glob files from %s"%args.input)
        args.input = sorted(glob.glob(args.input+"/*"))
        logger.info("found %d files"%(len(args.input)))
        
    source_maker = SourceMaker.type(cfg["source_type"],
                                    cfg["approximant"],
                                    catalogs=args.input,
                                    logger=logger)
    if cfg["source_type"] == "GB" and 'random_frequency' in cfg.keys():
        cat, units = source_maker.draw_random_catalog(**cfg)      
    else:
        cat, units = source_maker.choose_from_catalog(**cfg)
    #np.save(args.output, cat)
    try:
        i = np.where(cat["Name"]=='CD–30o11223')[0][0]
        cat[i]["Name"] = 'CD-30o11223'
    except:
        pass

    if cfg["source_type"] == "MBHB":
        isort = np.argsort(cat["CoalescenceTime"])
        cat = cat[isort]

    # remove unused parameters
    GW = HpHc.type('test', cfg["source_type"], cfg["approximant"])
    default_units = GW.info()
    for k in cat.dtype.names:
        if k not in default_units.keys():
            units.pop(k)
        
    h5io.save_array(args.output, cat, **units)
    
    close_logger(logger)
    
