#!/usr/bin/env python
"""Report memory usage during chunking of a large file
"""
# from pyaddons import utils
from merge_sort import chunks
import gzip
import csv
import os
import resource
import tempfile

infile = "/data/mapping_files/b38/v20220402/gwas_norm.v20220402.common.biallelic.vep.b38.vcf.gz"
delimiter = "\t"

with gzip.open(infile, 'rt') as infile:
    line = '##'
    while line.startswith('##'):
        line = infile.readline()

    with open(os.path.join(os.environ['HOME'], "merge_mem_usage.txt"), 'wt') as memhf:
        idx = 1
        reader = csv.reader(infile, delimiter=delimiter)
        chunk_dir = tempfile.mkdtemp(prefix='merge_sort_test_')
        chunksize = 100000
        with chunks.CsvSortedChunks(
                chunk_dir, lambda x: (x[0], int(x[1]), x[3], x[4]),
                chunksize=chunksize, delimiter=delimiter
        ) as chunker:
            for row in reader:
                if (idx % 100000) == 0:
                    mem = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss
                    print(f"{idx}: main max mem usage: {round(mem/1024, 2)} MB")
                idx += 1
                chunker.add_row(row)
