add_annotations.py

#!/usr/bin/env python
import argparse
import hail as hl
import logging
import re
import sys
import numpy as np

import sys, os
sys.path.append('/home/jupyter/')

from collections import Counter
from textwrap import dedent

from gnomad.utils.annotations import age_hists_expr
from gnomad.utils.reference_genome import add_reference_sequence
from gnomad.utils.slack import slack_notifications
from gnomad.utils.vep import vep_struct_to_csq
from gnomad_qc.v3.resources.meta import meta  # pylint: disable=import-error
from gnomad.resources.grch38.gnomad import POPS
from gnomad.resources.grch38.reference_data import dbsnp, _import_dbsnp
from gnomad_mitochondria.pipeline.annotation_descriptions import (
    add_descriptions,
    adjust_descriptions,
)

from generate_mtdna_call_mt.merging_constants import *

# Github repo locations for imports:
# gnomad: https://github.com/broadinstitute/gnomad_methods
# gnomad_qc: https://github.com/broadinstitute/gnomad_qc

# Include NA in POPS to account for cases where population annotations are missing
POPS = POPS["v3"]["genomes"]
POPS.append("NA")

#hl.init(tmp_dir="file:///tmp")
#hl.init(tmp_dir=f"{os.environ['WORKSPACE_BUCKET']}/tmp",
#        local_tmpdir="file:///tmp",
#        spark_conf={"spark.local.dir": "file:///tmp"})

RESOURCE_PATH = 'gcp-public-data--gnomad/resources/mitochondria'
RESOURCES = {
    "variant_context": f"gs://{RESOURCE_PATH}/variant_context/chrM_pos_ref_alt_context_categories.txt",
    "phylotree": f"gs://{RESOURCE_PATH}/phylotree/rCRS-centered_phylo_vars_final_update.txt",
    "pon_mt_trna": f"gs://{RESOURCE_PATH}/trna_predictions/pon_mt_trna_predictions_08_27_2020.txt",
    "mitotip": f"gs://{RESOURCE_PATH}/trna_predictions/mitotip_scores_08_27_2020.txt",
}


logging.basicConfig(
    format="%(asctime)s (%(name)s %(lineno)s): %(message)s",
    datefmt="%m/%d/%Y %I:%M:%S %p",
)
logger = logging.getLogger("add annotations")
logger.setLevel(logging.INFO)

_HAIL_VERSION_MINOR = int(hl.version().split('-')[0].split('.')[2])


def add_genotype(mt_path: str, min_hom_threshold: float = 0.95) -> hl.MatrixTable:
    """
    STABLE
    Add in genotype annotation based on heteroplasmy level.

    If the heteroplasmy level is above the min_hom_threshold, set the genotype to 1/1.
    If the heteroplasmy level is less than the min_hom_threshold, but greater than 0, set the genotype to 0/1.
    Otherwise set the genotype to 0/0.

    :param mt_path: Path to the MatrixTable (this MatrixTable can be generated by running combine_vcfs.py)
    :param min_hom_threshold: Minimum heteroplasmy level to define a variant as homoplasmic
    :return: MatrixTable with GT field added
    """
    logger.info("Reading in MT...")
    mt = hl.read_matrix_table(mt_path)

    # Add in genotype (GT) based on min_hom_threshold
    mt = mt.annotate_entries(
        GT=(
            hl.case()
            .when((mt.HL < min_hom_threshold) & (mt.HL > 0.0), hl.parse_call("0/1"))
            .when(mt.HL >= min_hom_threshold, hl.parse_call("1/1"))
            .when(mt.HL == 0, hl.parse_call("0/0"))
            .default(hl.missing(hl.tcall))
        ),
    )

    return mt


def add_variant_context(input_mt: hl.MatrixTable) -> hl.MatrixTable:
    """
    STABLE
    Add variant context annotations to the MatrixTable.

    This fucntion adds in information on regions/strand for SNPs that can be useful for determining mutational signatures.

    :param input_mt: MatrixTable
    :return: MatrixTable with variant context information added
    """
    # Read in variant context data
    vc_ht = hl.import_table(RESOURCES["variant_context"], impute=True)

    # Split columns into separate annotations
    vc_ht = vc_ht.annotate(
        ref=vc_ht["POS.REF.ALT"].split(r"\.")[1],
        alt=vc_ht["POS.REF.ALT"].split(r"\.")[2],
        strand=vc_ht.Context_category.split("_")[-1],
        variant=vc_ht.Context_category.split("_")[0],
    )

    # Rename and select certain columns
    vc_ht = vc_ht.rename({"MT_POS": "pos", "Annotation": "region"})
    vc_ht = vc_ht.select("pos", "ref", "alt", "strand", "region", "variant")

    # Key by locus and allele
    vc_ht = vc_ht.key_by(
        locus=hl.locus("MT", vc_ht.pos, reference_genome="GRCh37"),
        alleles=[vc_ht.ref, vc_ht.alt],
    )

    # Annotate original mt with variant context information
    input_mt = input_mt.annotate_rows(**vc_ht[input_mt.locus, input_mt.alleles])
    input_mt = input_mt.annotate_rows(
        variant_context=hl.str(input_mt.variant) + "_" + hl.str(input_mt.strand)
    )
    input_mt = input_mt.drop("pos", "ref", "alt", "strand", "variant")

    return input_mt


def add_gnomad_metadata(input_mt: hl.MatrixTable) -> hl.MatrixTable:
    """
    Add select gnomAD metadata to the MatrixTable.

    :param input_mt: MatrixTable
    :return: MatrixTable with select gnomAD metadata added
    """
    # TODO: Add option here to accomodate non-gnomAD metadata
    genome_meta_ht = meta.versions["3.1"].ht()

    genome_meta_struct = genome_meta_ht[input_mt.s]

    input_mt = input_mt.annotate_cols(
        release=genome_meta_struct.release,
        hard_filters=genome_meta_struct.sample_filters.hard_filters,
        research_project=genome_meta_struct.project_meta.research_project,
        project_id=genome_meta_struct.project_meta.project_id,
        product=genome_meta_struct.project_meta.product,
        sample_pi=genome_meta_struct.project_meta.sample_pi,
        sex_karyotype=genome_meta_struct.sex_imputation.sex_karyotype,
        age=hl.if_else(
            hl.is_defined(genome_meta_struct.project_meta.age),
            genome_meta_struct.project_meta.age,
            genome_meta_struct.project_meta.age_alt,
        ),
        broad_external=genome_meta_struct.project_meta.broad_external,
        pop=genome_meta_struct.population_inference.pop,
    )

    return input_mt


def add_age_and_pop(input_mt: hl.MatrixTable, participant_data: str) -> hl.MatrixTable:
    """
    STABLE
    Add sample-level metadata for age and pop to `input_mt`.

    :param input_mt: MatrixTable
    :param participant_data: Path to metadata file downloaded from Terra that contains sample age and pop information
    :return: MatrixTable with select age and pop annotations added
    """
    ht = hl.import_table(
        participant_data,
        types={"age": hl.tint32, "pop": hl.tstr},
        missing=["", "NA", "NaN", "nan"],
    ).key_by("s")

    ht = ht.select("age", "pop")

    input_mt = input_mt.annotate_cols(**ht[input_mt.col_key])

    # If a sample doesn't have an annotated population, set it to the string "NA"
    input_mt = input_mt.annotate_cols(
        pop=hl.if_else(hl.is_missing(input_mt.pop), "NA", input_mt.pop)
    )

    return input_mt


def filter_by_hom_overlap(input_mt: hl.MatrixTable, keep_all_samples: bool, sample_stats: str):
    stat_ht = hl.import_table(sample_stats, impute=True, key='s', types={'s':hl.tstr})
    input_mt = input_mt.annotate_cols(num_mt_overlaps = stat_ht[input_mt.col_key].mtdna_consensus_overlaps)
    n_removed = input_mt.aggregate_cols(hl.agg.count_where(input_mt.num_mt_overlaps > 0))
    if not keep_all_samples:
        input_mt = input_mt.filter_cols(input_mt.num_mt_overlaps == 0)
    
    return input_mt.drop('num_mt_overlaps'), n_removed


def filter_by_copy_number(
    input_mt: hl.MatrixTable, keep_all_samples: bool = False, max_cn: int = 500
) -> hl.MatrixTable:
    """
    Calculate the mitochondrial copy number based on mean mitochondrial coverage and median nuclear coverage. Filter out samples with more extreme copy numbers.

    Note that median and mean coverage for mitochondria are very similar. Mean mitochondria coverage was used based on metrics available at the time, but releases will switch to using median mitochondria coverage.

    :param hl.MatrixTable input_mt: MatrixTable
    :param keep_all_samples: If True, keep all samples (calculate mitochondrial copy number, but do not filter any samples based on this metric)
    :return: MatrixTable filtered to samples with a copy number of at least 50 and less than 500, number samples below 50 removed, number samples above 500 removed
    """
    # Calculate mitochondrial copy number, if median autosomal coverage is not present default to a wgs_median_coverage of 30x
    input_mt = input_mt.annotate_cols(
        mito_cn=2
        * input_mt.mt_mean_coverage
        / hl.if_else(
            hl.is_missing(input_mt.wgs_median_coverage),
            30,
            input_mt.wgs_median_coverage,
        )
    )
    n_removed_below_cn = input_mt.aggregate_cols(
        hl.agg.count_where(input_mt.mito_cn < 50)
    )
    n_removed_above_cn = input_mt.aggregate_cols(
        hl.agg.count_where(input_mt.mito_cn > max_cn)
    )

    if not keep_all_samples:
        # Remove samples with a mitochondrial copy number below 50 or greater than 500
        input_mt = input_mt.filter_cols(
            (input_mt.mito_cn >= 50) & (input_mt.mito_cn <= max_cn)
        )
    input_mt = input_mt.filter_rows(hl.agg.any(input_mt.HL > 0))

    return input_mt, n_removed_below_cn, n_removed_above_cn


def filter_by_contamination(
    input_mt: hl.MatrixTable, output_dir: str, keep_all_samples: bool = False
) -> hl.MatrixTable:
    """
    Calculate contamination based on internal algorithm and filter out samples with contamination above 2%.

    Contamination takes into account:
    a) mitochondria contamination output by HaploCheck
    b) nuclear contamination (freemix) output by VerifyBamID
    c) an internal algorithm with utilizes the PASS haplogroup-defining variants which should be homoplasmic (100% alternate alleles), but in contaminated samples show multiple alleles with heteroplasmy 85-99.8%

    :param input_mt: MatrixTable
    :param output_dir: Output directory to which results should be written
    :param keep_all_samples: If True, keep all samples (calculate contamination, but do not filter any samples based on this metric)
    :return: MatrixTable filtered to samples without contamination, number of contaminated samples removed
    """
    # Generate expression for genotypes with >= 85% heteroplasmy and no FT filters at haplogroup-defining sites that are not filtered as artifact-prone sites
    over_85_expr = (
        (input_mt.HL >= 0.85)
        & (input_mt.FT == {"PASS"})
        & input_mt.hap_defining_variant
        & ~hl.str(input_mt.filters).contains("artifact_prone_site")
    )

    input_mt = input_mt.annotate_cols(
        over_85_mean=hl.agg.filter(over_85_expr, hl.agg.mean(input_mt.HL)),
        over_85_count=hl.agg.filter(
            over_85_expr, hl.agg.count_where(hl.is_defined(input_mt.HL))
        ),
        bt_85_and_99_mean=hl.agg.filter(
            over_85_expr & (input_mt.HL <= 0.998), hl.agg.mean(input_mt.HL)
        ),
        bt_85_and_99_count=hl.agg.filter(
            over_85_expr & (input_mt.HL <= 0.998),
            hl.agg.count_where(hl.is_defined(input_mt.HL)),
        ),
    )

    input_mt = input_mt.annotate_cols(
        contam_high_het=hl.if_else(
            input_mt.bt_85_and_99_count >= 3,
            1 - input_mt.bt_85_and_99_mean,
            1 - input_mt.over_85_mean,
        )
    )

    # If contam_high_het is nan, set to 0 (to avoid filtering out missing values which would be more common with haplogroups closer to the reference haplogroup)
    input_mt = input_mt.annotate_cols(
        contam_high_het=hl.if_else(
            hl.is_nan(input_mt.contam_high_het), 0, input_mt.contam_high_het
        )
    )

    # Find samples on border of .02 that may flip between < 0.02 and > 0.02 from issues with floating point precision and mark these samples for removal
    epsilon = 0.000001
    border_samples = input_mt.aggregate_cols(
        hl.agg.filter(
            (input_mt.contam_high_het > (0.02 - epsilon))
            & (input_mt.contam_high_het < (0.02 + epsilon)),
            hl.agg.collect((input_mt.s)),
        )
    )

    border_samples = (
        hl.literal(border_samples) if border_samples else hl.empty_array(hl.tstr)
    )

    # Add annotation to keep only samples with a contamination less than 2%
    input_mt = input_mt.annotate_cols(freemix_percentage_imp = hl.if_else(hl.is_missing(input_mt.freemix_percentage), 0, input_mt.freemix_percentage))
    input_mt = input_mt.annotate_cols(
        keep=(input_mt.contamination < 0.02)
        & (input_mt.freemix_percentage_imp < 2)
        & (input_mt.contam_high_het < 0.02)
        & ~border_samples.contains(input_mt.s)
    )
    # Save sample contamination information to separate file
    n_contaminated = input_mt.aggregate_cols(hl.agg.count_where(~input_mt.keep))

    sample_data = input_mt.select_cols(
        "contamination",
        "freemix_percentage",
        "contam_high_het",
        "over_85_mean",
        "over_85_count",
        "bt_85_and_99_mean",
        "bt_85_and_99_count",
        "keep",
    )
    data_export = sample_data.cols()
    data_export.export(f"{output_dir}/sample_contamination.tsv")

    if not keep_all_samples:
        logger.info(
            "Removing %d samples with contamination above 2 percent", n_contaminated
        )
        input_mt = input_mt.filter_cols(input_mt.keep)
    input_mt = input_mt.drop("keep")

    input_mt = input_mt.filter_rows(hl.agg.any(input_mt.HL > 0))

    return input_mt, n_contaminated


