<commit_before>package org.broadinstitute.hellbender.tools.walkers.gnarlyGenotyper;

import com.google.common.annotations.VisibleForTesting;
import com.google.common.primitives.Ints;
import htsjdk.variant.variantcontext.*;
import htsjdk.variant.vcf.*;
import org.apache.commons.lang.StringUtils;
import org.apache.commons.lang3.tuple.Pair;
import org.broadinstitute.hellbender.exceptions.UserException;
import org.broadinstitute.hellbender.tools.walkers.annotator.*;
import org.broadinstitute.hellbender.tools.walkers.annotator.allelespecific.*;
import org.broadinstitute.hellbender.tools.walkers.genotyper.*;
import org.broadinstitute.hellbender.utils.*;
import org.broadinstitute.hellbender.utils.variant.GATKVCFConstants;
import org.broadinstitute.hellbender.utils.variant.GATKVariantContextUtils;
import org.broadinstitute.hellbender.utils.variant.HomoSapiensConstants;
import org.broadinstitute.hellbender.utils.variant.VariantContextGetters;
import org.reflections.Reflections;

import java.util.*;
import java.util.stream.Collectors;
import java.util.stream.IntStream;

/**
 * Guts of the GnarlyGenotyper
 *
 *
 */

public final class GnarlyGenotyperEngine {

    private final static double INDEL_QUAL_THRESHOLD = GenotypeCalculationArgumentCollection.DEFAULT_STANDARD_CONFIDENCE_FOR_CALLING - 10 * Math.log10(HomoSapiensConstants.INDEL_HETEROZYGOSITY);
    private final static double SNP_QUAL_THRESHOLD = GenotypeCalculationArgumentCollection.DEFAULT_STANDARD_CONFIDENCE_FOR_CALLING - 10 * Math.log10(HomoSapiensConstants.SNP_HETEROZYGOSITY);

    private static final int ASSUMED_PLOIDY = GATKVariantContextUtils.DEFAULT_PLOIDY;

    private static final RMSMappingQuality mqCalculator = RMSMappingQuality.getInstance();

    // cache the ploidy 2 PL array sizes for increasing numbers of alts up to the maximum of maxAltAllelesToOutput
    private int[] likelihoodSizeCache;
    private final ArrayList<GenotypeLikelihoodCalculator> glcCache = new ArrayList<>();
    private Set<Class<? extends InfoFieldAnnotation>> allASAnnotations;

    private final int maxAltAllelesToOutput;
    private final boolean summarizePls;  //for very large numbers of samples, save on space and hail import time by summarizing PLs with genotype quality metrics
    private final boolean keepAllSites;
    private final boolean stripASAnnotations;


    public GnarlyGenotyperEngine(final boolean keepAllSites, final int maxAltAllelesToOutput, final boolean summarizePls, final boolean stripASAnnotations) {
        this.maxAltAllelesToOutput = maxAltAllelesToOutput;
        this.summarizePls = summarizePls;
        this.keepAllSites = keepAllSites;
        this.stripASAnnotations = stripASAnnotations;

        if (!summarizePls) {
            final GenotypeLikelihoodCalculators GLCprovider = new GenotypeLikelihoodCalculators();

            //initialize PL size cache -- HTSJDK cache only goes up to 4 alts, but I need 6
            likelihoodSizeCache = new int[maxAltAllelesToOutput + 1 + 1]; //+1 for ref and +1 so index == numAlleles
            glcCache.add(null); //add a null at index zero because zero alleles (incl. ref) makes no sense
            for (final int numAlleles : IntStream.rangeClosed(1, maxAltAllelesToOutput + 1).boxed().collect(Collectors.toList())) {
                likelihoodSizeCache[numAlleles] = GenotypeLikelihoods.numLikelihoods(numAlleles, ASSUMED_PLOIDY);
                //GL calculator cache is indexed by the total number of alleles, including ref
                glcCache.add(numAlleles, GLCprovider.getInstance(ASSUMED_PLOIDY, numAlleles));
            }
        }

        //TODO: fix weird reflection logging?
        final Reflections reflections = new Reflections("org.broadinstitute.hellbender.tools.walkers.annotator.allelespecific");
        allASAnnotations = reflections.getSubTypesOf(InfoFieldAnnotation.class);
        allASAnnotations.addAll(reflections.getSubTypesOf(AS_StrandBiasTest.class));
        allASAnnotations.addAll(reflections.getSubTypesOf(AS_RankSumTest.class));
    }

    @SuppressWarnings({"unchecked", "rawtypes"})
    public VariantContext finalizeGenotype(final VariantContext variant) {
        return finalizeGenotype(variant, null);
    }

