#!/usr/bin/env python

"""
A file splitter.

A split-like utility given an input file, an output basename and the number
of lines required in each split output file. Here the input file is expected
to have a one-line header, which is replicated in each output file.

Alan Christie
July 2018
"""

import argparse


def header_slit(input_file, output_base, output_size, extension='.smi'):
    """Splits the lines in an input file (including a header)
    into a series of output files.

    :param input_file: The input file
    :type input_file: ``str``
    :param output_base: The basename of the output files. The files are written
                        to the current working directory with a numerical
                        suffix and .smi. If the base is `x` the first output
                        file will be named `x_1[extension]`.
    :type output_base: ``str``
    :param output_size: The number of lines in each output file.
    :type output_size: ``int``
    :param extension: The extension for the output files.
    :type extension: ``str``
    """

    with open(input_file) as smiles_file:

        # Get the input file's header
        smiles_file.seek(0, 0)
        header = smiles_file.readline()

        file_line_count = 0
        file_number = 1
        output_file = None
        line = smiles_file.readline()
        while line and line.strip():

            if file_line_count == 0:
                # Start a new file and write the header
                name = output_base + "_" + str(file_number) + extension
                output_file = open(name, 'w')
                output_file.write(header)
                file_line_count = 0

            output_file.write(line)
            file_line_count += 1

            if file_line_count == output_size:
                # Close file
                output_file.close()
                output_file = None
                file_line_count = 0
                file_number += 1

            # Next line...
            line = smiles_file.readline()

        # Done last input line.
        # Do we have anything un-written?
        if output_file:
            output_file.close()


if __name__ == "__main__":

    PARSER = argparse.ArgumentParser(description='File splitter.'
                                                 ' Given a file this utility'
                                                 ' splits the lines'
                                                 ' replicating its header'
                                                 ' to new files.')
    PARSER.add_argument('-i', '--input_file',
                        help='The name of the input file.',
                        required=True)
    PARSER.add_argument('-o', '--output_base',
                        help='The basename you want to use for your output'
                             ' files. Each output file will use this as a base'
                             ' name and will append a unique decimal number to'
                             ' the end before adding a .smi extension.',
                        required=True)
    PARSER.add_argument('-s', '--output_size',
                        help='The (maximum) number of lines in each output'
                             ' file.',
                        required=True,
                        type=int)
    ARGS = PARSER.parse_args()

    header_slit(ARGS.input_file, ARGS.output_base, int(ARGS.output_size))