def add_terra_metadata(
    input_mt: hl.MatrixTable, participant_data: str
) -> hl.MatrixTable:
    """
    STABLE
    Add Terra metadata to the MatrixTable.

    The participant_data file can be obtained by downloading the participant data after running Mutect2 in Terra. This file should contain the following columns:
        - entity:participant_id: Participant ID uploaded to Terra by user
        - s: Sample ID uploaded to Terra by user
        - contamination: Output by Mutect2, gives the estimate of mitochondrial contamination
        - freemix_percentage: Uploaded to Terra by user, can be calculated with VerifyBamID
        - major_haplogroup: Output by Mutect2 which utilizes Haplogrep
        - wgs_median_coverage: Uploaded to Terra by user, can be calculated with Picard's CollectWgsMetrics
        - mt_mean_coverage: Output by Mutect2, gives the mean mitochondrial coverage

    :param input_mt: MatrixTable
    :param participant_data: Path to metadata file downloaded from Terra
    :return: MatrixTable with Terra metadata annotations added
    """
    # Add haplogroup and Mutect2/Terra output annotations
    ht = hl.import_table(
        participant_data,
        types={
            "contamination": hl.tfloat64,
            "freemix_percentage": hl.tfloat64,
            "mt_mean_coverage": hl.tfloat64,
            "wgs_median_coverage": hl.tfloat64,
        },
        missing=["","NA"],
    ).key_by("s")
    ht = ht.rename({"entity:participant_id": "participant_id"})

    ht = ht.select(
        "participant_id",
        "contamination",
        "freemix_percentage",
        "major_haplogroup",
        "wgs_median_coverage",
        "mt_mean_coverage",
    )

    input_mt = input_mt.annotate_cols(**ht[input_mt.s])

    # Annotate the high level haplogroup by taking the first letter, with the exception of H and L haplogroups which are more commonly referred to using the first two letters
    input_mt = input_mt.annotate_cols(
        hap=hl.if_else(
            input_mt.major_haplogroup.startswith("HV")
            | input_mt.major_haplogroup.startswith("L"),
            input_mt.major_haplogroup[0:2],
            input_mt.major_haplogroup[0],
        )
    )

    return input_mt


def fill_missing_mt_mean_coverage_from_covdb(
    input_mt: hl.MatrixTable,
    coverage_h5_path: str | None,
    position_block_size: int = 1024,
    logger: logging.Logger | None = None,
) -> hl.MatrixTable:
    """Fill missing mt_mean_coverage using coverage.h5 (covdb) block reads.

    This reuses the same covdb read strategy as finalize and avoids loading
    the full (sample × position) matrix into Hail.
    """
    if not coverage_h5_path:
        return input_mt

    log = logger if logger is not None else logging.getLogger(__name__)

    ht_cols = input_mt.cols()
    ht_cols = ht_cols.select(mt_mean_coverage=ht_cols.mt_mean_coverage)
    col_rows = ht_cols.collect()
    missing_samples = [r.s for r in col_rows if r.mt_mean_coverage is None]
    if not missing_samples:
        return input_mt

    from generate_mtdna_call_mt.covdb_utils import open_covdb_index, read_covdb_block

    log.info(
        "Filling mt_mean_coverage for %d samples using covdb %s",
        len(missing_samples),
        coverage_h5_path,
    )

    idx = open_covdb_index(coverage_h5_path)
    missing_indices = []
    missing_lookup = []
    for s in missing_samples:
        cov_i = idx.sample_to_index.get(s)
        if cov_i is None:
            raise ValueError(
                f"Sample {s} missing from coverage.h5 /sample_id"
            )
        missing_indices.append(int(cov_i))
        missing_lookup.append(s)

    ht_rows = input_mt.rows()
    pos_ht = (
        ht_rows.select(pos=hl.int32(ht_rows.locus.position))
        .distinct()
        .order_by("pos")
    )
    positions = [int(r.pos) for r in pos_ht.collect()]
    if not positions:
        return input_mt

    pos_to_covdb_col = {p: idx.pos_to_index.get(p) for p in positions}
    missing_pos = [p for p, ci in pos_to_covdb_col.items() if ci is None]
    if missing_pos:
        raise ValueError(
            f"{len(missing_pos)} MT positions missing from coverage.h5 /pos (first 5: {missing_pos[:5]})"
        )

    sample_indices = np.array(missing_indices, dtype=np.int64)
    sums = np.zeros(len(sample_indices), dtype=np.float64)
    total_positions = 0
    for start in range(0, len(positions), position_block_size):
        block = positions[start : start + position_block_size]
        covdb_cols_block = np.array(
            [pos_to_covdb_col[p] for p in block], dtype=np.int64
        )
        cov_block = read_covdb_block(
            h5_path=coverage_h5_path,
            sample_indices=sample_indices,
            pos_indices=covdb_cols_block,
        )
        sums += cov_block.sum(axis=1)
        total_positions += cov_block.shape[1]

    if total_positions == 0:
        return input_mt

    mean_cov = (sums / float(total_positions)).tolist()
    mt_mean_by_sample = {
        s: float(m) for s, m in zip(missing_lookup, mean_cov)
    }
    mean_cov_hl = hl.literal(mt_mean_by_sample)
    input_mt = input_mt.annotate_cols(
        mt_mean_coverage=hl.if_else(
            hl.is_missing(input_mt.mt_mean_coverage),
            mean_cov_hl.get(input_mt.s),
            input_mt.mt_mean_coverage,
        )
    )

    return input_mt


def add_hap_defining(input_mt: hl.MatrixTable) -> hl.MatrixTable:
    """
    STABLE
    Add bool on whether or not a variant is a haplogroup-defining variant to the MatrixTable.

    Haplogroup-defining annotations were obtained from PhyloTree Build 17.

    :param input_mt: MatrixTable
    :return: MatrixTable with annotation on whether or not the variant is haplogroup-defining added
    """
    # TODO: move dataset location
    hap_defining_variants = hl.import_table(RESOURCES["phylotree"])

    hap_defining = hl.literal(set(hap_defining_variants.variant.collect()))
    input_mt = input_mt.annotate_rows(
        variant_collapsed=input_mt.alleles[0]
        + hl.str(input_mt.locus.position)
        + input_mt.alleles[1]
    )
    input_mt = input_mt.annotate_rows(
        hap_defining_variant=hap_defining.contains(input_mt.variant_collapsed)
    )  # set hap_defining_variant to True or False

    return input_mt


def add_trna_predictions(input_mt: hl.MatrixTable, avoid_fasta_workaround: bool) -> hl.MatrixTable:
    """
    STABLE
    Add tRNA predictions on pathogenicity from PON-mt-tRNA and MitoTIP to the MatrixTable.

    :param input_mt: MatrixTable
    :param avoid_fasta_workaround: bool
    :return: MatrixTable with tRNA predictions of pathogenicity added
    """
    # Add PON-mt-tRNA predictions
    pon_predictions = hl.import_table(RESOURCES["pon_mt_trna"])

    # If reference allele from fasta doesn't match Reference_nucleotide, PON-mt-tRNA is reporting the allele of opposite strand and need to get reverse complement for ref and alt

    if not avoid_fasta_workaround:
        # as a workaround for some hail weirdness, we are going to manually grab the GRCh37 sequence
        tab_workaround = hl.import_table(RESOURCES['variant_context'])
        tab_workaround = tab_workaround.select(pos = hl.int(tab_workaround.MT_POS), 
                                               REF = tab_workaround["POS.REF.ALT"].split('\.')[1])
        tab_workaround = tab_workaround.key_by('pos','REF').distinct()
        tab_workaround = tab_workaround.key_by('pos')
        pon_predictions = pon_predictions.annotate(
            ref=tab_workaround[hl.int(pon_predictions.mtDNA_position)].REF
        )

        # confirm that the workaround produces the correct results
        #add_reference_sequence(hl.get_reference("GRCh37"))
        #tab_workaround = tab_workaround.annotate(REF_corr = hl.get_sequence('MT', tab_workaround.pos, reference_genome='GRCh37'))
        #tab_workaround.filter(tab_workaround.REF_corr != tab_workaround.REF).count() # > 0
    else:
        add_reference_sequence(hl.get_reference("GRCh37"))
        pon_predictions = pon_predictions.annotate(
            ref=hl.get_sequence(
                "MT", hl.int(pon_predictions.mtDNA_position), reference_genome='GRCh37'
            )
        )

    n_miss_ref = pon_predictions.filter(~hl.is_defined(pon_predictions.ref)).count()
    if n_miss_ref != 0:
        raise ValueError('ERROR: pon predictions table has undefined reference sequences.')
    
    pon_predictions = pon_predictions.annotate(
        alt=hl.if_else(
            pon_predictions.Reference_nucleotide == pon_predictions.ref,
            pon_predictions.New_nucleotide,
            hl.reverse_complement(pon_predictions.New_nucleotide),
        )
    )
    pon_predictions = pon_predictions.key_by(
        variant_id=pon_predictions.ref
        + hl.str(pon_predictions.mtDNA_position)
        + pon_predictions.alt
    )
    input_mt = input_mt.annotate_rows(
        pon_mt_trna_prediction=pon_predictions[input_mt.variant_collapsed]
        .Classification.lower()
        .replace(" ", "_"),
        pon_ml_probability_of_pathogenicity=hl.float(
            pon_predictions[input_mt.variant_collapsed].ML_probability_of_pathogenicity
        ),
    )

    # Add MitoTIP predictions
    mitotip_predictions = hl.import_table(RESOURCES["mitotip"])
    mitotip_predictions = mitotip_predictions.key_by(
        variant_id=mitotip_predictions.rCRS
        + hl.str(mitotip_predictions.Position)
        + mitotip_predictions.Alt
    )
    input_mt = input_mt.annotate_rows(
        mitotip_score=hl.float(
            mitotip_predictions[input_mt.variant_collapsed].MitoTIP_Score
        )
    )
    # Set pathogenicity based on MitoTIP scores, classifications obtained from MitoTIP's website
    input_mt = input_mt.annotate_rows(
        mitotip_trna_prediction=(
            hl.case()
            .when(input_mt.mitotip_score > 16.25, "likely_pathogenic")
            .when(
                (input_mt.mitotip_score <= 16.25) & (input_mt.mitotip_score > 12.66),
                "possibly_pathogenic",
            )
            .when(
                (input_mt.mitotip_score <= 12.66) & (input_mt.mitotip_score >= 8.44),
                "possibly_benign",
            )
            .when((input_mt.mitotip_score < 8.44), "likely_benign")
            .or_missing()
        )
    )

    return input_mt


def get_indel_expr(input_mt: hl.MatrixTable) -> hl.expr.BooleanExpression:
    """
    Generate expression for filtering to indels that should be used to evaluate indel stacks.

    To be considered a variant to be used to evaluate indel stacks, the variant should:
    a) be an indel
    b) have a heteroplasmy level >= 0.01 and <= 0.95
    c) have a PASS genotype

    :param input_mt: MatrixTable
    :return: Expression to be used for determining if a variant is an indel that should to be used to evaluate indel stacks
    """
    indel_expr = (
        hl.is_indel(input_mt.alleles[0], input_mt.alleles[1])
        & (input_mt.HL <= 0.95)
        & (input_mt.HL >= 0.01)
        & (input_mt.FT == {"PASS"})
    )

    return indel_expr


def generate_expressions(
    input_mt: hl.MatrixTable, min_hom_threshold: float = 0.95
) -> hl.MatrixTable:
    """
    STABLE
    Create expressions to use for annotating the MatrixTable.

    The expressions include AC, AN, AF, filtering allele frequency (FAF) split by homplasmic/heteroplasmic, haplgroup, and population.
    Also includes calcuations of mean DP, MQ, and TLOD.

    :param input_mt: MatrixTable
    :param min_hom_threshold: Minimum heteroplasmy level to define a variant as homoplasmic
    :return: Tuple of hail expressions
    """
    # Calculate AC and AN
    AC = hl.agg.count_where((input_mt.HL > 0.0))
    AN = hl.agg.count_where(hl.is_defined(input_mt.HL))
    # Note: if AN is zero, AFs will evaluate to NaN, which may need to be converted to zero for downstream tools
    AF = AC / AN

    # Calculate AC for het and hom variants, and histogram for HL
    AC_hom = hl.agg.count_where(input_mt.HL >= min_hom_threshold)
    AC_het = hl.agg.count_where((input_mt.HL < min_hom_threshold) & (input_mt.HL > 0.0))
    HL_hist = hl.agg.filter(input_mt.HL > 0, hl.agg.hist(input_mt.HL, 0, 1, 10))
    DP_hist_alt = hl.agg.filter(
        input_mt.GT.is_non_ref(), hl.agg.hist(input_mt.DP, 0, 2000, 10)
    )
    DP_hist_all = hl.agg.hist(input_mt.DP, 0, 2000, 10)
    DP_mean = hl.agg.mean(input_mt.DP)
    MQ_mean = hl.agg.mean(input_mt.MQ)
    TLOD_mean = hl.agg.mean(input_mt.TLOD)

    # Calculate AF
    # Note: if AN is zero, AFs will evaluate to NaN, which may need to be converted to zero for downstream tools
    AF_hom = AC_hom / AN
    AF_het = AC_het / AN

    # Calculate max individual heteroplasmy
    max_HL = hl.agg.max(input_mt.HL)

    # Haplogroup annotations
    pre_hap_AC = hl.agg.group_by(input_mt.hap, AC)
    pre_hap_AN = hl.agg.group_by(input_mt.hap, AN)
    pre_hap_AF = hl.agg.group_by(input_mt.hap, AF)
    pre_hap_AC_het = hl.agg.group_by(input_mt.hap, AC_het)
    pre_hap_AC_hom = hl.agg.group_by(input_mt.hap, AC_hom)
    pre_hap_AF_hom = hl.agg.group_by(input_mt.hap, AF_hom)
    pre_hap_AF_het = hl.agg.group_by(input_mt.hap, AF_het)
    pre_hap_HL_hist = hl.agg.group_by(input_mt.hap, HL_hist.bin_freq)
    pre_hap_FAF = hl.agg.group_by(
        input_mt.hap,
        hl.experimental.filtering_allele_frequency(hl.int32(AC), hl.int32(AN), 0.95),
    )
    pre_hap_FAF_hom = hl.agg.group_by(
        input_mt.hap,
        hl.experimental.filtering_allele_frequency(
            hl.int32(AC_hom), hl.int32(AN), 0.95
        ),
    )

    # population annotations
    pre_pop_AC = hl.agg.group_by(input_mt.pop, AC)
    pre_pop_AN = hl.agg.group_by(input_mt.pop, AN)
    pre_pop_AF = hl.agg.group_by(input_mt.pop, AF)
    pre_pop_AC_het = hl.agg.group_by(input_mt.pop, AC_het)
    pre_pop_AC_hom = hl.agg.group_by(input_mt.pop, AC_hom)
    pre_pop_AF_hom = hl.agg.group_by(input_mt.pop, AF_hom)
    pre_pop_AF_het = hl.agg.group_by(input_mt.pop, AF_het)
    pre_pop_HL_hist = hl.agg.group_by(input_mt.pop, HL_hist.bin_freq)

    return hl.struct(
        AC=AC,
        AN=AN,
        AF=AF,
        AC_hom=AC_hom,
        AC_het=AC_het,
        hl_hist=HL_hist,
        dp_hist_all=DP_hist_all,
        dp_hist_alt=DP_hist_alt,
        dp_mean=DP_mean,
        mq_mean=MQ_mean,
        tlod_mean=TLOD_mean,
        AF_hom=AF_hom,
        AF_het=AF_het,
        max_hl=max_HL,
        pre_hap_AC=pre_hap_AC,
        pre_hap_AN=pre_hap_AN,
        pre_hap_AF=pre_hap_AF,
        pre_hap_AC_het=pre_hap_AC_het,
        pre_hap_AF_het=pre_hap_AF_het,
        pre_hap_AC_hom=pre_hap_AC_hom,
        pre_hap_AF_hom=pre_hap_AF_hom,
        pre_hap_hl_hist=pre_hap_HL_hist,
        pre_hap_faf=pre_hap_FAF,
        pre_hap_faf_hom=pre_hap_FAF_hom,
        pre_pop_AN=pre_pop_AN,
        pre_pop_AC_het=pre_pop_AC_het,
        pre_pop_AF_het=pre_pop_AF_het,
        pre_pop_AC_hom=pre_pop_AC_hom,
        pre_pop_AF_hom=pre_pop_AF_hom,
        pre_pop_hl_hist=pre_pop_HL_hist,
    )


