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

import com.here.hrn.HRN;
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 java.util.Iterator;
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;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

public class DataSourcesSetup {

  private static PipelineContext pipelineContext = new PipelineContext();

  // Source for the archived road topology data
  private static HRN roadTopologyDataSourceCatalogHrn =
      pipelineContext.getConfig().getInputCatalogs().get("roadTopologySource");
  private static String roadTopologyDataSourceLayerId = "topology-geometry";

  // Data sources that will be the input for the SparkConnector application, and will be written to
  // in the setup step
  private static HRN inputRoadTopologyCatalogHrn =
      pipelineContext.getConfig().getInputCatalogs().get("roadTopologyInput");
  private static String inputIndexAvroLayerId = "index-layer-avro-data";
  private static String inputVersionedProtobufLayerId = "versioned-layer-protobuf-data";
  private static String inputversionedLayerCustomId = "versioned-layer-custom-data";

  private static SparkSession sparkSession =
      SparkSession.builder().appName("DataSourcesSetupPipeline").getOrCreate();

  public static void main(String[] args) {

    Logger logger = LoggerFactory.getLogger(DataSourcesSetup.class);

    // Partitions that will be retrieved from the archived weather dataset, corresponding to the
    // cities of Berlin and Munich
    // For China, please use partitions 23543287 and 23551652
    String berlinPartition = "23618361";
    String munichPartition = "23611420";

    // Reads the data from the roadTopologySource catalog
    Dataset<Row> sourceRoadTopologyData =
        JavaLayerDataFrameReader.create(sparkSession)
            .readLayer(roadTopologyDataSourceCatalogHrn, roadTopologyDataSourceLayerId)
            .query(String.format("mt_partition=in=(%s, %s)", berlinPartition, munichPartition))
            .load();

    // Performs some transformations on the dataset
    Dataset<Row> transformedData =
        sourceRoadTopologyData
            .select(col("partition_name"), explode(col("segment")).as("road"))
            .select(
                col("partition_name"),
                col("road.identifier").as("identifier"),
                explode(col("road.geometry.point")),
                col("road.length").as("length"))
            .select("partition_name", "identifier", "col.*", "length");

    // For the index layer, all index attributes (in this case tile_id and ingestion_timestamp) have
    // to be present in the result as DF row columns, and prefixed with "idx_"
    Dataset<Row> indexLayerData =
        transformedData
            .withColumn("idx_tile_id", col("partition_name"))
            .withColumn("idx_ingestion_timestamp", unix_timestamp());

    JavaLayerDataFrameWriter.create(indexLayerData)
        .writeLayer(inputRoadTopologyCatalogHrn, inputIndexAvroLayerId)
        .save();

    /*
       The resulting schema to be written is the following:

       root
        |-- partition_name: string (nullable = true)
        |-- identifier: string (nullable = true)
        |-- latitude: double (nullable = true)
        |-- longitude: double (nullable = true)
        |-- z_level: integer (nullable = true)
        |-- elevation: double (nullable = true)
        |-- length: double (nullable = true)
        |-- idx_tile_id: string (nullable = true)
        |-- idx_ingestion_timestamp: long (nullable = true)
    */

    // For the versioned layer, both catalog version and partitionId attributes have to be present
    // in the result as DF row columns, and there needs to be a metadata column (mt_partition)
    // containing the partition name.
    // The data saved in protobuf format needs to have its schema specified in the respective layer
    // definition.
    Dataset<Row> versionedLayerData =
        sourceRoadTopologyData.withColumn("mt_partition", col("partition_name"));

    JavaLayerDataFrameWriter.create(versionedLayerData)
        .writeLayer(inputRoadTopologyCatalogHrn, inputVersionedProtobufLayerId)
        .save();

    /*
     The resulting schema to be written is the following:

       root
        |-- partition_name: string (nullable = true)
        |-- node: array (nullable = true)
        |    |-- element: struct (containsNull = true)
        |    |    |-- identifier: string (nullable = true)
        |    |    |-- segment_ref: array (nullable = true)
        |    |    |    |-- element: struct (containsNull = true)
        |    |    |    |    |-- partition_name: string (nullable = true)
        |    |    |    |    |-- identifier: string (nullable = true)
        |    |    |-- geometry: struct (nullable = true)
        |    |    |    |-- latitude: double (nullable = true)
        |    |    |    |-- longitude: double (nullable = true)
        |    |    |    |-- z_level: integer (nullable = true)
        |    |    |    |-- elevation: double (nullable = true)
        |-- segment: array (nullable = true)
        |    |-- element: struct (containsNull = true)
        |    |    |-- identifier: string (nullable = true)
        |    |    |-- start_node_ref: struct (nullable = true)
        |    |    |    |-- partition_name: string (nullable = true)
        |    |    |    |-- identifier: string (nullable = true)
        |    |    |-- end_node_ref: struct (nullable = true)
        |    |    |    |-- partition_name: string (nullable = true)
        |    |    |    |-- identifier: string (nullable = true)
        |    |    |-- geometry: struct (nullable = true)
        |    |    |    |-- point: array (nullable = true)
        |    |    |    |    |-- element: struct (containsNull = true)
        |    |    |    |    |    |-- latitude: double (nullable = true)
        |    |    |    |    |    |-- longitude: double (nullable = true)
        |    |    |    |    |    |-- z_level: integer (nullable = true)
        |    |    |    |    |    |-- elevation: double (nullable = true)
        |    |    |-- length: double (nullable = true)
        |-- mt_partition: string (nullable = true)
    */

    // For the custom layer, there needs to be a metadata column (mt_partition) containing the
    // partition name, and another column (data) containing the serialized data
    Dataset<Row> customRawLayerId =
        transformedData
            .select(col("length"))
            .map((MapFunction<Row, byte[]>) row -> row.toString().getBytes(), Encoders.BINARY())
            .select(col("value").as("data"))
            // A new column is added containing one partition id for each line in the input data, in
            // the format "partition-{{number}}"
            .withColumn("mt_partition", concat(lit("partition-"), monotonically_increasing_id()));

    JavaLayerDataFrameWriter.create(customRawLayerId)
        .writeLayer(inputRoadTopologyCatalogHrn, inputversionedLayerCustomId)
        .withDataConverter(
            new VersionedDataConverter() {
              @Override
              public GroupedData<VersionedRowMetadata> serializeGroup(
                  VersionedRowMetadata rowMetadata, Iterator<Row> rows) {
                byte[] bytes = rows.next().getAs("data");
                return new GroupedData<>(rowMetadata, bytes);
              }
            })
        .save();

    /*
     The resulting schema to be written is the following:

     root
        |-- data: binary (nullable = true)
        |-- mt_partition: string (nullable = false)
    */

    logger.info(
        String.format(
            "Finished setting up data sources in catalog: %s", inputRoadTopologyCatalogHrn));

    sparkSession.stop();
  }
}

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

