/*
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you 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.apache.asterix.optimizer.rules;

import java.util.ArrayList;
import java.util.BitSet;
import java.util.Iterator;
import java.util.LinkedHashMap;
import java.util.List;
import java.util.Map;

import org.apache.asterix.om.functions.BuiltinFunctions;
import org.apache.asterix.om.util.container.IObjectPool;
import org.apache.asterix.om.util.container.ListObjectPool;
import org.apache.asterix.om.util.container.ObjectFactories;
import org.apache.commons.lang3.mutable.Mutable;
import org.apache.commons.lang3.mutable.MutableObject;
import org.apache.hyracks.algebricks.common.exceptions.AlgebricksException;
import org.apache.hyracks.algebricks.core.algebra.base.ILogicalExpression;
import org.apache.hyracks.algebricks.core.algebra.base.ILogicalOperator;
import org.apache.hyracks.algebricks.core.algebra.base.ILogicalPlan;
import org.apache.hyracks.algebricks.core.algebra.base.IOptimizationContext;
import org.apache.hyracks.algebricks.core.algebra.base.LogicalExpressionTag;
import org.apache.hyracks.algebricks.core.algebra.base.LogicalOperatorTag;
import org.apache.hyracks.algebricks.core.algebra.base.LogicalVariable;
import org.apache.hyracks.algebricks.core.algebra.expressions.AbstractFunctionCallExpression;
import org.apache.hyracks.algebricks.core.algebra.expressions.ConstantExpression;
import org.apache.hyracks.algebricks.core.algebra.expressions.VariableReferenceExpression;
import org.apache.hyracks.algebricks.core.algebra.functions.FunctionIdentifier;
import org.apache.hyracks.algebricks.core.algebra.operators.logical.AggregateOperator;
import org.apache.hyracks.algebricks.core.algebra.operators.logical.AssignOperator;
import org.apache.hyracks.algebricks.core.algebra.operators.logical.DistinctOperator;
import org.apache.hyracks.algebricks.core.algebra.operators.logical.GroupByOperator;
import org.apache.hyracks.algebricks.core.algebra.operators.logical.InnerJoinOperator;
import org.apache.hyracks.algebricks.core.algebra.plan.ALogicalPlanImpl;
import org.apache.hyracks.algebricks.core.algebra.util.OperatorManipulationUtil;
import org.apache.hyracks.algebricks.core.rewriter.base.IAlgebraicRewriteRule;

/**
 * This rule rewrites agg-distinct() function calls inside AGGREGATE operator into into regular agg() function calls
 * over DISTINCT operator. <br/>
 * In the simplest case:
 * <pre>
 * AGGREGATE [$x] <- [agg-distinct($y)]
 * INPUT_OP
 * --->
 * AGGREGATE [$x] <- [agg($y)]
 * DISTINCT [$y]
 * INPUT_OP
 * </pre>
 *
 * In the general case a single AGGREGATE operator could contain regular aggregate function calls and agg-distinct()
 * function calls that work on different arguments:
 * <pre>
 * AGGREGATE [...] <- [avg($a), sum($b), avg-distinct($c), sum-distinct($c), avg-distinct($d), sum-distinct($d)]
 * </pre>
 * In this case distinct aggregate function calls are divided into groups based on their arguments, with one additional
 * group introduced for all non-distinct aggregates. Each distinct aggregate is the rewritten into regular aggregate
 * over DISTINCT operator as described above. Then all groups are combined together as follows.
 * If the original AGGREGATE operator was in a nested plan of a GROUPBY operator then each aggregate group becomes
 * a separate nested plan, otherwise (if the AGGREGATE operator was not inside GROUPBY) these groups are joined
 * together by a trivial join:
 *
 * <pre>
 * 1. AGGREGATE inside GROUPBY:
 *
 *   GROUPBY (
 *     {
 *       AGGREGATE [...] <- [avg($a), sum($b), avg-distinct($c), sum-distinct($c), avg-distinct($d), sum-distinct($d)]
 *       NESTED_TUPLE_SOURCE
 *     }
 *   )
 *   -->
 *   GROUPBY (
 *     {
 *       AGGREGATE [...] <- [avg($a), sum($b)]
 *       NESTED_TUPLE_SOURCE
 *     },
 *     {
 *       AGGREGATE [...] <- [avg($c), sum($c)]
 *       DISTINCT [$c]
 *       NESTED_TUPLE_SOURCE
 *     },
 *     {
 *       AGGREGATE [...] <- [avg($d), sum($d)]
 *       DISTINCT [$d]
 *       NESTED_TUPLE_SOURCE
 *     }
 *   )
 *
 *  2. AGGREGATE not inside GROUPBY:
 *
 *    AGGREGATE [...] <- [avg($a), sum($b), avg-distinct($c), sum-distinct($c), avg-distinct($d), sum-distinct($d)]
 *    INPUT_OP
 *    -->
 *    JOIN (TRUE) {
 *      JOIN (TRUE) {
 *        AGGREGATE [...] <- [avg($a), sum($b)]
 *          INPUT_OP
 *        AGGREGATE [...] <- [avg($c), sum($c)]
 *          DISTINCT [$c]
 *           INPUT_OP
 *      },
 *      AGGREGATE [...] <- [avg($d), sum($d)]
 *        DISTINCT [$d]
 *          INPUT_OP
 *    }
 *
 *    Note that in this case the input operator (INPUT_OP) is cloned for each aggregate group.
 * </pre>
 *
 * Notes:
 * <ul>
 * <li>
 * If distinct aggregate function argument is not a variable reference expression then an additional ASSIGN
 * operator will be introduced below the DISTINCT</li>
 * <li>
 * This rule must run before {@link AsterixIntroduceGroupByCombinerRule} because distinct aggregates must be converted
 * into regular before the local/global split.
 * </li>
 * <li>
 * Scalar distinct aggregate function calls are not handled by this rule. If they were not rewritten into regular
 * aggregates then they will be evaluated at runtime.
 * </li>
 * </ul>
 */