def standardize_haps(
    input_mt: hl.MatrixTable, annotation: str, haplogroup_order: list
) -> list:
    """
    STABLE
    Convert the dictionary of haplogroup annotations into an array of values in a predefined haplogroup order.

    :param input_mt: MatrixTable
    :param annotation: Annotation to convert and sort
    :param haplogroup_order: Order in which to sort the haplogroups
    :return: Sorted list of haplogroup annotations (the values of the dictionary)
    """
    # Converts haplogroup dictionary to sorted array
    value = [input_mt[annotation][x] for x in haplogroup_order]

    return value


def standardize_pops(
    input_mt: hl.MatrixTable, annotation: str, population_order: list
) -> list:
    """
    STABLE
    Convert the dictionary of population annotations into an array of values in a predefined population order.

    :param input_mt: MatrixTable
    :param annotation: Annotation to convert and sort
    :param population_order: Order in which to sort the populations
    :return: Sorted list of population annotations (the values of the dictionary)
    """
    # Converts haplogroup dictionary to sorted array
    value = [input_mt[annotation][x] for x in population_order]

    return value


def add_quality_histograms(input_mt: hl.MatrixTable) -> hl.MatrixTable:
    """
    STABLE
    Add histogram annotations for quality metrics to the MatrixTable.

    :param input_mt: MatrixTable
    :return: MatrixTable annotated with quality metric histograms
    """
    # Generate histogram for site quality metrics across all variants
    # TODO: decide on bin edges
    dp_hist_all_variants = input_mt.aggregate_rows(
        hl.agg.hist(input_mt.dp_mean, 0, 4000, 40)
    )
    input_mt = input_mt.annotate_globals(
        dp_hist_all_variants_bin_freq=dp_hist_all_variants.bin_freq,
        dp_hist_all_variants_n_larger=dp_hist_all_variants.n_larger,
        dp_hist_all_variants_bin_edges=dp_hist_all_variants.bin_edges,
    )

    mq_hist_all_variants = input_mt.aggregate_rows(
        hl.agg.hist(input_mt.mq_mean, 0, 80, 40)
    )  # is 80 the actual max value here?
    input_mt = input_mt.annotate_globals(
        mq_hist_all_variants_bin_freq=mq_hist_all_variants.bin_freq,
        mq_hist_all_variants_n_larger=mq_hist_all_variants.n_larger,
        mq_hist_all_variants_bin_edges=mq_hist_all_variants.bin_edges,
    )

    tlod_hist_all_variants = input_mt.aggregate_rows(
        hl.agg.hist(input_mt.tlod_mean, 0, 40000, 40)
    )
    input_mt = input_mt.annotate_globals(
        tlod_hist_all_variants_bin_freq=tlod_hist_all_variants.bin_freq,
        tlod_hist_all_variants_n_larger=tlod_hist_all_variants.n_larger,
        tlod_hist_all_variants_bin_edges=tlod_hist_all_variants.bin_edges,
    )

    # Generate histogram for overall age distribution
    age_hist_all_samples = input_mt.aggregate_cols(
        hl.agg.hist(input_mt.age, 30, 80, 10)
    )
    input_mt = input_mt.annotate_globals(
        age_hist_all_samples_bin_freq=age_hist_all_samples.bin_freq,
        age_hist_all_samples_n_larger=age_hist_all_samples.n_larger,
        age_hist_all_samples_n_smaller=age_hist_all_samples.n_smaller,
        age_hist_all_samples_bin_edges=age_hist_all_samples.bin_edges,
    )

    # Add age histograms per variant type (heteroplasmic or homoplasmic)
    age_data = age_hists_expr(True, input_mt.GT, input_mt.age)
    input_mt = input_mt.annotate_rows(
        age_hist_hom=age_data.age_hist_hom, age_hist_het=age_data.age_hist_het
    )

    return input_mt


def add_annotations_by_hap_and_pop(input_mt: hl.MatrixTable, temp_dir) -> hl.MatrixTable:
    """
    Add variant annotations (such as AC, AN, AF, heteroplasmy histogram, and filtering allele frequency) split by haplogroup and population.

    :param input_mt: MatrixTable
    :return: MatrixTable with variant annotations
    """
    # Order the haplogroup-specific annotations
    list_hap_order = list(set(input_mt.hap.collect()))
    input_mt = input_mt.annotate_globals(hap_order=sorted(list_hap_order))

    # Sanity check for haplogroups (make sure that they at least start with a letter)
    for i in list_hap_order:
        if not re.match("^[A-Z]", i):
            sys.exit(f"Invalid haplogroup {i}, does not start with a letter")

    pre_hap_annotation_labels_1 = [
        "pre_hap_AC",
        "pre_hap_AN",
        "pre_hap_AF",
        "pre_hap_AC_het",
        "pre_hap_AC_hom",
    ]
    for_annot = {re.sub("pre_", "", i): standardize_haps(input_mt, i, sorted(list_hap_order)) for i in pre_hap_annotation_labels_1}
    input_mt = input_mt.annotate_rows(**for_annot)
    input_mt = input_mt.checkpoint(f"{temp_dir}/temp3.mt", overwrite=True)

    pre_hap_annotation_labels_2 = [
        "pre_hap_AF_hom",
        "pre_hap_AF_het",
        "pre_hap_hl_hist",
        "pre_hap_faf",
        "pre_hap_faf_hom",
    ]
    for_annot = {re.sub("pre_", "", i): standardize_haps(input_mt, i, sorted(list_hap_order)) for i in pre_hap_annotation_labels_2}
    input_mt = input_mt.annotate_rows(**for_annot)
    input_mt = input_mt.checkpoint(f"{temp_dir}/temp4.mt", overwrite=True)

    # Get a list of indexes where AC of the haplogroup is greater than 0, then get the list of haplogroups with that index
    input_mt = input_mt.annotate_rows(
        alt_haps=hl.enumerate(input_mt.hap_AC)
        .filter(lambda x: x[1] > 0)
        .map(lambda x: input_mt.hap_order[x[0]])
    )
    # Count number of haplogroups containing an alt allele
    input_mt = input_mt.annotate_rows(n_alt_haps=hl.len(input_mt.alt_haps))

    # Calculate hapmax
    input_mt = input_mt.annotate_rows(
        hapmax_AF_hom=input_mt.hap_order[(hl.argmax(input_mt.hap_AF_hom, unique=True))],
        hapmax_AF_het=input_mt.hap_order[(hl.argmax(input_mt.hap_AF_het, unique=True))],
    )

    # Calculate faf hapmax
    input_mt = input_mt.annotate_rows(
        faf_hapmax=hl.max(input_mt.hap_faf), faf_hapmax_hom=hl.max(input_mt.hap_faf_hom)
    )

    # Add populatation annotations
    final_pops = set(input_mt.pop.collect())

    # Order according to POPS
    #final_pops = [x for x in POPS if x in found_pops]
    #if len(found_pops - set(POPS)) > 0:
    #    sys.exit("Invalid population found")
    
    input_mt = input_mt.annotate_globals(pop_order=final_pops)

    pre_pop_annotation_labels = [
        "pre_pop_AN",
        "pre_pop_AC_het",
        "pre_pop_AC_hom",
        "pre_pop_AF_hom",
        "pre_pop_AF_het",
        "pre_pop_hl_hist",
    ]

    for i in pre_pop_annotation_labels:
        # Remove "pre" prefix for final annotations
        final_annotation = re.sub("pre_", "", i)
        input_mt = input_mt.annotate_rows(
            **{final_annotation: standardize_pops(input_mt, i, final_pops)}
        )

    # Drop intermediate annotations
    annotations_to_drop = [
        "pre_hap_AC",
        "pre_hap_AN",
        "pre_hap_AF",
        "pre_hap_AC_het",
        "pre_hap_AC_hom",
        "pre_hap_AF_hom",
        "pre_hap_AF_het",
        "pre_hap_hl_hist",
        "pre_hap_faf",
        "pre_hap_faf_hom",
        "AC_mid_het",
        "AF_mid_het",
        "pre_pop_AN",
        "pre_pop_AC_het",
        "pre_pop_AC_hom",
        "pre_pop_AF_hom",
        "pre_pop_AF_het",
        "pre_pop_hl_hist",
    ]

    input_mt = input_mt.drop(*annotations_to_drop)
    # Last-minute drops (ever add back in?)
    input_mt = input_mt.drop(
        "AC",
        "AF",
        "hap_AC",
        "hap_AF",
        "hap_faf",
        "faf_hapmax",
        "alt_haps",
        "n_alt_haps",
    )

    input_mt = input_mt.annotate_rows(filters=input_mt.filters.difference({"PASS"}))

    return input_mt


def apply_common_low_het_flag(input_mt: hl.MatrixTable) -> hl.MatrixTable:
    """
    Apply the common_low_heteroplasmy flag to the MatrixTable.

    The common_low_heteroplasmy flag marks variants where the overall frequency is > 0.001 for samples with a heteroplasmy level > 0 and < 0.50 and either "low_allele_frac" or "PASS" for the genotype filter

    NOTE: The "low_allele_frac" is applied by Mutect2 to variants with a heteroplasmy level below the supplied vaf_filter_threshold

    :param input_mt: MatrixTable
    :return: MatrixTable with the common_low_heteroplasmy flag added
    """
    input_mt = format_filters(input_mt)
    input_mt = input_mt.annotate_rows(
        AC_mid_het=hl.agg.count_where(
            (input_mt.HL < 0.50)
            & (input_mt.HL > 0.0)
            & ((input_mt.FT == {"PASS"}) | (input_mt.FT == {"low_allele_frac"}))
        )
    )
    input_mt = input_mt.annotate_rows(
        AF_mid_het=input_mt.AC_mid_het
        / hl.agg.count_where(
            hl.is_defined(input_mt.HL)
            & ((input_mt.FT == {"PASS"}) | (input_mt.FT == {"low_allele_frac"}))
        )
    )
    input_mt = input_mt.annotate_rows(
        common_low_heteroplasmy=input_mt.AF_mid_het > 0.001
    )

    return format_filters(input_mt)


def remove_low_allele_frac_genotypes(
    input_mt: hl.MatrixTable, vaf_filter_threshold: float = 0.01
) -> hl.MatrixTable:
    """
    STABLE
    Remove low_allele_frac genotypes and sets the call to homoplasmic reference.

    NOTE: vaf_filter_threshold should match what was supplied to the vaf_filter_threshold when running Mutect2, variants below this value will be set to homoplasmic reference after calculating the common_low_heteroplasmy filter, Mutect2 will have flagged these variants as "low_allele_frac"

    :param input_mt: MatrixTable
    :param vaf_filter_threshold: Should match vaf_filter_threshold supplied to Mutect2, variants below this value will be set to homoplasmic reference after calculating the common_low_heteroplasmy filter
    :return: MatrixTable with genotypes below the vaf_filter_threshold set to homoplasmic reference
    """
    # Set HL to 0 if < vaf_filter_threshold and remove variants that no longer have at least one alt call
    input_mt = input_mt.annotate_entries(
        HL=hl.if_else(
            (input_mt.HL > 0) & (input_mt.HL < vaf_filter_threshold), 0, input_mt.HL
        )
    )
    # Filter field for all variants with a heteroplasmy of 0 should be set to PASS
    # This step is needed to prevent homref calls that are filtered
    input_mt = input_mt.annotate_entries(
        FT=hl.if_else(input_mt.HL < vaf_filter_threshold, {"PASS"}, input_mt.FT)
    )
    input_mt = input_mt.annotate_entries(
        GT=hl.if_else(
            input_mt.HL < vaf_filter_threshold, hl.parse_call("0/0"), input_mt.GT
        )
    )

    # Check that variants no longer contain the "low_allele_frac" filter (vaf_filter_threshold should be set to appropriate level to remove these variants)
    laf_rows = input_mt.filter_rows(
        hl.agg.any(hl.str(input_mt.FT).contains("low_allele_frac"))
    )
    n_laf_rows = laf_rows.count_rows()
    if n_laf_rows > 0:
        sys.exit(
            "low_allele_frac filter should no longer be present after applying vaf_filter_threshold (vaf_filter_threshold should equal the vaf_filter_threshold supplied to Mutect2)"
        )
    input_mt = input_mt.filter_rows(hl.agg.any(input_mt.HL > 0))

    return input_mt


