import subprocess
import argparse
import os


def stratify_vcf(input_vcf, output_vcf, data_source):
    """Stratify a VCF file using the VarStrat tool.

    Args:
        input_vcf (str): Path to the input VCF file.
        output_vcf (str): Path to the output VCF file.
        data_source (str): Path to the data source file.

    Returns:
        None
    """
   
    # Get the directory of the current script
    script_dir = os.path.dirname(os.path.realpath(__file__))
    stratify_script = os.path.join(script_dir, 'stratify.sh')

    # Run the VarStrat tool
    subprocess.run(['bash', stratify_script, input_vcf, output_vcf, data_source])


def main():
    parser = argparse.ArgumentParser(description="A tool for stratifying genetic variants.")
    parser.add_argument("--input_vcf", type=str, help="Path to the input VCF file.")
    parser.add_argument("--output_vcf", type=str, help="Path to the output VCF file.")
    parser.add_argument("--data_source", type=str, help="Path to the data source directory.")

    args = parser.parse_args()

    stratify_vcf(args.input_vcf, args.output_vcf, args.data_source)


if __name__ == "__main__":
    main()