/*
 * 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 org.apache.spark.sql.functions.col;
import static org.apache.spark.sql.functions.collect_list;
import static org.apache.spark.sql.functions.udf;

import akka.actor.ActorSystem;
import akka.japi.Pair;
import com.here.hrn.HRN;
import com.here.platform.data.client.base.javadsl.BaseClient;
import com.here.platform.data.client.base.javadsl.BaseClientJava;
import com.here.platform.data.client.engine.javadsl.DataEngine;
import com.here.platform.data.client.engine.javadsl.WriteEngine;
import com.here.platform.data.client.javadsl.CommitPartition;
import com.here.platform.data.client.javadsl.DataClient;
import com.here.platform.data.client.javadsl.NewPartition;
import com.here.platform.data.client.javadsl.PublishApi;
import com.here.platform.data.client.javadsl.QueryApi;
import com.here.platform.data.client.model.BatchToken;
import com.here.platform.data.client.model.VersionDependency;
import com.here.platform.data.client.spark.DataClientSparkContextUtils;
import com.here.platform.data.client.spark.javadsl.JavaLayerDataFrameReader;
import com.here.platform.location.core.geospatial.GeoCoordinate;
import com.here.platform.location.core.mapmatching.MatchResult;
import com.here.platform.location.core.mapmatching.NoTransition;
import com.here.platform.location.core.mapmatching.OnRoad;
import com.here.platform.location.core.mapmatching.javadsl.MatchResults;
import com.here.platform.location.core.mapmatching.javadsl.PathMatcher;
import com.here.platform.location.inmemory.graph.Vertex;
import com.here.platform.location.integration.optimizedmap.OptimizedMapLayers;
import com.here.platform.location.integration.optimizedmap.dcl2.javadsl.OptimizedMapCatalog;
import com.here.platform.location.integration.optimizedmap.mapmatching.javadsl.PathMatchers;
import com.here.platform.pipeline.PipelineContext;
import java.io.Serializable;
import java.util.*;
import java.util.concurrent.ExecutionException;
import java.util.stream.Collectors;
import java.util.stream.StreamSupport;
import org.apache.spark.api.java.JavaRDD;
import org.apache.spark.sql.Dataset;
import org.apache.spark.sql.Encoder;
import org.apache.spark.sql.Encoders;
import org.apache.spark.sql.Row;
import org.apache.spark.sql.SparkSession;
import org.apache.spark.sql.api.java.UDF1;
import org.apache.spark.sql.expressions.UserDefinedFunction;
import org.apache.spark.sql.types.ArrayType;
import org.apache.spark.sql.types.DataTypes;
import scala.collection.JavaConverters;
import scala.collection.mutable.WrappedArray;

/**
 * Read sensor data from an SDII archive and output the map-matched paths to a GeoJSON output
 * catalog with the same tiling scheme.
 */
public class BatchPathMatchingJava {

  private static final PipelineContext pipelineContext = new PipelineContext();

  private static final HRN OUTPUT_CATALOG_HRN = pipelineContext.getConfig().getOutputCatalog();
  private static final String OUTPUT_LAYER = "matched-trips";

  private static final HRN SENSOR_DATA_ARCHIVE_HRN =
      pipelineContext.getConfig().getInputCatalogs().get("sensor");
  private static final String SENSOR_DATA_ARCHIVE_VERSION_LAYER_NAME = "sample-index-layer";

  private static final HRN LOCATION_CATALOG_HRN =
      pipelineContext.getConfig().getInputCatalogs().get("location");
  private static final Long LOCATION_CATALOG_VERSION =
      pipelineContext.getJob().get().getInputCatalogs().get("location").version();

  // Since it is not convenient to recreate an instance of BaseClient, it's connection pool
  // and OptimizedMapCatalog caches, they are stored in a singleton and called from it.
  private enum OptimizedMapLayersSingleton {
    INSTANCE;

    private final OptimizedMapLayers optimizedMap;

