// Copyright 2018 Yahoo Holdings. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
package com.yahoo.searchdefinition.processing;

import com.yahoo.config.application.api.DeployLogger;
import com.yahoo.search.query.profile.QueryProfileRegistry;
import com.yahoo.searchdefinition.MapEvaluationTypeContext;
import com.yahoo.searchdefinition.RankProfile;
import com.yahoo.searchdefinition.RankProfileRegistry;
import com.yahoo.searchdefinition.Search;
import com.yahoo.searchlib.rankingexpression.ExpressionFunction;
import com.yahoo.searchlib.rankingexpression.RankingExpression;
import com.yahoo.searchlib.rankingexpression.Reference;
import com.yahoo.searchlib.rankingexpression.rule.ExpressionNode;
import com.yahoo.tensor.TensorType;
import com.yahoo.tensor.evaluation.TypeContext;
import com.yahoo.vespa.model.container.search.QueryProfiles;

import java.util.Map;
import java.util.logging.Level;

/**
 * Resolves and assigns types to all functions in a ranking expression, and
 * validates the types of all ranking expressions under a search instance:
 * Some operators constrain the types of inputs, and first-and second-phase expressions
 * must return scalar values.
 *
 * In addition, the existence of all referred attribute, query and constant
 * features is ensured.
 *
 * @author bratseth
 */
public class RankingExpressionTypeResolver extends Processor {

    private final QueryProfileRegistry queryProfiles;

    public RankingExpressionTypeResolver(Search search,
                                         DeployLogger deployLogger,
                                         RankProfileRegistry rankProfileRegistry,
                                         QueryProfiles queryProfiles) {
        super(search, deployLogger, rankProfileRegistry, queryProfiles);
        this.queryProfiles = queryProfiles.getRegistry();
    }

    @Override
    public void process(boolean validate, boolean documentsOnly) {
        if (documentsOnly) return;

        for (RankProfile profile : rankProfileRegistry.rankProfilesOf(search)) {
            try {
                resolveTypesIn(profile, validate);
            }
            catch (IllegalArgumentException e) {
                throw new IllegalArgumentException("In " + search + ", " + profile, e);
            }
        }
    }

    /**
     * Resolves the types of all functions in the given profile
     *
     * @throws IllegalArgumentException if validate is true and the given rank profile does not produce valid types
     */
    private void resolveTypesIn(RankProfile profile, boolean validate) {
        MapEvaluationTypeContext context = profile.typeContext(queryProfiles);
        for (Map.Entry<String, RankProfile.RankingExpressionFunction> function : profile.getFunctions().entrySet()) {
            if (hasUntypedArguments(function.getValue().function())) continue;
            TensorType type = resolveType(function.getValue().function().getBody(),
                                          "function '" + function.getKey() + "'",
                                          context);
            function.getValue().setReturnType(type);
        }

        if (validate) {
            profile.getSummaryFeatures().forEach(f -> resolveType(f, "summary feature " + f, context));
            ensureValidDouble(profile.getFirstPhaseRanking(), "first-phase expression", context);
            ensureValidDouble(profile.getSecondPhaseRanking(), "second-phase expression", context);
            if ( context.tensorsAreUsed() && ! context.queryFeaturesNotDeclared().isEmpty()) {
                deployLogger.log(Level.WARNING, "The following query features are not declared in query profile " +
                                             "types and will be interpreted as scalars, not tensors: " +
                                             context.queryFeaturesNotDeclared());
            }
        }
    }

    private boolean hasUntypedArguments(ExpressionFunction function) {
        return function.arguments().size() > function.argumentTypes().size();
    }

    private TensorType resolveType(RankingExpression expression, String expressionDescription, TypeContext context) {
        if (expression == null) return null;
        return resolveType(expression.getRoot(), expressionDescription, context);
    }

    private TensorType resolveType(ExpressionNode expression, String expressionDescription, TypeContext context) {
        TensorType type;
        try {
            type = expression.type(context);
        }
        catch (IllegalArgumentException e) {
            throw new IllegalArgumentException("The " + expressionDescription + " is invalid", e);
        }
        if (type == null) // Not expected to happen
            throw new IllegalStateException("Could not determine the type produced by " + expressionDescription);
        return type;
    }

    private void ensureValidDouble(RankingExpression expression, String expressionDescription, TypeContext context) {
        if (expression == null) return;
        TensorType type = resolveType(expression, expressionDescription, context);
        if ( ! type.equals(TensorType.empty))
            throw new IllegalArgumentException("The " + expressionDescription + " must produce a double " +
                                               "(a tensor with no dimensions), but produces " + type);
    }

}

@@ -2256,16 +2256,34 @@
 (%22In %22 +
+ (search != null ?
  search 
@@ -2288,16 +2288,22 @@
 h + %22, %22
+ : %22%22)
  + profi
@@ -2830,62 +2830,134 @@
     
-if (hasUntypedArguments(function.getValue().function()
+ExpressionFunction expressionFunction = function.getValue().function();%0A            if (hasUntypedArguments(expressionFunction
 )) c
@@ -2965,16 +2965,18 @@
 ntinue;%0A
+%0A 
         
@@ -2982,162 +2982,369 @@
     
- TensorType
+// Add any missing inputs for
  type
- =
  resol
-veType(function.getValue().function().getBody(),%0A                                          %22function '%22 + function.getKey() + %22'%22,
+ution%0A            for (String argument : expressionFunction.arguments()) %7B%0A                Reference ref = Reference.fromIdentifier(argument);%0A                if (context.getType(ref).equals(TensorType.empty)) %7B%0A                    context.setType(ref, expressionFunction.argumentTypes().get(argument));%0A                %7D
 %0A   
@@ -3344,32 +3344,34 @@
   %7D%0A            
+%7D%0A
                 
@@ -3366,33 +3366,159 @@
         
-                 
+context.forgetResolvedTypes();%0A%0A            TensorType type = resolveType(expressionFunction.getBody(), %22function '%22 + function.getKey() + %22'%22,
  context

Revert "Revert "Lesters/resolve input types for stateless model evaluation 2""