def apply_indel_stack_filter(input_mt: hl.MatrixTable) -> hl.MatrixTable:
    """
    STABLE
    Apply the indel_stack filter to the MatrixTable.

    The indel_stack filter marks alleles where all samples with the variant call had at least 2 different indels called at the position

    :param input_mt: MatrixTable
    :return: MatrixTable with the indel_stack filter added
    """
    # Add variant-level indel_stack at any indel allele where all samples with a variant call had at least 2 different indels called at that position
    # If any sample had a solo indel at that position, do not filter
    indel_expr = get_indel_expr(input_mt)
    input_mt = input_mt.annotate_cols(
        indel_pos_counter=hl.agg.filter(
            indel_expr, hl.agg.counter(input_mt.locus.position)
        )
    )
    indel_expr = get_indel_expr(input_mt)
    input_mt = input_mt.annotate_entries(
        indel_occurences=(
            hl.case()
            .when(
                (
                    indel_expr
                    & (input_mt.indel_pos_counter.get(input_mt.locus.position) >= 2)
                ),
                "stack",
            )
            .when(
                (
                    indel_expr
                    & (input_mt.indel_pos_counter.get(input_mt.locus.position) == 1)
                ),
                "solo",
            )
            .or_missing()
        )
    )

    # If stack is true and solo is false, the indel is stack only and should be filtered out
    input_mt = input_mt.annotate_rows(
        filters=hl.if_else(
            hl.agg.any(input_mt.indel_occurences == "stack")
            & ~hl.agg.any(input_mt.indel_occurences == "solo"),
            input_mt.filters.add("indel_stack"),
            input_mt.filters,
        )
    )

    return format_filters(input_mt)


def filter_genotypes_below_min_het_threshold(
    input_mt: hl.MatrixTable, min_het_threshold: float = 0.10
) -> hl.MatrixTable:
    """
    Filter out genotypes with a heteroplasmy below the min_het_threshold.

    This filter is a genotype level filter to remove variants with a heteroplasmy level below the specified min_het_threshold
    NOTE: Should later parameterize this function to allow other heteroplasmy cutoffs?

    :param input_mt: MatrixTable
    :param min_het_threshold: Minimum heteroplasmy level to define a variant as a PASS heteroplasmic variant, genotypes below this threshold will count towards the heteroplasmy_below_min_het_threshold filter and be set to missing
    :return: MatrixTable with the heteroplasmy_below_min_het_threshold in the FT field added where applicable
    """
    input_mt = input_mt.annotate_entries(
        FT=hl.if_else(
            (input_mt.HL < min_het_threshold) & (input_mt.GT.is_het()),
            input_mt.FT.add("heteroplasmy_below_min_het_threshold"),
            input_mt.FT,
        )
    )


    return format_filters(input_mt)


def apply_npg_filter(input_mt: hl.MatrixTable) -> hl.MatrixTable:
    """
    Apply the npg filter to the MatrixTable.

    The npg (no pass genotypes) filter marks sites that don't have at least one pass alt call

    :param input_mt: MatrixTable
    :return: MatrixTable with the npg filter added
    """
    input_mt = format_filters(input_mt)
    input_mt = input_mt.annotate_rows(
        filters=hl.if_else(
            ~(hl.agg.any((input_mt.HL > 0.0) & (input_mt.FT == {"PASS"}))),
            input_mt.filters.add("npg"),
            input_mt.filters,
        )
    )

    return format_filters(input_mt)


def generate_filter_histogram(
    input_mt: hl.MatrixTable, filter_name: str
) -> hl.ArrayExpression:
    """
    STABLE
    Generate histogram for number of indiviudals with the specified sample-level filter at different heteroplasmy levels.

    :param input_mt: MatrixTable
    :param filter_name: Name of sample-filter for which to generate a histogram
    :return: Histogram containing the counts of individuals with a variant filtered by the specified filter name across binned heteroplasmy levels
    """
    filter_histogram = hl.agg.filter(
        hl.str(input_mt.FT).contains(filter_name), hl.agg.hist(input_mt.HL, 0, 1, 10)
    ).bin_freq

    return filter_histogram


def format_filters(mt, row_f = ['filters'], entry_f = ['FT','FT_LIFT']):
    """ 
    Ensures that for all relevant filter fields, PASS is only present when there are no filters present.
    If a field is hl.missing(), then len and if_else produces a missing value. Thus it remains missing.
    If a field is length 0, then it is processed here and does not become missing.
    This function prevents any filters of length 0.
    """
    mt = mt.annotate_rows(**{x: mt[x].difference({'PASS'}) for x in row_f})
    mt = mt.annotate_rows(**{x: hl.if_else(hl.len(mt[x]) == 0, {'PASS'}, mt[x]) for x in row_f})
    
    mt = mt.annotate_entries(**{x: mt[x].difference({'PASS'}) for x in entry_f})
    mt = mt.annotate_entries(**{x: hl.if_else(hl.len(mt[x]) == 0, {'PASS'}, mt[x]) for x in entry_f})
    return mt


def add_filter_annotations(
    input_mt: hl.MatrixTable, 
    vaf_filter_threshold: float = 0.01,
    min_het_threshold: float = 0.10,
    temp_dir: str = ''
) -> hl.MatrixTable:
    """
    Generate histogram for number of individuals with the specified sample-level filter at different heteroplasmy levels.

    :param input_mt: MatrixTable
    :param vaf_filter_threshold: Should match vaf_filter_threshold supplied to Mutect2, variants below this value will be set to homoplasmic reference after calculating the common_low_heteroplasmy filter
    :param min_het_threshold: Minimum heteroplasmy level to define a variant as a PASS heteroplasmic variant, genotypes below this threshold will count towards the heteroplasmy_below_min_het_threshold filter and be set to missing
    :return: MatrixTable with added annotations for sample and variant level filters and number of genotypes with heteroplasmy_below_min_het_threshold
    """
    # TODO: pull these from header instead?
    filters = [
        "base_qual",
        "map_qual",
        "position",
        "strand_bias",
        "weak_evidence",
        "contamination",
        "heteroplasmy_below_min_het_threshold",
    ]

    logger.info("Applying common low heteroplasmy flag...")
    input_mt = apply_common_low_het_flag(input_mt)

    logger.info("Removing low_allele_frac genotypes...")
    input_mt = remove_low_allele_frac_genotypes(input_mt, vaf_filter_threshold)
    input_mt = input_mt.checkpoint(f"{temp_dir}/temp_low_frac.mt", overwrite=True)

    logger.info("Applying indel_stack filter...")
    input_mt = apply_indel_stack_filter(input_mt)

    logger.info(
        "Filtering genotypes below with heteroplasmy below the min_het_threshold..."
    )
    input_mt = filter_genotypes_below_min_het_threshold(input_mt, min_het_threshold)
    input_mt = input_mt.checkpoint(f"{temp_dir}/temp_het_filt.mt", overwrite=True)
    
    n_het_below_min_het_threshold = input_mt.aggregate_entries(
        hl.agg.count_where(
            hl.str(input_mt.FT).contains("heteroplasmy_below_min_het_threshold")
        )
    )

    logger.info("Applying npg filter...")
    input_mt = apply_npg_filter(input_mt)

    logger.info("Generating filter histograms and calculating excluded_AC...")
    for i in filters:
        annotation_name = i + "_hist"
        input_mt = input_mt.annotate_rows(
            **{annotation_name: generate_filter_histogram(input_mt, i)}
        )
    input_mt = input_mt.annotate_rows(
        excluded_AC=hl.agg.count_where(input_mt.FT != {"PASS"})
    )

    return format_filters(input_mt), n_het_below_min_het_threshold


def filter_genotypes(input_mt: hl.MatrixTable, pass_set: set = {}) -> hl.MatrixTable:
    """
    Set all genotype field values to missing if the variant is not "PASS" for that sample.

    :param input_mt: MatrixTable
    :return: MatrixTable with filtered genotype fields set to missing
    """
    if len(pass_set) > 0:
        pass_expr = (input_mt.FT == {"PASS"}) | input_mt.FT.is_subset(hl.literal(pass_set))
    else:
        pass_expr = input_mt.FT == {"PASS"}
    
    input_mt = input_mt.annotate_entries(
        GT=hl.or_missing(pass_expr, input_mt.GT),
        DP=hl.or_missing(pass_expr, input_mt.DP),
        HL=hl.or_missing(pass_expr, input_mt.HL),
        #FT=hl.or_missing(pass_expr, input_mt.FT),
        #FT_LIFT=hl.or_missing(pass_expr, input_mt.FT_LIFT),
        MQ=hl.or_missing(pass_expr, input_mt.MQ),
        TLOD=hl.or_missing(pass_expr, input_mt.TLOD),
    )

    input_mt = input_mt.annotate_entries(FT = hl.if_else(input_mt.FT == {'PASS'}, input_mt.FT.union({'GT_PASS'}), input_mt.FT))
    input_mt = input_mt.annotate_entries(FT = input_mt.FT.difference({'PASS'}), FT_LIFT = input_mt.FT_LIFT.difference({'PASS'}))

    return input_mt


def add_sample_annotations(
    input_mt: hl.MatrixTable, min_hom_threshold: float = 0.95
) -> hl.MatrixTable:
    """
    STABLE
    Add sample annotations to the MatrixTable.

    These sample annotations include the callrate, number of heteroplasmic/homoplasmic SNPs/indels, and the number of singletons.
    Filtered variants (such as artifact_prone_site, npg, and indel_stack) are excluded from these calculations.

    :param input_mt: MatrixTable
    :param min_hom_threshold: Minimum heteroplasmy level to define a variant as homoplasmic
    :return: MatrixTable with sample annotations added
    """
    # Count number of variants
    num_rows = input_mt.count_rows()

    # Add sample qc annotations
    filter_expr = hl.len(input_mt.filters) == 0
    input_mt = input_mt.annotate_cols(
        callrate=hl.agg.filter(
            filter_expr, (hl.agg.count_where(hl.is_defined(input_mt.HL))) / num_rows
        ),
        n_singletons_het=hl.agg.filter(
            filter_expr,
            hl.agg.count_where(
                (input_mt.AC_het == 1)
                & ((input_mt.HL < min_hom_threshold) & (input_mt.HL > 0.0))
            ),
        ),
        n_singletons_hom=hl.agg.filter(
            filter_expr,
            hl.agg.count_where(
                (input_mt.AC_hom == 1) & (input_mt.HL >= min_hom_threshold)
            ),
        ),
        n_snp_het=hl.agg.filter(
            filter_expr,
            hl.agg.count_where(
                (input_mt.HL < min_hom_threshold)
                & (input_mt.HL > 0.0)
                & hl.is_snp(input_mt.alleles[0], input_mt.alleles[1])
            ),
        ),
        n_snp_hom=hl.agg.filter(
            filter_expr,
            hl.agg.count_where(
                (input_mt.HL >= min_hom_threshold)
                & hl.is_snp(input_mt.alleles[0], input_mt.alleles[1])
            ),
        ),
        n_indel_het=hl.agg.filter(
            filter_expr,
            hl.agg.count_where(
                (input_mt.HL < min_hom_threshold)
                & (input_mt.HL > 0.0)
                & (~hl.is_snp(input_mt.alleles[0], input_mt.alleles[1]))
            ),
        ),
        n_indel_hom=hl.agg.filter(
            filter_expr,
            hl.agg.count_where(
                (input_mt.HL >= min_hom_threshold)
                & (~hl.is_snp(input_mt.alleles[0], input_mt.alleles[1]))
            ),
        ),
    )

    return input_mt


def add_vep(input_mt: hl.MatrixTable, run_vep: bool, vep_output: str) -> hl.MatrixTable:
    """
    STABLE
    Add vep annotations to the MatrixTable.

    :param input_mt: MatrixTable
    :param run_vep: Whether or not to run vep
    :param vep_output: Path to the MatrixTable output vep results (either the existing results or where to ouput new vep results)
    :return: MatrixTable with vep annotations
    """
    if run_vep:
        vep_mt = hl.vep(input_mt)
        vep_mt = vep_mt.checkpoint(vep_output, overwrite=True)
    else:
        vep_mt = hl.read_matrix_table(vep_output)

    input_mt = input_mt.annotate_rows(
        vep=vep_mt.index_rows(input_mt.locus, input_mt.alleles).vep
    )
    # TODO: get vep version directly from config file
    input_mt = input_mt.annotate_globals(vep_version="v101")

    # If only filter is END_TRUNC, change lof for LC to HC and remove the END_TRUNC filter
    # Remove SINGLE_EXON flags because all exons are single exon in the mitochondria
    input_mt = input_mt.annotate_rows(
        vep=input_mt.vep.annotate(
            transcript_consequences=input_mt.vep.transcript_consequences.map(
                lambda x: x.annotate(
                    lof=hl.if_else(x.lof_filter == "END_TRUNC", "HC", x.lof),
                    lof_filter=hl.if_else(
                        x.lof_filter == "END_TRUNC", hl.missing(hl.tstr), x.lof_filter
                    ),
                    lof_flags=hl.if_else(
                        x.lof_flags == "SINGLE_EXON", hl.missing(hl.tstr), x.lof_flags
                    ),
                )
            )
        )
    )

    end_trunc_count = input_mt.filter_rows(
        hl.str(input_mt.vep.transcript_consequences[0].lof_filter).contains("END_TRUNC")
    ).count_rows()
    if end_trunc_count > 0:
        sys.exit(
            f"END_TRUNC filter should no longer be present but was found for {end_trunc_count} variants"
        )

    single_exon_count = input_mt.filter_rows(
        hl.str(input_mt.vep.transcript_consequences[0].lof_flags).contains(
            "SINGLE_EXON"
        )
    ).count_rows()
    if single_exon_count > 0:
        sys.exit(
            f"SINGLE_EXON flag should no longer be present but was found for {single_exon_count} variants"
        )

    return input_mt