    OptimizedMapLayersSingleton() {
      BaseClient baseClient = BaseClientJava.instance();
      optimizedMap =
          OptimizedMapCatalog.from(LOCATION_CATALOG_HRN)
              .usingBaseClient(baseClient)
              .newInstance()
              .version(LOCATION_CATALOG_VERSION);
    }

    private OptimizedMapLayers optimizedMap() {
      return optimizedMap;
    }
  }

  public static void main(String[] args) throws InterruptedException, ExecutionException {
    SparkSession sparkSession =
        SparkSession.builder().appName("BatchPathMatchingPipeline").getOrCreate();

    Dataset<SensorData> sensorData = getSensorData(sparkSession);
    JavaRDD<Pair<Long, String>> matchedPaths = matchPaths(sensorData);
    publish(matchedPaths);

    sparkSession.stop();
  }

  /** Retrieves the contents of the sensor data archive partitions. */
  private static Dataset<SensorData> getSensorData(SparkSession sparkSession) {

    // Extracting input data.
    // There are two different options for the query parameter: `hour_from` and `tile_id`.
    // You can either make it `hour_from>0` to get all available messages or make it
    // `tile_id==$tile_id` to get messages by the specific partition.
    // Possible tile_ids = [377893756, 377894443, 377894442, ...].
    // For CN environment = [389695267, 389696772, 389695401, ...].
    // You can implement more complex queries with RSQL
    // (https://here-tech.skawa.fun/documentation/data-client-library/dev_guide/client/rsql.html)
    Dataset<Row> sdiiMessages =
        JavaLayerDataFrameReader.create(sparkSession)
            .readLayer(SENSOR_DATA_ARCHIVE_HRN, SENSOR_DATA_ARCHIVE_VERSION_LAYER_NAME)
            .query("hour_from>0")
            .load();

    // Define structure of List of coordinate pairs - Spark uses it to encode data
    ArrayType schema = DataTypes.createArrayType(DataTypes.StringType);

    // UDF to extract only coordinates from the complex structure of path
    UserDefinedFunction positionEstimateToCoordinatePair =
        udf(
            (UDF1<WrappedArray<Row>, List<String>>)
                positionEstimateList ->
                    JavaConverters.seqAsJavaList(positionEstimateList)
                        .stream()
                        .map(
                            positionEstimate -> {
                              Double lat = positionEstimate.getAs("latitude_deg");
                              Double lon = positionEstimate.getAs("longitude_deg");
                              return lat + "," + lon;
                            })
                        .collect(Collectors.toList()),
            schema);

    Encoder<SensorData> sensorDataEncoder = Encoders.bean(SensorData.class);
    // Extract list of paths (each path represented as a list of coordinates)
    return sdiiMessages
        .select(
            col("idx_tile_id").as("tileId"),
            col("path.positionEstimate").as("fullPositionEstimate"))
        .withColumn(
            "positionEstimate", positionEstimateToCoordinatePair.apply(col("fullPositionEstimate")))
        .groupBy("tileId")
        .agg(collect_list("positionEstimate").as("positionEstimateList"))
        .as(sensorDataEncoder);
  }

  /**
   * Map match each path inside each sensor archive tile
   * (https://en.wikipedia.org/wiki/Map_matching) and return a list of lists of (lon,lat) string
   * pairs. Each element in the top-level list represents a path as a list of coordinates. Convert
   * each path to GeoJson.
   */
  private static JavaRDD<Pair<Long, String>> matchPaths(Dataset<SensorData> sensorData) {
    JavaRDD<Pair<Long, List<List<GeoCoordinate>>>> matchedTrips =
        sensorData
            .javaRDD()
            .mapPartitions(
                tiles -> {
                  OptimizedMapLayers optimizedMap =
                      OptimizedMapLayersSingleton.INSTANCE.optimizedMap();
                  PathMatcher<GeoCoordinate, Vertex, NoTransition> pathMatcher =
                      new PathMatchers(optimizedMap).carPathMatcherWithoutTransitions();
                  Iterable<SensorData> tileIterable = () -> tiles;
                  return StreamSupport.stream(tileIterable.spliterator(), false)
                      .map(
                          sensorDataElement ->
                              new Pair<>(
                                  sensorDataElement.getTileId(),
                                  getCoordinatesMatchedToPath(
                                      pathMatcher, sensorDataElement.getPositionEstimateList())))
                      .iterator();
                });

    return convertToGeoJSON(matchedTrips);
  }

