<commit_before>package org.opencb.opencga.storage.core.variant.query;

import org.apache.commons.collections.CollectionUtils;
import org.apache.commons.lang3.EnumUtils;
import org.apache.commons.lang3.StringUtils;
import org.apache.commons.lang3.tuple.Pair;
import org.opencb.biodata.models.feature.Genotype;
import org.opencb.biodata.models.variant.Variant;
import org.opencb.biodata.models.variant.annotation.ConsequenceTypeMappings;
import org.opencb.biodata.models.variant.avro.ClinicalSignificance;
import org.opencb.biodata.models.variant.avro.VariantType;
import org.opencb.biodata.models.variant.metadata.VariantFileHeaderComplexLine;
import org.opencb.cellbase.core.variant.annotation.VariantAnnotationUtils;
import org.opencb.commons.datastore.core.Query;
import org.opencb.commons.datastore.core.QueryOptions;
import org.opencb.commons.datastore.core.QueryParam;
import org.opencb.opencga.storage.core.metadata.VariantStorageMetadataManager;
import org.opencb.opencga.storage.core.metadata.models.SampleMetadata;
import org.opencb.opencga.storage.core.metadata.models.StudyMetadata;
import org.opencb.opencga.storage.core.metadata.models.TaskMetadata;
import org.opencb.opencga.storage.core.utils.CellBaseUtils;
import org.opencb.opencga.storage.core.variant.adaptors.*;

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

import static org.opencb.opencga.storage.core.variant.VariantStorageEngine.Options.EXCLUDE_GENOTYPES;
import static org.opencb.opencga.storage.core.variant.VariantStorageEngine.Options.LOADED_GENOTYPES;
import static org.opencb.opencga.storage.core.variant.adaptors.VariantQueryParam.*;
import static org.opencb.opencga.storage.core.variant.adaptors.VariantQueryUtils.*;

/**
 * Created on 01/04/19.
 *
 * @author <NAME> &lt;<EMAIL>&gt;
 */
public class VariantQueryParser {

    protected final CellBaseUtils cellBaseUtils;
    protected final VariantStorageMetadataManager metadataManager;

    public VariantQueryParser(CellBaseUtils cellBaseUtils, VariantStorageMetadataManager metadataManager) {
        this.cellBaseUtils = cellBaseUtils;
        this.metadataManager = metadataManager;
    }

    //    public VariantQuery parseQuery(Query query, QueryOptions options) throws StorageEngineException {};

    public Query preProcessQuery(Query originalQuery, QueryOptions options) {
        // Copy input query! Do not modify original query!
        Query query = originalQuery == null ? new Query() : new Query(originalQuery);

        preProcessAnnotationParams(query);

        preProcessStudyParams(options, query);

        return query;
    }