def add_rsids(input_mt: hl.MatrixTable, band_aid_fix: bool) -> hl.MatrixTable:
    """
    STABLE
    Add rsid annotations to the MatrixTable.

    :param input_mt: MatrixTable
    :return: MatrixTable with rsid annotations added
    """
    dbsnp_import_args = dbsnp.versions["b154"].import_args
    # Replace the contig recoding with just the chrM mapping
    dbsnp_import_args.update({"contig_recoding": {"NC_012920.1": "chrM"}})
    # If enabled, fix paths
    if band_aid_fix:
        band_aid_fun = lambda x: re.sub('gnomad-public-requester-pays', 'gcp-public-data--gnomad', x)
        dbsnp_import_args.update({"path": band_aid_fun(dbsnp_import_args['path']),
                                  'header_file': band_aid_fun(dbsnp_import_args['header_file'])})
    dbsnp_ht = _import_dbsnp(**dbsnp_import_args)

    input_mt = input_mt.annotate_rows(
        rsid=dbsnp_ht[input_mt.locus, input_mt.alleles].rsid
    )
    input_mt = input_mt.annotate_globals(dbsnp_version="b154")

    return input_mt


def export_simplified_variants(input_ht: hl.Table, output_dir: str) -> None:
    """
    Export a text file containing only several high-level variant annotations.

    :param input_ht: Hail Table of variants
    :param output_dir: Output directory to which results should be output
    :return: None
    """
    reduced_ht = (
        input_ht.key_by(
            chromosome=input_ht.locus.contig,
            position=input_ht.locus.position,
            ref=input_ht.alleles[0],
            alt=input_ht.alleles[1],
        )
        .select("filters", "AC_hom", "AC_het", "AF_hom", "AF_het", "AN", "max_hl")
        .rename({"max_hl": "max_observed_heteroplasmy"})
    )
    reduced_ht = reduced_ht.annotate(
        filters=hl.if_else(
            hl.len(reduced_ht.filters) == 0,
            "PASS",
            hl.str(",").join(hl.array(reduced_ht.filters)),
        )
    )

    reduced_ht.export(f"{output_dir}/reduced_annotations.txt")


def generate_output_paths(
    output_dir: str, file_name: str, subset_name: str, extension: str
) -> list:
    """
    Generate output paths for results files based on the given output directory, file name, subset name, and extension.

    :param output_dir: Output directory to which results should be output
    :param file_name: Name of the file, preceeds subset_name
    :param subset_name: Name that should be appended to output file names
    :param extension: Extension for the output file
    :return: Path for the file
    """
    # set up output paths for callset
    file_path = f"{output_dir}/{file_name}{subset_name}.{extension}"

    return file_path


def report_stats(
    input_mt: hl.MatrixTable,
    output_dir: str,
    pass_only: bool,
    n_samples_below_cn: int,
    n_samples_above_cn: int,
    n_samples_contam: int,
    n_het_below_min_het_threshold: int,
    n_removed_overlap: int,
    min_het_threshold: float = 0.10,
    min_hom_threshold: float = 0.95,
    max_cn: int = 500
) -> None:
    """
    Generate output report with basic stats.

    :param input_mt: MatrixTable
    :param output_dir: Output directory to which results should be output
    :param pass_only: Whether or not directory should be filtered to pass_only variants
    :param n_samples_below_cn: Number of samples removed because mitochondrial number is less than 50
    :param n_samples_above_cn: Number of samples removed because mitochondrial number is above 500
    :param n_samples_contam: Number of samples removed because of contamination
    :param n_het_below_min_het_threshold: Number of genotypes filtered because the heteroplasmy levels was below the min_het_threshold
    :param n_removed_overlap: Number of samples filtered because there were overlapping homoplasmies
    :param min_het_threshold: Minimum heteroplasmy level to define a variant as a PASS heteroplasmic variant, genotypes below this threshold will count towards the heteroplasmy_below_min_het_threshold filter and be set to missing
    :param min_hom_threshold: Minimum heteroplasmy level to define a variant as homoplasmic
    :return: None
    """
    if pass_only:
        suffix = "_pass"
        input_mt = input_mt.filter_rows(hl.len(input_mt.filters) == 0)
    else:
        suffix = ""
    out_stats = hl.hadoop_open(f"{output_dir}/stats{suffix}.txt", "w")

    if pass_only:
        out_stats.write("Below metrics are for PASS-only variants\n\n")

    # Report numbers of filtered samples/genotypes
    out_stats.write(
        f"Number of samples removed because of overlapping homoplasmies: {n_removed_overlap}\n"
    )
    out_stats.write(
        f"Number of samples removed because contamination above 2%: {n_samples_contam}\n"
    )
    out_stats.write(
        f"Number of samples removed because mitochondrial copy number below 50: {n_samples_below_cn}\n"
    )
    out_stats.write(
        f"Number of samples removed because mitochondrial copy number above {str(max_cn)}: {n_samples_above_cn}\n"
    )
    out_stats.write(
        f'Number of genotypes filtered because "heteroplasmy_below_min_het_threshold": {n_het_below_min_het_threshold}\n\n'
    )

    # Count variant, samples, bases
    unique_variants, samples = input_mt.count()
    out_stats.write(f"Number of samples: {samples}\n")
    out_stats.write(f"Number of unique variants: {unique_variants}\n")
    bases_w_variant = len(set(input_mt.locus.position.collect()))
    out_stats.write(f"Number of bases with variation: {bases_w_variant}\n\n")

    # Count number of filters
    for filter_name, filter_count in Counter(
        [i for sublist in input_mt.filters.collect() for i in sublist]
    ).items():
        out_stats.write(
            f'Number of variants with "{filter_name}" filter: {filter_count} variants\n'
        )

    # Calculate row stats
    row_stats = input_mt.aggregate_rows(
        hl.struct(
            common_low_het_count=hl.agg.count_where(input_mt.common_low_heteroplasmy),
            het_only_sites=hl.agg.count_where(
                (input_mt.AC_het > 0) & (input_mt.AC_hom == 0)
            ),
            hom_only_sites=hl.agg.count_where(
                (input_mt.AC_hom > 0) & (input_mt.AC_het == 0)
            ),
            het_and_hom_sites=hl.agg.count_where(
                (input_mt.AC_hom > 0) & (input_mt.AC_het > 0)
            ),
            hap_defining_sites=hl.agg.count_where(input_mt.hap_defining_variant),
            snps=hl.agg.count_where(
                hl.is_snp(input_mt.alleles[0], input_mt.alleles[1])
            ),
            indels=hl.agg.count_where(
                hl.is_indel(input_mt.alleles[0], input_mt.alleles[1])
            ),
            transitions=hl.agg.count_where(
                hl.is_transition(input_mt.alleles[0], input_mt.alleles[1])
            ),
            transversions=hl.agg.count_where(
                hl.is_transversion(input_mt.alleles[0], input_mt.alleles[1])
            ),
        )
    )

    # Calculate col stats
    col_stats = input_mt.aggregate_cols(
        hl.struct(
            unique_haplogroups=hl.len(hl.agg.collect_as_set(input_mt.major_haplogroup)),
            unique_top_level_haplogroups=hl.len(hl.agg.collect_as_set(input_mt.hap)),
        )
    )

    # Calculate entry stats
    entry_stats = input_mt.aggregate_entries(
        hl.struct(
            total_variants=hl.agg.count_where(input_mt.HL > 0),
            total_hom_variants=hl.agg.count_where(input_mt.HL >= min_hom_threshold),
            total_het_variants=hl.agg.count_where(
                (input_mt.HL < min_hom_threshold) & (input_mt.HL >= min_het_threshold)
            ),
            min_hl=hl.agg.filter(input_mt.HL > 0, hl.agg.min(input_mt.HL)),
            max_hl=hl.agg.filter(input_mt.HL > 0, hl.agg.max(input_mt.HL)),
        )
    )

    # Count number of flags
    out_stats.write(
        f'Number of variants with "common_low_heteroplasmy" flag: {row_stats["common_low_het_count"]} variants\n\n'
    )

    # Count variants
    out_stats.write(f'Total number of variants: {entry_stats["total_variants"]}\n')

    # Count of homoplasmic/heteroplasmic variants
    out_stats.write(
        f'Number of homoplasmic-only sites: {row_stats["hom_only_sites"]}\n'
    )
    out_stats.write(
        f'Number of heteroplasmic-only sites: {row_stats["het_only_sites"]}\n'
    )
    out_stats.write(f'Number of het and hom sites: {row_stats["het_and_hom_sites"]}\n')

    percent_hom = round(
        entry_stats["total_hom_variants"] / entry_stats["total_variants"], 2
    )
    percent_het = round(
        entry_stats["total_het_variants"] / entry_stats["total_variants"], 2
    )
    out_stats.write(
        f'Total number of homoplasmic variants: {entry_stats["total_hom_variants"]}\n'
    )
    out_stats.write(f"Percent homoplasmic variants: {percent_hom}\n")
    out_stats.write(
        f'Total number of heteroplasmic variants: {entry_stats["total_het_variants"]}\n'
    )
    out_stats.write(f"Percent heteroplasmic variants: {percent_het}\n\n")

    out_stats.write(f'Minimum heteroplasmy detected: {entry_stats["min_hl"]}\n')
    out_stats.write(f'Maximum heteroplasmy detected: {entry_stats["max_hl"]}\n\n')

    # Count number of snps and indels
    out_stats.write(f'Number of SNPs: {row_stats["snps"]}\n')
    out_stats.write(f'Number of indels: {row_stats["indels"]}\n')

    # Count number of transitions and transversions
    out_stats.write(f'Number of transitions: {row_stats["transitions"]}\n')
    out_stats.write(f'Number of transversions: {row_stats["transversions"]}\n')
    out_stats.write(
        f'Number of haplogroup defining variants: {row_stats["hap_defining_sites"]}\n\n'
    )

    # Count number of haplogroups
    out_stats.write(
        f'Number of unique haplogroups: {col_stats["unique_haplogroups"]}\n'
    )
    out_stats.write(
        f'Number of top-level haplogroups: {col_stats["unique_top_level_haplogroups"]}\n'
    )

    out_stats.close()


def change_to_grch38_chrm(input_mt: hl.MatrixTable) -> None:
    """
    Change build to GRCh38 and filters reference genome to chrM.

    :param input_mt: MatrixTable
    :return: MatrixTable with GRCh38 reference genome subsetted to just chrM (so that extraneous contigs will be excluded from VCF output)
    """
    ref = hl.get_reference("GRCh38")
    my_ref = hl.ReferenceGenome(
        "GRCh38_chrM", contigs=["chrM"], lengths={"chrM": ref.lengths["chrM"]}
    )
    assert "chrM" in ref.contigs
    input_mt = input_mt.key_rows_by(
        locus=hl.locus("chrM", input_mt.locus.position, reference_genome="GRCh38_chrM"),
        alleles=input_mt.alleles,
    )

    return input_mt


