#!/bin/bash
##############################################################################
# A bash script for generating a job array file of row numbers in an input   #
# file that can be subset and processed in a parallel job array.             #
# Arguments:                                                                 #
# infile : The path to the file that will have row number subsets,           #
#          can be gzip/bzip2 compressed                                      #
# chunk_size : The number of rows in each subset.                            #
# split_prefix: A path and prefix that will make an output file name.        #
#                                                                            #
# Output is to STDOUT                                                        #
##############################################################################
test_file="$1"
chunk_size="$2"
split_prefix="$3"

if [[ -z $split_prefix ]]; then
cat <<EOF
A bash script for generating a job array file of row numbers in an input
file that can be subset and processed in a parallel job array.

Arguments:
infile : The path to the file that will have row number subsets,
         can be gzip/bzip2 compressed.
chunk_size : The number of rows in each subset.
split_prefix: A path and prefix that will make an output file name.

Output is to STDOUT
EOF
exit 1
fi
set -eu

# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
# Test if a file is uncompressed,gzip or bzip2 compressed open and stream to
#  STDOUT.
#
# Globals:
#   None
# Arguments:
#   infile, the path to the input file to test.
# Outputs:
#   File contents streamed to STDOUT
# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
test_file() {
	local infile="$1"
	if file -b "$infile" | grep -q "gzip compressed"; then
	    zcat "$infile"
	elif file -b "$infile" | grep -q "bzip2 compressed"; then
		bzip2 -dc "$infile"
	elif file -b "$infile" | grep -q "compressed"; then
	    echo "Can only handle gzip/bzip2 compression" 1>&2
	    exit 1
	else
		cat "$infile"
	fi
}


test_file "$test_file" |
 tail -n+2 |
 tqdm --desc "[info] processing..." --unit ' row(s)' |
 awk -vchunk=$chunk_size -vinfile="$test_file" -voutprefix="$split_prefix" '


BEGIN{
	OFS="\t";
	rowidx=1;
	start=1;
	outrows=0;
	print "rowidx", "infile", "outfile", "startrow", "endrow"
}
{
	if (outrows==chunk) {
		print rowidx, infile, outprefix rowidx"-"start"-"NR".txt.gz", start, NR;
		start=NR;
		rowidx++;
		outrows=0;
	}
	outrows++
} END {
	if (NR>=start) {
		print rowidx, infile, outprefix rowidx"-"start"-"NR".txt.gz", start, NR+1;
	}
}'