    protected void preProcessAnnotationParams(Query query) {
        convertGoToGeneQuery(query, cellBaseUtils);
        convertExpressionToGeneQuery(query, cellBaseUtils);

        VariantQueryXref xrefs = parseXrefs(query);
        List<String> allIds = new ArrayList<>(xrefs.getIds().size() + xrefs.getVariants().size());
        allIds.addAll(xrefs.getIds());
        for (Variant variant : xrefs.getVariants()) {
            allIds.add(variant.toString());
        }
        query.put(ID.key(), allIds);
        query.put(GENE.key(), xrefs.getGenes());
        query.put(ANNOT_XREF.key(), xrefs.getOtherXrefs());
        query.remove(ANNOT_CLINVAR.key());
        query.remove(ANNOT_COSMIC.key());

        if (VariantQueryUtils.isValidParam(query, TYPE)) {
            Set<String> types = new HashSet<>();
            if (query.getString(TYPE.key()).contains(NOT)) {
                // Invert negations
                for (VariantType value : VariantType.values()) {
                    types.add(value.name());
                }
                for (String type : query.getAsStringList(TYPE.key())) {
                    if (isNegated(type)) {
                        type = removeNegation(type);
                    } else {
                        throw VariantQueryException.malformedParam(TYPE, "Can not mix negated and no negated values");
                    }
                    // Expand types to subtypes
                    type = type.toUpperCase();
                    Set<VariantType> subTypes = Variant.subTypes(VariantType.valueOf(type));
                    types.remove(type);
                    subTypes.forEach(subType -> types.remove(subType.toString()));
                }
            } else {
                // Expand types to subtypes
                for (String type : query.getAsStringList(TYPE.key())) {
                    type = type.toUpperCase();
                    Set<VariantType> subTypes = Variant.subTypes(VariantType.valueOf(type));
                    types.add(type);
                    subTypes.forEach(subType -> types.add(subType.toString()));
                }
            }
            query.put(TYPE.key(), new ArrayList<>(types));
        }

        if (VariantQueryUtils.isValidParam(query, ANNOT_CLINICAL_SIGNIFICANCE)) {
            String v = query.getString(ANNOT_CLINICAL_SIGNIFICANCE.key());
            QueryOperation operator = VariantQueryUtils.checkOperator(v);
            List<String> values = VariantQueryUtils.splitValue(v, operator);
            List<String> clinicalSignificanceList = new ArrayList<>(values.size());
            for (String clinicalSignificance : values) {
                ClinicalSignificance enumValue = EnumUtils.getEnum(ClinicalSignificance.class, clinicalSignificance);
                if (enumValue == null) {
                    String key = clinicalSignificance.toLowerCase().replace(' ', '_');
                    enumValue = EnumUtils.getEnum(ClinicalSignificance.class, key);
                }
                if (enumValue == null) {
                    String key = clinicalSignificance.toLowerCase();
                    if (VariantAnnotationUtils.CLINVAR_CLINSIG_TO_ACMG.containsKey(key)) {
                        // No value set
                        enumValue = VariantAnnotationUtils.CLINVAR_CLINSIG_TO_ACMG.get(key);
                    }
                }
                if (enumValue != null) {
                    clinicalSignificance = enumValue.toString();
                } // else should throw exception?

                clinicalSignificanceList.add(clinicalSignificance);
            }
            query.put(ANNOT_CLINICAL_SIGNIFICANCE.key(), clinicalSignificanceList);
        }

        if (isValidParam(query, ANNOT_SIFT)) {
            String sift = query.getString(ANNOT_SIFT.key());
            String[] split = splitOperator(sift);
            if (StringUtils.isNotEmpty(split[0])) {
                throw VariantQueryException.malformedParam(ANNOT_SIFT, sift);
            }
            if (isValidParam(query, ANNOT_PROTEIN_SUBSTITUTION)) {
                String proteinSubstitution = query.getString(ANNOT_PROTEIN_SUBSTITUTION.key());
                if (proteinSubstitution.contains("sift")) {
                    throw VariantQueryException.malformedParam(ANNOT_SIFT,
                            "Conflict with parameter \"" + ANNOT_PROTEIN_SUBSTITUTION.key() + "\"");
                }
                query.put(ANNOT_PROTEIN_SUBSTITUTION.key(), proteinSubstitution + AND + "sift" + split[1] + split[2]);
            } else {
                query.put(ANNOT_PROTEIN_SUBSTITUTION.key(), "sift" + split[1] + split[2]);
            }
            query.remove(ANNOT_SIFT.key());
        }

        if (isValidParam(query, ANNOT_POLYPHEN)) {
            String polyphen = query.getString(ANNOT_POLYPHEN.key());
            String[] split = splitOperator(polyphen);
            if (StringUtils.isNotEmpty(split[0])) {
                throw VariantQueryException.malformedParam(ANNOT_POLYPHEN, polyphen);
            }
            if (isValidParam(query, ANNOT_PROTEIN_SUBSTITUTION)) {
                String proteinSubstitution = query.getString(ANNOT_PROTEIN_SUBSTITUTION.key());
                if (proteinSubstitution.contains("sift")) {
                    throw VariantQueryException.malformedParam(ANNOT_SIFT,
                            "Conflict with parameter \"" + ANNOT_PROTEIN_SUBSTITUTION.key() + "\"");
                }
                query.put(ANNOT_PROTEIN_SUBSTITUTION.key(), proteinSubstitution + AND + "polyphen" + split[1] + split[2]);
            } else {
                query.put(ANNOT_PROTEIN_SUBSTITUTION.key(), "polyphen" + split[1] + split[2]);
            }
            query.remove(ANNOT_POLYPHEN.key());
        }

        if (isValidParam(query, ANNOT_CONSEQUENCE_TYPE)) {
            List<String> cts = query.getAsStringList(ANNOT_CONSEQUENCE_TYPE.key());
            List<String> parsedCts = new ArrayList<>(cts.size());
            for (String ct : cts) {
                if (ct.equalsIgnoreCase("lof")) {
                    parsedCts.addAll(VariantQueryUtils.LOF_SET);
                } else {
                    parsedCts.add(ConsequenceTypeMappings.accessionToTerm.get(VariantQueryUtils.parseConsequenceType(ct)));
                }
            }
            query.put(ANNOT_CONSEQUENCE_TYPE.key(), parsedCts);
        }
    }

