/*
 * Copyright (c) 2018-2023 HERE Europe B.V.
 *
 * Licensed 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.
 */

import com.here.hrn.HRN;
import com.here.olp.util.geo.BoundingBox;
import com.here.olp.util.quad.factory.HereQuadFactory;
import com.here.platform.data.client.spark.javadsl.JavaLayerDataFrameReader;
import com.here.platform.data.client.spark.javadsl.JavaLayerDataFrameWriter;
import com.here.platform.data.client.spark.javadsl.VersionedDataConverter;
import com.here.platform.data.client.spark.scaladsl.GroupedData;
import com.here.platform.data.client.spark.scaladsl.VersionedRowMetadata;
import com.here.platform.pipeline.PipelineContext;
import com.here.schema.rib.v2.TopologyGeometry.Node;
import com.here.schema.rib.v2.TopologyGeometryPartitionOuterClass;
import com.here.schema.rib.v2.TopologyGeometryPartitionOuterClass.TopologyGeometryPartition;
import com.typesafe.config.ConfigBeanFactory;
import java.io.Serializable;
import java.util.Iterator;
import java.util.List;
import java.util.stream.Collectors;
import org.apache.spark.api.java.function.MapFunction;
import org.apache.spark.sql.Dataset;
import org.apache.spark.sql.Encoders;
import org.apache.spark.sql.Row;
import org.apache.spark.sql.SparkSession;

public class SparkConnectorLocal {

  // Class representing the number of nodes in a level 14 sub-tile
  public static class NodeCount {
    private long tileId;
    private int count;

    public NodeCount(long tileId, int count) {
      this.tileId = tileId;
      this.count = count;
    }

    public String toGeoJson() {
      BoundingBox box = HereQuadFactory.INSTANCE.getMapQuadByLongKey(tileId).getBoundingBox();
      double e = box.getEast();
      double w = box.getWest();
      double n = box.getNorth();
      double s = box.getSouth();

      // from black (0 nodes) to red (500 nodes)
      String color = String.format("rgb(%d,0,0)", (int) (Math.min(count, 500) / 500.0 * 255));

      return String.format(
          "{\"type\":\"Feature\","
              + "\"geometry\":{\"type\":\"Polygon\",\"coordinates\":[[[%f, %f], [%f, %f], [%f, %f], [%f, %f], [%f, %f]]]},"
              + "\"properties\":{\"description\":{\"tileId\":%d,\"count\":%d},\"style\":{\"color\":\"%s\"}}}",
          w, s, e, s, e, n, w, n, w, s, tileId, count, color);
    }
  }

  // Class representing the decoded content of an output partition
  public static class NodeCountPartition {
    private String partitionName;
    private List<NodeCount> counts;

    public NodeCountPartition(String partitionName, List<NodeCount> counts) {
      this.partitionName = partitionName;
      this.counts = counts;
    }

    public String toGeoJson() {
      return String.format(
          "{\"type\":\"FeatureCollection\",\"features\":[%s]}",
          counts.stream().map(NodeCount::toGeoJson).collect(Collectors.joining(",")));
    }

    public String getPartitionName() {
      return partitionName;
    }
  }

  // Class representing an encoded output partition
  public static class OutputData implements Serializable {
    private String mt_partition;
    private byte[] data;

    public OutputData(String mt_partition, byte[] data) {
      this.mt_partition = mt_partition;
      this.data = data;
    }

    // Minimal bean interface - required to use Encoders.bean
    public String getMt_partition() {
      return mt_partition;
    }

    public void setMt_partition(String mt_partition) {
      this.mt_partition = mt_partition;
    }

    public byte[] getData() {
      return data;
    }

    public void setData(byte[] data) {
      this.data = data;
    }
  }

  // We expose two static methods, `main` and `run`. `main` reads the pipeline context and the
  // configuration from `pipeline-config.conf` and `application.conf`, while `run` accepts them
  // as parameters. This makes it convenient to test the logic later using `run` with different
  // configuration parameters, without having to pass them through the classpath.
  public static void main(String[] args) {
    // Read the default pipeline context
    PipelineContext pipelineContext = new PipelineContext();

    // Read the bounding box configured in `application.conf`
    BoundingBox boundingBox =
        ConfigBeanFactory.create(
            pipelineContext.applicationConfig().getConfig("tutorial.boundingBox"),
            BoundingBox.class);

    run(pipelineContext, boundingBox);
  }

