/*
 * 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.sdk.transforms;

import static org.apache.beam.sdk.transforms.display.DisplayDataMatchers.hasDisplayItem;
import static org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions.checkArgument;
import static org.hamcrest.MatcherAssert.assertThat;
import static org.hamcrest.Matchers.allOf;
import static org.hamcrest.Matchers.everyItem;
import static org.hamcrest.Matchers.in;
import static org.hamcrest.core.Is.is;
import static org.junit.Assert.assertEquals;
import static org.junit.Assert.assertTrue;

import java.io.IOException;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collection;
import java.util.Collections;
import java.util.HashSet;
import java.util.List;
import java.util.Set;
import java.util.TreeSet;
import org.apache.beam.sdk.TestUtils;
import org.apache.beam.sdk.coders.BigEndianIntegerCoder;
import org.apache.beam.sdk.coders.StringUtf8Coder;
import org.apache.beam.sdk.testing.CombineFnTester;
import org.apache.beam.sdk.testing.NeedsRunner;
import org.apache.beam.sdk.testing.PAssert;
import org.apache.beam.sdk.testing.TestPipeline;
import org.apache.beam.sdk.transforms.display.DisplayData;
import org.apache.beam.sdk.transforms.windowing.FixedWindows;
import org.apache.beam.sdk.transforms.windowing.IntervalWindow;
import org.apache.beam.sdk.transforms.windowing.Window;
import org.apache.beam.sdk.values.PCollection;
import org.apache.beam.sdk.values.TimestampedValue;
import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Joiner;
import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableList;
import org.hamcrest.Matchers;
import org.joda.time.Duration;
import org.joda.time.Instant;
import org.junit.Rule;
import org.junit.Test;
import org.junit.experimental.categories.Category;
import org.junit.runner.RunWith;
import org.junit.runners.JUnit4;
import org.junit.runners.Parameterized;

/** Tests for Sample transform. */
public class SampleTest {
  private static final Integer[] EMPTY = new Integer[] {};
  private static final Integer[] DATA = new Integer[] {1, 2, 3, 4, 5};
  private static final Integer[] REPEATED_DATA = new Integer[] {1, 1, 2, 2, 3, 3, 4, 4, 5, 5};

  /** Test variations for Sample transform. */
  @RunWith(Parameterized.class)
  public static class PickAnyTest {
    @Rule public final transient TestPipeline p = TestPipeline.create();

    @Parameterized.Parameters(name = "limit_{1}")
    public static Iterable<Object[]> data() throws IOException {
      return ImmutableList.<Object[]>builder()
          .add(
              new Object[] {TestUtils.NO_LINES, 0},
              new Object[] {TestUtils.NO_LINES, 1},
              new Object[] {TestUtils.LINES, 1},
              new Object[] {TestUtils.LINES, TestUtils.LINES.size() / 2},
              new Object[] {TestUtils.LINES, TestUtils.LINES.size() * 2},
              new Object[] {TestUtils.LINES, TestUtils.LINES.size() - 1},
              new Object[] {TestUtils.LINES, TestUtils.LINES.size()},
              new Object[] {TestUtils.LINES, TestUtils.LINES.size() + 1})
          .build();
    }

    @SuppressWarnings("DefaultAnnotationParam")
    @Parameterized.Parameter(0)
    public List<String> lines;

    @Parameterized.Parameter(1)
    public int limit;

    private static class VerifyAnySample implements SerializableFunction<Iterable<String>, Void> {
      private final List<String> lines;
      private final int limit;

      private VerifyAnySample(List<String> lines, int limit) {
        this.lines = lines;
        this.limit = limit;
      }

      @Override
      public Void apply(Iterable<String> actualIter) {
        final int expectedSize = Math.min(limit, lines.size());

        // Make sure actual is the right length, and is a
        // subset of expected.
        List<String> actual = new ArrayList<>();
        for (String s : actualIter) {
          actual.add(s);
        }
        assertEquals(expectedSize, actual.size());
        Set<String> actualAsSet = new TreeSet<>(actual);
        Set<String> linesAsSet = new TreeSet<>(lines);
        assertEquals(actual.size(), actualAsSet.size());
        assertEquals(lines.size(), linesAsSet.size());
        assertTrue(linesAsSet.containsAll(actualAsSet));
        return null;
      }
    }

    void runPickAnyTest(final List<String> lines, int limit) {
      checkArgument(new HashSet<>(lines).size() == lines.size(), "Duplicates are unsupported.");

      PCollection<String> input = p.apply(Create.of(lines).withCoder(StringUtf8Coder.of()));

      PCollection<String> output = input.apply(Sample.any(limit));

      PAssert.that(output).satisfies(new VerifyAnySample(lines, limit));

      p.run();
    }