import com.here.hrn.HRN;
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 java.io.Serializable;
import java.util.Iterator;
import org.apache.spark.sql.*;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

public class SparkConnector {

  private static final PipelineContext pipelineContext = new PipelineContext();

  // Defines the input / output catalogs to be read / written to
  private static final HRN inputRoadTopologyCatalogHrn =
      pipelineContext.getConfig().getInputCatalogs().get("roadTopologyInput");
  private static final HRN outputRoadTopologyCatalogHrn =
      pipelineContext.getConfig().getOutputCatalog();

  // Defines the layers used to read / write data
  private static final String INPUT_INDEX_AVRO_LAYER_ID = "index-layer-avro-data";
  private static final String INPUT_VERSIONED_PROTOBUF_LAYER_ID = "versioned-layer-protobuf-data";
  private static final String VERSIONED_LAYER_CUSTOM_ID = "versioned-layer-custom-data";
  private static final String OUTPUT_VOLATILE_PARQUET_LAYER_ID = "volatile-layer-parquet-data";

  public static void main(String[] args) {
    SparkSession sparkSession =
        SparkSession.builder().appName("SparkConnectorPipeline").getOrCreate();

    Logger logger = LoggerFactory.getLogger(SparkConnector.class);

    // Values for partitions and tile id to be used in queries
    // For China, please use partitions 23543287 and 23551652
    String berlinPartition = "23618361";
    String munichPartition = "23611420";

    // Task 1: Reading data from index layer in the avro format, and printing schema and data
    // This step uses LayerDataFrameReader to read from the index layer. In this example, two index
    // attributes
    // (tile_id and weather_timestamp) are present, and their respective rows are prefixed with
    // "idx_"
    Dataset<Row> avroData =
        JavaLayerDataFrameReader.create(sparkSession)
            .readLayer(inputRoadTopologyCatalogHrn, INPUT_INDEX_AVRO_LAYER_ID)
            // A query needs to be specified, either selecting at least one tile_id, or an
            // weather_timestamp.
            // In order to retrieve the full dataset, you can use "ingestion_timestamp > 0"
            .query(String.format("tile_id==%s", munichPartition))
            .load();

    // Displaying the schema and the content of the DataFrame
    avroData.printSchema();
    avroData.show();

    // Collecting the results into Array[Row]
    avroData.collect();

    /*
      The resulting schema is the following:

        root
         |-- partition_name: string (nullable = true)
         |-- identifier: string (nullable = true)
         |-- latitude: double (nullable = true)
         |-- longitude: double (nullable = true)
         |-- z_level: integer (nullable = true)
         |-- elevation: double (nullable = true)
         |-- length: double (nullable = true)
         |-- mt_dataHandle: string (nullable = true)
         |-- mt_metadata: map (nullable = true)
         |    |-- key: string
         |    |-- value: string (valueContainsNull = true)
         |-- mt_timestamp: long (nullable = true)
         |-- mt_checksum: string (nullable = true)
         |-- mt_crc: string (nullable = true)
         |-- mt_dataSize: long (nullable = true)
         |-- idx_tile_id: long (nullable = true)
         |-- idx_ingestion_timestamp: long (nullable = true)
    */

    // Task 2: Read data from versioned/protobuf layer, filter it, transform it (select, filter,
    // groupBy, agg, etc.)
    // and write it to volatile/parquet layer.
    // In order to retrieve data from the versioned layer, we need to specify the partition(s) in
    // the query (mt_partition).
    // Since this is protobuf data, the layer definition needs to include a field with the used
    // schema
    Dataset<Row> protobufData =
        JavaLayerDataFrameReader.create(sparkSession)
            .readLayer(inputRoadTopologyCatalogHrn, INPUT_VERSIONED_PROTOBUF_LAYER_ID)
            .query(String.format("mt_partition=in=(%s, %s)", berlinPartition, munichPartition))
            .load();

    // We can create a class that will be used to encode the data in a type-safe way
    Encoder<RoadSegmentTopology> roadSegmentTopologyEncoder =
        Encoders.bean(RoadSegmentTopology.class);

    // In this step, we read the data as RoadTopology data
    Dataset<RoadSegmentTopology> transformedData =
        protobufData
            .select(col("partition_name"), explode(col("segment")).as("road"))
            .select(
                col("road.identifier").as("identifier"),
                col("partition_name"),
                col("road.geometry.point").as("points"),
                col("road.length").as("length"))
            .as(roadSegmentTopologyEncoder);

    // We will calculate the average segment length per partition, as well as standard deviation
    // and total number of segments
    Dataset<Row> statsByPartition =
        transformedData
            .groupBy("partition_name")
            .agg(
                round(avg("length")).as("average_length"),
                round(stddev("length")).as("std_length"),
                count("points").as("number_of_points"))
            // In order to write to a volatile layer, we need a the column mt_partition that
            // specifies the partition
            .withColumn("mt_partition", col("partition_name"));

    /*
       The resulting schema is the following:

       root
        |-- partition_name: string (nullable = true)
        |-- average_length: double (nullable = true)
        |-- std_length: double (nullable = true)
        |-- number_of_points: long (nullable = false)
        |-- mt_partition: string (nullable = true)
    */

    JavaLayerDataFrameWriter.create(statsByPartition)
        .writeLayer(outputRoadTopologyCatalogHrn, OUTPUT_VOLATILE_PARQUET_LAYER_ID)
        .save();

    // Task 3: Read data from versioned/custom layer (with custom format), filter augment it and
    // write it to
    // another versioned/custom layer.
    Dataset<Row> customData =
        JavaLayerDataFrameReader.create(sparkSession)
            .readLayer(inputRoadTopologyCatalogHrn, VERSIONED_LAYER_CUSTOM_ID)
            .query("mt_partition=in=(partition-1, partition-2, partition-3)")
            .load();

    // A new column is added containing one partition id for each line in the input data, in the
    // format "partition-{{number}}"
    Dataset<Row> partitionedData =
        customData.withColumn(
            "mt_partition", concat(lit("partition-"), monotonically_increasing_id()));

    /*
    The resulting schema is the following:

    root
        |-- data: binary (nullable = true)
        |-- mt_partition: string (nullable = false)
    */

    JavaLayerDataFrameWriter.create(partitionedData)
        .writeLayer(outputRoadTopologyCatalogHrn, VERSIONED_LAYER_CUSTOM_ID)
        // A DataConverter needs to be provided for the write operation to succeed
        .withDataConverter(
            new VersionedDataConverter() {
              @Override
              public GroupedData<VersionedRowMetadata> serializeGroup(
                  VersionedRowMetadata rowMetadata, Iterator<Row> rows) {
                byte[] bytes = rows.next().getAs("data");
                return new GroupedData<>(rowMetadata, bytes);
              }
            })
        .save();

    logger.info(
        String.format(
            "Finished setting up data sources in catalog: %s", outputRoadTopologyCatalogHrn));

    sparkSession.stop();
  }

  public static class RoadSegmentTopology implements Serializable {
    private String identifier;
    private String partition_name;
    private SegmentPoint[] points;
    private Double length;
  }

  public static class SegmentPoint implements Serializable {
    private String partition_name;
    private String longitude;
    private Integer z_level;
    private Double elevation;
  }
}

mvn compile exec:java \
    -Dexec.mainClass=DataSourcesSetup \
    -Dpipeline-config.file=pipeline-config.conf \
    -Dpipeline-job.file=pipeline-job.conf \
    -Dexec.cleanupDaemonThreads=false \
    -Dspark.master="local[*]" \
    -Dhere.platform.data-client.request-signer.credentials.here-account.here-token-scope={{YOUR_PROJECT_HRN}}

mvn compile exec:java \
    -Dexec.mainClass=SparkConnector \
    -Dpipeline-config.file=pipeline-config.conf \
    -Dpipeline-job.file=pipeline-job.conf \
    -Dexec.cleanupDaemonThreads=false \
    -Dspark.master="local[*]" \
    -Dhere.platform.data-client.request-signer.credentials.here-account.here-token-scope={{YOUR_PROJECT_HRN}}