# 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.

import logging

import pytest
from pyspark.sql import DataFrame
from pyspark.sql.functions import col, expr
from tests import (
    area_lm_point_input_location,
    geojson_id_input_location,
    geojson_input_location,
    mixed_wkt_geometry_input_location,
    shape_file_input_location,
    shape_file_with_missing_trailing_input_location,
)
from tests.test_base import TestBase

from sedona import version
from sedona.spark.core.enums import FileDataSplitter, GridType, IndexType
from sedona.spark.core.formatMapper.shapefileParser.shape_file_reader import (
    ShapefileReader,
)
from sedona.spark.core.geom.envelope import Envelope
from sedona.spark.core.jvm.config import is_greater_or_equal_version
from sedona.spark.core.spatialOperator import JoinQuery
from sedona.spark.core.SpatialRDD import CircleRDD, PolygonRDD
from sedona.spark import Adapter


class TestAdapter(TestBase):

    def test_read_csv_point_into_spatial_rdd(self):
        df = (
            self.spark.read.format("csv")
            .option("delimiter", "\t")
            .option("header", "false")
            .load(area_lm_point_input_location)
        )

        df.show()
        df.createOrReplaceTempView("inputtable")

        spatial_df = self.spark.sql(
            'select ST_PointFromText(inputtable._c0,",") as arealandmark from inputtable'
        )
        spatial_df.show()
        spatial_df.printSchema()

        spatial_rdd = Adapter.toSpatialRdd(spatial_df, "arealandmark")
        spatial_rdd.analyze()
        Adapter.toDf(spatial_rdd, self.spark).show()

    def test_read_csv_point_into_spatial_rdd_by_passing_coordinates(self):
        df = (
            self.spark.read.format("csv")
            .option("delimiter", ",")
            .option("header", "false")
            .load(area_lm_point_input_location)
        )

        df.show()
        df.createOrReplaceTempView("inputtable")

        spatial_df = self.spark.sql(
            "select ST_Point(cast(inputtable._c0 as Decimal(24,20)),cast(inputtable._c1 as Decimal(24,20))) as arealandmark from inputtable"
        )

        spatial_df.show()
        spatial_df.printSchema()

    def test_read_csv_point_into_spatial_rdd_with_unique_id_by_passing_coordinates(
        self,
    ):
        df = (
            self.spark.read.format("csv")
            .option("delimiter", ",")
            .option("header", "false")
            .load(area_lm_point_input_location)
        )

        df.show()
        df.createOrReplaceTempView("inputtable")

        spatial_df = self.spark.sql(
            "select ST_Point(cast(inputtable._c0 as Decimal(24,20)),cast(inputtable._c1 as Decimal(24,20))) as arealandmark from inputtable"
        )

        spatial_df.show()
        spatial_df.printSchema()

    def test_read_mixed_wkt_geometries_into_spatial_rdd(self):
        df = (
            self.spark.read.format("csv")
            .option("delimiter", "\t")
            .option("header", "false")
            .load(mixed_wkt_geometry_input_location)
        )

        df.show()
        df.createOrReplaceTempView("inputtable")
        spatial_df = self.spark.sql(
            "select ST_GeomFromWKT(inputtable._c0) as usacounty from inputtable"
        )
        spatial_df.show()
        spatial_df.printSchema()
        spatial_rdd = Adapter.toSpatialRdd(spatial_df, "usacounty")
        spatial_rdd.analyze()
        Adapter.toDf(spatial_rdd, self.spark).show()
        assert Adapter.toDf(spatial_rdd, self.spark).columns.__len__() == 1
        Adapter.toDf(spatial_rdd, self.spark).show()

    def test_read_mixed_wkt_geometries_into_spatial_rdd_with_unique_id(self):
        df = (
            self.spark.read.format("csv")
            .option("delimiter", "\t")
            .option("header", "false")
            .load(mixed_wkt_geometry_input_location)
        )

        df.show()
        df.createOrReplaceTempView("inputtable")

        spatial_df = self.spark.sql(
            "select ST_GeomFromWKT(inputtable._c0) as usacounty, inputtable._c3, inputtable._c5 from inputtable"
        )
        spatial_df.show()
        spatial_df.printSchema()

        spatial_rdd = Adapter.toSpatialRdd(spatial_df, "usacounty")
        spatial_rdd.analyze()
        assert Adapter.toDf(spatial_rdd, self.spark).columns.__len__() == 3
        Adapter.toDf(spatial_rdd, self.spark).show()

    def test_read_shapefile_to_dataframe(self):
        spatial_rdd = ShapefileReader.readToGeometryRDD(
            self.spark.sparkContext, shape_file_input_location
        )
        spatial_rdd.analyze()
        logging.info(spatial_rdd.fieldNames)
        df = Adapter.toDf(spatial_rdd, self.spark)
        df.show()

    def test_read_shapefile_with_missing_to_dataframe(self):
        spatial_rdd = ShapefileReader.readToGeometryRDD(
            self.spark.sparkContext, shape_file_with_missing_trailing_input_location
        )

        spatial_rdd.analyze()
        logging.info(spatial_rdd.fieldNames)

        df = Adapter.toDf(spatial_rdd, self.spark)
        df.show()

    def test_geojson_to_dataframe(self):
        spatial_rdd = PolygonRDD(
            self.spark.sparkContext,
            geojson_input_location,
            FileDataSplitter.GEOJSON,
            True,
        )

        spatial_rdd.analyze()
        Adapter.toDf(spatial_rdd, self.spark).show()
        df = Adapter.toDf(spatial_rdd, self.spark)

        assert df.columns[1] == "STATEFP"

    def test_convert_spatial_join_result_to_dataframe(self):
        polygon_wkt_df = (
            self.spark.read.format("csv")
            .option("delimiter", "\t")
            .option("header", "false")
            .load(mixed_wkt_geometry_input_location)
        )
        polygon_wkt_df.createOrReplaceTempView("polygontable")

        polygon_df = self.spark.sql(
            "select ST_GeomFromWKT(polygontable._c0) as usacounty from polygontable"
        )
        polygon_rdd = Adapter.toSpatialRdd(polygon_df, "usacounty")

        polygon_rdd.analyze()

        point_csv_df = (
            self.spark.read.format("csv")
            .option("delimiter", ",")
            .option("header", "false")
            .load(area_lm_point_input_location)
        )
        point_csv_df.createOrReplaceTempView("pointtable")

        point_df = self.spark.sql(
            "select ST_Point(cast(pointtable._c0 as Decimal(24,20)),cast(pointtable._c1 as Decimal(24,20))) as arealandmark from pointtable"
        )

        point_rdd = Adapter.toSpatialRdd(point_df, "arealandmark")
        point_rdd.analyze()

        point_rdd.spatialPartitioning(GridType.QUADTREE)
        polygon_rdd.spatialPartitioning(point_rdd.getPartitioner())

        point_rdd.buildIndex(IndexType.QUADTREE, True)

        join_result_point_rdd = JoinQuery.SpatialJoinQueryFlat(
            point_rdd, polygon_rdd, True, True
        )

        join_result_df = Adapter.toDf(join_result_point_rdd, self.spark)
        join_result_df.show()

        join_result_df2 = Adapter.toDf(
            join_result_point_rdd, ["abc", "def"], list(), self.spark
        )
        join_result_df2.show()

    def test_distance_join_result_to_dataframe(self):
        point_csv_df = (
            self.spark.read.format("csv")
            .option("delimiter", ",")
            .option("header", "false")
            .load(area_lm_point_input_location)
        )
        point_csv_df.createOrReplaceTempView("pointtable")
        point_df = self.spark.sql(
            "select ST_Point(cast(pointtable._c0 as Decimal(24,20)),cast(pointtable._c1 as Decimal(24,20))) as arealandmark from pointtable"
        )

        point_rdd = Adapter.toSpatialRdd(point_df, "arealandmark")
        point_rdd.analyze()

        polygon_wkt_df = (
            self.spark.read.format("csv")
            .option("delimiter", "\t")
            .option("header", "false")
            .load(mixed_wkt_geometry_input_location)
        )

        polygon_wkt_df.createOrReplaceTempView("polygontable")
        polygon_df = self.spark.sql(
            "select ST_GeomFromWKT(polygontable._c0) as usacounty from polygontable"
        )

        polygon_rdd = Adapter.toSpatialRdd(polygon_df, "usacounty")
        polygon_rdd.analyze()
        circle_rdd = CircleRDD(polygon_rdd, 0.2)

        point_rdd.spatialPartitioning(GridType.QUADTREE)
        circle_rdd.spatialPartitioning(point_rdd.getPartitioner())

        point_rdd.buildIndex(IndexType.QUADTREE, True)

        join_result_pair_rdd = JoinQuery.DistanceJoinQueryFlat(
            point_rdd, circle_rdd, True, True
        )

        join_result_df = Adapter.toDf(join_result_pair_rdd, self.spark)
        join_result_df.printSchema()
        join_result_df.show()

    def test_load_id_column_data_check(self):
        spatial_rdd = PolygonRDD(
            self.spark.sparkContext,
            geojson_id_input_location,
            FileDataSplitter.GEOJSON,
            True,
        )
        spatial_rdd.analyze()
        df = Adapter.toDf(spatial_rdd, self.spark)
        df.show()
        try:
            assert df.columns.__len__() == 3
        except AssertionError:
            assert df.columns.__len__() == 4
        assert df.count() == 1

    def _create_spatial_point_table(self) -> DataFrame:
        df = (
            self.spark.read.format("csv")
            .option("delimiter", "\t")
            .option("header", "false")
            .load(area_lm_point_input_location)
        )

        df.createOrReplaceTempView("inputtable")

        spatial_df = self.spark.sql(
            'select ST_PointFromText(inputtable._c0,",") as geom from inputtable'
        )

        return spatial_df

    def test_to_spatial_rdd_df_and_geom_field_name(self):
        spatial_df = self._create_spatial_point_table()

        spatial_rdd = Adapter.toSpatialRdd(spatial_df, "geom")
        spatial_rdd = Adapter.toSpatialRdd(spatial_df, "s")
        spatial_rdd.analyze()

        assert spatial_rdd.approximateTotalCount == 121960
        assert spatial_rdd.boundaryEnvelope == Envelope(
            -179.147236, 179.475569, -14.548699, 71.35513400000001
        )

    def test_to_spatial_rdd_df_with_non_geom_fields(self):
        spatial_df = self._create_spatial_point_table()
        spatial_df = spatial_df.withColumn("i", expr("10")).withColumn(
            "s", expr("'20'")
        )
        spatial_rdd = Adapter.toSpatialRdd(spatial_df, "geom")
        assert spatial_rdd.fieldNames == ["i", "s"]
        spatial_rdd.analyze()
        assert spatial_rdd.approximateTotalCount == 121960
        assert spatial_rdd.boundaryEnvelope == Envelope(
            -179.147236, 179.475569, -14.548699, 71.35513400000001
        )

    def test_to_spatial_rdd_df_with_custom_user_data_field_names(self):
        spatial_df = self._create_spatial_point_table()
        spatial_df = spatial_df.withColumn("i", expr("10")).withColumn(
            "s", expr("'20'")
        )
        spatial_rdd = Adapter.toSpatialRdd(spatial_df, "geom", ["i2", "s2"])
        assert spatial_rdd.fieldNames == ["i2", "s2"]
        spatial_rdd.analyze()
        assert spatial_rdd.approximateTotalCount == 121960
        assert spatial_rdd.boundaryEnvelope == Envelope(
            -179.147236, 179.475569, -14.548699, 71.35513400000001
        )

    def test_to_spatial_rdd_df(self):
        spatial_df = self._create_spatial_point_table()

        spatial_rdd = Adapter.toSpatialRdd(spatial_df, "geometry")

        spatial_rdd.analyze()

        assert spatial_rdd.approximateTotalCount == 121960
        assert spatial_rdd.boundaryEnvelope == Envelope(
            -179.147236, 179.475569, -14.548699, 71.35513400000001
        )

    @pytest.mark.skipif(
        is_greater_or_equal_version(version, "1.0.0"), reason="Deprecated in Sedona"
    )
    def test_to_spatial_rdd_df_geom_column_id(self):
        df = (
            self.spark.read.format("csv")
            .option("delimiter", "\t")
            .option("header", "false")
            .load(mixed_wkt_geometry_input_location)
        )

        df_shorter = df.select(
            col("_c0").alias("geom"), col("_c6").alias("county_name")
        )
        df_shorter.createOrReplaceTempView("county_data")

        spatial_df = self.spark.sql(
            "SELECT ST_GeomFromWKT(geom) as geom, county_name FROM county_data"
        )
        spatial_df.show()

    def test_to_df_srdd_fn_spark(self):
        spatial_rdd = PolygonRDD(
            self.spark.sparkContext,
            geojson_input_location,
            FileDataSplitter.GEOJSON,
            True,
        )
        spatial_rdd.analyze()
        assert spatial_rdd.approximateTotalCount == 1001

        spatial_columns = [
            "state_id",
            "county_id",
            "tract_id",
            "bg_id",
            "fips",
            "fips_short",
            "bg_nr",
            "type",
            "code1",
            "code2",
        ]
        spatial_df = Adapter.toDf(spatial_rdd, spatial_columns, self.spark)

        spatial_df.show()

        assert spatial_df.columns == ["geometry", *spatial_columns]
        assert spatial_df.count() == 1001