    @Test
    @Category(NeedsRunner.class)
    public void testPickAny() {
      runPickAnyTest(lines, limit);
    }

    @Test
    public void testCombineFn() {
      CombineFnTester.testCombineFn(
          Sample.combineFn(limit),
          lines,
          allOf(
              Matchers.<String>iterableWithSize(Math.min(lines.size(), limit)),
              everyItem(is(in(lines)))));
    }
  }

  /** Further tests for Sample transform. */
  @RunWith(JUnit4.class)
  public static class MiscTest {

    @Rule public final transient TestPipeline pipeline = TestPipeline.create();

    /**
     * Verifies that the result of a Sample operation contains the expected number of elements, and
     * that those elements are a subset of the items in expected.
     */
    @SuppressWarnings("rawtypes")
    public static class VerifyCorrectSample<T extends Comparable>
        implements SerializableFunction<Iterable<T>, Void> {
      private Object[] expectedValues;
      private int expectedSize;

      /**
       * expectedSize is the number of elements that the Sample should contain. expected is the set
       * of elements that the sample may contain.
       */
      @SafeVarargs
      VerifyCorrectSample(int expectedSize, T... expected) {
        this.expectedValues = expected;
        this.expectedSize = expectedSize;
      }

      /**
       * expectedSize is the number of elements that the Sample should contain. expected is the set
       * of elements that the sample may contain.
       */
      VerifyCorrectSample(int expectedSize, Collection<T> expected) {
        this.expectedValues = expected.toArray();
        this.expectedSize = expectedSize;
      }

      @Override
      @SuppressWarnings("unchecked")
      public Void apply(Iterable<T> in) {
        List<T> actual = new ArrayList<>();
        for (T elem : in) {
          actual.add(elem);
        }

        assertEquals(expectedSize, actual.size());

        Collections.sort(actual); // We assume that @expected is already sorted.
        int i = 0; // Index into @expected
        for (T s : actual) {
          boolean matchFound = false;
          for (; i < expectedValues.length; i++) {
            if (s.equals(expectedValues[i])) {
              matchFound = true;
              break;
            }
          }
          assertTrue("Invalid sample: " + Joiner.on(',').join(actual), matchFound);
          i++; // Don't match the same element again.
        }
        return null;
      }
    }

    private static TimestampedValue<Integer> tv(int i) {
      return TimestampedValue.of(i, new Instant(i * 1000));
    }

    @Test
    @Category(NeedsRunner.class)
    public void testSampleAny() {
      PCollection<Integer> input =
          pipeline
              .apply(
                  Create.timestamped(ImmutableList.of(tv(0), tv(1), tv(2), tv(3), tv(4), tv(5)))
                      .withCoder(BigEndianIntegerCoder.of()))
              .apply(Window.into(FixedWindows.of(Duration.standardSeconds(3))));
      PCollection<Integer> output = input.apply(Sample.any(2));

      PAssert.that(output)
          .inWindow(new IntervalWindow(new Instant(0), Duration.standardSeconds(3)))
          .satisfies(new VerifyCorrectSample<>(2, Arrays.asList(0, 1, 2)));
      PAssert.that(output)
          .inWindow(new IntervalWindow(new Instant(3000), Duration.standardSeconds(3)))
          .satisfies(new VerifyCorrectSample<>(2, Arrays.asList(3, 4, 5)));
      pipeline.run();
    }

    @Test
    @Category(NeedsRunner.class)
    public void testSampleAnyEmpty() {
      PCollection<Integer> input = pipeline.apply(Create.empty(BigEndianIntegerCoder.of()));
      PCollection<Integer> output =
          input
              .apply(Window.into(FixedWindows.of(Duration.standardSeconds(3))))
              .apply(Sample.any(2));

      PAssert.that(output).satisfies(new VerifyCorrectSample<>(0, EMPTY));
      pipeline.run();
    }

    @Test
    @Category(NeedsRunner.class)
    public void testSampleAnyZero() {
      PCollection<Integer> input =
          pipeline.apply(
              Create.timestamped(ImmutableList.of(tv(0), tv(1), tv(2), tv(3), tv(4), tv(5)))
                  .withCoder(BigEndianIntegerCoder.of()));
      PCollection<Integer> output =
          input
              .apply(Window.into(FixedWindows.of(Duration.standardSeconds(3))))
              .apply(Sample.any(0));

      PAssert.that(output)
          .inWindow(new IntervalWindow(new Instant(0), Duration.standardSeconds(3)))
          .satisfies(new VerifyCorrectSample<>(0, EMPTY));
      PAssert.that(output)
          .inWindow(new IntervalWindow(new Instant(3000), Duration.standardSeconds(3)))
          .satisfies(new VerifyCorrectSample<>(0, EMPTY));
      pipeline.run();
    }