def format_vcf(
    input_mt: hl.MatrixTable,
    output_dir: str,
    min_hom_threshold: float = 0.95,
    vaf_filter_threshold: float = 0.01,
    min_het_threshold: float = 0.10,
    skip_vep: bool = False,
) -> dict:
    """
    Generate dictionary for VCF header annotations.

    :param input_mt: MatrixTable
    :param min_hom_threshold: Minimum heteroplasmy level to define a variant as homoplasmic
    :param vaf_filter_threshold: Should match vaf_filter_threshold supplied to Mutect2, variants below this value will be set to homoplasmic reference after calculating the common_low_heteroplasmy filter
    :param min_het_threshold: Minimum heteroplasmy level to define a variant as a PASS heteroplasmic variant, genotypes below this threshold will count towards the heteroplasmy_below_min_het_threshold filter and be set to missing
    :param output_dir: Output directory to which appended header info should be written
    :return: MatrixTable with VCF annotations in the info field and dictionary of filter, info, and format fields to be output in the VCF header; path of VCF headers to append
    """
    input_mt = change_to_grch38_chrm(input_mt)

    haplogroup_order = hl.eval(input_mt.hap_order)
    population_order = hl.eval(input_mt.pop_order)

    age_hist_hom_bin_edges = input_mt.age_hist_hom.bin_edges.take(1)[0]
    age_hist_het_bin_edges = input_mt.age_hist_het.bin_edges.take(1)[0]
    hl_hist_bin_edges = input_mt.hl_hist.bin_edges.take(1)[0]
    dp_hist_all_bin_edges = input_mt.dp_hist_all.bin_edges.take(1)[0]
    dp_hist_alt_bin_edges = input_mt.dp_hist_alt.bin_edges.take(1)[0]

    input_mt = input_mt.annotate_rows(
        hl_hist=input_mt.hl_hist.bin_freq,
        age_hist_hom_bin_freq=input_mt.age_hist_hom.bin_freq,
        age_hist_hom_n_smaller=input_mt.age_hist_hom.n_smaller,
        age_hist_hom_n_larger=input_mt.age_hist_hom.n_larger,
        age_hist_het_bin_freq=input_mt.age_hist_het.bin_freq,
        age_hist_het_n_smaller=input_mt.age_hist_het.n_smaller,
        age_hist_het_n_larger=input_mt.age_hist_het.n_larger,
        dp_hist_all_n_larger=input_mt.dp_hist_all.n_larger,
        dp_hist_alt_n_larger=input_mt.dp_hist_alt.n_larger,
        dp_hist_all_bin_freq=input_mt.dp_hist_all.bin_freq,
        dp_hist_alt_bin_freq=input_mt.dp_hist_alt.bin_freq,
    )

    if not skip_vep:
        input_mt = input_mt.annotate_rows(vep=vep_struct_to_csq(input_mt.vep))

    # Get length of annotations to use in Number fields in the VCF where necessary
    len_hap_hl_hist = len(input_mt.hap_hl_hist.take(1)[0])
    len_pop_hl_hist = len(input_mt.pop_hl_hist.take(1)[0])

    # Output appended header info to file
    vcf_header_file = output_dir + "/extra_fields_for_header.tsv"
    if not skip_vep:
        appended_vcf_header = dedent(
            f"""
        ##VEP version={hl.eval(input_mt.vep_version)}
        ##dbSNP version={hl.eval(input_mt.dbsnp_version)}
        ##age distributions=bin_edges:{hl.eval(input_mt.age_hist_all_samples_bin_edges)}, bin_freq:{hl.eval(input_mt.age_hist_all_samples_bin_freq)}, n_smaller:{hl.eval(input_mt.age_hist_all_samples_n_smaller)}, n_larger:{hl.eval(input_mt.age_hist_all_samples_n_larger)}

        """
        )
    else:
        appended_vcf_header = dedent(
            f"""
        ##dbSNP version={hl.eval(input_mt.dbsnp_version)}
        ##age distributions=bin_edges:{hl.eval(input_mt.age_hist_all_samples_bin_edges)}, bin_freq:{hl.eval(input_mt.age_hist_all_samples_bin_freq)}, n_smaller:{hl.eval(input_mt.age_hist_all_samples_n_smaller)}, n_larger:{hl.eval(input_mt.age_hist_all_samples_n_larger)}

        """
        )
    with hl.hadoop_open(vcf_header_file, "w") as out:
        out.write(appended_vcf_header)

    # Drop intermediate annotations
    input_mt = input_mt.drop(
        "region",
        "variant_context",
        "age_hist_het",
        "age_hist_hom",
        "dp_hist_all",
        "dp_hist_alt",
    )

    ht = input_mt.rows()
    # Move row annotations into info struct
    input_mt = input_mt.annotate_rows(info=hl.struct())
    input_mt = input_mt.annotate_rows(
        info=input_mt.info.annotate(**ht[input_mt.row_key]).drop("rsid")
    ).select_rows(
        "rsid", "filters", "info"
    )  # create info annotation

    # Convert "rsid" array to str for VCF output
    input_mt = input_mt.annotate_rows(rsid=hl.str(";").join(input_mt.rsid))

    # Convert "," to "|" for array annotations
    for key, value in input_mt.row_value.info.items():
        if str(value).startswith("<Array"):
            if str(value.dtype).startswith("array<array"):
                # If value is an array of arrays, only replace the commas within each individual array
                input_mt = input_mt.annotate_rows(
                    info=input_mt.info.annotate(
                        **{
                            key: hl.map(
                                lambda x: hl.delimit(x, delimiter="|"),
                                input_mt.info[key],
                            )
                        }
                    )
                )
            else:
                input_mt = input_mt.annotate_rows(
                    info=input_mt.info.annotate(
                        **{key: hl.delimit(input_mt.info[key], delimiter="|")}
                    )
                )

    meta_dict = {
        "filter": {
            "artifact_prone_site": {
                "Description": "Variant overlaps site that is commonly reported in literature to be artifact prone"
            },
            "npg": {
                "Description": "No-pass-genotypes site (no individuals were PASS for the variant)"
            },
            "indel_stack": {
                "Description": "Allele where all samples with the variant call had at least 2 different heteroplasmic indels called at the position"
            },
        },
        "info": {
            "variant_collapsed": {
                "Description": "Variant in format of RefPosAlt",
                "Number": "1",
                "Type": "String",
            },
            "hap_defining_variant": {
                "Description": "Present if variant is present as a haplogroup defining variant in PhyloTree build 17",
                "Number": "0",
                "Type": "Flag",
            },
            "common_low_heteroplasmy": {
                "Description": f"Present if variant is found at an overall frequency of .001 across all samples with a heteroplasmy level > 0 and < 0.50 (includes variants <{vaf_filter_threshold} heteroplasmy which are subsequently filtered)",
                "Number": "0",
                "Type": "Flag",
            },
            "AN": {
                "Description": "Overall allele number (number of samples with non-missing genotype)",
                "Number": "1",
                "Type": "Integer",
            },
            "AC_hom": {
                "Description": f"Allele count restricted to variants with a heteroplasmy level >= {min_hom_threshold}",
                "Number": "1",
                "Type": "Integer",
            },  # should put in threshold variable
            "AC_het": {
                "Description": f"Allele count restricted to variants with a heteroplasmy level >= 0.10 and < {min_hom_threshold}",
                "Number": "1",
                "Type": "Integer",
            },
            "hl_hist": {
                "Description": f"Histogram of heteroplasmy levels; bin edges are: {hl_hist_bin_edges}",
                "Number": "1",
                "Type": "String",
            },
            "hap_hl_hist": {
                "Description": f"Histogram of heteroplasmy levels for each haplogroup; bin edges are: {hl_hist_bin_edges}, haplogroup order: {haplogroup_order}",
                "Number": f"{len_hap_hl_hist}",
                "Type": "String",
            },
            "AF_hom": {
                "Description": f"Allele frequency restricted to variants with a heteroplasmy level >= {min_hom_threshold}",
                "Number": "1",
                "Type": "Float",
            },
            "AF_het": {
                "Description": f"Allele frequency restricted to variants with a heteroplasmy level >= 0.10 and < {min_hom_threshold}",
                "Number": "1",
                "Type": "Float",
            },
            "max_hl": {
                "Description": "Maximum heteroplasmy level observed among all samples for that variant",
                "Number": "1",
                "Type": "Float",
            },
            "hap_AN": {
                "Description": f"List of overall allele number for each haplogroup, haplogroup order: {haplogroup_order}",
                "Number": "1",
                "Type": "String",
            },
            "hap_AC_het": {
                "Description": f"List of AC_het for each haplogroup, haplogroup order: {haplogroup_order}",
                "Number": "1",
                "Type": "String",
            },
            "hap_AC_hom": {
                "Description": f"List of AC_hom for each haplogroup, haplogroup order: {haplogroup_order}",
                "Number": "1",
                "Type": "String",
            },
            "hap_AF_hom": {
                "Description": f"List of AF_hom for each haplogroup, haplogroup order: {haplogroup_order}",
                "Number": "1",
                "Type": "String",
            },
            "hap_AF_het": {
                "Description": f"List of AF_het for each haplogroup, haplogroup order: {haplogroup_order}",
                "Number": "1",
                "Type": "String",
            },
            "hap_faf_hom": {
                "Description": f"List of filtering allele frequency for each haplogroup restricted to homoplasmic variants, haplogroup order: {haplogroup_order}",
                "Number": "1",
                "Type": "String",
            },
            "hapmax_AF_hom": {
                "Description": "Haplogroup with maximum AF_hom",
                "Number": "1",
                "Type": "String",
            },
            "hapmax_AF_het": {
                "Description": "Haplogroup with maximum AF_het",
                "Number": "1",
                "Type": "String",
            },
            "faf_hapmax_hom": {
                "Description": "Maximum filtering allele frequency across haplogroups restricted to homoplasmic variants",
                "Number": "1",
                "Type": "Float",
            },
            "bin_edges_hl_hist": {
                "Description": "Bin edges for histogram of heteroplasmy levels",
                "Number": "1",
                "Type": "String",
            },
            "pop_AN": {
                "Description": f"List of overall allele number for each population, population order: {population_order}",
                "Number": "1",
                "Type": "String",
            },
            "pop_AC_het": {
                "Description": f"List of AC_het for each population, population order: {population_order}",
                "Number": "1",
                "Type": "String",
            },
            "pop_AC_hom": {
                "Description": f"List of AC_hom for each population, population order: {population_order}",
                "Number": "1",
                "Type": "String",
            },
            "pop_AF_hom": {
                "Description": f"List of AF_hom for each population, population order: {population_order}",
                "Number": "1",
                "Type": "String",
            },
            "pop_AF_het": {
                "Description": f"List of AF_het for each population, population order: {population_order}",
                "Number": "1",
                "Type": "String",
            },
            "pop_hl_hist": {
                "Description": f"Histogram of heteroplasmy levels for each population; bin edges are: {hl_hist_bin_edges}, population order: {population_order}",
                "Number": f"{len_pop_hl_hist}",
                "Type": "String",
            },
            "age_hist_hom_bin_freq": {
                "Description": f"Histogram of ages of individuals with a homoplasmic variant; bin edges are: {age_hist_hom_bin_edges}",
                "Number": "1",
                "Type": "String",
            },
            "age_hist_hom_n_smaller": {
                "Description": "Count of age values falling below lowest histogram bin edge for individuals with a homoplasmic variant",
                "Number": "1",
                "Type": "String",
            },
            "age_hist_hom_n_larger": {
                "Description": "Count of age values falling above highest histogram bin edge for individuals with a homoplasmic variant",
                "Number": "1",
                "Type": "String",
            },
            "age_hist_het_bin_freq": {
                "Description": f"Histogram of ages of individuals with a heteroplasmic variant; bin edges are: {age_hist_het_bin_edges}",
                "Number": "1",
                "Type": "String",
            },
            "age_hist_het_n_smaller": {
                "Description": "Count of age values falling below lowest histogram bin edge for individuals with a heteroplasmic variant",
                "Number": "1",
                "Type": "String",
            },
            "age_hist_het_n_larger": {
                "Description": "Count of age values falling above highest histogram bin edge for individuals with a heteroplasmic variant",
                "Number": "1",
                "Type": "String",
            },
            "dp_hist_all_n_larger": {
                "Description": "Count of dp values falling above highest histogram bin edge for all individuals",
                "Number": "1",
                "Type": "String",
            },
            "dp_hist_alt_n_larger": {
                "Description": "Count of dp values falling above highest histogram bin edge for individuals with the alternative allele",
                "Number": "1",
                "Type": "String",
            },
            "dp_hist_all_bin_freq": {
                "Description": f"Histogram of dp values for all individuals; bin edges are: {dp_hist_all_bin_edges}",
                "Number": "1",
                "Type": "String",
            },
            "dp_hist_alt_bin_freq": {
                "Description": f"Histogram of dp values for individuals with the alternative allele; bin edges are: {dp_hist_alt_bin_edges}",
                "Number": "1",
                "Type": "String",
            },
            "dp_mean": {
                "Description": "Mean depth across all individuals for the site",
                "Number": "1",
                "Type": "Float",
            },
            "mq_mean": {
                "Description": "Mean MMQ (median mapping quality) across individuals with a variant for the site",
                "Number": "1",
                "Type": "Float",
            },
            "tlod_mean": {
                "Description": "Mean TLOD (Log 10 likelihood ratio score of variant existing versus not existing) across individuals with a variant for the site",
                "Number": "1",
                "Type": "Float",
            },
            "pon_mt_trna_prediction": {
                "Description": "tRNA pathogenicity classification from PON-mt-tRNA",
                "Number": "1",
                "Type": "String",
            },
            "pon_ml_probability_of_pathogenicity": {
                "Description": "tRNA ML_probability_of_pathogenicity from PON-mt-tRNA",
                "Number": "1",
                "Type": "Float",
            },
            "mitotip_score": {
                "Description": "MitoTip raw score",
                "Number": "1",
                "Type": "Float",
            },
            "mitotip_trna_prediction": {
                "Description": "MitoTip score interpretation",
                "Number": "1",
                "Type": "String",
            },
            "vep": {
                "Description": "Consequence annotations from Ensembl VEP; note that the SINGLE_EXON flag and END_TRUNC filters have been removed from the LOFTEE annotations to avoid misinterpretation in context of the mitochondrial genome. Format: Allele|Consequence|IMPACT|SYMBOL|Gene|Feature_type|Feature|BIOTYPE|EXON|INTRON|HGVSc|HGVSp|cDNA_position|CDS_position|Protein_position|Amino_acids|Codons|ALLELE_NUM|DISTANCE|STRAND|VARIANT_CLASS|MINIMISED|SYMBOL_SOURCE|HGNC_ID|CANONICAL|TSL|APPRIS|CCDS|ENSP|SWISSPROT|TREMBL|UNIPARC|GENE_PHENO|SIFT|PolyPhen|DOMAINS|HGVS_OFFSET|MOTIF_NAME|MOTIF_POS|HIGH_INF_POS|MOTIF_SCORE_CHANGE|LoF|LoF_filter|LoF_flags|LoF_info",
                "Number": ".",
                "Type": "String",
            },
            "filters": {
                "Description": "Site-level filters",
                "Number": ".",
                "Type": "String",
            },
            "base_qual_hist": {
                "Description": f"Histogram of number of individuals failing the base_qual filter (alternate allele median base quality) across heteroplasmy levels, bin edges are: {hl_hist_bin_edges}",
                "Number": "1",
                "Type": "String",
            },
            "heteroplasmy_below_min_het_threshold_hist": {
                "Description": f"Histogram of number of individuals with a heteroplasmy level below {min_het_threshold}, bin edges are: {hl_hist_bin_edges}",
                "Number": "1",
                "Type": "String",
            },
            "position_hist": {
                "Description": f"Histogram of number of individuals failing the position filter (median distance of alternate variants from end of reads) across heteroplasmy levels, bin edges are: {hl_hist_bin_edges}",
                "Number": "1",
                "Type": "String",
            },
            "strand_bias_hist": {
                "Description": f"Histogram of number of individuals failing the strand_bias filter (evidence for alternate allele comes from one read direction only) across heteroplasmy levels, bin edges are: {hl_hist_bin_edges}",
                "Number": "1",
                "Type": "String",
            },
            "weak_evidence_hist": {
                "Description": f"Histogram of number of individuals failing the weak_evidence filter (mutation does not meet likelihood threshold) across heteroplasmy levels, bin edges are: {hl_hist_bin_edges}",
                "Number": "1",
                "Type": "String",
            },
            "contamination_hist": {
                "Description": f"Histogram of number of individuals failing the contamination filter across heteroplasmy levels, bin edges are: {hl_hist_bin_edges}",
                "Number": "1",
                "Type": "String",
            },
            "excluded_AC": {
                "Description": "Excluded allele count (number of individuals in which the variant was filtered out)",
                "Number": "1",
                "Type": "String",
            },
        },
        "format": {
            "GT": {
                "Description": f"Genotype, 1/1 if heteroplasmy level >= {min_hom_threshold}, and 0/1 if heteroplasmy level < {min_hom_threshold}",
                "Number": "1",
                "Type": "String",
            },
            "DP": {
                "Description": "Depth of coverage",
                "Number": "1",
                "Type": "Integer",
            },
            "FT": {
                "Description": "Sample-level filters",
                "Number": ".",
                "Type": "String",
            },
            "FT_LIFT": {
                "Description": "Sample-level liftover tags",
                "Number": ".",
                "Type": "String",
            },
            "HL": {"Description": "Heteroplasmy level", "Number": "1", "Type": "Float"},
            "MQ": {"Description": "Mapping quality", "Number": "1", "Type": "Float"},
            "TLOD": {
                "Description": "Log 10 likelihood ratio score of variant existing versus not existing",
                "Number": "1",
                "Type": "Float",
            },
        },
    }

    return input_mt, meta_dict, vcf_header_file