  public static void run(PipelineContext pipelineContext, BoundingBox bbox) {
    // Defines the input / output catalogs to be read / written to
    HRN inputHrn = pipelineContext.getConfig().getInputCatalogs().get("here-map-content");
    HRN outputHrn = pipelineContext.getConfig().getOutputCatalog();

    // Input / output layers
    String topologyGeometryLayer = "topology-geometry";
    String outputLayer = "node-count";

    SparkSession sparkSession =
        SparkSession.builder().appName("SparkTopologyNodeCount").getOrCreate();

    // Read the input data as a Dataset of topology partitions
    Dataset<TopologyGeometryPartitionOuterClass.TopologyGeometryPartition> topologyGeometryData =
        JavaLayerDataFrameReader.create(sparkSession)
            .readLayer(inputHrn, topologyGeometryLayer)
            .query(
                String.format(
                    "mt_partition=inboundingbox=(%f,%f,%f,%f)",
                    bbox.getNorth(), bbox.getSouth(), bbox.getEast(), bbox.getWest()))
            .option("olp.connector.force-raw-data", true)
            .load()
            .select("data")
            .as(Encoders.BINARY())
            .map(
                (MapFunction<byte[], TopologyGeometryPartition>)
                    TopologyGeometryPartition::parseFrom,
                Encoders.kryo(TopologyGeometryPartitionOuterClass.TopologyGeometryPartition.class));

    // Compute the output partitions
    Dataset<NodeCountPartition> nodeCountPartitions =
        topologyGeometryData.map(
            (MapFunction<TopologyGeometryPartition, NodeCountPartition>)
                partition -> {
                  List<NodeCount> counts =
                      partition
                          .getNodeList()
                          .stream()
                          .map(Node::getGeometry)
                          .map(
                              g ->
                                  HereQuadFactory.INSTANCE
                                      .getMapQuadByLocation(g.getLatitude(), g.getLongitude(), 14)
                                      .getLongKey())
                          .collect(Collectors.groupingBy(t -> t))
                          .entrySet()
                          .stream()
                          .map(e -> new NodeCount(e.getKey(), e.getValue().size()))
                          .collect(Collectors.toList());
                  return new NodeCountPartition(partition.getPartitionName(), counts);
                },
            Encoders.kryo(NodeCountPartition.class));

    // Encode the output partitions, convert to DataFrame and publish the output data
    JavaLayerDataFrameWriter.create(
            nodeCountPartitions
                .map(
                    (MapFunction<NodeCountPartition, OutputData>)
                        p -> new OutputData(p.getPartitionName(), p.toGeoJson().getBytes()),
                    Encoders.bean(OutputData.class))
                .toDF())
        .writeLayer(outputHrn, outputLayer)
        .withDataConverter(
            new VersionedDataConverter() {
              @Override
              public GroupedData<VersionedRowMetadata> serializeGroup(
                  VersionedRowMetadata rowMetadata, Iterator<Row> rows) {
                return new GroupedData<>(rowMetadata, rows.next().getAs("data"));
              }
            })
        .save();

    sparkSession.stop();
  }
}

mvn compile exec:java \
    -Dexec.mainClass=SparkConnectorLocal \
    -Dpipeline-config.file=pipeline-config-cn.conf \
    -Dconfig.file=application-cn.conf \
    -Dspark.master="local[*]"

mvn compile exec:java \
    -Dexec.mainClass=SparkConnectorLocal \
    -Dpipeline-config.file=pipeline-config-local.conf \
    -Dconfig.file=application-cn.conf \
    -Dspark.master="local[*]" \
    -Dhere.platform.data-client.endpoint-locator.discovery-service-env=local

/*
 * Copyright (c) 2018-2023 HERE Europe B.V.
 *
 * Licensed 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.
 */

import static junit.framework.TestCase.assertEquals;

import com.here.hrn.HRN;
import com.here.olp.util.geo.BoundingBox;
import com.here.olp.util.quad.factory.HereQuadFactory;
import com.here.platform.data.client.engine.javadsl.DataEngine;
import com.here.platform.data.client.engine.javadsl.ReadEngine;
import com.here.platform.data.client.javadsl.DataClient;
import com.here.platform.data.client.javadsl.NewPartition;
import com.here.platform.data.client.javadsl.Partition;
import com.here.platform.data.client.javadsl.QueryApi;
import com.here.platform.data.client.model.*;
import com.here.platform.data.client.spark.DataClientSparkContextUtils;
import com.here.platform.pipeline.PipelineConfig;
import com.here.platform.pipeline.PipelineContext;
import com.here.schema.geometry.v2.GeometryOuterClass;
import com.here.schema.rib.v2.TopologyGeometry;
import com.here.schema.rib.v2.TopologyGeometryPartitionOuterClass;
import java.util.*;
import java.util.concurrent.CompletionStage;
import java.util.regex.Matcher;
import java.util.regex.Pattern;
import java.util.stream.Collectors;
import java.util.stream.IntStream;
import java.util.stream.StreamSupport;
import org.junit.Test;

