#!/grid/gillis/home/nfox/software/miniconda3/bin/python

import h5py
import os
import sys


def main(networks_file, outfile=sys.stdout):
    """Count the number of contacts in 1bp and 10kbp networks.

    Count the total number of Hi-C total in the 1bp and 10kbp
    networks in a "networks.hdf5" Run networks file. Validate
    that they are equal then print the Run ID and the
    number of contacts.

    Args:
        networks_file: str. Path to the "networks.hdf5" file
            that will be queried.
        outfile: Open File object. Default is standard out.
            Run IDs and corresponding contact counts will
            be written here.
    """
    if type(networks_file) is str and networks_file.startswith('./'):
        networks_file = f'{os.getcwd()}/{networks_file[2:]}'
    run = os.path.basename(os.path.dirname(networks_file))
    sum_1bp = 0
    sum_10kbp = 0
    with h5py.File(networks_file, 'r') as hfile:
        for chrom in hfile['/networks/1bp_raw/'].keys():
            if '_vs_' not in chrom:
                continue
            sum_1bp += hfile[f'/networks/1bp_raw/{chrom}/data'][:].sum()
        for chrom in hfile['/networks/10kbp_raw/'].keys():
            if '_vs_' not in chrom:
                continue
            sum_10kbp += hfile[f'/networks/10kbp_raw/{chrom}/data'][:].sum()
    assert sum_1bp == sum_10kbp
    sum_1bp = int(sum_1bp)
    sum_10kbp = int(sum_10kbp)
    try:
        outfile.write(f'{run:12s} : {sum_1bp}\n')
        outfile.close()
    except AttributeError:
        if type(outfile) is str:
            with open(outfile, 'a') as f:
                f.write(f'{run:12s} : {sum_1bp}\n')
        else:
            print(f'{outfile} is an invalid value for outfile.')
            return


if __name__ == '__main__':
    import argparse
    parser = argparse.ArgumentParser()
    parser.add_argument('networks_file')
    parser.add_argument('outfile')
    args = parser.parse_args()
    main(args.networks_file, args.outfile)