def modify_ft_liftover(mt):
    these_f = hl.literal(CUSTOMLIFTOVERFILTERS.union(LIFTOVERFILTERS))
    mt = mt.annotate_entries(FT_LIFT = mt.FT.intersection(these_f))
    mt = mt.annotate_entries(FT = mt.FT.difference(these_f))    
    #mt = mt.annotate_entries(FT = mt.FT.difference({'PASS'}))
    #mt = mt.annotate_entries(FT = hl.if_else(hl.len(mt.FT) == 0, {"PASS"}, mt.FT))
    return format_filters(mt)


def make_comma_delim(expr):
    return hl.literal(',').join(hl.map(hl.str, expr))


def process_mt_for_flat_file_analysis(mt, skip_vep, allow_gt_fail):
    """ 
    Function to format the MT into high-yield fields used for downstream analysis.
    
    NOTE: We now filter to sites that have HL missing or > 0 because /most sites/ are hom ref
    with HL = 0. The sites with missing HL are uncertain - these sites have either low coverage 
    or have a variant that was filtered. This way the set of true zeros can be obtained
    efficiently as the set of samples not present in this dataset.
    """
    base_row_set = ['rsid', 'common_low_heteroplasmy', 'filters', 
                    'hap_defining_variant', 'pon_mt_trna_prediction',
                    'pon_ml_probability_of_pathogenicity','mitotip_score',
                    'mitotip_trna_prediction','region', 'variant_context',
                    'dp_mean', 'mq_mean', 'tlod_mean', 'AF_hom', 'AF_het', 
                    'AC_hom', 'AC_het', 'max_hl','excluded_AC']
    if not skip_vep:
        mt = mt.annotate_rows(ancestral = mt.vep.ancestral, 
                              most_severe_csq=mt.vep.most_severe_consequence)
        base_row_set = base_row_set + ['ancestral', 'most_severe_csq']
    base_col_set = ['batch']
    base_col_set = [x for x in base_col_set if x in mt.col]
    addl_col_set = ['contamination', 'freemix_percentage', 
                    'contam_high_het', 'freemix_percentage_imp',
                    'major_haplogroup', 'hap', 'wgs_median_coverage',
                    'mt_mean_coverage', 'mito_cn', 'age', 'pop', 
                    'over_85_mean', 'over_85_count']
    mt = mt.select_globals().select_rows(*base_row_set
                           ).select_cols(*base_col_set, *addl_col_set)
    ht = mt.filter_entries(hl.is_missing(mt.HL) | (mt.HL > 0)).entries()

    # there should not be any empty filters
    if ht.aggregate(hl.agg.count_where(hl.len(ht.FT) == 0)) > 0:
        raise ValueError('There should be no empty FT entries.')
    
    # there may be HL that are missing a call but do not have a reason for failure; these should have dp < 100
    htf = ht.filter(hl.is_missing(ht.HL) & hl.is_missing(ht.FT))
    if htf.aggregate(hl.agg.any(hl.is_defined(htf.DP))):
        raise ValueError('Any instances of missing HL and missing FT should also be missing DP.')

    htf = ht.filter(ht.HL == 0)
    if htf.aggregate(hl.agg.count_where(htf.FT != {'PASS'})) > 0:
        raise ValueError('No entries with HL = 0 should have failed a genotype filter.')
    
    # thus any records with DP > 100 and missing HL should have a reason for failure
    htf = ht.filter(ht.DP > 100)
    if htf.aggregate(hl.agg.count_where(hl.is_missing(htf.FT))) > 0:
        raise ValueError('There should be no missing filters when DP > 100.')
    if htf.aggregate(hl.agg.count_where(hl.is_missing(htf.HL) & htf.FT.contains('GT_PASS'))) > 0:
        raise ValueError('Any missing heteroplasmies at DP > 100 should not be passing in terms of FT.')

    ht = ht.annotate(AD_ref = ht.AD[0], AD_alt = ht.AD[1], FT = ht.FT.union(ht.filters)).drop('AD','filters')
    if 'F1R2' in ht.row and 'F2R1' in ht.row:
        ht = ht.annotate(F2R1_ref = ht.F2R1[0], F2R1_alt = ht.F2R1[1], F1R2_ref = ht.F1R2[0], F1R2_alt = ht.F1R2[1]).drop('F2R1','F1R2')
    if 'RPA' in ht.row:
        ht = ht.annotate(RPA_ref = ht.RPA[0], RPA_alt = ht.RPA[1]).drop('RPA')
    if 'AS_SB_TABLE' in ht.row:
        ht = ht.annotate(AS_SB_SPLIT = ht.AS_SB_TABLE.split('\\|'))
        ht = ht.annotate(FWD_SB_ref = hl.if_else(hl.is_defined(ht.AS_SB_TABLE), hl.int32(ht.AS_SB_SPLIT[0].split(',')[0]), hl.missing(hl.tint32)),
                         FWD_SB_alt = hl.if_else(hl.is_defined(ht.AS_SB_TABLE), hl.int32(ht.AS_SB_SPLIT[1].split(',')[0]), hl.missing(hl.tint32)),
                         REV_SB_ref = hl.if_else(hl.is_defined(ht.AS_SB_TABLE), hl.int32(ht.AS_SB_SPLIT[0].split(',')[1]), hl.missing(hl.tint32)),
                         REV_SB_alt = hl.if_else(hl.is_defined(ht.AS_SB_TABLE), hl.int32(ht.AS_SB_SPLIT[1].split(',')[1]), hl.missing(hl.tint32)))
        ht = ht.drop('AS_SB_TABLE', 'AS_SB_SPLIT')
    ht = ht.annotate(FT = ht.FT.difference({'PASS'}), FT_LIFT = ht.FT_LIFT.difference({'PASS'}))
    
    # fail_gt should be a subset of missing_call
    # case/control should be computed with controls excluding missing_call samples
    ht = ht.annotate(artifact_prone = ht.FT.contains('artifact_prone_site'),
                     lifted = hl.len(ht.FT_LIFT) > 0,
                     fail_gt = ~ht.FT.contains('GT_PASS'),
                     missing_call = hl.is_missing(ht.HL)).key_by()

    # confirm that all fail_gt are missing a call
    if not allow_gt_fail and ht.aggregate(hl.agg.count_where(ht.fail_gt & (~ht.missing_call))) > 0:
        raise ValueError('All samples where "GT_PASS" is not found should be missing a heteroplasmy call.')

    ht = ht.annotate(**{x: make_comma_delim(ht[x]) for x in ['FT','FT_LIFT','OriginalSelfRefAlleles','alleles','rsid']})
    ht = ht.annotate(locus = ht.locus.contig.replace('MT','chrM') + ':' + hl.str(ht.locus.position))
    ht = ht.annotate(variant = ht.locus + ':' + ht.alleles).key_by('locus','alleles','s')
    return ht.key_by('locus', 'alleles', 's'), base_col_set + addl_col_set, base_row_set


def process_flat_ht_slim(ht, mt_full, row_keep, col_keep, skip_vep):
    mt = mt_full.select_globals()

    if not skip_vep:
        mt = mt.annotate_rows(ancestral = mt.vep.ancestral, 
                              most_severe_csq=mt.vep.most_severe_consequence)

    # organize column data
    sample_table = mt.cols()
    sample_table = sample_table.select(*[x for x in col_keep if x not in sample_table.key])

    # organize row data
    row_table = mt.rows()
    row_table = row_table.select(*[x for x in row_keep if x not in row_table.key]).drop('filters')

    # organize entry data
    cols_to_keep = [x for x in ht.row if x not in sample_table.row and x not in row_table.row and x not in ht.key]
    entry_table = ht.select(*cols_to_keep)

    # entry to drop
    entry_drop = ['F2R1_ref', 'F2R1_alt', 'F1R2_ref', 'F1R2_alt',
                  'FWD_SB_ref', 'FWD_SB_alt', 'REV_SB_ref', 'REV_SB_alt',
                  'RPA_ref', 'RPA_alt', 'MPOS', 'STR', 'STRQ']
    entry_table = entry_table.drop(*[x for x in entry_drop if x in entry_table.row])

    return entry_table, sample_table, row_table


