/**
 * 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.hadoop.hive.ql.optimizer.calcite.rules;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

import org.apache.calcite.plan.RelOptPredicateList;
import org.apache.calcite.plan.RelOptRule;
import org.apache.calcite.plan.RelOptRuleCall;
import org.apache.calcite.plan.RelOptUtil;
import org.apache.calcite.plan.hep.HepRelVertex;
import org.apache.calcite.rel.RelCollations;
import org.apache.calcite.rel.RelFieldCollation;
import org.apache.calcite.rel.RelNode;
import org.apache.calcite.rel.core.Sort;
import org.apache.calcite.rel.metadata.RelMetadataQuery;
import org.apache.calcite.rel.type.RelDataTypeField;
import org.apache.calcite.rex.RexBuilder;
import org.apache.calcite.rex.RexLiteral;
import org.apache.calcite.rex.RexNode;
import org.apache.calcite.rex.RexUtil;
import org.apache.calcite.tools.RelBuilder;
import org.apache.calcite.tools.RelBuilderFactory;
import org.apache.calcite.util.Pair;
import org.apache.calcite.util.mapping.Mappings;
import org.apache.hadoop.hive.ql.optimizer.calcite.HiveRelFactories;
import org.apache.hadoop.hive.ql.optimizer.calcite.reloperators.HiveSortLimit;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import com.google.common.collect.ImmutableList;

/**
 * Planner rule that pulls up constant keys through a SortLimit operator.
 * 
 * This rule is only applied on SortLimit operators that are not the root
 * of the plan tree. This is done because the interaction of this rule
 * with the AST conversion may cause some optimizations to not kick in
 * e.g. SimpleFetchOptimizer. Nevertheless, this will not have any
 * performance impact in the resulting plans.
 */
public class HiveSortLimitPullUpConstantsRule extends RelOptRule {

  protected static final Logger LOG = LoggerFactory.getLogger(HiveSortLimitPullUpConstantsRule.class);


  public static final HiveSortLimitPullUpConstantsRule INSTANCE =
          new HiveSortLimitPullUpConstantsRule(HiveSortLimit.class,
                  HiveRelFactories.HIVE_BUILDER);

  private HiveSortLimitPullUpConstantsRule(Class<? extends Sort> sortClass,
      RelBuilderFactory relBuilderFactory) {
    super(operand(RelNode.class, unordered(operand(sortClass, any()))), relBuilderFactory, null);
  }

  @Override
  public void onMatch(RelOptRuleCall call) {
    final RelNode parent = call.rel(0);
    final Sort sort = call.rel(1);

    final int count = sort.getInput().getRowType().getFieldCount();
    if (count == 1) {
      // No room for optimization since we cannot convert to an empty
      // Project operator.
      return;
    }

    final RexBuilder rexBuilder = sort.getCluster().getRexBuilder();
    final RelMetadataQuery mq = RelMetadataQuery.instance();
    final RelOptPredicateList predicates = mq.getPulledUpPredicates(sort.getInput());
    if (predicates == null) {
      return;
    }

    Map<RexNode, RexNode> conditionsExtracted = HiveReduceExpressionsRule.predicateConstants(
            RexNode.class, rexBuilder, predicates);
    Map<RexNode, RexNode> constants = new HashMap<>();
    for (int i = 0; i < count ; i++) {
      RexNode expr = rexBuilder.makeInputRef(sort.getInput(), i);
      if (conditionsExtracted.containsKey(expr)) {
        constants.put(expr, conditionsExtracted.get(expr));
      }
    }

    // None of the expressions are constant. Nothing to do.
    if (constants.isEmpty()) {
      return;
    }

    if (count == constants.size()) {
      // At least a single item in project is required.
      constants.remove(constants.keySet().iterator().next());
    }

    // Create expressions for Project operators before and after the Sort
    List<RelDataTypeField> fields = sort.getInput().getRowType().getFieldList();
    List<Pair<RexNode, String>> newChildExprs = new ArrayList<>();
    List<RexNode> topChildExprs = new ArrayList<>();
    List<String> topChildExprsFields = new ArrayList<>();
    for (int i = 0; i < count ; i++) {
      RexNode expr = rexBuilder.makeInputRef(sort.getInput(), i);
      RelDataTypeField field = fields.get(i);
      if (constants.containsKey(expr)) {
        topChildExprs.add(constants.get(expr));
        topChildExprsFields.add(field.getName());
      } else {
        newChildExprs.add(Pair.<RexNode,String>of(expr, field.getName()));
        topChildExprs.add(expr);
        topChildExprsFields.add(field.getName());
      }
    }

    // Update field collations
    final Mappings.TargetMapping mapping =
            RelOptUtil.permutation(Pair.left(newChildExprs), sort.getInput().getRowType()).inverse();
    List<RelFieldCollation> fieldCollations = new ArrayList<>();
    for (RelFieldCollation fc : sort.getCollation().getFieldCollations()) {
      final int target = mapping.getTargetOpt(fc.getFieldIndex());
      if (target < 0) {
        // It is a constant, we can ignore it
        continue;
      }
      fieldCollations.add(fc.copy(target));
    }

    // Update top Project positions
    topChildExprs = ImmutableList.copyOf(RexUtil.apply(mapping, topChildExprs));

    // Create new Project-Sort-Project sequence
    final RelBuilder relBuilder = call.builder();
    relBuilder.push(sort.getInput());
    relBuilder.project(Pair.left(newChildExprs), Pair.right(newChildExprs));
    final ImmutableList<RexNode> sortFields =
            relBuilder.fields(RelCollations.of(fieldCollations));
    relBuilder.sortLimit(sort.offset == null ? -1 : RexLiteral.intValue(sort.offset),
            sort.fetch == null ? -1 : RexLiteral.intValue(sort.fetch), sortFields);
    // Create top Project fixing nullability of fields
    relBuilder.project(topChildExprs, topChildExprsFields);
    relBuilder.convert(sort.getRowType(), false);

    List<RelNode> inputs = new ArrayList<>();
    for (RelNode child : parent.getInputs()) {
      if (!((HepRelVertex) child).getCurrentRel().equals(sort)) {
        inputs.add(child);
      } else {
        inputs.add(relBuilder.build());
      }
    }
    call.transformTo(parent.copy(parent.getTraitSet(), inputs));
  }

}