public class SparkConnectorLocalTest {

  // Simple Pair class used to collect maps with Java streams.
  static class Pair<K, V> {
    private K key;
    private V value;

    public Pair(K key, V value) {
      this.key = key;
      this.value = value;
    }

    public K getKey() {
      return key;
    }

    public V getValue() {
      return value;
    }
  }

  // Specify the bounding box used by our pipeline under test.
  private BoundingBox testBoundingBox = new BoundingBox(52.67551, 52.33826, 13.76116, 13.08835);

  // Utilities to create a random topology node within a bounding box.
  private Random random = new Random(42);

  private double randomInRange(double min, double max) {
    return min + (max - min) * random.nextDouble();
  }

  private TopologyGeometry.Node randomNode(BoundingBox boundingBox) {
    return TopologyGeometry.Node.newBuilder()
        .setGeometry(
            GeometryOuterClass.Point.newBuilder()
                .setLatitude(randomInRange(boundingBox.getSouth(), boundingBox.getNorth()))
                .setLongitude(randomInRange(boundingBox.getWest(), boundingBox.getEast()))
                .build())
        .build();
  }

  // Utility to create a catalog with a single HERE Tile versioned layer. The catalog will be local
  // because we use `here.platform.data-client.endpoint-locator.discovery-service-env=local` in
  // src/test/resources/application.conf. It will be kept in memory only, because we also use
  // `here.platform.data-local.memory-mode=true`.
  private HRN createCatalog(String catalogIdPrefix, String layerId) {
    return await(
        DataClient.get(DataClientSparkContextUtils.context().actorSystem())
            .adminApi()
            .createCatalog(
                new WritableCatalogConfiguration.Builder()
                    .withId(catalogIdPrefix + "-" + UUID.randomUUID().toString())
                    .withName("name")
                    .withSummary("summary")
                    .withDescription("description")
                    .withLayers(
                        Collections.singletonList(
                            new WritableLayer.Builder()
                                .withId(layerId)
                                .withName("name")
                                .withSummary("summary")
                                .withDescription("description")
                                .withLayerType(LayerTypes.Versioned())
                                .withPartitioning(
                                    Partitioning.HereTile()
                                        .withTileLevels(Collections.singletonList(12))
                                        .build())
                                .withVolume(Volumes.Durable())
                                .withContentType("application/octet-stream")))
                    .build()));
  }

  // Utility to publish some test data in a catalog.
  private void writeData(HRN hrn, String layer, Map<String, byte[]> data) {
    Iterator<NewPartition> partitions =
        data.entrySet()
            .stream()
            .map(
                e ->
                    new NewPartition.Builder()
                        .withPartition(e.getKey())
                        .withLayer(layer)
                        .withData(e.getValue())
                        .build())
            .collect(Collectors.toList())
            .iterator();

    await(
        DataEngine.get(DataClientSparkContextUtils.context().actorSystem())
            .writeEngine(hrn)
            .publishBatch2(
                4,
                Optional.of(Collections.singletonList(layer)),
                Collections.emptyList(),
                partitions));
  }

  // Utility to read the data from the latest version of a catalog.
  private Map<String, byte[]> readData(HRN hrn, String layer) {
    ReadEngine readEngine =
        DataEngine.get(DataClientSparkContextUtils.context().actorSystem()).readEngine(hrn);
    QueryApi queryApi =
        DataClient.get(DataClientSparkContextUtils.context().actorSystem()).queryApi(hrn);

    Iterable<Partition> partitions =
        () ->
            await(
                queryApi
                    .getLatestVersion(OptionalLong.empty())
                    .thenApply(
                        v -> v.orElseThrow(() -> new RuntimeException("Output catalog is empty")))
                    .thenCompose(
                        v -> queryApi.getPartitionsAsIterator(v, layer, Collections.emptySet())));

    return StreamSupport.stream(partitions.spliterator(), false)
        .map(p -> new Pair<>(p.getPartition(), await(readEngine.getDataAsBytes(p))))
        .collect(Collectors.toMap(Pair::getKey, Pair::getValue));
  }

