/*
 * 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.calcite.adapter.enumerable;

import org.apache.calcite.adapter.enumerable.impl.AggResultContextImpl;
import org.apache.calcite.adapter.java.JavaTypeFactory;
import org.apache.calcite.linq4j.Ord;
import org.apache.calcite.linq4j.function.Function0;
import org.apache.calcite.linq4j.function.Function2;
import org.apache.calcite.linq4j.tree.BlockBuilder;
import org.apache.calcite.linq4j.tree.Expression;
import org.apache.calcite.linq4j.tree.Expressions;
import org.apache.calcite.linq4j.tree.ParameterExpression;
import org.apache.calcite.plan.RelOptCluster;
import org.apache.calcite.plan.RelTraitSet;
import org.apache.calcite.rel.RelCollation;
import org.apache.calcite.rel.RelCollations;
import org.apache.calcite.rel.RelFieldCollation;
import org.apache.calcite.rel.RelNode;
import org.apache.calcite.rel.core.Aggregate;
import org.apache.calcite.rel.core.AggregateCall;
import org.apache.calcite.rex.RexUtil;
import org.apache.calcite.util.BuiltInMethod;
import org.apache.calcite.util.ImmutableBitSet;
import org.apache.calcite.util.Pair;
import org.apache.calcite.util.Util;
import org.apache.calcite.util.mapping.Mappings;

import com.google.common.collect.ImmutableList;

import org.checkerframework.checker.nullness.qual.Nullable;

import java.lang.reflect.Type;
import java.util.ArrayList;
import java.util.List;

import static java.util.Objects.requireNonNull;

/** Sort based physical implementation of {@link Aggregate} in
 * {@link EnumerableConvention enumerable calling convention}. */
public class EnumerableSortedAggregate extends EnumerableAggregateBase implements EnumerableRel {
  public EnumerableSortedAggregate(
      RelOptCluster cluster,
      RelTraitSet traitSet,
      RelNode input,
      ImmutableBitSet groupSet,
      @Nullable List<ImmutableBitSet> groupSets,
      List<AggregateCall> aggCalls) {
    super(cluster, traitSet, ImmutableList.of(), input, groupSet, groupSets, aggCalls);
    assert getConvention() instanceof EnumerableConvention;
  }

  @Override public EnumerableSortedAggregate copy(RelTraitSet traitSet, RelNode input,
      ImmutableBitSet groupSet,
      @Nullable List<ImmutableBitSet> groupSets, List<AggregateCall> aggCalls) {
    return new EnumerableSortedAggregate(getCluster(), traitSet, input,
        groupSet, groupSets, aggCalls);
  }

  @Override public @Nullable Pair<RelTraitSet, List<RelTraitSet>> passThroughTraits(
      final RelTraitSet required) {
    if (!isSimple(this)) {
      return null;
    }

    RelTraitSet inputTraits = getInput().getTraitSet();
    RelCollation collation =
        requireNonNull(required.getCollation(),
            () -> "collation trait is null, required traits are " + required);
    ImmutableBitSet requiredKeys = ImmutableBitSet.of(RelCollations.ordinals(collation));
    ImmutableBitSet groupKeys = ImmutableBitSet.range(groupSet.cardinality());

    Mappings.TargetMapping mapping =
        Mappings.source(groupSet.toList(), input.getRowType().getFieldCount());

    if (requiredKeys.equals(groupKeys)) {
      RelCollation inputCollation = RexUtil.apply(mapping, collation);
      return Pair.of(required, ImmutableList.of(inputTraits.replace(inputCollation)));
    } else if (groupKeys.contains(requiredKeys)) {
      // group by a,b,c order by c,b
      List<RelFieldCollation> list = new ArrayList<>(collation.getFieldCollations());
      groupKeys.except(requiredKeys).forEachInt(k ->
          list.add(new RelFieldCollation(k)));
      RelCollation aggCollation = RelCollations.of(list);
      RelCollation inputCollation = RexUtil.apply(mapping, aggCollation);
      return Pair.of(traitSet.replace(aggCollation),
          ImmutableList.of(inputTraits.replace(inputCollation)));
    }

    // Group keys doesn't contain all the required keys, e.g.
    // group by a,b order by a,b,c
    // nothing we can do to propagate traits to child nodes.
    return null;
  }