    @Test
    @Category(NeedsRunner.class)
    public void testSampleAnyInsufficientElements() {
      PCollection<Integer> input = pipeline.apply(Create.empty(BigEndianIntegerCoder.of()));
      PCollection<Integer> output =
          input
              .apply(Window.into(FixedWindows.of(Duration.standardSeconds(3))))
              .apply(Sample.any(10));

      PAssert.that(output)
          .inWindow(new IntervalWindow(new Instant(0), Duration.standardSeconds(3)))
          .satisfies(new VerifyCorrectSample<>(0, EMPTY));
      pipeline.run();
    }

    @Test(expected = IllegalArgumentException.class)
    public void testSampleAnyNegative() {
      pipeline.enableAbandonedNodeEnforcement(false);
      pipeline.apply(Create.empty(BigEndianIntegerCoder.of())).apply(Sample.any(-10));
    }

    @Test
    @Category(NeedsRunner.class)
    public void testSample() {

      PCollection<Integer> input =
          pipeline.apply(
              Create.of(ImmutableList.copyOf(DATA)).withCoder(BigEndianIntegerCoder.of()));
      PCollection<Iterable<Integer>> output = input.apply(Sample.fixedSizeGlobally(3));

      PAssert.thatSingletonIterable(output).satisfies(new VerifyCorrectSample<>(3, DATA));
      pipeline.run();
    }

    @Test
    @Category(NeedsRunner.class)
    public void testSampleEmpty() {

      PCollection<Integer> input = pipeline.apply(Create.empty(BigEndianIntegerCoder.of()));
      PCollection<Iterable<Integer>> output = input.apply(Sample.fixedSizeGlobally(3));

      PAssert.thatSingletonIterable(output).satisfies(new VerifyCorrectSample<>(0, EMPTY));
      pipeline.run();
    }

    @Test
    @Category(NeedsRunner.class)
    public void testSampleZero() {

      PCollection<Integer> input =
          pipeline.apply(
              Create.of(ImmutableList.copyOf(DATA)).withCoder(BigEndianIntegerCoder.of()));
      PCollection<Iterable<Integer>> output = input.apply(Sample.fixedSizeGlobally(0));

      PAssert.thatSingletonIterable(output).satisfies(new VerifyCorrectSample<>(0, DATA));
      pipeline.run();
    }

    @Test
    @Category(NeedsRunner.class)
    public void testSampleInsufficientElements() {

      PCollection<Integer> input =
          pipeline.apply(
              Create.of(ImmutableList.copyOf(DATA)).withCoder(BigEndianIntegerCoder.of()));
      PCollection<Iterable<Integer>> output = input.apply(Sample.fixedSizeGlobally(10));

      PAssert.thatSingletonIterable(output).satisfies(new VerifyCorrectSample<>(5, DATA));
      pipeline.run();
    }

    @Test(expected = IllegalArgumentException.class)
    public void testSampleNegative() {
      pipeline.enableAbandonedNodeEnforcement(false);

      PCollection<Integer> input =
          pipeline.apply(
              Create.of(ImmutableList.copyOf(DATA)).withCoder(BigEndianIntegerCoder.of()));
      input.apply(Sample.fixedSizeGlobally(-1));
    }

    @Test
    @Category(NeedsRunner.class)
    public void testSampleMultiplicity() {

      PCollection<Integer> input =
          pipeline.apply(
              Create.of(ImmutableList.copyOf(REPEATED_DATA)).withCoder(BigEndianIntegerCoder.of()));
      // At least one value must be selected with multiplicity.
      PCollection<Iterable<Integer>> output = input.apply(Sample.fixedSizeGlobally(6));

      PAssert.thatSingletonIterable(output).satisfies(new VerifyCorrectSample<>(6, REPEATED_DATA));
      pipeline.run();
    }

    @Test
    public void testSampleGetName() {
      assertEquals("Sample.Any", Sample.<String>any(1).getName());
    }

    @Test
    public void testDisplayData() {
      PTransform<?, ?> sampleAny = Sample.any(1234);
      DisplayData sampleAnyDisplayData = DisplayData.from(sampleAny);
      assertThat(sampleAnyDisplayData, hasDisplayItem("sampleSize", 1234));

      PTransform<?, ?> samplePerKey = Sample.fixedSizePerKey(2345);
      DisplayData perKeyDisplayData = DisplayData.from(samplePerKey);
      assertThat(perKeyDisplayData, hasDisplayItem("sampleSize", 2345));
    }
  }
}
