#!/usr/bin/env python
# Copyright (C) The Arvados Authors. All rights reserved.
#
# SPDX-License-Identifier: Apache-2.0

import arvados
import os
import re
import sys
import pyrtg

this_job = arvados.current_job()
this_task = arvados.current_task()
ref_dir = arvados.util.collection_extract(
    collection = this_job['script_parameters']['reference'],
    path = 'reference',
    decompress = False)
input_dir = arvados.util.collection_extract(
    collection = this_job['script_parameters']['input'],
    path = 'input')
bam_files = map(lambda f: os.path.join(input_dir, f),
                filter(lambda f: re.search(r'^(.*/)?alignments.bam$', f),
                       arvados.util.listdir_recursive(input_dir)))
out_dir = os.path.join(arvados.current_task().tmpdir, 'out')
arvados.util.run_command(['rm', '-rf', out_dir], stderr=sys.stderr)

# call sequence variants
pyrtg.run_rtg('snp', out_dir,
              ['-t', ref_dir] + bam_files)

# store output
out = arvados.CollectionWriter()
out.write_directory_tree(out_dir, max_manifest_depth=0)
this_task.set_output(out.finish())