  @Override public Result implement(EnumerableRelImplementor implementor, Prefer pref) {
    if (!Aggregate.isSimple(this)) {
      throw Util.needToImplement("EnumerableSortedAggregate");
    }

    final JavaTypeFactory typeFactory = implementor.getTypeFactory();
    final BlockBuilder builder = new BlockBuilder();
    final EnumerableRel child = (EnumerableRel) getInput();
    final Result result = implementor.visitChild(this, 0, child, pref);
    Expression childExp =
        builder.append(
            "child",
            result.block);

    final PhysType physType =
        PhysTypeImpl.of(
            typeFactory, getRowType(), pref.preferCustom());

    final PhysType inputPhysType = result.physType;

    ParameterExpression parameter =
        Expressions.parameter(inputPhysType.getJavaRowType(), "a0");

    final PhysType keyPhysType =
        inputPhysType.project(groupSet.asList(), getGroupType() != Group.SIMPLE,
            JavaRowFormat.LIST);
    final int groupCount = getGroupCount();

    final List<AggImpState> aggs = new ArrayList<>(aggCalls.size());
    for (Ord<AggregateCall> call : Ord.zip(aggCalls)) {
      aggs.add(new AggImpState(call.i, call.e, false));
    }

    // Function0<Object[]> accumulatorInitializer =
    //     new Function0<Object[]>() {
    //         public Object[] apply() {
    //             return new Object[] {0, 0};
    //         }
    //     };
    final List<Expression> initExpressions = new ArrayList<>();
    final BlockBuilder initBlock = new BlockBuilder();

    final List<Type> aggStateTypes =
        createAggStateTypes(initExpressions, initBlock, aggs, typeFactory);

    final PhysType accPhysType =
        PhysTypeImpl.of(typeFactory,
            typeFactory.createSyntheticType(aggStateTypes));

    declareParentAccumulator(initExpressions, initBlock, accPhysType);

    final Expression accumulatorInitializer =
        builder.append("accumulatorInitializer",
            Expressions.lambda(
                Function0.class,
                initBlock.toBlock()));

    // Function2<Object[], Employee, Object[]> accumulatorAdder =
    //     new Function2<Object[], Employee, Object[]>() {
    //         public Object[] apply(Object[] acc, Employee in) {
    //              acc[0] = ((Integer) acc[0]) + 1;
    //              acc[1] = ((Integer) acc[1]) + in.salary;
    //             return acc;
    //         }
    //     };
    final ParameterExpression inParameter =
        Expressions.parameter(inputPhysType.getJavaRowType(), "in");
    final ParameterExpression acc_ =
        Expressions.parameter(accPhysType.getJavaRowType(), "acc");

    createAccumulatorAdders(
        inParameter, aggs, accPhysType, acc_, inputPhysType, builder, implementor, typeFactory);

    final ParameterExpression lambdaFactory =
        Expressions.parameter(AggregateLambdaFactory.class,
            builder.newName("lambdaFactory"));

    implementLambdaFactory(builder, inputPhysType, aggs, accumulatorInitializer,
        false, lambdaFactory);

    final BlockBuilder resultBlock = new BlockBuilder();
    final List<Expression> results = Expressions.list();
    final ParameterExpression key_;
    final Type keyType = keyPhysType.getJavaRowType();
    key_ = Expressions.parameter(keyType, "key");
    for (int j = 0; j < groupCount; j++) {
      final Expression ref = keyPhysType.fieldReference(key_, j);
      results.add(ref);
    }

    for (final AggImpState agg : aggs) {
      results.add(
          agg.implementor.implementResult(
              requireNonNull(agg.context, () -> "agg.context is null for " + agg),
              new AggResultContextImpl(resultBlock, agg.call,
                  requireNonNull(agg.state, () -> "agg.state is null for " + agg),
                  key_,
                  keyPhysType)));
    }
    resultBlock.add(physType.record(results));

    final Expression keySelector_ =
        builder.append("keySelector",
            inputPhysType.generateSelector(parameter, groupSet.asList(),
                keyPhysType.getFormat()));
    // Generate the appropriate key Comparator. In the case of NULL values
    // in group keys, the comparator must be able to support NULL values by giving a
    // consistent sort ordering.
    final Expression comparator =
        keyPhysType.generateComparator(
            requireNonNull(getTraitSet().getCollation(),
                () -> "getTraitSet().getCollation() is null; traits are "
                    + getTraitSet()));

    final Expression resultSelector_ =
        builder.append("resultSelector",
            Expressions.lambda(Function2.class, resultBlock.toBlock(), key_,
                acc_));

    builder.add(
        Expressions.return_(null,
            Expressions.call(childExp,
                BuiltInMethod.SORTED_GROUP_BY.method,
                Expressions.list(keySelector_,
                    Expressions.call(lambdaFactory,
                        BuiltInMethod.AGG_LAMBDA_FACTORY_ACC_INITIALIZER.method),
                    Expressions.call(lambdaFactory,
                        BuiltInMethod.AGG_LAMBDA_FACTORY_ACC_ADDER.method),
                    Expressions.call(lambdaFactory,
                        BuiltInMethod.AGG_LAMBDA_FACTORY_ACC_RESULT_SELECTOR.method,
                        resultSelector_), comparator))));

    return implementor.result(physType, builder.toBlock());
  }
}
