/*
 * 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.lucene.search.grouping;

import java.util.ArrayList;
import java.util.Comparator;
import java.util.List;
import java.util.NavigableSet;
import java.util.TreeSet;
import org.apache.lucene.search.FieldComparator;
import org.apache.lucene.search.Pruning;
import org.apache.lucene.search.ScoreDoc;
import org.apache.lucene.search.Sort;
import org.apache.lucene.search.SortField;
import org.apache.lucene.search.TopDocs;
import org.apache.lucene.search.TopFieldDocs;
import org.apache.lucene.search.TotalHits;
import org.apache.lucene.search.TotalHits.Relation;

/**
 * Represents result returned by a grouping search.
 *
 * @lucene.experimental
 */
public class TopGroups<T> {
  /** Number of documents matching the search */
  public final int totalHitCount;

  /** Number of documents grouped into the topN groups */
  public final int totalGroupedHitCount;

  /** The total number of unique groups. If <code>null</code> this value is not computed. */
  public final Integer totalGroupCount;

  /** Group results in groupSort order */
  public final GroupDocs<T>[] groups;

  /** How groups are sorted against each other */
  public final SortField[] groupSort;

  /** How docs are sorted within each group */
  public final SortField[] withinGroupSort;

  /** Highest score across all hits, or <code>Float.NaN</code> if scores were not computed. */
  public final float maxScore;

  public TopGroups(
      SortField[] groupSort,
      SortField[] withinGroupSort,
      int totalHitCount,
      int totalGroupedHitCount,
      GroupDocs<T>[] groups,
      float maxScore) {
    this.groupSort = groupSort;
    this.withinGroupSort = withinGroupSort;
    this.totalHitCount = totalHitCount;
    this.totalGroupedHitCount = totalGroupedHitCount;
    this.groups = groups;
    this.totalGroupCount = null;
    this.maxScore = maxScore;
  }

  public TopGroups(TopGroups<T> oldTopGroups, Integer totalGroupCount) {
    this.groupSort = oldTopGroups.groupSort;
    this.withinGroupSort = oldTopGroups.withinGroupSort;
    this.totalHitCount = oldTopGroups.totalHitCount;
    this.totalGroupedHitCount = oldTopGroups.totalGroupedHitCount;
    this.groups = oldTopGroups.groups;
    this.maxScore = oldTopGroups.maxScore;
    this.totalGroupCount = totalGroupCount;
  }

  /** How the GroupDocs score (if any) should be merged. */
  public enum ScoreMergeMode {
    /** Set score to Float.NaN */
    None,
    /* Sum score across all shards for this group. */
    Total,
    /* Avg score across all shards for this group. */
    Avg,
  }

  /**
   * If either value is NaN then return the other value, otherwise return the greater of the two
   * values by calling Math.max.
   *
   * @param a - one value
   * @param b - another value
   * @return ignoring any NaN return the greater of a and b
   */
  static float nonNANmax(float a, float b) {
    if (Float.isNaN(a)) return b;
    if (Float.isNaN(b)) return a;
    return Math.max(a, b);
  }

