<commit_before>/*
 * Copyright 2015-2020 OpenCB
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package org.opencb.opencga.analysis.variant.knockout;

import com.google.common.collect.Iterables;
import org.apache.commons.lang3.StringUtils;
import org.apache.commons.lang3.time.StopWatch;
import org.opencb.biodata.models.core.Region;
import org.opencb.biodata.models.variant.Genotype;
import org.opencb.biodata.models.variant.StudyEntry;
import org.opencb.biodata.models.variant.Variant;
import org.opencb.biodata.models.variant.avro.*;
import org.opencb.commons.datastore.core.DataResult;
import org.opencb.commons.datastore.core.Query;
import org.opencb.commons.datastore.core.QueryOptions;
import org.opencb.opencga.analysis.variant.manager.VariantCatalogQueryUtils;
import org.opencb.opencga.analysis.variant.manager.VariantStorageManager;
import org.opencb.opencga.analysis.variant.manager.VariantStorageToolExecutor;
import org.opencb.opencga.core.common.TimeUtils;
import org.opencb.opencga.core.models.analysis.knockout.KnockoutByGene;
import org.opencb.opencga.core.models.analysis.knockout.KnockoutByIndividual;
import org.opencb.opencga.core.models.analysis.knockout.KnockoutByIndividual.KnockoutGene;
import org.opencb.opencga.core.models.analysis.knockout.KnockoutTranscript;
import org.opencb.opencga.core.models.analysis.knockout.KnockoutVariant;
import org.opencb.opencga.core.tools.annotations.ToolExecutor;
import org.opencb.opencga.storage.core.metadata.models.Trio;
import org.opencb.opencga.storage.core.variant.adaptors.GenotypeClass;
import org.opencb.opencga.storage.core.variant.adaptors.VariantField;
import org.opencb.opencga.storage.core.variant.adaptors.VariantQueryParam;
import org.opencb.opencga.storage.core.variant.adaptors.iterators.VariantDBIterator;
import org.opencb.opencga.storage.core.variant.query.VariantQueryUtils;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.IOException;
import java.nio.file.Files;
import java.nio.file.Path;
import java.util.*;
import java.util.function.Predicate;

import static org.opencb.biodata.models.clinical.ClinicalProperty.ModeOfInheritance.COMPOUND_HETEROZYGOUS;
import static org.opencb.opencga.storage.core.variant.annotation.VariantAnnotationConstants.PROTEIN_CODING;
import static org.opencb.opencga.storage.core.variant.query.VariantQueryUtils.ALL;
import static org.opencb.opencga.storage.core.variant.query.VariantQueryUtils.IS;

@ToolExecutor(id = "opencga-local",
        tool = KnockoutAnalysis.ID,
        source = ToolExecutor.Source.STORAGE,
        framework = ToolExecutor.Framework.LOCAL)
public class KnockoutLocalAnalysisExecutor extends KnockoutAnalysisExecutor implements VariantStorageToolExecutor {

    private Logger logger = LoggerFactory.getLogger(this.getClass());
    private VariantStorageManager variantStorageManager;
    private boolean allProteinCoding;
    private List<String> biotype;

    @Override
    protected void run() throws Exception {
        variantStorageManager = getVariantStorageManager();
        biotype = StringUtils.isEmpty(getBiotype()) ? null : Arrays.asList(getBiotype().split(","));
        if (biotype != null && biotype.contains(PROTEIN_CODING)) {
            biotype = new ArrayList<>(biotype);
            biotype.remove(PROTEIN_CODING);
        }
        allProteinCoding = getProteinCodingGenes().size() == 1 && getProteinCodingGenes().iterator().next().equals(ALL);

        String executionMethod = getExecutorParams().getString("executionMethod", "auto");
        boolean bySample;
        switch (executionMethod) {
            case "bySample":
                bySample = true;
                break;
            case "byGene":
                if (getSkipGenesFile()) {
                    throw new IllegalArgumentException("Unable to execute '" + executionMethod + "' "
                            + "when excluding genes file");
                }
                if (allProteinCoding) {
                    throw new IllegalArgumentException("Unable to execute '" + executionMethod + "' "
                            + "when analysing all protein coding genes");
                }
                bySample = false;
                break;
            case "auto":
            case "":
                if (allProteinCoding || getSkipGenesFile()) {
                    bySample = true;
                } else if (getSamples().size() < (getProteinCodingGenes().size() + getOtherGenes().size())) {
                    bySample = true;
                } else {
                    bySample = false;
                }
                break;
            default:
                throw new IllegalArgumentException("Unknown executionMethod '" + executionMethod + "'");
        }
//        if (bySample || (auto && (allProteinCoding || (getProteinCodingGenes().size() + getOtherGenes().size()) >= getSamples().size()))) {
        if (bySample) {
            logger.info("Execute knockout analysis by sample");
            addAttribute("executionMethod", "bySample");
            new KnockoutBySampleExecutor().run();
        } else {
            logger.info("Execute knockout analysis by gene");
            addAttribute("executionMethod", "byGene");
            new KnockoutByGeneExecutor().run();
        }
    }

    private class KnockoutBySampleExecutor {
        public void run() throws Exception {
            Query baseQuery = new Query()
                    .append(VariantQueryParam.STUDY.key(), getStudy())
                    .append(VariantQueryParam.FILTER.key(), getFilter())
                    .append(VariantQueryParam.QUAL.key(), getQual());
            for (String sample : getSamples()) {
                StopWatch stopWatch = StopWatch.createStarted();
                logger.info("Processing sample {}", sample);
                Map<String, KnockoutGene> knockoutGenes = new LinkedHashMap<>();
                Trio trio = getTrios().get(sample);

                // Protein coding genes (if any)
                if (allProteinCoding) {
                    // All protein coding genes
                    Query query = new Query(baseQuery)
                            .append(VariantQueryParam.ANNOT_BIOTYPE.key(), PROTEIN_CODING)
                            .append(VariantQueryParam.INCLUDE_SAMPLE.key(), sample)
                            .append(VariantQueryParam.ANNOT_CONSEQUENCE_TYPE.key(), getCt());
                    knockouts(query, sample, trio, knockoutGenes,
                            getCts()::contains,
                            b -> b.equals(PROTEIN_CODING),
                            g -> true);
                } else if (!getProteinCodingGenes().isEmpty()) {
                    // Set of protein coding genes
                    Query query = new Query(baseQuery)
                            .append(VariantQueryParam.GENE.key(), getProteinCodingGenes())
                            .append(VariantQueryParam.ANNOT_BIOTYPE.key(), PROTEIN_CODING)
                            .append(VariantQueryParam.ANNOT_CONSEQUENCE_TYPE.key(), getCt())
                            .append(VariantQueryParam.INCLUDE_SAMPLE.key(), sample);
                    knockouts(query, sample, trio, knockoutGenes,
                            getCts()::contains,
                            b -> b.equals(PROTEIN_CODING),
                            getProteinCodingGenes()::contains);
                }

                // Other genes (if any)
                if (!getOtherGenes().isEmpty()) {
                    Query query = new Query(baseQuery)
                            .append(VariantQueryParam.ANNOT_BIOTYPE.key(), biotype)
                            .append(VariantQueryParam.GENE.key(), getOtherGenes())
                            .append(VariantQueryParam.INCLUDE_SAMPLE.key(), sample);
                    knockouts(query, sample, trio, knockoutGenes,
                            ct -> true,  // Accept any CT
                            biotype == null ? (b -> !b.equals(PROTEIN_CODING)) : new HashSet<>(biotype)::contains,
                            getOtherGenes()::contains);
                }

                if (knockoutGenes.isEmpty()) {
                    logger.info("No results for sample {}", sample);
                } else {
                    KnockoutByIndividual.GeneKnockoutByIndividualStats stats = getGeneKnockoutBySampleStats(knockoutGenes.values());
                    writeSampleFile(new KnockoutByIndividual()
                            .setSampleId(sample)
                            .setStats(stats)
                            .setGenes(knockoutGenes.values()));
                }
                logger.info("Sample {} processed in {}", sample, TimeUtils.durationToString(stopWatch));
                logger.info("-----------------------------------------------------------");
            }

            if (getSkipGenesFile()) {
                logger.info("Skip genes file transpose");
            } else {
                transposeSampleToGeneOutputFiles();
            }
        }

        private void transposeSampleToGeneOutputFiles() throws IOException {
            int samplesBatchSize = 200;// Maths.max(1, 2000 - getSamples().size());
            int numBatches = (int) Math.ceil((float) getSamples().size() / samplesBatchSize);

            for (int batch = 0; batch < numBatches; batch++) {
                Map<String, KnockoutByGene> byGeneMap = new HashMap<>();
                List<String> samplesBatch = getSamples().subList(batch * samplesBatchSize,<commit_msg>analysis: Add parallelism to KnockoutAnalysis<commit_after>/*
 * Copyright 2015-2020 OpenCB
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package org.opencb.opencga.analysis.variant.knockout;

import com.google.common.collect.Iterables;
import org.apache.commons.lang3.StringUtils;
import org.apache.commons.lang3.time.StopWatch;
import org.opencb.biodata.models.core.Region;
import org.opencb.biodata.models.variant.Genotype;
import org.opencb.biodata.models.variant.StudyEntry;
import org.opencb.biodata.models.variant.Variant;
import org.opencb.biodata.models.variant.avro.*;
import org.opencb.commons.datastore.core.DataResult;
import org.opencb.commons.datastore.core.Query;
import org.opencb.commons.datastore.core.QueryOptions;
import org.opencb.commons.run.ParallelTaskRunner;
import org.opencb.opencga.analysis.variant.manager.VariantCatalogQueryUtils;
import org.opencb.opencga.analysis.variant.manager.VariantStorageManager;
import org.opencb.opencga.analysis.variant.manager.VariantStorageToolExecutor;
import org.opencb.opencga.core.common.TimeUtils;
import org.opencb.opencga.core.models.analysis.knockout.KnockoutByGene;
import org.opencb.opencga.core.models.analysis.knockout.KnockoutByIndividual;
import org.opencb.opencga.core.models.analysis.knockout.KnockoutByIndividual.KnockoutGene;
import org.opencb.opencga.core.models.analysis.knockout.KnockoutTranscript;
import org.opencb.opencga.core.models.analysis.knockout.KnockoutVariant;
import org.opencb.opencga.core.tools.annotations.ToolExecutor;
import org.opencb.opencga.storage.core.io.plain.StringDataReader;
import org.opencb.opencga.storage.core.metadata.models.Trio;
import org.opencb.opencga.storage.core.variant.adaptors.GenotypeClass;
import org.opencb.opencga.storage.core.variant.adaptors.VariantField;
import org.opencb.opencga.storage.core.variant.adaptors.VariantQueryParam;
import org.opencb.opencga.storage.core.variant.adaptors.iterators.VariantDBIterator;
import org.opencb.opencga.storage.core.variant.query.VariantQueryUtils;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.ByteArrayInputStream;
import java.io.IOException;
import java.nio.file.Files;
import java.nio.file.Path;
import java.util.*;
import java.util.function.Predicate;

import static org.opencb.biodata.models.clinical.ClinicalProperty.ModeOfInheritance.COMPOUND_HETEROZYGOUS;
import static org.opencb.opencga.storage.core.variant.annotation.VariantAnnotationConstants.PROTEIN_CODING;
import static org.opencb.opencga.storage.core.variant.query.VariantQueryUtils.ALL;
import static org.opencb.opencga.storage.core.variant.query.VariantQueryUtils.IS;

@ToolExecutor(id = "opencga-local",
        tool = KnockoutAnalysis.ID,
        source = ToolExecutor.Source.STORAGE,
        framework = ToolExecutor.Framework.LOCAL)
public class KnockoutLocalAnalysisExecutor extends KnockoutAnalysisExecutor implements VariantStorageToolExecutor {

    private Logger logger = LoggerFactory.getLogger(this.getClass());
    private VariantStorageManager variantStorageManager;
    private boolean allProteinCoding;
    private List<String> biotype;

    @Override
    protected void run() throws Exception {
        variantStorageManager = getVariantStorageManager();
        biotype = StringUtils.isEmpty(getBiotype()) ? null : Arrays.asList(getBiotype().split(","));
        if (biotype != null && biotype.contains(PROTEIN_CODING)) {
            biotype = new ArrayList<>(biotype);
            biotype.remove(PROTEIN_CODING);
        }
        allProteinCoding = getProteinCodingGenes().size() == 1 && getProteinCodingGenes().iterator().next().equals(ALL);

        String executionMethod = getExecutorParams().getString("executionMethod", "auto");
        boolean bySample;
        switch (executionMethod) {
            case "bySample":
                bySample = true;
                break;
            case "byGene":
                if (getSkipGenesFile()) {
                    throw new IllegalArgumentException("Unable to execute '" + executionMethod + "' "
                            + "when excluding genes file");
                }
                if (allProteinCoding) {
                    throw new IllegalArgumentException("Unable to execute '" + executionMethod + "' "
                            + "when analysing all protein coding genes");
                }
                bySample = false;
                break;
            case "auto":
            case "":
                if (allProteinCoding || getSkipGenesFile()) {
                    bySample = true;
                } else if (getSamples().size() < (getProteinCodingGenes().size() + getOtherGenes().size())) {
                    bySample = true;
                } else {
                    bySample = false;
                }
                break;
            default:
                throw new IllegalArgumentException("Unknown executionMethod '" + executionMethod + "'");
        }
//        if (bySample || (auto && (allProteinCoding || (getProteinCodingGenes().size() + getOtherGenes().size()) >= getSamples().size()))) {
        if (bySample) {
            logger.info("Execute knockout analysis by sample");
            addAttribute("executionMethod", "bySample");
            new KnockoutBySampleExecutor().run();
        } else {
            logger.info("Execute knockout analysis by gene");
            addAttribute("executionMethod", "byGene");
            new KnockoutByGeneExecutor().run();
        }
    }

    private class KnockoutBySampleExecutor {
        public void run() throws Exception {
            Iterator<String> samplesIterator = getSamples().iterator();
            ParallelTaskRunner<String, Void> ptr = new ParallelTaskRunner<>(
                    batchSize -> {
                        List<String> samples = new ArrayList<>(batchSize);
                        for (int i = 0; i < batchSize && samplesIterator.hasNext(); i++) {
                            samples.add(samplesIterator.next());
                        }
                        return samples;
                    },
                    samples -> {
                        for (String sample : samples) {
                            processSample(sample);
                        }
                        return null;
                    }, null,
                    ParallelTaskRunner.Config.builder()
                            .setBatchSize(1)
                            .setNumTasks(Runtime.getRuntime().availableProcessors())
                            .build()
            );

            ptr.run();

            if (getSkipGenesFile()) {
                logger.info("Skip genes file transpose");
            } else {
                transposeSampleToGeneOutputFiles();
            }
        }

        protected void processSample(String sample) throws Exception {
            Query baseQuery = new Query()
                    .append(VariantQueryParam.STUDY.key(), getStudy())
                    .append(VariantQueryParam.FILTER.key(), getFilter())
                    .append(VariantQueryParam.QUAL.key(), getQual());
            StopWatch stopWatch = StopWatch.createStarted();
            logger.info("Processing sample {}", sample);
            Map<String, KnockoutGene> knockoutGenes = new LinkedHashMap<>();
            Trio trio = getTrios().get(sample);

            // Protein coding genes (if any)
            if (allProteinCoding) {
                // All protein coding genes
                Query query = new Query(baseQuery)
                        .append(VariantQueryParam.ANNOT_BIOTYPE.key(), PROTEIN_CODING)
                        .append(VariantQueryParam.INCLUDE_SAMPLE.key(), sample)
                        .append(VariantQueryParam.ANNOT_CONSEQUENCE_TYPE.key(), getCt());
                knockouts(query, sample, trio, knockoutGenes,
                        getCts()::contains,
                        b -> b.equals(PROTEIN_CODING),
                        g -> true);
            } else if (!getProteinCodingGenes().isEmpty()) {
                // Set of protein coding genes
                Query query = new Query(baseQuery)
                        .append(VariantQueryParam.GENE.key(), getProteinCodingGenes())
                        .append(VariantQueryParam.ANNOT_BIOTYPE.key(), PROTEIN_CODING)
                        .append(VariantQueryParam.ANNOT_CONSEQUENCE_TYPE.key(), getCt())
                        .append(VariantQueryParam.INCLUDE_SAMPLE.key(), sample);
                knockouts(query, sample, trio, knockoutGenes,
                        getCts()::contains,
                        b -> b.equals(PROTEIN_CODING),
                        getProteinCodingGenes()::contains);
            }

            // Other genes (if any)
            if (!getOtherGenes().isEmpty()) {
                Query query = new Query(baseQuery)
                        .append(VariantQueryParam.ANNOT_BIOTYPE.key(), biotype)
                        .append(VariantQueryParam.GENE.key(), getOtherGenes())
                        .append(VariantQueryParam.INCLUDE_SAMPLE.key(), sample);
                knockouts(query, sample, trio, knockoutGenes,
                        ct -> true,  // Accept any CT
                        biotype == null ? (b -> !b.equals(PROTEIN_CODING)) : new HashSet<>(biotype)::contains,
                        getOtherGenes()::contains);
            }

            if (knockoutGenes.isEmpty()) {
                logger.info("No results for sample {}", sample);
            } else {
                KnockoutByIndividual.GeneKnockoutByIndividualStats stats = getGeneKnockoutBySampleStats(knockoutGenes.values());
                writeSampleFile(new KnockoutByIndividual()
                        .setSampleId(sample)
                        .setStats(stats)
                        .setGenes(knockoutGenes.values()));
            }
            logger.info("Sample {} processed in {}", sample, TimeUtils.durationToString(stopWatch));
            logger.info("-----------------------------------------------------------");
        }

        private void transposeSampleToGeneOutputFiles() throws IOException {
            int samplesBatchSize = 200;// Maths.max(1, 2000 - getSamples().size());
            int numBatches = (int) Math.ceil((float) getSamples().size() / samplesBatchSize);

            for (int batch = 0; batch < numBatches; batch++) {
                Map<String, KnockoutByGene> byGeneMap = new HashMap<>();
                List<String> samplesBatch = getSamples().subList(batch * samplesBatchSize,
{'fraction_non_alphanumeric': 0.10942393470627314, 'fraction_numerical': 0.0025623991648476797, 'mean_word_length': 2.1994838317898893, 'pattern_counts': {'":': 8, '<': 29, '<?xml version=': 0, '>': 38, 'https://': 0, 'lorem ipsum': 0, 'www.': 2, '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': '2376903', 'n_tokens_mistral': 5777, 'n_tokens_neox': 5385, 'n_words': 1053}