/*
 * 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.beam.runners.spark.structuredstreaming.translation.batch;

import static org.apache.beam.runners.spark.structuredstreaming.translation.helpers.EncoderHelpers.encoderFor;
import static org.apache.beam.runners.spark.structuredstreaming.translation.helpers.EncoderHelpers.windowedValueEncoder;
import static org.hamcrest.MatcherAssert.assertThat;
import static org.hamcrest.Matchers.equalTo;
import static org.joda.time.Duration.standardMinutes;

import java.util.Arrays;
import java.util.Collection;
import java.util.Map;
import java.util.TreeMap;
import java.util.stream.Collectors;
import org.apache.beam.runners.spark.structuredstreaming.translation.helpers.EncoderHelpers;
import org.apache.beam.sdk.coders.Coder;
import org.apache.beam.sdk.transforms.Combine;
import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
import org.apache.beam.sdk.transforms.windowing.IntervalWindow;
import org.apache.beam.sdk.transforms.windowing.PaneInfo;
import org.apache.beam.sdk.transforms.windowing.Sessions;
import org.apache.beam.sdk.transforms.windowing.SlidingWindows;
import org.apache.beam.sdk.transforms.windowing.TimestampCombiner;
import org.apache.beam.sdk.transforms.windowing.WindowFn;
import org.apache.beam.sdk.transforms.windowing.WindowMappingFn;
import org.apache.beam.sdk.values.KV;
import org.apache.beam.sdk.values.WindowedValue;
import org.apache.beam.sdk.values.WindowedValues;
import org.apache.beam.sdk.values.WindowingStrategy;
import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Streams;
import org.apache.spark.sql.Encoder;
import org.apache.spark.sql.expressions.Aggregator;
import org.apache.spark.util.MutablePair;
import org.hamcrest.Matcher;
import org.joda.time.Duration;
import org.joda.time.Instant;
import org.junit.Test;
import org.junit.experimental.runners.Enclosed;
import org.junit.runner.RunWith;

@RunWith(Enclosed.class)
public class AggregatorsTest {

  // just something easy readable
  private static final Instant NOW = Instant.parse("2000-01-01T00:00Z");

  /** Tests for NonMergingWindowedAggregator in {@link Aggregators}. */
  public static class NonMergingWindowedAggregatorTest {

    private SlidingWindows sliding =
        SlidingWindows.of(standardMinutes(15)).every(standardMinutes(5));

    private Aggregator<
            WindowedValue<Integer>,
            Map<IntervalWindow, MutablePair<Instant, Integer>>,
            Collection<WindowedValue<Integer>>>
        agg = windowedAgg(sliding);

    @Test
    public void testReduce() {
      Map<IntervalWindow, MutablePair<Instant, Integer>> acc;

      acc = agg.reduce(agg.zero(), windowedValue(1, at(10)));
      assertThat(
          acc,
          equalsToMap(
              KV.of(intervalWindow(0, 15), pair(at(10), 1)),
              KV.of(intervalWindow(5, 20), pair(at(10), 1)),
              KV.of(intervalWindow(10, 25), pair(at(10), 1))));

      acc = agg.reduce(acc, windowedValue(2, at(16)));
      assertThat(
          acc,
          equalsToMap(
              KV.of(intervalWindow(0, 15), pair(at(10), 1)),
              KV.of(intervalWindow(5, 20), pair(at(16), 3)),
              KV.of(intervalWindow(10, 25), pair(at(16), 3)),
              KV.of(intervalWindow(15, 30), pair(at(16), 2))));
    }

    @Test
    public void testMerge() {
      Map<IntervalWindow, MutablePair<Instant, Integer>> acc;

      assertThat(agg.merge(agg.zero(), agg.zero()), equalTo(agg.zero()));

      acc = mapOf(KV.of(intervalWindow(0, 15), pair(at(0), 1)));

      assertThat(agg.merge(acc, agg.zero()), equalTo(acc));
      assertThat(agg.merge(agg.zero(), acc), equalTo(acc));

      acc = agg.merge(acc, acc);
      assertThat(acc, equalsToMap(KV.of(intervalWindow(0, 15), pair(at(0), 1 + 1))));

      acc = agg.merge(acc, mapOf(KV.of(intervalWindow(5, 20), pair(at(5), 3))));
      assertThat(
          acc,
          equalsToMap(
              KV.of(intervalWindow(0, 15), pair(at(0), 1 + 1)),
              KV.of(intervalWindow(5, 20), pair(at(5), 3))));

      acc = agg.merge(mapOf(KV.of(intervalWindow(10, 25), pair(at(10), 4))), acc);
      assertThat(
          acc,
          equalsToMap(
              KV.of(intervalWindow(0, 15), pair(at(0), 1 + 1)),
              KV.of(intervalWindow(5, 20), pair(at(5), 3)),
              KV.of(intervalWindow(10, 25), pair(at(10), 4))));
    }

    private WindowedValue<Integer> windowedValue(Integer value, Instant ts) {
      return WindowedValues.of(value, ts, sliding.assignWindows(ts), PaneInfo.NO_FIRING);
    }
  }

  /**
   * Shared implementation of tests for SessionsAggregator and MergingWindowedAggregator in {@link
   * Aggregators}.
   */
  public abstract static class AbstractSessionsTest<
      AccT extends Map<IntervalWindow, MutablePair<Instant, Integer>>> {

    static final Duration SESSIONS_GAP = standardMinutes(15);

    final Aggregator<WindowedValue<Integer>, AccT, Collection<WindowedValue<Integer>>> agg;

    AbstractSessionsTest(WindowFn<?, ?> windowFn) {
      agg = windowedAgg(windowFn);
    }

    abstract AccT accOf(KV<IntervalWindow, MutablePair<Instant, Integer>>... entries);

    @Test
    public void testReduce() {
      AccT acc;

      acc = agg.reduce(agg.zero(), sessionValue(10, at(0)));
      assertThat(acc, equalsToMap(KV.of(sessionWindow(0), pair(at(0), 10))));

      // 2nd session after 1st
      acc = agg.reduce(acc, sessionValue(7, at(20)));
      assertThat(
          acc,
          equalsToMap(
              KV.of(sessionWindow(0), pair(at(0), 10)), KV.of(sessionWindow(20), pair(at(20), 7))));

      // merge into 2nd session
      acc = agg.reduce(acc, sessionValue(6, at(18)));
      assertThat(
          acc,
          equalsToMap(
              KV.of(sessionWindow(0), pair(at(0), 10)),
              KV.of(sessionWindow(18, 35), pair(at(20), 7 + 6))));

      // merge into 2nd session
      acc = agg.reduce(acc, sessionValue(5, at(21)));
      assertThat(
          acc,
          equalsToMap(
              KV.of(sessionWindow(0), pair(at(0), 10)),
              KV.of(sessionWindow(18, 36), pair(at(21), 7 + 6 + 5))));

      // 3rd session after 2nd
      acc = agg.reduce(acc, sessionValue(2, NOW.plus(standardMinutes(45))));
      assertThat(
          acc,
          equalsToMap(
              KV.of(sessionWindow(0), pair(at(0), 10)),
              KV.of(sessionWindow(18, 36), pair(at(21), 7 + 6 + 5)),
              KV.of(sessionWindow(45), pair(at(45), 2))));

      // merge with 1st and 2nd
      acc = agg.reduce(acc, sessionValue(1, at(10)));
      assertThat(
          acc,
          equalsToMap(
              KV.of(sessionWindow(0, 36), pair(at(21), 10 + 7 + 6 + 5 + 1)),
              KV.of(sessionWindow(45), pair(at(45), 2))));
    }

    @Test
    public void testMerge() {
      AccT acc;

      assertThat(agg.merge(agg.zero(), agg.zero()), equalTo(agg.zero()));

      acc = accOf(KV.of(sessionWindow(0), pair(at(0), 1)));

      assertThat(agg.merge(acc, agg.zero()), equalTo(acc));
      assertThat(agg.merge(agg.zero(), acc), equalTo(acc));

      acc = agg.merge(acc, acc);
      assertThat(acc, equalsToMap(KV.of(sessionWindow(0), pair(at(0), 1 + 1))));

      acc = agg.merge(acc, accOf(KV.of(sessionWindow(20), pair(at(20), 2))));
      assertThat(
          acc,
          equalsToMap(
              KV.of(sessionWindow(0), pair(at(0), 1 + 1)),
              KV.of(sessionWindow(20), pair(at(20), 2))));

      acc = agg.merge(accOf(KV.of(sessionWindow(40), pair(at(40), 3))), acc);
      assertThat(
          acc,
          equalsToMap(
              KV.of(sessionWindow(0), pair(at(0), 1 + 1)),
              KV.of(sessionWindow(20), pair(at(20), 2)),
              KV.of(sessionWindow(40), pair(at(40), 3))));

      acc = agg.merge(acc, accOf(KV.of(sessionWindow(10), pair(at(10), 4))));
      assertThat(
          acc,
          equalsToMap(
              KV.of(sessionWindow(0, 35), pair(at(20), 1 + 1 + 2 + 4)),
              KV.of(sessionWindow(40), pair(at(40), 3))));

      acc = agg.merge(accOf(KV.of(sessionWindow(5, 45), pair(at(30), 5))), acc);
      assertThat(
          acc, equalsToMap(KV.of(sessionWindow(0, 55), pair(at(40), 1 + 1 + 2 + 4 + 3 + 5))));
    }

    private WindowedValue<Integer> sessionValue(Integer value, Instant ts) {
      return WindowedValues.of(value, ts, new IntervalWindow(ts, SESSIONS_GAP), PaneInfo.NO_FIRING);
    }

    private IntervalWindow sessionWindow(int fromMinutes) {
      return new IntervalWindow(at(fromMinutes), SESSIONS_GAP);
    }

    private static IntervalWindow sessionWindow(int fromMinutes, int toMinutes) {
      return intervalWindow(fromMinutes, toMinutes);
    }
  }

  /** Tests for specialized SessionsAggregator in {@link Aggregators}. */
  public static class SessionsAggregatorTest
      extends AbstractSessionsTest<TreeMap<IntervalWindow, MutablePair<Instant, Integer>>> {

    public SessionsAggregatorTest() {
      super(Sessions.withGapDuration(SESSIONS_GAP));
    }

    @Override
    TreeMap<IntervalWindow, MutablePair<Instant, Integer>> accOf(
        KV<IntervalWindow, MutablePair<Instant, Integer>>... entries) {
      return new TreeMap<>(mapOf(entries));
    }
  }

  /** Tests for MergingWindowedAggregator in {@link Aggregators}. */
  public static class MergingWindowedAggregatorTest
      extends AbstractSessionsTest<Map<IntervalWindow, MutablePair<Instant, Integer>>> {

    public MergingWindowedAggregatorTest() {
      super(new CustomSessions<>());
    }

    @Override
    Map<IntervalWindow, MutablePair<Instant, Integer>> accOf(
        KV<IntervalWindow, MutablePair<Instant, Integer>>... entries) {
      return mapOf(entries);
    }

    /** Wrapper around {@link Sessions} to test the MergingWindowedAggregator. */
    private static class CustomSessions<T> extends WindowFn<T, IntervalWindow> {
      private final Sessions sessions = Sessions.withGapDuration(SESSIONS_GAP);

      @Override
      public Collection<IntervalWindow> assignWindows(WindowFn<T, IntervalWindow>.AssignContext c) {
        return sessions.assignWindows((WindowFn.AssignContext) c);
      }

      @Override
      public void mergeWindows(WindowFn<T, IntervalWindow>.MergeContext c) throws Exception {
        sessions.mergeWindows((WindowFn<Object, IntervalWindow>.MergeContext) c);
      }

      @Override
      public boolean isCompatible(WindowFn<?, ?> other) {
        return sessions.isCompatible(other);
      }

      @Override
      public Coder<IntervalWindow> windowCoder() {
        return sessions.windowCoder();
      }

      @Override
      public WindowMappingFn<IntervalWindow> getDefaultWindowMappingFn() {
        return sessions.getDefaultWindowMappingFn();
      }
    }
  }

  private static IntervalWindow intervalWindow(int fromMinutes, int toMinutes) {
    return new IntervalWindow(at(fromMinutes), at(toMinutes));
  }

  private static Instant at(int minutes) {
    return NOW.plus(standardMinutes(minutes));
  }

  private static Matcher<Map<IntervalWindow, MutablePair<Instant, Integer>>> equalsToMap(
      KV<IntervalWindow, MutablePair<Instant, Integer>>... entries) {
    return equalTo(mapOf(entries));
  }

  private static Map<IntervalWindow, MutablePair<Instant, Integer>> mapOf(
      KV<IntervalWindow, MutablePair<Instant, Integer>>... entries) {
    return Arrays.asList(entries).stream().collect(Collectors.toMap(KV::getKey, KV::getValue));
  }

  private static MutablePair<Instant, Integer> pair(Instant ts, int value) {
    return new MutablePair<>(ts, value);
  }

  private static <AccT>
      Aggregator<WindowedValue<Integer>, AccT, Collection<WindowedValue<Integer>>> windowedAgg(
          WindowFn<?, ?> windowFn) {
    Encoder<Integer> intEnc = EncoderHelpers.encoderOf(Integer.class);
    Encoder<BoundedWindow> windowEnc = encoderFor((Coder) IntervalWindow.getCoder());
    Encoder<WindowedValue<Integer>> outputEnc = windowedValueEncoder(intEnc, windowEnc);

    WindowingStrategy<?, ?> windowing =
        WindowingStrategy.of(windowFn).withTimestampCombiner(TimestampCombiner.LATEST);

    Aggregator<WindowedValue<Integer>, ?, Collection<WindowedValue<Integer>>> agg =
        Aggregators.windowedValue(
            new SimpleSum(), WindowedValue::getValue, windowing, windowEnc, intEnc, outputEnc);
    return (Aggregator) agg;
  }

  private static class SimpleSum extends Combine.CombineFn<Integer, Integer, Integer> {

    @Override
    public Integer createAccumulator() {
      return 0;
    }

    @Override
    public Integer addInput(Integer acc, Integer input) {
      return acc + input;
    }

    @Override
    public Integer mergeAccumulators(Iterable<Integer> accs) {
      return Streams.stream(accs.iterator()).reduce((a, b) -> a + b).orElseGet(() -> 0);
    }

    @Override
    public Integer extractOutput(Integer acc) {
      return acc;
    }
  }
}