  /**
   * Merges an array of TopGroups, for example obtained from the second-pass collector across
   * multiple shards. Each TopGroups must have been sorted by the same groupSort and docSort, and
   * the top groups passed to all second-pass collectors must be the same.
   *
   * <p><b>NOTE</b>: We can't always compute an exact totalGroupCount. Documents belonging to a
   * group may occur on more than one shard and thus the merged totalGroupCount can be higher than
   * the actual totalGroupCount. In this case the totalGroupCount represents a upper bound. If the
   * documents of one group do only reside in one shard then the totalGroupCount is exact.
   *
   * <p><b>NOTE</b>: the topDocs in each GroupDocs is actually an instance of TopDocsAndShards
   */
  public static <T> TopGroups<T> merge(
      TopGroups<T>[] shardGroups,
      Sort groupSort,
      Sort docSort,
      int docOffset,
      int docTopN,
      ScoreMergeMode scoreMergeMode) {

    // System.out.println("TopGroups.merge");

    if (shardGroups.length == 0) {
      return null;
    }

    int totalHitCount = 0;
    int totalGroupedHitCount = 0;
    // Optionally merge the totalGroupCount.
    Integer totalGroupCount = null;

    final int numGroups = shardGroups[0].groups.length;
    for (TopGroups<T> shard : shardGroups) {
      if (numGroups != shard.groups.length) {
        throw new IllegalArgumentException(
            "number of groups differs across shards; you must pass same top groups to all shards' second-pass collector");
      }
      totalHitCount += shard.totalHitCount;
      totalGroupedHitCount += shard.totalGroupedHitCount;
      if (shard.totalGroupCount != null) {
        if (totalGroupCount == null) {
          totalGroupCount = 0;
        }

        totalGroupCount += shard.totalGroupCount;
      }
    }

    @SuppressWarnings({"unchecked", "rawtypes"})
    final GroupDocs<T>[] mergedGroupDocs = new GroupDocs[numGroups];

    final TopDocs[] shardTopDocs;
    final boolean sortByRelevance = docSort.equals(Sort.RELEVANCE);
    if (sortByRelevance) {
      shardTopDocs = new TopDocs[shardGroups.length];
    } else {
      shardTopDocs = new TopFieldDocs[shardGroups.length];
    }
    float totalMaxScore = Float.NaN;

    for (int groupIDX = 0; groupIDX < numGroups; groupIDX++) {
      final T groupValue = shardGroups[0].groups[groupIDX].groupValue();
      // System.out.println("  merge groupValue=" + groupValue + " sortValues=" +
      // Arrays.toString(shardGroups[0].groups[groupIDX].groupSortValues));
      float maxScore = Float.NaN;
      int totalHits = 0;
      double scoreSum = 0.0;
      for (int shardIDX = 0; shardIDX < shardGroups.length; shardIDX++) {
        // System.out.println("    shard=" + shardIDX);
        final TopGroups<T> shard = shardGroups[shardIDX];
        final GroupDocs<?> shardGroupDocs = shard.groups[groupIDX];
        if (groupValue == null) {
          if (shardGroupDocs.groupValue() != null) {
            throw new IllegalArgumentException(
                "group values differ across shards; you must pass same top groups to all shards' second-pass collector");
          }
        } else if (!groupValue.equals(shardGroupDocs.groupValue())) {
          throw new IllegalArgumentException(
              "group values differ across shards; you must pass same top groups to all shards' second-pass collector");
        }

        /*
        for(ScoreDoc sd : shardGroupDocs.scoreDocs) {
          System.out.println("      doc=" + sd.doc);
        }
        */

        if (sortByRelevance) {
          shardTopDocs[shardIDX] =
              new TopDocs(shardGroupDocs.totalHits(), shardGroupDocs.scoreDocs());
        } else {
          shardTopDocs[shardIDX] =
              new TopFieldDocs(
                  shardGroupDocs.totalHits(), shardGroupDocs.scoreDocs(), docSort.getSort());
        }

        for (int i = 0; i < shardTopDocs[shardIDX].scoreDocs.length; i++) {
          shardTopDocs[shardIDX].scoreDocs[i].shardIndex = shardIDX;
        }

        if (!sortByRelevance) {
          maxScore = nonNANmax(maxScore, shardGroupDocs.maxScore());
        }
        assert shardGroupDocs.totalHits().relation() == Relation.EQUAL_TO;
        totalHits += shardGroupDocs.totalHits().value();
        scoreSum += shardGroupDocs.score();
      }

      final TopDocs mergedTopDocs;
      if (sortByRelevance) {
        mergedTopDocs = TopDocs.merge(docOffset + docTopN, shardTopDocs);
        // When sorting by relevance, the highest-scoring doc is first, so we can
        // derive maxScore directly instead of accumulating across shards.
        maxScore =
            mergedTopDocs.scoreDocs.length == 0 ? Float.NaN : mergedTopDocs.scoreDocs[0].score;
      } else {
        mergedTopDocs = TopDocs.merge(docSort, docOffset + docTopN, (TopFieldDocs[]) shardTopDocs);
      }

      // Slice;
      final ScoreDoc[] mergedScoreDocs;
      if (docOffset == 0) {
        mergedScoreDocs = mergedTopDocs.scoreDocs;
      } else if (docOffset >= mergedTopDocs.scoreDocs.length) {
        mergedScoreDocs = new ScoreDoc[0];
      } else {
        mergedScoreDocs = new ScoreDoc[mergedTopDocs.scoreDocs.length - docOffset];
        System.arraycopy(
            mergedTopDocs.scoreDocs,
            docOffset,
            mergedScoreDocs,
            0,
            mergedTopDocs.scoreDocs.length - docOffset);
      }

      final float groupScore;
      switch (scoreMergeMode) {
        case None:
          groupScore = Float.NaN;
          break;
        case Avg:
          if (totalHits > 0) {
            groupScore = (float) (scoreSum / totalHits);
          } else {
            groupScore = Float.NaN;
          }
          break;
        case Total:
          groupScore = (float) scoreSum;
          break;
        default:
          throw new IllegalArgumentException("can't handle ScoreMergeMode " + scoreMergeMode);
      }

      // System.out.println("SHARDS=" + Arrays.toString(mergedTopDocs.shardIndex));
      mergedGroupDocs[groupIDX] =
          new GroupDocs<>(
              groupScore,
              maxScore,
              new TotalHits(totalHits, TotalHits.Relation.EQUAL_TO),
              mergedScoreDocs,
              groupValue,
              shardGroups[0].groups[groupIDX].groupSortValues());
      totalMaxScore = nonNANmax(totalMaxScore, maxScore);
    }

    if (totalGroupCount != null) {
      TopGroups<T> result =
          new TopGroups<>(
              groupSort.getSort(),
              docSort.getSort(),
              totalHitCount,
              totalGroupedHitCount,
              mergedGroupDocs,
              totalMaxScore);
      return new TopGroups<>(result, totalGroupCount);
    } else {
      return new TopGroups<>(
          groupSort.getSort(),
          docSort.getSort(),
          totalHitCount,
          totalGroupedHitCount,
          mergedGroupDocs,
          totalMaxScore);
    }
  }