public final class RewriteDistinctAggregateRule implements IAlgebraicRewriteRule {

    private final IObjectPool<BitSet, Void> bitSetAllocator = new ListObjectPool<>(ObjectFactories.BIT_SET_FACTORY);

    private final Map<ILogicalExpression, BitSet> aggGroups = new LinkedHashMap<>();

    private final List<AggregateOperator> newAggOps = new ArrayList<>();

    @Override
    public boolean rewritePre(Mutable<ILogicalOperator> opRef, IOptimizationContext context)
            throws AlgebricksException {
        ILogicalOperator op = opRef.getValue();
        switch (op.getOperatorTag()) {
            case AGGREGATE:
                if (context.checkIfInDontApplySet(this, op)) {
                    return false;
                }
                newAggOps.clear();
                if (rewriteAggregate((AggregateOperator) op, newAggOps, context)) {
                    ILogicalOperator newOp = join(newAggOps, context);
                    opRef.setValue(newOp);
                    return true;
                }
                return false;

            case GROUP:
                if (context.checkIfInDontApplySet(this, op)) {
                    return false;
                }
                GroupByOperator gbyOp = (GroupByOperator) op;
                List<ILogicalPlan> nestedPlans = gbyOp.getNestedPlans();
                boolean applied = false;
                for (int i = nestedPlans.size() - 1; i >= 0; i--) {
                    ILogicalPlan nestedPlan = nestedPlans.get(i);
                    for (Mutable<ILogicalOperator> rootOpRef : nestedPlan.getRoots()) {
                        ILogicalOperator rootOp = rootOpRef.getValue();
                        if (rootOp.getOperatorTag() == LogicalOperatorTag.AGGREGATE) {
                            newAggOps.clear();
                            if (rewriteAggregate((AggregateOperator) rootOp, newAggOps, context)) {
                                applied = true;
                                rootOpRef.setValue(newAggOps.get(0));
                                for (int j = 1, ln = newAggOps.size(); j < ln; j++) {
                                    nestedPlans.add(new ALogicalPlanImpl(new MutableObject<>(newAggOps.get(j))));
                                }
                            }
                        }
                    }
                }
                if (applied) {
                    context.computeAndSetTypeEnvironmentForOperator(gbyOp);
                    context.addToDontApplySet(this, gbyOp);
                    return true;
                }
                return false;

            default:
                return false;
        }
    }

    private ILogicalOperator join(List<? extends ILogicalOperator> ops, IOptimizationContext context)
            throws AlgebricksException {
        ILogicalOperator resultOp = ops.get(0);
        for (int i = 1, n = ops.size(); i < n; i++) {
            ILogicalOperator op = ops.get(i);
            InnerJoinOperator joinOp = new InnerJoinOperator(new MutableObject<>(ConstantExpression.TRUE));
            joinOp.setSourceLocation(resultOp.getSourceLocation());
            joinOp.getInputs().add(new MutableObject<>(resultOp));
            joinOp.getInputs().add(new MutableObject<>(op));
            context.computeAndSetTypeEnvironmentForOperator(joinOp);
            resultOp = joinOp;
        }
        return resultOp;
    }