def main(args):  # noqa: D103
    mt_path = args.mt_path
    output_dir = args.output_dir
    temp_dir = args.temp_dir
    hl.init(
        log='annotations_logging.log',
        tmp_dir=f"file://{os.path.abspath(temp_dir)}",
        local_tmpdir=f"file://{os.path.abspath(temp_dir)}",
        spark_conf={"spark.local.dir": os.path.abspath(temp_dir)},
    )
    if _HAIL_VERSION_MINOR >= 75:  # only use this if using hail 0.2.75 or greater
        logger.info("Setting hail flag to avoid array index out of bounds error...")
        # Setting this flag isn't generally recommended, but is needed (since at least Hail version 0.2.75) to avoid an array index out of bounds error until changes are made in future versions of Hail
        # TODO: reassess if this flag is still needed for future versions of Hail
        hl._set_flags(no_whole_stage_codegen="1")
    participant_data = args.participant_data
    vep_results = args.vep_results
    min_hom_threshold = args.min_hom_threshold
    vaf_filter_threshold = args.vaf_filter_threshold
    min_het_threshold = args.min_het_threshold
    gnomad_subset = args.subset_to_gnomad_release
    keep_all_samples = args.keep_all_samples
    run_vep = args.run_vep
    max_cn = args.max_cn
    avoid_fasta_workaround = args.avoid_fasta_workaround
    coverage_h5_path = args.coverage_h5_path
    covdb_position_block_size = args.covdb_position_block_size

    logger.info("Cutoff for homoplasmic variants is set to %.2f...", min_hom_threshold)

    if coverage_h5_path:
        if os.path.exists(coverage_h5_path):
            logger.info("Using coverage DB at %s", coverage_h5_path)
        else:
            logger.warning("coverage_h5_path not found: %s", coverage_h5_path)

    # Define mt path, output directory, subset name
    subset_name = "_gnomad" if gnomad_subset else ""

    logger.info("Setting up output paths...")
    annotated_mt_path = generate_output_paths(
        output_dir, "annotated_combined", subset_name, "mt"
    )
    sites_ht_path = generate_output_paths(
        output_dir, "combined_sites_only", subset_name, "ht"
    )
    sites_txt_path = generate_output_paths(
        output_dir, "combined_sites_only", subset_name, "txt"
    )
    sites_vcf_path = generate_output_paths(
        output_dir, "combined_sites_only", subset_name, "vcf.bgz"
    )
    samples_txt_path = generate_output_paths(
        output_dir, "sample_annotations", subset_name, "txt"
    )
    samples_vcf_path = generate_output_paths(
        output_dir, "sample_vcf", subset_name, "vcf.bgz"
    )

    logger.info("Results will be output to the following files:")
    print(
        "\n".join(
            [
                annotated_mt_path,
                sites_ht_path,
                sites_txt_path,
                sites_vcf_path,
                samples_txt_path,
                samples_vcf_path,
            ]
        )
    )

    run_full = True
    if args.debug_allow_intermediate_read:
        logger.info("INTERMEDIATE READ ENABLED... checking for intermediate mt...")
        if hl.hadoop_exists(f'{annotated_mt_path}/_SUCCESS'):
            logger.info("INTERMEDIATE MT FOUND. NOT EXPORTING STATS.")
            mt = hl.read_matrix_table(annotated_mt_path)
            run_full = False
        else:
            logger.info("INTERMEDIATE MT NOT FOUND.")

    if run_full:
        if not args.overwrite and hl.hadoop_exists(f"{output_dir}/prior_to_sample_filt.mt"):
            mt = hl.read_matrix_table(f"{output_dir}/prior_to_sample_filt.mt")
        
        else:
            logger.info("Adding genotype annotation...")
            # NOTE: on import, there are no instances of hl.len(FT) == 0. Missing FT implies no HL measured with confidence.
            mt = add_genotype(mt_path, min_hom_threshold)

            logger.info("Moving Liftover FT fields to a new entry...")
            mt = modify_ft_liftover(mt)

            logger.info("Adding annotations from Terra...")
            mt = add_terra_metadata(mt, participant_data)

            mt = fill_missing_mt_mean_coverage_from_covdb(
                mt,
                coverage_h5_path,
                position_block_size=covdb_position_block_size,
                logger=logger,
            )

            logger.info("Annotating haplogroup-defining variants...")
            mt = add_hap_defining(mt)

            mt = mt.checkpoint(
                f"{temp_dir}/temp_prior_to_trna.mt", overwrite=True
            )

            logger.info("Annotating tRNA predictions...")
            mt = add_trna_predictions(mt, avoid_fasta_workaround)

            # If 'subset-to-gnomad-release' is set, 'age' and 'pop' are added by the add_gnomad_metadata function.
            # If 'subset-to-gnomad-release' is not set, the user should include an 'age' and 'pop' column in the file supplied to `participant-data`.
            if gnomad_subset:
                logger.info("Adding gnomAD metadata sample annotations...")
                mt = add_gnomad_metadata(mt)
            else:
                logger.info("Checking for and adding age and pop annotations...")
                mt = add_age_and_pop(mt, participant_data)

            logger.info("Adding variant context annotations...")
            mt = add_variant_context(mt)

            # If specified, subet to only the gnomAD samples in the current release
            if gnomad_subset:
                logger.warning("Subsetting results to gnomAD release samples...")

                # Subset to release samples and filter out rows that no longer have at least one alt call
                mt = mt.filter_cols(mt.release)  # Filter to cols where release is true
                mt = mt.filter_rows(hl.agg.any(mt.HL > 0))

            mt = mt.checkpoint(
                f"{output_dir}/prior_to_sample_filt.mt", overwrite=args.overwrite
            )


        if not args.overwrite and hl.hadoop_exists(f"{output_dir}/prior_to_vep.mt"):
            mt = hl.read_matrix_table(f"{output_dir}/prior_to_vep.mt")
        
        else:
            logger.info('Removing samples which show overlapping homoplasmies in self-reference construction...')
            mt, n_removed_overlap = filter_by_hom_overlap(
                mt, keep_all_samples, args.sample_stats
            )

            logger.info("Checking for samples with low/high mitochondrial copy number...")
            mt, n_removed_below_cn, n_removed_above_cn = filter_by_copy_number(
                mt, keep_all_samples, max_cn
            )

            logger.info("Checking for contaminated samples...")
            mt, n_contaminated = filter_by_contamination(mt, output_dir, keep_all_samples)

            logger.info("Switch build and checkpoint...")
            # Switch build 37 to build 38
            mt = mt.key_rows_by(
                locus=hl.locus("chrM", mt.locus.position, reference_genome="GRCh38"),
                alleles=mt.alleles,
            )
            # NOTE: at this stage there should still be no instances of hl.len(FT) == 0. Missing FT implies HL not called.
            # NOTE: all missing HL entries have missing FT. These are entries with low DP so cannot be called hom ref.
            mt = mt.checkpoint(f"{output_dir}/prior_to_vep.mt", overwrite=args.overwrite)


        if not args.overwrite and hl.hadoop_exists(f"{output_dir}/prior_to_filter_genotypes.mt"):
            mt = hl.read_matrix_table(f"{output_dir}/prior_to_filter_genotypes.mt")
        
        else:
            if not args.fully_skip_vep:
                logger.info("Adding vep annotations...")
                mt = add_vep(mt, run_vep, vep_results)

            logger.info("Adding dbsnp annotations...")
            mt = add_rsids(mt, args.band_aid_dbsnp_path_fix)

            logger.info("Annotating MT...")
            mt, n_het_below_min_het_threshold = add_filter_annotations(
                mt, vaf_filter_threshold, min_het_threshold, temp_dir
            )
            mt = mt.checkpoint(
                f"{output_dir}/prior_to_filter_genotypes.mt", overwrite=args.overwrite
            )

        # Some checks
        # NOTE: at this stage there should still be no instances of hl.len(FT) == 0. Missing FT implies HL not called.
        if mt.aggregate_entries(hl.agg.count_where(hl.len(mt.FT) == 0)) > 0:
            raise ValueError('Before filtering genotypes, there should be no entries with FT of length 0.')
        # NOTE: anything with HL == 0 should have no genotype filters
        mtf = mt.filter_entries(mt.HL == 0)
        if mtf.aggregate_entries(hl.agg.count_where(mtf.FT != {'PASS'})) > 0:
            raise ValueError('No entries with HL = 0 should have failed a genotype filter.')
        # NOTE: all missing HL entries have missing FT. These are entries with low DP so cannot be called hom ref.
        mtf = mt.filter_entries(hl.is_missing(mt.FT))
        if mtf.aggregate_entries(hl.agg.count_where(hl.is_defined(mtf.HL))) > 0:
            raise ValueError('No entries with missing FT should have defined HL.')
        mtf = mt.filter_entries(hl.is_missing(mt.HL))
        if mtf.aggregate_entries(hl.agg.count_where(hl.is_defined(mtf.FT))) > 0:
            raise ValueError('No entries with missing HL should have defined FT.')
        # NOTE: at this stage, there should be no missing filters and no filters of length 0
        if mt.filter_rows(hl.is_missing(mt.filters)).count_rows() > 0:
            raise ValueError('There should be no rows with missing variant-level filters.')
        if mt.filter_rows(hl.len(mt.filters) == 0).count_rows() > 0:
            raise ValueError('There should be no rows with variant-level filters of length 0.')


        # After this, passing GTs have "GT_PASS" rather than PASS in FT. Failing GTs have a reason for failure.
        # FT_LIFT also no longer has "PASS" and can have length 0.
        pass_set_filter_genotypes = {'strand_bias'} if args.allow_strand_bias else {}
        mt = filter_genotypes(mt, pass_set=pass_set_filter_genotypes)
        # Add variant annotations such as AC, AF, and AN
        mt = mt.annotate_rows(**dict(generate_expressions(mt, min_hom_threshold)))
        # Checkpoint to help avoid Hail errors from large queries
        mt = mt.checkpoint(f"{temp_dir}/temp.mt", overwrite=True)
        
        mt = add_quality_histograms(mt)
        mt = mt.checkpoint(f"{temp_dir}/temp2.mt", overwrite=True)
        
        # After this, filters no longer contain "PASS"
        # FT and FT_LIFT do not contain PASS as of filter_genotypes
        # FT instead contains "GT_PASS"; FT_LIFT can be empty
        mt = add_annotations_by_hap_and_pop(mt, temp_dir=temp_dir)
        
        mt = add_descriptions(
            mt, min_hom_threshold, vaf_filter_threshold, min_het_threshold
        )

        mt = mt.checkpoint(
            annotated_mt_path, overwrite=args.overwrite
        )  # Full matrix table for internal use

        logger.info("Generating summary statistics reports...")
        report_stats(
            mt,
            output_dir,
            False,
            n_removed_below_cn,
            n_removed_above_cn,
            n_contaminated,
            n_het_below_min_het_threshold,
            n_removed_overlap,
            max_cn=max_cn
        )
        report_stats(
            mt,
            output_dir,
            True,
            n_removed_below_cn,
            n_removed_above_cn,
            n_contaminated,
            n_het_below_min_het_threshold,
            n_removed_overlap,
            max_cn=max_cn
        )

    logger.info('Writing variants flat file for internal use...')
    ht_for_output, col_set, row_set = process_mt_for_flat_file_analysis(mt, args.fully_skip_vep, args.allow_strand_bias)
    ht_for_output.export(annotated_mt_path.replace('.mt','_processed_flat.tsv.bgz'))

    logger.info('Writing slimmed flat file for internal use...')
    ht_for_output_slim, col_data, row_data = process_flat_ht_slim(ht_for_output, mt_full=mt, row_keep=row_set, col_keep=col_set, skip_vep=args.fully_skip_vep)
    ht_for_output_slim.export(annotated_mt_path.replace('.mt','_processed_flat_slim.tsv.bgz'))
    col_data.export(annotated_mt_path.replace('.mt','_processed_flat_slim_sample_data.tsv.bgz'))
    row_data.export(annotated_mt_path.replace('.mt','_processed_flat_slim_variant_data.tsv.bgz'))

    logger.info('Generate non-missing HL file for internal use...')
    ht_per_var_N = ht_for_output_slim.group_by(ht_for_output_slim.locus, ht_for_output_slim.alleles
                                    ).aggregate(N_failed = hl.agg.count_where(hl.is_missing(ht_for_output_slim.HL)), 
                                                N_variant_pass = hl.agg.count_where(~hl.is_missing(ht_for_output_slim.HL)))
    ht_for_output_slim_def = ht_for_output_slim.filter(hl.is_defined(ht_for_output_slim.HL))
    ht_per_var_N.export(annotated_mt_path.replace('.mt','_processed_flat_slim_per_variant_N_missing.tsv.bgz'))
    ht_for_output_slim_def.export(annotated_mt_path.replace('.mt','_processed_flat_slim_qc_pass_HL_only.tsv.bgz'))

    logger.info("Writing ht...")
    variant_ht = mt.rows()
    variant_ht = variant_ht.drop("region", "variant_context")
    variant_ht = adjust_descriptions(variant_ht)
    variant_ht.export(sites_txt_path)  # Sites-only txt file for external use
    variant_ht.write(
        sites_ht_path, overwrite=args.overwrite
    )  # Sites-only ht for external use

    logger.info("Writing sample annotations...")
    mt = add_sample_annotations(mt, min_hom_threshold)
    sample_ht = mt.cols().persist()
    sample_ht.group_by(sample_ht.hap).aggregate(n=hl.agg.count()).export(
        f"{output_dir}/haplogroup_counts.txt"
    )  # Counts of top level haplogroups
    sample_ht.export(samples_txt_path)  # Sample annotations txt file for internal use

    logger.info("Formatting and writing VCF...")
    rows_ht = mt.rows()
    export_simplified_variants(rows_ht, output_dir)

    if not args.skip_vcf:
        vcf_mt, vcf_meta, vcf_header_file = format_vcf(mt, output_dir, min_hom_threshold, skip_vep=args.fully_skip_vep)
        hl.export_vcf(
            vcf_mt,
            samples_vcf_path,
            metadata=vcf_meta,
            append_to_header=vcf_header_file,
            tabix=True,
        )  # Full VCF for internal use
        vcf_variant_ht = vcf_mt.rows()
        rows_mt = hl.MatrixTable.from_rows_table(vcf_variant_ht).key_cols_by(s="foo")
        hl.export_vcf(
            rows_mt,
            sites_vcf_path,
            metadata=vcf_meta,
            append_to_header=vcf_header_file,
            tabix=True,
        )  # Sites-only VCF for external use
    else:
        logger.info('Skipping VCF output.')

    logger.info("All annotation steps are completed")


if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        description="This script adds variant annotations to the mitochondria VCF/MT"
    )
    parser.add_argument("-m", "--mt-path", help="Path to combined mt", required=True)
    parser.add_argument(
        "-d",
        "--output-dir",
        help="Path to directory to which output should be written",
        required=True,
    )
    parser.add_argument(
        "-a",
        "--participant-data",
        help="Output file that results from Terra data download",
        required=True,
    )
    parser.add_argument(
        "-v",
        "--vep-results",
        help="MatrixTable path to output vep results (either the existing results or where to ouput new vep results if also setting run_vep)",
        required=True,
    )
    parser.add_argument(
        "--slack-token", help="Slack token that allows integration with slack",
    )
    parser.add_argument(
        "--slack-channel", help="Slack channel to post results and notifications to",
    )
    parser.add_argument(
        "--min-het-threshold",
        help="Minimum heteroplasmy level to define a variant as a PASS heteroplasmic variant, genotypes below this threshold will count towards the heteroplasmy_below_min_het_threshold filter and be set to missing",
        type=float,
        default=0.10,
    )
    parser.add_argument(
        "--min-hom-threshold",
        help="Minimum heteroplasmy level to define a variant as homoplasmic",
        type=float,
        default=0.95,
    )
    parser.add_argument(
        "--vaf-filter-threshold",
        help="Should match vaf_filter_threshold supplied to Mutect2, variants below this value will be set to homoplasmic reference after calculating the common_low_heteroplasmy filter",
        type=float,
        default=0.01,
    )
    parser.add_argument(
        "--subset-to-gnomad-release",
        help="Set to True to only include released gnomAD samples",
        action="store_true",
    )
    parser.add_argument(
        "--keep-all-samples",
        help="Set to True to keep all samples (will skip steps that filter samples because of contamination, overlapping homoplasmies, and/or mitochondrial copy number)",
        action="store_true",
    )
    parser.add_argument(
        "--run-vep", help="Set to True to run/rerun vep", action="store_true"
    )
    parser.add_argument('--fully-skip-vep', action='store_true', help='If true, will skip VEP entirely.')
    parser.add_argument(
        "--overwrite", help="Overwrites existing files", action="store_true"
    )
    parser.add_argument(
        '--band-aid-dbsnp-path-fix', action='store_true', help='If enabled, uses a regex replace to fix the path to the dbSNP database.'
    )
    parser.add_argument(
        '--debug-allow-intermediate-read', action='store_true', help='If true, will read intermediates from disk. If found, assumes stats have been exported.'
    )
    parser.add_argument('--avoid-fasta-workaround',
                        action='store_true',
                        help='If True, will NOT use a workaround for loading in the GRCh37 FASTA sequence. The workaround involves avoiding hail for sequence data and instead using the all sites variant context file.' + \
                             'We have verified that this file has the correct reference sequence. This workaround does not change the tRNA annotations.')
    parser.add_argument(
        '--sample-stats', required=True, help='Path to sample statistics file. Used to remove samples that show overlapping mtDNA homoplasmies.'
    )
    parser.add_argument(
        '--coverage-h5-path',
        required=False,
        help='Optional path to coverage.h5 from the v2 coverage DB for filling missing mt_mean_coverage.'
    )
    parser.add_argument(
        '--covdb-position-block-size',
        type=int,
        default=1024,
        help='Position block size to use when reading coverage.h5 for missing mt_mean_coverage values.'
    )
    parser.add_argument(
        '--max-cn', default=500, type=int, help='Allows specification of the maximum copy number filter. Only operates if keep-samples if false.'
    )
    parser.add_argument(
        '--allow-strand-bias', action='store_true', help='In some cases, one may want to allow strand bias calls to persist in the final callset.'
    )
    parser.add_argument(
        '--skip-vcf', action='store_true', help='Will skip VCF output. Recommend enabling when working with N > 200k.'
    )
    parser.add_argument(
        '--temp-dir', type=str, help='Required temporary directory.', required=True
    )

    args = parser.parse_args()

    # Both a slack token and slack channel must be supplied to receive notifications on slack
    if args.slack_channel and args.slack_token:
        with slack_notifications(args.slack_token, args.slack_channel):
            main(args)
    else:
        main(args)