  // Utility to create an input test catalog given the expected outcome. `expectedCounts` is a map
  // containing, for each level 14 tile ID, the number of nodes contained in that sub-tile.
  // This method creates the requested number of topology nodes in each level-14 tile, encodes them
  // in a topology partition and publishes them in a freshly created test catalog.
  private HRN createInputCatalog(Map<Long, Integer> expectedNodeCount) {
    HRN hrn = createCatalog("input", "topology-geometry");

    Map<String, byte[]> partitions =
        expectedNodeCount
            .entrySet()
            .stream()
            .map(
                e -> {
                  BoundingBox boundingBox =
                      HereQuadFactory.INSTANCE.getMapQuadByLongKey(e.getKey()).getBoundingBox();
                  List<TopologyGeometry.Node> nodes =
                      IntStream.rangeClosed(1, e.getValue())
                          .mapToObj(notused -> randomNode(boundingBox))
                          .collect(Collectors.toList());
                  return new Pair<>(e.getKey(), nodes);
                })
            .collect(
                Collectors.groupingBy(
                    p ->
                        HereQuadFactory.INSTANCE
                            .getMapQuadByLongKey(p.getKey())
                            .getAncestor(12)
                            .getLongKey()))
            .entrySet()
            .stream()
            .map(
                e -> {
                  byte[] data =
                      TopologyGeometryPartitionOuterClass.TopologyGeometryPartition.newBuilder()
                          .setPartitionName(String.valueOf(e.getKey()))
                          .addAllNode(
                              e.getValue()
                                  .stream()
                                  .flatMap(p -> p.getValue().stream())
                                  .collect(Collectors.toList()))
                          .build()
                          .toByteArray();

                  return new Pair<>(String.valueOf(e.getKey()), data);
                })
            .collect(Collectors.toMap(Pair::getKey, Pair::getValue));

    writeData(hrn, "topology-geometry", partitions);

    return hrn;
  }

  // Utility to create an output catalog for the pipeline.
  private HRN createOutputCatalog() {
    return createCatalog("output", "node-count");
  }

  private <T> T await(CompletionStage<T> stage) {
    try {
      return stage.toCompletableFuture().get();
    } catch (Exception e) {
      throw new RuntimeException(e);
    }
  }

  // Utility to retrieve the actual node count per level 14 tile in the output catalog. It reads the
  // latest version of the catalog, retrieves the data, decodes it and returns the node count
  // encoded in each level 14 tile. This is used, after a run of the pipeline under test, to verify
  // that the node count is the expected on.
  private Map<Long, Integer> readActualNodeCount(HRN hrn) {

    // We use regular expressions to extract the tileId of the level 14 tiles and their node count.
    // Using a proper JSON parser is left as an exercise.
    Pattern descriptionRegex =
        Pattern.compile("\"description\":\\s*\\{\"tileId\":\\s*(\\d+),\\s*\"count\":\\s*(\\d+)}");

    return readData(hrn, "node-count")
        .entrySet()
        .stream()
        .flatMap(
            e -> {
              String geoJson = new String(e.getValue());
              Matcher matcher = descriptionRegex.matcher(geoJson);
              List<Pair<Long, Integer>> matches = new ArrayList<>();
              while (matcher.find()) {
                matches.add(
                    new Pair<>(
                        Long.parseLong(matcher.group(1)), Integer.parseInt(matcher.group(2))));
              }
              return matches.stream();
            })
        .collect(Collectors.toMap(Pair::getKey, Pair::getValue));
  }

  private Map<Long, Integer> initializeExpectedNodeCount() {
    return StreamSupport.stream(
            HereQuadFactory.INSTANCE
                .iterableBoundingBoxToMapQuad(testBoundingBox, 14)
                .spliterator(),
            false)
        .map(quad -> new Pair<>(quad.getLongKey(), Math.abs(random.nextInt()) % 100))
        .filter(pair -> pair.getValue() != 0)
        .collect(Collectors.toMap(Pair::getKey, Pair::getValue));
  }

  // Randomly generated expected result. For each level 14 tile in `testBoundingBox` we generate a
  // random non-negative node count. We filter out zeros because those are by design not encoded by
  // the pipeline under test.
  private Map<Long, Integer> expectedNodeCount = initializeExpectedNodeCount();

  @Test
  public void correctNodeCountTest() {
    // Create the input catalog and publish the test data.
    HRN inputHrn = createInputCatalog(expectedNodeCount);
    // Create an empty output catalog.
    HRN outputHrn = createOutputCatalog();

    PipelineContext pipelineContext =
        new PipelineContext(
            new PipelineConfig(outputHrn, Collections.singletonMap("here-map-content", inputHrn)),
            Optional.empty());

    // Run the pipeline under test.
    SparkConnectorLocal.run(pipelineContext, testBoundingBox);

    // Check the output data.
    assertEquals(readActualNodeCount(outputHrn), expectedNodeCount);
  }
}