    private boolean rewriteAggregate(AggregateOperator aggOp, List<AggregateOperator> outAggOps,
            IOptimizationContext context) throws AlgebricksException {
        aggGroups.clear();
        bitSetAllocator.reset();

        List<Mutable<ILogicalExpression>> aggExprs = aggOp.getExpressions();
        int aggCount = aggExprs.size();
        int distinctAggCount = 0;

        for (int i = 0; i < aggCount; i++) {
            ILogicalExpression aggExpr = aggExprs.get(i).getValue();
            ILogicalExpression distinctValueExpr = null;
            if (aggExpr.getExpressionTag() == LogicalExpressionTag.FUNCTION_CALL) {
                AbstractFunctionCallExpression callExpr = (AbstractFunctionCallExpression) aggExpr;
                FunctionIdentifier aggFn = callExpr.getFunctionIdentifier();
                if (BuiltinFunctions.getAggregateFunctionForDistinct(aggFn) != null) {
                    distinctValueExpr = callExpr.getArguments().get(0).getValue();
                    distinctAggCount++;
                }
            }

            BitSet argIndexes = aggGroups.get(distinctValueExpr);
            if (argIndexes == null) {
                argIndexes = bitSetAllocator.allocate(null);
                argIndexes.clear();
                aggGroups.put(distinctValueExpr, argIndexes);
            }
            argIndexes.set(i);
        }

        if (distinctAggCount == 0) {
            return false;
        }

        for (Iterator<Map.Entry<ILogicalExpression, BitSet>> i = aggGroups.entrySet().iterator(); i.hasNext();) {
            Map.Entry<ILogicalExpression, BitSet> me = i.next();
            boolean isDistinct = me.getKey() != null;
            BitSet argIndexes = me.getValue();

            AggregateOperator newAggOp;
            if (i.hasNext()) {
                newAggOp = (AggregateOperator) OperatorManipulationUtil.deepCopyWithNewVars(aggOp, context).first;
                newAggOp.getVariables().clear();
                newAggOp.getVariables().addAll(aggOp.getVariables());
            } else {
                newAggOp = aggOp;
            }

            OperatorManipulationUtil.retainAssignVariablesAndExpressions(newAggOp.getVariables(),
                    newAggOp.getExpressions(), argIndexes);

            if (isDistinct) {
                rewriteDistinctAggregate(newAggOp, context);
            }

            context.computeAndSetTypeEnvironmentForOperator(newAggOp);
            context.addToDontApplySet(this, newAggOp);
            outAggOps.add(newAggOp);
        }

        return true;
    }

    private void rewriteDistinctAggregate(AggregateOperator aggOp, IOptimizationContext context)
            throws AlgebricksException {
        ILogicalOperator inputOp = aggOp.getInputs().get(0).getValue();

        // Introduce ASSIGN operator if distinct expression is not a variable
        ILogicalExpression distinctExpr = ((AbstractFunctionCallExpression) aggOp.getExpressions().get(0).getValue())
                .getArguments().get(0).getValue();
        AssignOperator assignOp = null;
        LogicalVariable distinctVar;
        if (distinctExpr.getExpressionTag() == LogicalExpressionTag.VARIABLE) {
            distinctVar = ((VariableReferenceExpression) distinctExpr).getVariableReference();
        } else {
            distinctVar = context.newVar();
            assignOp = new AssignOperator(distinctVar, new MutableObject<>(distinctExpr));
            assignOp.setSourceLocation(aggOp.getSourceLocation());
            assignOp.getInputs().add(new MutableObject<>(inputOp));
            assignOp.setExecutionMode(inputOp.getExecutionMode());
            context.computeAndSetTypeEnvironmentForOperator(assignOp);
            inputOp = assignOp;
        }

        // Introduce DISTINCT operator
        VariableReferenceExpression distinctVarRef = new VariableReferenceExpression(distinctVar);
        distinctVarRef.setSourceLocation(distinctExpr.getSourceLocation());

        List<Mutable<ILogicalExpression>> distinctExprs = new ArrayList<>(1);
        distinctExprs.add(new MutableObject<>(distinctVarRef));
        DistinctOperator distinctOp = new DistinctOperator(distinctExprs);
        distinctOp.setSourceLocation(aggOp.getSourceLocation());
        distinctOp.getInputs().add(new MutableObject<>(inputOp));
        distinctOp.setExecutionMode(inputOp.getExecutionMode());
        context.computeAndSetTypeEnvironmentForOperator(distinctOp);

        // Change function identifiers from agg-distinct(expr) to agg(expr),
        // replace 'expr' with variable reference if ASSIGN was introduced
        for (Mutable<ILogicalExpression> aggExprRef : aggOp.getExpressions()) {
            AbstractFunctionCallExpression callExpr = (AbstractFunctionCallExpression) aggExprRef.getValue();
            FunctionIdentifier regularAggForDistinct =
                    BuiltinFunctions.getAggregateFunctionForDistinct(callExpr.getFunctionIdentifier());
            if (regularAggForDistinct == null) {
                throw new IllegalStateException(String.valueOf(callExpr.getFunctionIdentifier()));
            }
            callExpr.setFunctionInfo(BuiltinFunctions.getBuiltinFunctionInfo(regularAggForDistinct));

            if (assignOp != null) {
                callExpr.getArguments().get(0).setValue(distinctVarRef.cloneExpression());
            }
        }

        aggOp.getInputs().clear();
        aggOp.getInputs().add(new MutableObject<>(distinctOp));
        context.computeAndSetTypeEnvironmentForOperator(aggOp);
    }

    @Override
    public boolean rewritePost(Mutable<ILogicalOperator> opRef, IOptimizationContext context) {
        return false;
    }
}
