#!/usr/bin/env python3
import numpy as np
import re
import os
import sys
import datetime
import h5py

import ldc.io.hdf5 as hdf5io
import ldc.io.npz as npzio
import ldc.io.yml as ymlio
from ldc.utils.logging import init_logger, close_logger

def from_extension(filename):
    """ Return IO manager corresponding to filename
    """
    extension = filename.split(".")[-1]
    if extension in ["npz", "npy"]:
        return npzio
    elif extension in ["hdf5", "h5"]:
        return hdf5io
    

if __name__ == "__main__":

    import argparse
    parser = argparse.ArgumentParser()
    parser.add_argument('--data-files', nargs="+", help= "Path to input file(s)")
    parser.add_argument('--dset-names', nargs="+", default=[], help=
                        "Data set name for each input file(s)")
    parser.add_argument('--config-files', nargs='+', default=[],
                        help= "Path to configuration file(s)")
    parser.add_argument('--config-names', nargs="+", default=[],
                        help= "Config name for each config file(s)")
    parser.add_argument('-o', '--out', default="./output.h5", help= "Output hdf5 file")
    parser.add_argument('-l', '--log', type=str, default="", help="Log file")

    args = parser.parse_args()
    logger = init_logger(args.log)		

    if not args.dset_names:
        args.dset_names = [os.path.basename(f) for f in args.data_files]

    # create and init hdf5 file
    hdf5 = hdf5io.save_config(args.out,
                              {'date':datetime.datetime.now().ctime()},
                              mode='w', name="header")

    # add all array
    for name, dfile in zip(args.dset_names, args.data_files):

        extension = dfile.split(".")[-1]

        if extension in ["hdf5", "h5"]:
            #with h5py.File(dfile, "r") as fid:
            #    dsets = list(fid.keys())
            # if len(dsets) > 1 and 'tdi' in name:
            #     arr = []
            #     names = []
            #     for i,dset in enumerate(dsets):
            #         subarr, descr = hdf5io.load_array(dfile, name=dset)
            #         if i==0:
            #             ineg = np.where(subarr[:,0]>=0)[0][0]
            #             arr.append(subarr[ineg:,0])
            #             arr.append(subarr[ineg:,1])
            #             names.append("time")
            #             names.append(dset)
            #         else:
            #             arr.append(subarr[ineg:,1])
            #             names.append(dset)
            #     arr = np.rec.fromarrays(arr, names=names)
            #     del descr["name"]
            #     hdf5io.save_array(args.out, arr, name=name, mode='a', **descr)
            # else:
            arr, descr = hdf5io.load_array(dfile)#, name=dsets[0])
            hdf5io.save_array(args.out, arr, name=name, mode='a', **descr)
        elif extension in ["npz", "npy"]:
            arr, descr = npzio.load_array(dfile)
            hdf5io.save_array(args.out, arr, name=name, mode='a', **descr)
                
    # add all config
    for name, dfile in zip(args.config_names, args.config_files):
        
        extension = dfile.split(".")[-1]
        
        if extension == "yml":
            cfg = ymlio.load_config(dfile)
            if 'sources' in cfg.keys():
                del cfg['sources']
            hdf5io.save_config(args.out, cfg, name=name, mode='a')
        elif extension == "py":
            pcode = open(dfile, "r").read() 
            cfg = {}
            exec(pcode, cfg)
            cfg_ = {}
            for k,v in cfg.items():
                if v is None:
                    v = 'None'
                if k[0]!='_':
                    cfg_[k] = v
            hdf5io.save_config(args.out, cfg_, name=name, mode='a')
            