  /** Using path matcher get new list of lists of GeoCoordinates, each list representing a path. */
  private static List<List<GeoCoordinate>> getCoordinatesMatchedToPath(
      PathMatcher<GeoCoordinate, Vertex, NoTransition> pathMatcher,
      List<List<String>> positionEstimatePathList) {
    return positionEstimatePathList
        .stream()
        .map(BatchPathMatchingJava::convertToGeoCoordinates)
        .map(
            positionEstimatePath ->
                pathMatcher
                    .matchPath(positionEstimatePath)
                    .results()
                    .stream()
                    .map(BatchPathMatchingJava::matchResultOnRoad)
                    .filter(Objects::nonNull)
                    .collect(Collectors.toList()))
        .collect(Collectors.toList());
  }

  /** Convert list of string pairs, representing coordinates to actual GeoCoordinates */
  private static List<GeoCoordinate> convertToGeoCoordinates(List<String> coordinatePairs) {
    return coordinatePairs
        .stream()
        .map(
            coordinatePair -> {
              String[] coordinates = coordinatePair.split(",", 2);
              double latitude = Double.parseDouble(coordinates[0]);
              double longitude = Double.parseDouble(coordinates[1]);
              return new GeoCoordinate(latitude, longitude);
            })
        .collect(Collectors.toList());
  }

  /** Construct new GeoCoordinates for on-road matches. */
  private static GeoCoordinate matchResultOnRoad(MatchResult<Vertex> matchResult) {
    if (MatchResults.isOnRoad(matchResult)) {
      OnRoad<Vertex> onRoad = (OnRoad<Vertex>) matchResult;
      GeoCoordinate nearest = onRoad.elementProjection().nearest();
      return new GeoCoordinate(nearest.latitude(), nearest.longitude());
    }
    return null;
  }

  /** Convert the trips to GeoJson; each trip is converted to a linestring. */
  private static JavaRDD<Pair<Long, String>> convertToGeoJSON(
      JavaRDD<Pair<Long, List<List<GeoCoordinate>>>> matchedTrips) {
    return matchedTrips.map(
        matchedTrip -> {
          String featureCollection =
              matchedTrip
                  .second()
                  .stream()
                  .map(BatchPathMatchingJava::mapToFeature)
                  .collect(
                      Collectors.joining(
                          ",", "{ \"type\": \"FeatureCollection\", \"features\": [", "]}"));
          return new Pair<>(matchedTrip.first(), featureCollection);
        });
  }

  /** Convert path to GeoJson Feature */
  private static String mapToFeature(List<GeoCoordinate> coordinates) {
    StringJoiner coordJoiner = new StringJoiner(",", "[", "]");
    coordinates
        .stream()
        .map(c -> "[" + c.getLongitude() + ", " + c.getLatitude() + "]")
        .forEach(coordJoiner::add);
    return "{ \"type\": \"Feature\", \"geometry\": "
        + "{ \"type\": \"LineString\", \"coordinates\": "
        + coordJoiner
        + "}"
        + ", \"properties\": {} }";
  }

