/*
 * 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.spark.shuffle.celeborn;

import java.util.*;

import org.apache.commons.lang3.tuple.Pair;
import org.junit.Assert;
import org.junit.Test;

import org.apache.celeborn.common.protocol.PartitionLocation;
import org.apache.celeborn.common.protocol.StorageInfo;

public class CelebornPartitionUtilSuiteJ {
  @Test
  public void testSkewPartitionSplitIgnoreEmpty1() {
    ArrayList<PartitionLocation> locations = new ArrayList<>();
    PartitionLocation loc = genPartitionLocation(0, new Long[] {0L, 20L, 30L});
    List<Long> offsets = loc.getStorageInfo().getChunkOffsets();
    locations.add(loc);
    int partitionNumber = 10;
    Long sum = 0L;
    for (int i = 0; i < partitionNumber; i++) {
      Map<String, Pair<Integer, Integer>> res =
          CelebornPartitionUtil.splitSkewedPartitionLocations(locations, partitionNumber, i);
      if (!res.isEmpty()) {
        int l = res.get(loc.getUniqueId()).getLeft();
        int r = res.get(loc.getUniqueId()).getRight();
        Assert.assertTrue(r >= l);
        for (int j = l; j <= r; j++) {
          sum += offsets.get(j + 1) - offsets.get(j);
        }
      }
    }
    Assert.assertEquals(offsets.get(offsets.size() - 1), sum);
  }

  @Test
  public void testSkewPartitionSplit() {

    ArrayList<PartitionLocation> locations = new ArrayList<>();
    for (int i = 0; i < 13; i++) {
      locations.add(genPartitionLocation(i, new Long[] {0L, 100L, 200L, 300L, 500L, 1000L}));
    }
    locations.add(genPartitionLocation(91, new Long[] {0L, 1L}));

    int subPartitionSize = 3;

    Map<String, Pair<Integer, Integer>> result1 =
        CelebornPartitionUtil.splitSkewedPartitionLocations(locations, subPartitionSize, 0);
    Map<String, Pair<Integer, Integer>> expectResult1 =
        genRanges(
            new Object[][] {
              {"0-0", 0, 4},
              {"0-1", 0, 4},
              {"0-10", 0, 4},
              {"0-11", 0, 4},
              {"0-12", 0, 2}
            });
    Assert.assertEquals(expectResult1, result1);

    Map<String, Pair<Integer, Integer>> result2 =
        CelebornPartitionUtil.splitSkewedPartitionLocations(locations, subPartitionSize, 1);
    Map<String, Pair<Integer, Integer>> expectResult2 =
        genRanges(
            new Object[][] {
              {"0-12", 3, 4},
              {"0-2", 0, 4},
              {"0-3", 0, 4},
              {"0-4", 0, 4},
              {"0-5", 0, 3}
            });
    Assert.assertEquals(expectResult2, result2);

    Map<String, Pair<Integer, Integer>> result3 =
        CelebornPartitionUtil.splitSkewedPartitionLocations(locations, subPartitionSize, 2);
    Map<String, Pair<Integer, Integer>> expectResult3 =
        genRanges(
            new Object[][] {
              {"0-5", 4, 4},
              {"0-6", 0, 4},
              {"0-7", 0, 4},
              {"0-8", 0, 4},
              {"0-9", 0, 4},
              {"0-91", 0, 0}
            });
    Assert.assertEquals(expectResult3, result3);
  }

  @Test
  public void testBoundary() {
    ArrayList<PartitionLocation> locations = new ArrayList<>();
    locations.add(genPartitionLocation(0, new Long[] {0L, 100L, 200L, 300L, 400L, 500L}));

    for (int i = 0; i < 5; i++) {
      Map<String, Pair<Integer, Integer>> result =
          CelebornPartitionUtil.splitSkewedPartitionLocations(locations, 5, i);
      Map<String, Pair<Integer, Integer>> expectResult = genRanges(new Object[][] {{"0-0", i, i}});
      Assert.assertEquals(expectResult, result);
    }
  }

  @Test
  public void testSplitStable() {
    ArrayList<PartitionLocation> locations = new ArrayList<>();
    for (int i = 0; i < 13; i++) {
      locations.add(genPartitionLocation(i, new Long[] {0L, 100L, 200L, 300L, 500L, 1000L}));
    }
    locations.add(genPartitionLocation(91, new Long[] {0L, 1L}));

    Collections.shuffle(locations);

    Map<String, Pair<Integer, Integer>> result =
        CelebornPartitionUtil.splitSkewedPartitionLocations(locations, 3, 0);
    Map<String, Pair<Integer, Integer>> expectResult =
        genRanges(
            new Object[][] {
              {"0-0", 0, 4},
              {"0-1", 0, 4},
              {"0-10", 0, 4},
              {"0-11", 0, 4},
              {"0-12", 0, 2}
            });
    Assert.assertEquals(expectResult, result);
  }

  private ArrayList<PartitionLocation> genPartitionLocations(Map<Integer, Long[]> epochToOffsets) {
    ArrayList<PartitionLocation> locations = new ArrayList<>();
    epochToOffsets.forEach(
        (epoch, offsets) -> {
          PartitionLocation location =
              new PartitionLocation(
                  0, epoch, "localhost", 0, 0, 0, 0, PartitionLocation.Mode.PRIMARY);
          StorageInfo storageInfo =
              new StorageInfo(
                  StorageInfo.Type.HDD,
                  "mountPoint",
                  false,
                  "filePath",
                  StorageInfo.LOCAL_DISK_MASK,
                  1,
                  Arrays.asList(offsets));
          location.setStorageInfo(storageInfo);
          locations.add(location);
        });
    return locations;
  }

  private PartitionLocation genPartitionLocation(int epoch, Long[] offsets) {
    PartitionLocation location =
        new PartitionLocation(0, epoch, "localhost", 0, 0, 0, 0, PartitionLocation.Mode.PRIMARY);
    StorageInfo storageInfo =
        new StorageInfo(
            StorageInfo.Type.HDD,
            "mountPoint",
            false,
            "filePath",
            StorageInfo.LOCAL_DISK_MASK,
            offsets[offsets.length - 1],
            Arrays.asList(offsets));
    location.setStorageInfo(storageInfo);
    return location;
  }

  private Map<String, Pair<Integer, Integer>> genRanges(Object[][] inputs) {
    Map<String, Pair<Integer, Integer>> ranges = new HashMap<>();
    for (Object[] idToChunkRange : inputs) {
      String uid = (String) idToChunkRange[0];
      Pair<Integer, Integer> range = Pair.of((int) idToChunkRange[1], (int) idToChunkRange[2]);
      ranges.put(uid, range);
    }
    return ranges;
  }
}