    /**
     * Calculates final annotation value for VC and calls genotypes, if necessary
     * @param variant   expected to have <NON_REF>
     * @param annotationDBBuilder   may be null, otherwise modified to retain the combined form of the annotations
     * @return may be null if variant doesn't exceed the quality threshold
     */
    @SuppressWarnings({"unchecked", "rawtypes"})
    public VariantContext finalizeGenotype(final VariantContext variant, final VariantContextBuilder annotationDBBuilder) {
        //GenomicsDB or Evoquer merged all the annotations, but we still need to finalize MQ and QD annotations
        //return a VC with the finalized annotations and dbBuilder gets the raw annotations for the database

        final double QUALapprox = variant.getAttributeAsInt(GATKVCFConstants.RAW_QUAL_APPROX_KEY, 0);
        //Don't apply the indel prior to mixed sites if there's a SNP
        final boolean hasSnpAllele = variant.getAlternateAlleles().stream().anyMatch(allele -> allele.length() == variant.getReference().length());
        final boolean isIndel = !hasSnpAllele;
        final double sitePrior = isIndel ? HomoSapiensConstants.INDEL_HETEROZYGOSITY : HomoSapiensConstants.SNP_HETEROZYGOSITY;
        if((isIndel && QUALapprox < INDEL_QUAL_THRESHOLD) || (!isIndel && QUALapprox < SNP_QUAL_THRESHOLD)) {
            if (keepAllSites) {
                final VariantContextBuilder builder = new VariantContextBuilder(mqCalculator.finalizeRawMQ(variant));
                //we don't need the rest of the annotations since this is filtered and won't go to VQSR
                builder.filter(GATKVCFConstants.LOW_QUAL_FILTER_NAME);
                builder.attribute(GATKVCFConstants.AC_ADJUSTED_KEY, 0);
                return builder.make();
            }
            return null;
        }

        //GenomicsDB merged all the annotations, but we still need to finalize MQ and QD annotations
        //vcfBuilder gets the finalized annotations and annotationDBBuilder (if present) gets the raw annotations for the database
        final VariantContext vcWithMQ = mqCalculator.finalizeRawMQ(variant);
        final VariantContextBuilder vcfBuilder = new VariantContextBuilder(vcWithMQ);

        //Because AS_StrandBias annotations both use and return the raw key
        final Map<String, Object> annotationsToBeModified = new HashMap<>(vcWithMQ.getAttributes());
        for (final Class c : allASAnnotations) {
            try {
                final InfoFieldAnnotation annotation = (InfoFieldAnnotation) c.getDeclaredConstructor().newInstance();
                if (annotation instanceof AS_StandardAnnotation && annotation instanceof ReducibleAnnotation) {
                    final ReducibleAnnotation ann = (ReducibleAnnotation) annotation;
                    if (variant.hasAttribute(ann.getPrimaryRawKey())) {
                        if (!stripASAnnotations) {
                            //here we still have the non-ref
                            final Map<String, Object> finalValue = ann.finalizeRawData(vcfBuilder.make(), variant);
                            finalValue.forEach((key, value) -> annotationsToBeModified.put(key, value));
                            if (annotationDBBuilder != null) {
                                annotationDBBuilder.attribute(ann.getPrimaryRawKey(), variant.getAttribute(ann.getPrimaryRawKey()));
                            }
                        }
                    }
                }
            }
            catch (final Exception e) {
                throw new IllegalStateException("Something went wrong: ", e);
            }
        }
        vcfBuilder.attributes(annotationsToBeModified);

        final int variantDP = variant.getAttributeAsInt(GATKVCFConstants.VARIANT_DEPTH_KEY, 0);
        final double QD = QUALapprox / (double)variantDP;
        vcfBuilder.attribute(GATKVCFConstants.QUAL_BY_DEPTH_KEY, QD).log10PError(QUALapprox/-10.0-Math.log10(sitePrior));<commit_msg>bugfix for gnarly qualapprox calculation, tolerate missing VarDP, support AS_QUALapprox if QUALapprox is missing<commit_after>package org.broadinstitute.hellbender.tools.walkers.gnarlyGenotyper;

import com.google.common.annotations.VisibleForTesting;
import com.google.common.primitives.Ints;
import htsjdk.variant.variantcontext.*;
import htsjdk.variant.vcf.*;
import org.apache.commons.lang.StringUtils;
import org.apache.commons.lang3.tuple.Pair;
import org.broadinstitute.hellbender.exceptions.UserException;
import org.broadinstitute.hellbender.tools.walkers.annotator.*;
import org.broadinstitute.hellbender.tools.walkers.annotator.allelespecific.*;
import org.broadinstitute.hellbender.tools.walkers.genotyper.*;
import org.broadinstitute.hellbender.utils.*;
import org.broadinstitute.hellbender.utils.variant.GATKVCFConstants;
import org.broadinstitute.hellbender.utils.variant.GATKVariantContextUtils;
import org.broadinstitute.hellbender.utils.variant.HomoSapiensConstants;
import org.broadinstitute.hellbender.utils.variant.VariantContextGetters;
import org.reflections.Reflections;

import java.util.*;
import java.util.stream.Collectors;
import java.util.stream.IntStream;

/**
 * Guts of the GnarlyGenotyper
 *
 *
 */

public final class GnarlyGenotyperEngine {

    private final static double INDEL_QUAL_THRESHOLD = GenotypeCalculationArgumentCollection.DEFAULT_STANDARD_CONFIDENCE_FOR_CALLING - 10 * Math.log10(HomoSapiensConstants.INDEL_HETEROZYGOSITY);
    private final static double SNP_QUAL_THRESHOLD = GenotypeCalculationArgumentCollection.DEFAULT_STANDARD_CONFIDENCE_FOR_CALLING - 10 * Math.log10(HomoSapiensConstants.SNP_HETEROZYGOSITY);

    private static final int ASSUMED_PLOIDY = GATKVariantContextUtils.DEFAULT_PLOIDY;

    private static final RMSMappingQuality mqCalculator = RMSMappingQuality.getInstance();

    // cache the ploidy 2 PL array sizes for increasing numbers of alts up to the maximum of maxAltAllelesToOutput
    private int[] likelihoodSizeCache;
    private final ArrayList<GenotypeLikelihoodCalculator> glcCache = new ArrayList<>();
    private Set<Class<? extends InfoFieldAnnotation>> allASAnnotations;

    private final int maxAltAllelesToOutput;
    private final boolean summarizePls;  //for very large numbers of samples, save on space and hail import time by summarizing PLs with genotype quality metrics
    private final boolean keepAllSites;
    private final boolean stripASAnnotations;


    public GnarlyGenotyperEngine(final boolean keepAllSites, final int maxAltAllelesToOutput, final boolean summarizePls, final boolean stripASAnnotations) {
        this.maxAltAllelesToOutput = maxAltAllelesToOutput;
        this.summarizePls = summarizePls;
        this.keepAllSites = keepAllSites;
        this.stripASAnnotations = stripASAnnotations;

        if (!summarizePls) {
            final GenotypeLikelihoodCalculators GLCprovider = new GenotypeLikelihoodCalculators();

            //initialize PL size cache -- HTSJDK cache only goes up to 4 alts, but I need 6
            likelihoodSizeCache = new int[maxAltAllelesToOutput + 1 + 1]; //+1 for ref and +1 so index == numAlleles
            glcCache.add(null); //add a null at index zero because zero alleles (incl. ref) makes no sense
            for (final int numAlleles : IntStream.rangeClosed(1, maxAltAllelesToOutput + 1).boxed().collect(Collectors.toList())) {
                likelihoodSizeCache[numAlleles] = GenotypeLikelihoods.numLikelihoods(numAlleles, ASSUMED_PLOIDY);
                //GL calculator cache is indexed by the total number of alleles, including ref
                glcCache.add(numAlleles, GLCprovider.getInstance(ASSUMED_PLOIDY, numAlleles));
            }
        }

        //TODO: fix weird reflection logging?
        final Reflections reflections = new Reflections("org.broadinstitute.hellbender.tools.walkers.annotator.allelespecific");
        allASAnnotations = reflections.getSubTypesOf(InfoFieldAnnotation.class);
        allASAnnotations.addAll(reflections.getSubTypesOf(AS_StrandBiasTest.class));
        allASAnnotations.addAll(reflections.getSubTypesOf(AS_RankSumTest.class));
    }

    @SuppressWarnings({"unchecked", "rawtypes"})
    public VariantContext finalizeGenotype(final VariantContext variant) {
        return finalizeGenotype(variant, null);
    }

    /**
     * Calculates final annotation value for VC and calls genotypes, if necessary
     * @param variant   expected to have <NON_REF>
     * @param annotationDBBuilder   may be null, otherwise modified to retain the combined form of the annotations
     * @return may be null if variant doesn't exceed the quality threshold
     */
    @SuppressWarnings({"unchecked", "rawtypes"})
    public VariantContext finalizeGenotype(final VariantContext variant, final VariantContextBuilder annotationDBBuilder) {
        //GenomicsDB or Evoquer merged all the annotations, but we still need to finalize MQ and QD annotations
        //return a VC with the finalized annotations and dbBuilder gets the raw annotations for the database

        // Prefer non-AS QUALapprox, but if we only have AS QUALapprox sum the values
        final double QUALapprox;
        if (variant.hasAttribute(GATKVCFConstants.RAW_QUAL_APPROX_KEY)) {
            QUALapprox = variant.getAttributeAsInt(GATKVCFConstants.RAW_QUAL_APPROX_KEY, 0);
        }
        else if (variant.hasAttribute(GATKVCFConstants.AS_RAW_QUAL_APPROX_KEY)) {
            List<Integer> alleleSpecificQualList = AS_QualByDepth.parseQualList(variant);
            QUALapprox = alleleSpecificQualList.stream()
                    .mapToInt(Integer::intValue)
                    .sum();
        }
        else {
            QUALapprox = 0;
        }

        //Don't apply the indel prior to mixed sites if there's a SNP, but don't count a '*' as a SNP
        final boolean hasSnpAllele = variant.getAlternateAlleles().stream().anyMatch(allele -> allele != Allele.SPAN_DEL && allele.length() == variant.getReference().length());
        final boolean isIndel = !hasSnpAllele;
        final double sitePrior = isIndel ? HomoSapiensConstants.INDEL_HETEROZYGOSITY : HomoSapiensConstants.SNP_HETEROZYGOSITY;
        if((isIndel && QUALapprox < INDEL_QUAL_THRESHOLD) || (!isIndel && QUALapprox < SNP_QUAL_THRESHOLD)) {
            if (keepAllSites) {
                final VariantContextBuilder builder = new VariantContextBuilder(mqCalculator.finalizeRawMQ(variant));
                //we don't need the rest of the annotations since this is filtered and won't go to VQSR
                builder.filter(GATKVCFConstants.LOW_QUAL_FILTER_NAME);
                builder.attribute(GATKVCFConstants.AC_ADJUSTED_KEY, 0);
                return builder.make();
            }
            return null;
        }

        //GenomicsDB merged all the annotations, but we still need to finalize MQ and QD annotations
        //vcfBuilder gets the finalized annotations and annotationDBBuilder (if present) gets the raw annotations for the database
        final VariantContext vcWithMQ = mqCalculator.finalizeRawMQ(variant);
        final VariantContextBuilder vcfBuilder = new VariantContextBuilder(vcWithMQ);

        //Because AS_StrandBias annotations both use and return the raw key
        final Map<String, Object> annotationsToBeModified = new HashMap<>(vcWithMQ.getAttributes());
        for (final Class c : allASAnnotations) {
            try {
                final InfoFieldAnnotation annotation = (InfoFieldAnnotation) c.getDeclaredConstructor().newInstance();
                if (annotation instanceof AS_StandardAnnotation && annotation instanceof ReducibleAnnotation) {
                    final ReducibleAnnotation ann = (ReducibleAnnotation) annotation;
                    if (variant.hasAttribute(ann.getPrimaryRawKey())) {
                        if (!stripASAnnotations) {
                            //here we still have the non-ref
                            final Map<String, Object> finalValue = ann.finalizeRawData(vcfBuilder.make(), variant);
                            finalValue.forEach((key, value) -> annotationsToBeModified.put(key, value));
                            if (annotationDBBuilder != null) {
                                annotationDBBuilder.attribute(ann.getPrimaryRawKey(), variant.getAttribute(ann.getPrimaryRawKey()));
                            }
                        }
                    }
                }
            }
            catch (final Exception e) {
                throw new IllegalStateException("Something went wrong: ", e);
            }
        }
        vcfBuilder.attributes(annotationsToBeModified);

        // tolerate lack of VarDP annotation
        if ( variant.hasAttribute(GATKVCFConstants.VARIANT_DEPTH_KEY) ) {
            final int variantDP = variant.getAttributeAsInt(GATKVCFConstants.VARIANT_DEPTH_KEY, 0);
            final double QD = QUALapprox / (double) variantDP;
            vcfBuilder.attribute(GATKVCFConstants.QUAL_BY_DEPTH_KEY, QD).log10PError(QUALapprox / -10.0 - Math.log10(sitePrior));
        }

{'fraction_non_alphanumeric': 0.07526507908917086, 'fraction_numerical': 0.0033026247175386753, 'mean_word_length': 3.295669487307118, 'pattern_counts': {'":': 0, '<': 24, '<?xml version=': 0, '>': 24, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '3408675', 'n_tokens_mistral': 4977, 'n_tokens_neox': 4648, 'n_words': 1158}