#!/usr/bin/env python

import argparse
import gzip
import os
import logging

from Bio import SeqIO
import numpy as np
import h5py


def parse_args():
    parser = argparse.ArgumentParser()

    parser.add_argument('--input', default='/dev/stdin')
    parser.add_argument('--hdf5', required=True)
    parser.add_argument('--format', default='fasta')
    parser.add_argument('--gzip', default=False, action='store_true', help='sequences are gzipped?')
    parser.add_argument('--init-size', default=10, type=int)
    parser.add_argument('--growth-size', default=10, type=int)
    parser.add_argument('--dataset', default='sequences')
    parser.add_argument('--log', default='/dev/stderr')

    return parser.parse_args()

def load_seqs(dset, records, growth_size=10000):
    for i, record in enumerate(records):
        try:
            dset[i] = str(record.seq)
        except ValueError:
            dset.resize(dset.size + growth_size, axis=0)
            dset[i] = str(record.seq)

        yield i


def main():
    args = parse_args()

    h5_file = h5py.File(args.hdf5, 'w')

    logging.basicConfig(level=logging.INFO, filename=args.log)

    logging.info('initializing db with size %s' % args.init_size)
    dset = h5_file.create_dataset(args.dataset,
                                  (args.init_size,),
                                  dtype=h5py.special_dtype(vlen=str),
                                  maxshape=(None, ),
                                  chunks=True)


    if args.gzip:
        opener = gzip.open
    else:
        opener = open

    size = float(os.path.getsize(args.input))

    with opener(args.input) as handle:
        records = SeqIO.parse(handle, args.format)
        try:
            for i in load_seqs(dset, records, growth_size=args.growth_size):
                if i%args.growth_size == 0:
                    logging.info('%s sequences' % i)
        except KeyboardInterrupt:
            logging.info('\nFinishing up...')
        finally:
            logging.info('wrote %s sequences' % dset.size)
            h5_file.close()


if __name__ == '__main__':
    main()