    protected void preProcessStudyParams(QueryOptions options, Query query) {
        StudyMetadata defaultStudy = getDefaultStudy(query, options, metadataManager);
        QueryOperation formatOperator = null;
        if (isValidParam(query, FORMAT)) {
            extractGenotypeFromFormatFilter(query);

            Pair<QueryOperation, Map<String, String>> pair = parseFormat(query);
            formatOperator = pair.getKey();

            for (Map.Entry<String, String> entry : pair.getValue().entrySet()) {
                String sampleName = entry.getKey();
                if (defaultStudy == null) {
                    throw VariantQueryException.missingStudyForSample(sampleName, metadataManager.getStudyNames());
                }
                Integer sampleId = metadataManager.getSampleId(defaultStudy.getId(), sampleName, true);
                if (sampleId == null) {
                    throw VariantQueryException.sampleNotFound(sampleName, defaultStudy.getName());
                }
                List<String> formats = splitValue(entry.getValue()).getValue();
                for (String format : formats) {<commit_msg>storage: Ensure preserve operator in ConsequenceType filter. #1262<commit_after>package org.opencb.opencga.storage.core.variant.query;

import org.apache.commons.collections.CollectionUtils;
import org.apache.commons.lang3.EnumUtils;
import org.apache.commons.lang3.StringUtils;
import org.apache.commons.lang3.tuple.Pair;
import org.opencb.biodata.models.feature.Genotype;
import org.opencb.biodata.models.variant.Variant;
import org.opencb.biodata.models.variant.annotation.ConsequenceTypeMappings;
import org.opencb.biodata.models.variant.avro.ClinicalSignificance;
import org.opencb.biodata.models.variant.avro.VariantType;
import org.opencb.biodata.models.variant.metadata.VariantFileHeaderComplexLine;
import org.opencb.cellbase.core.variant.annotation.VariantAnnotationUtils;
import org.opencb.commons.datastore.core.Query;
import org.opencb.commons.datastore.core.QueryOptions;
import org.opencb.commons.datastore.core.QueryParam;
import org.opencb.opencga.storage.core.metadata.VariantStorageMetadataManager;
import org.opencb.opencga.storage.core.metadata.models.SampleMetadata;
import org.opencb.opencga.storage.core.metadata.models.StudyMetadata;
import org.opencb.opencga.storage.core.metadata.models.TaskMetadata;
import org.opencb.opencga.storage.core.utils.CellBaseUtils;
import org.opencb.opencga.storage.core.variant.adaptors.*;

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

import static org.opencb.opencga.storage.core.variant.VariantStorageEngine.Options.EXCLUDE_GENOTYPES;
import static org.opencb.opencga.storage.core.variant.VariantStorageEngine.Options.LOADED_GENOTYPES;
import static org.opencb.opencga.storage.core.variant.adaptors.VariantQueryParam.*;
import static org.opencb.opencga.storage.core.variant.adaptors.VariantQueryUtils.*;

/**
 * Created on 01/04/19.
 *
 * @author <NAME> &lt;<EMAIL>&gt;
 */
public class VariantQueryParser {

    protected final CellBaseUtils cellBaseUtils;
    protected final VariantStorageMetadataManager metadataManager;

    public VariantQueryParser(CellBaseUtils cellBaseUtils, VariantStorageMetadataManager metadataManager) {
        this.cellBaseUtils = cellBaseUtils;
        this.metadataManager = metadataManager;
    }

    //    public VariantQuery parseQuery(Query query, QueryOptions options) throws StorageEngineException {};

    public Query preProcessQuery(Query originalQuery, QueryOptions options) {
        // Copy input query! Do not modify original query!
        Query query = originalQuery == null ? new Query() : new Query(originalQuery);

        preProcessAnnotationParams(query);

        preProcessStudyParams(options, query);

        return query;
    }

    protected void preProcessAnnotationParams(Query query) {
        convertGoToGeneQuery(query, cellBaseUtils);
        convertExpressionToGeneQuery(query, cellBaseUtils);

        VariantQueryXref xrefs = parseXrefs(query);
        List<String> allIds = new ArrayList<>(xrefs.getIds().size() + xrefs.getVariants().size());
        allIds.addAll(xrefs.getIds());
        for (Variant variant : xrefs.getVariants()) {
            allIds.add(variant.toString());
        }
        query.put(ID.key(), allIds);
        query.put(GENE.key(), xrefs.getGenes());
        query.put(ANNOT_XREF.key(), xrefs.getOtherXrefs());
        query.remove(ANNOT_CLINVAR.key());
        query.remove(ANNOT_COSMIC.key());

        if (VariantQueryUtils.isValidParam(query, TYPE)) {
            Set<String> types = new HashSet<>();
            if (query.getString(TYPE.key()).contains(NOT)) {
                // Invert negations
                for (VariantType value : VariantType.values()) {
                    types.add(value.name());
                }
                for (String type : query.getAsStringList(TYPE.key())) {
                    if (isNegated(type)) {
                        type = removeNegation(type);
                    } else {
                        throw VariantQueryException.malformedParam(TYPE, "Can not mix negated and no negated values");
                    }
                    // Expand types to subtypes
                    type = type.toUpperCase();
                    Set<VariantType> subTypes = Variant.subTypes(VariantType.valueOf(type));
                    types.remove(type);
                    subTypes.forEach(subType -> types.remove(subType.toString()));
                }
            } else {
                // Expand types to subtypes
                for (String type : query.getAsStringList(TYPE.key())) {
                    type = type.toUpperCase();
                    Set<VariantType> subTypes = Variant.subTypes(VariantType.valueOf(type));
                    types.add(type);
                    subTypes.forEach(subType -> types.add(subType.toString()));
                }
            }
            query.put(TYPE.key(), new ArrayList<>(types));
        }

        if (VariantQueryUtils.isValidParam(query, ANNOT_CLINICAL_SIGNIFICANCE)) {
            String v = query.getString(ANNOT_CLINICAL_SIGNIFICANCE.key());
            QueryOperation operator = VariantQueryUtils.checkOperator(v);
            List<String> values = VariantQueryUtils.splitValue(v, operator);
            List<String> clinicalSignificanceList = new ArrayList<>(values.size());
            for (String clinicalSignificance : values) {
                ClinicalSignificance enumValue = EnumUtils.getEnum(ClinicalSignificance.class, clinicalSignificance);
                if (enumValue == null) {
                    String key = clinicalSignificance.toLowerCase().replace(' ', '_');
                    enumValue = EnumUtils.getEnum(ClinicalSignificance.class, key);
                }
                if (enumValue == null) {
                    String key = clinicalSignificance.toLowerCase();
                    if (VariantAnnotationUtils.CLINVAR_CLINSIG_TO_ACMG.containsKey(key)) {
                        // No value set
                        enumValue = VariantAnnotationUtils.CLINVAR_CLINSIG_TO_ACMG.get(key);
                    }
                }
                if (enumValue != null) {
                    clinicalSignificance = enumValue.toString();
                } // else should throw exception?

                clinicalSignificanceList.add(clinicalSignificance);
            }
            query.put(ANNOT_CLINICAL_SIGNIFICANCE.key(), clinicalSignificanceList);
        }

        if (isValidParam(query, ANNOT_SIFT)) {
            String sift = query.getString(ANNOT_SIFT.key());
            String[] split = splitOperator(sift);
            if (StringUtils.isNotEmpty(split[0])) {
                throw VariantQueryException.malformedParam(ANNOT_SIFT, sift);
            }
            if (isValidParam(query, ANNOT_PROTEIN_SUBSTITUTION)) {
                String proteinSubstitution = query.getString(ANNOT_PROTEIN_SUBSTITUTION.key());
                if (proteinSubstitution.contains("sift")) {
                    throw VariantQueryException.malformedParam(ANNOT_SIFT,
                            "Conflict with parameter \"" + ANNOT_PROTEIN_SUBSTITUTION.key() + "\"");
                }
                query.put(ANNOT_PROTEIN_SUBSTITUTION.key(), proteinSubstitution + AND + "sift" + split[1] + split[2]);
            } else {
                query.put(ANNOT_PROTEIN_SUBSTITUTION.key(), "sift" + split[1] + split[2]);
            }
            query.remove(ANNOT_SIFT.key());
        }

        if (isValidParam(query, ANNOT_POLYPHEN)) {
            String polyphen = query.getString(ANNOT_POLYPHEN.key());
            String[] split = splitOperator(polyphen);
            if (StringUtils.isNotEmpty(split[0])) {
                throw VariantQueryException.malformedParam(ANNOT_POLYPHEN, polyphen);
            }
            if (isValidParam(query, ANNOT_PROTEIN_SUBSTITUTION)) {
                String proteinSubstitution = query.getString(ANNOT_PROTEIN_SUBSTITUTION.key());
                if (proteinSubstitution.contains("sift")) {
                    throw VariantQueryException.malformedParam(ANNOT_SIFT,
                            "Conflict with parameter \"" + ANNOT_PROTEIN_SUBSTITUTION.key() + "\"");
                }
                query.put(ANNOT_PROTEIN_SUBSTITUTION.key(), proteinSubstitution + AND + "polyphen" + split[1] + split[2]);
            } else {
                query.put(ANNOT_PROTEIN_SUBSTITUTION.key(), "polyphen" + split[1] + split[2]);
            }
            query.remove(ANNOT_POLYPHEN.key());
        }

        if (isValidParam(query, ANNOT_CONSEQUENCE_TYPE)) {
            Pair<QueryOperation, List<String>> pair = VariantQueryUtils.splitValue(query.getString(ANNOT_CONSEQUENCE_TYPE.key()));
            QueryOperation op = pair.getLeft();
            List<String> cts = pair.getRight();
            List<String> parsedCts = new ArrayList<>(cts.size());
            for (String ct : cts) {
                if (ct.equalsIgnoreCase("lof")) {
                    parsedCts.addAll(VariantQueryUtils.LOF_SET);
                } else {
                    parsedCts.add(ConsequenceTypeMappings.accessionToTerm.get(VariantQueryUtils.parseConsequenceType(ct)));
                }
            }
            query.put(ANNOT_CONSEQUENCE_TYPE.key(), op == null ? parsedCts : String.join(op.separator(), cts));
        }
    }

    protected void preProcessStudyParams(QueryOptions options, Query query) {
        StudyMetadata defaultStudy = getDefaultStudy(query, options, metadataManager);
        QueryOperation formatOperator = null;
        if (isValidParam(query, FORMAT)) {
            extractGenotypeFromFormatFilter(query);

            Pair<QueryOperation, Map<String, String>> pair = parseFormat(query);
            formatOperator = pair.getKey();

            for (Map.Entry<String, String> entry : pair.getValue().entrySet()) {
                String sampleName = entry.getKey();
                if (defaultStudy == null) {
                    throw VariantQueryException.missingStudyForSample(sampleName, metadataManager.getStudyNames());
                }
                Integer sampleId = metadataManager.getSampleId(defaultStudy.getId(), sampleName, true);
                if (sampleId == null) {
                    throw VariantQueryException.sampleNotFound(sampleName, defaultStudy.getName());
                }
                List<String> formats = splitValue(entry.getValue()).getValue();
                for (String format : formats) {
{'fraction_non_alphanumeric': 0.1063912704598597, 'fraction_numerical': 0.0020459859703819173, 'mean_word_length': 2.437541862022773, 'pattern_counts': {'":': 0, '<': 43, '<?xml version=': 0, '>': 47, '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': '5147941', 'n_tokens_mistral': 5805, 'n_tokens_neox': 5491, 'n_words': 884}