  private record MergedBlockGroup(Object[] topValues, int shardIndex, int groupIndex) {}

  private static class GroupComparator implements Comparator<MergedBlockGroup> {
    @SuppressWarnings("rawtypes")
    private final FieldComparator[] comparators;

    private final int[] reversed;

    @SuppressWarnings({"rawtypes"})
    public GroupComparator(Sort groupSort) {
      final SortField[] sortFields = groupSort.getSort();
      comparators = new FieldComparator[sortFields.length];
      reversed = new int[sortFields.length];
      for (int compIDX = 0; compIDX < sortFields.length; compIDX++) {
        final SortField sortField = sortFields[compIDX];
        comparators[compIDX] = sortField.getComparator(1, Pruning.NONE);
        reversed[compIDX] = sortField.getReverse() ? -1 : 1;
      }
    }

    @Override
    @SuppressWarnings({"unchecked"})
    public int compare(MergedBlockGroup group, MergedBlockGroup other) {
      if (group == other) {
        return 0;
      }
      final Object[] groupValues = group.topValues;
      final Object[] otherValues = other.topValues;
      for (int compIDX = 0; compIDX < comparators.length; compIDX++) {
        final int c =
            reversed[compIDX]
                * comparators[compIDX].compareValues(groupValues[compIDX], otherValues[compIDX]);
        if (c != 0) {
          return c;
        }
      }

      assert group.shardIndex != other.shardIndex;
      return group.shardIndex - other.shardIndex;
    }
  }