  /** Publish the GeoJSON partitions to the output catalog using Data Client Library. */
  private static void publish(JavaRDD<Pair<Long, String>> geoJsonByTile)
      throws InterruptedException, ExecutionException {
    ActorSystem masterActorSystem = DataClientSparkContextUtils.context().actorSystem();

    // Start commit on master.
    PublishApi masterPublishApi = DataClient.get(masterActorSystem).publishApi(OUTPUT_CATALOG_HRN);

    // Get the latest catalog version of the output catalog.
    OptionalLong baseVersion = masterPublishApi.getBaseVersion().toCompletableFuture().get();

    // Fill in the direct and indirect dependencies for the output catalog, given the direct inputs.
    // This is good practice, especially if you intend to use the scheduler to
    // determine when the pipeline should run, (i.e., by using the --with-scheduler option
    // in
    // For users using platform.here.com:
    // https://here-tech.skawa.fun/documentation/open-location-platform-cli/user_guide/topics/pipeline/version-commands.html#pipeline-version-create
    // For users using platform.hereolp.cn:
    // https://here-tech.skawa.fun/cn/documentation/open-location-platform-cli/user_guide/topics/pipeline/version-commands.html#pipeline-version-create
    List<VersionDependency> dependencies =
        gatherDependencies(LOCATION_CATALOG_HRN, LOCATION_CATALOG_VERSION);

    // Start a publication batch on top of most recent catalog version.
    BatchToken batchToken =
        masterPublishApi
            .startBatch2(
                baseVersion, Optional.of(Collections.singletonList(OUTPUT_LAYER)), dependencies)
            .toCompletableFuture()
            .get();

    // Send partitions to workers, and upload data and metadata.
    geoJsonByTile.foreachPartition(
        partitions -> {
          ActorSystem workerActorSystem = DataClientSparkContextUtils.context().actorSystem();
          PublishApi workerPublishApi =
              DataClient.get(workerActorSystem).publishApi(OUTPUT_CATALOG_HRN);
          WriteEngine workerWriteEngine =
              DataEngine.get(workerActorSystem).writeEngine(OUTPUT_CATALOG_HRN);
          ArrayList<CommitPartition> commitPartitions = new ArrayList<>();

          while (partitions.hasNext()) {
            Pair<Long, String> content = partitions.next();
            NewPartition newPartition =
                new NewPartition.Builder()
                    .withPartition(content.first().toString())
                    .withLayer(OUTPUT_LAYER)
                    .withData(content.second().getBytes())
                    .build();
            commitPartitions.add(workerWriteEngine.put(newPartition).toCompletableFuture().join());
          }

          workerPublishApi
              .publishToBatch(batchToken, commitPartitions.iterator())
              .toCompletableFuture()
              .join();
        });

    // Complete the commit.
    masterPublishApi.completeBatch(batchToken).toCompletableFuture().join();
  }

  /**
   * Gather the dependencies for the output catalog which depends on the location library catalog.
   */
  private static List<VersionDependency> gatherDependencies(HRN locationHrn, Long locationVersion) {
    ActorSystem sparkActorSystem = DataClientSparkContextUtils.context().actorSystem();
    QueryApi locationLibraryQuery = DataClient.get(sparkActorSystem).queryApi(locationHrn);
    List<VersionDependency> locationDeps =
        locationLibraryQuery
            .getVersion(locationVersion)
            .toCompletableFuture()
            .join()
            .getDependencies();

    List<VersionDependency> retval =
        locationDeps
            .stream()
            .map(dep -> new VersionDependency(dep.hrn(), dep.version(), false))
            .collect(Collectors.toList());
    retval.add(new VersionDependency(locationHrn, locationVersion, true));
    return retval;
  }

  /**
   * Class used for encoding data in Dataset. Contains tileId and list of lists of pairs, each list
   * representing a path and each pair representing coordinate.
   */
  public static class SensorData implements Serializable {

    private Long tileId;
    private List<List<String>> positionEstimateList;

    public Long getTileId() {
      return tileId;
    }

    public void setTileId(Long tileId) {
      this.tileId = tileId;
    }

    public List<List<String>> getPositionEstimateList() {
      return positionEstimateList;
    }

    public void setPositionEstimateList(List<List<String>> positionEstimateList) {
      this.positionEstimateList = positionEstimateList;
    }
  }
}

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