  /**
   * Merge TopGroups that are partitioned into blocks per shard. This method assumes that within
   * each shard, the groups are sorted according to the groupSort.
   *
   * @param shardGroups list of TopGroups, one per shard.
   * @param groupSort The {@link Sort} used to sort the groups. The top sorted document within each
   *     group according to groupSort, determines how that group sorts against other groups. This
   *     must be non-null, ie, if you want to groupSort by relevance use Sort.RELEVANCE.
   * @param groupOffset Which group to start from.
   * @param topNGroups How many top groups to keep.
   * @param docSort The sort to use within each group
   * @return TopGroups instance or null if there are no groups.
   */
  @SuppressWarnings("unchecked")
  public static <T> TopGroups<T> mergeBlockGroups(
      List<TopGroups<T>> shardGroups,
      Sort groupSort,
      int groupOffset,
      int topNGroups,
      Sort docSort) {
    if (shardGroups.isEmpty()) {
      return new TopGroups<>(
          groupSort.getSort(),
          docSort.getSort(),
          0,
          0,
          (GroupDocs<T>[]) new GroupDocs<?>[0],
          Float.NaN);
    }

    Integer totalGroupCount = null;
    int totalHitCount = 0;
    int totalGroupedHitCount = 0;
    for (TopGroups<T> sg : shardGroups) {
      totalHitCount += sg.totalHitCount;
      if (sg.totalGroupCount != null) {
        if (totalGroupCount == null) {
          totalGroupCount = 0;
        }
        totalGroupCount += sg.totalGroupCount;
      }
    }

    // k-way merge
    GroupComparator groupComp = new GroupComparator(groupSort);
    NavigableSet<MergedBlockGroup> queue = new TreeSet<>(groupComp);

    float totalMaxScore = Float.NaN;
    final boolean groupSortByRelevance = groupSort.equals(Sort.RELEVANCE);
    // init queue
    for (int idx = 0; idx < shardGroups.size(); idx++) {
      TopGroups<T> topGroups = shardGroups.get(idx);
      if (topGroups.groups.length == 0) {
        continue;
      }
      if (!groupSortByRelevance) {
        totalMaxScore = nonNANmax(totalMaxScore, topGroups.maxScore);
      }
      GroupDocs<T> firstGroupDocs = topGroups.groups[0];
      queue.add(new MergedBlockGroup(firstGroupDocs.groupSortValues(), idx, 0));
    }

    if (groupSortByRelevance && !queue.isEmpty()) {
      totalMaxScore = shardGroups.get(queue.first().shardIndex).maxScore;
    }

    final List<GroupDocs<T>> groupDocsList = new ArrayList<>();
    int count = 0;
    while (!queue.isEmpty()) {
      final MergedBlockGroup mergedBlockGroup = queue.pollFirst();
      TopGroups<T> shardGroup = shardGroups.get(mergedBlockGroup.shardIndex);

      int currentGroupIndex = mergedBlockGroup.groupIndex;
      GroupDocs<T> currentGroupDocs = shardGroup.groups[currentGroupIndex];
      if (count++ >= groupOffset) {
        groupDocsList.add(currentGroupDocs);
        totalGroupedHitCount += (int) currentGroupDocs.totalHits().value();
        if (groupDocsList.size() == topNGroups) {
          break;
        }
      }

      int nextGroupIndex = currentGroupIndex + 1;
      if (nextGroupIndex < shardGroup.groups.length) {
        GroupDocs<T> nextGroupDocs = shardGroup.groups[nextGroupIndex];
        queue.add(
            new MergedBlockGroup(
                nextGroupDocs.groupSortValues(), mergedBlockGroup.shardIndex, nextGroupIndex));
      }
    }

    @SuppressWarnings({"unchecked"})
    GroupDocs<T>[] groupDocs = (GroupDocs<T>[]) groupDocsList.toArray(GroupDocs[]::new);

    return new TopGroups<>(
        new TopGroups<>(
            groupSort.getSort(),
            docSort.getSort(),
            totalHitCount,
            totalGroupedHitCount,
            groupDocs,
            totalMaxScore),
        totalGroupCount);
  }
}
