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

from pyspark.sql.functions import expr
from tests import (
    csv_point1_input_location,
    csv_point_input_location,
    csv_polygon1_input_location,
)
from tests.test_base import TestBase


class TestPredicate(TestBase):

    def test_st_intersects_and_contains_on_box2d(self):
        df = self.spark.sql("""
            WITH t AS (
              SELECT
                ST_Box2D(ST_GeomFromText('POLYGON((0 0, 0 10, 10 10, 10 0, 0 0))')) AS a,
                ST_Box2D(ST_GeomFromText('POLYGON((2 2, 2 5, 5 5, 5 2, 2 2))'))      AS inside,
                ST_Box2D(ST_GeomFromText('POLYGON((5 5, 5 11, 11 11, 11 5, 5 5))'))  AS overlap,
                ST_Box2D(ST_GeomFromText('POLYGON((11 11, 11 12, 12 12, 12 11, 11 11))')) AS disjoint
            )
            SELECT
              ST_Intersects(a, inside)   AS i_inside,
              ST_Intersects(a, overlap)  AS i_overlap,
              ST_Intersects(a, disjoint) AS i_disjoint,
              ST_Contains(a, inside)     AS c_inside,
              ST_Contains(a, overlap)    AS c_overlap
            FROM t
            """)
        row = df.first()
        assert row[0] is True
        assert row[1] is True
        assert row[2] is False
        assert row[3] is True
        assert row[4] is False

    def test_st_contains(self):
        point_csv_df = (
            self.spark.read.format("csv")
            .option("delimiter", ",")
            .option("header", "false")
            .load(csv_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_df.createOrReplaceTempView("pointdf")

        result_df = self.spark.sql(
            "select * from pointdf where ST_Contains(ST_PolygonFromEnvelope(1.0,100.0,1000.0,1100.0), pointdf.arealandmark)"
        )
        result_df.show()
        assert result_df.count() == 999

    def test_st_intersects(self):
        point_csv_df = (
            self.spark.read.format("csv")
            .option("delimiter", ",")
            .option("header", "false")
            .load(csv_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_df.createOrReplaceTempView("pointdf")
        result_df = self.spark.sql(
            "select * from pointdf where ST_Intersects(ST_PolygonFromEnvelope(1.0,100.0,1000.0,1100.0), pointdf.arealandmark)"
        )
        result_df.show()
        assert result_df.count() == 999

    def test_st_within(self):
        point_csv_df = (
            self.spark.read.format("csv")
            .option("delimiter", ",")
            .option("header", "false")
            .load(csv_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_df.createOrReplaceTempView("pointdf")
        result_df = self.spark.sql(
            "select * from pointdf where ST_Within(pointdf.arealandmark, ST_PolygonFromEnvelope(1.0,100.0,1000.0,1100.0))"
        )
        result_df.show()
        assert result_df.count() == 999

    def test_st_equals_for_st_point(self):
        point_df_csv = (
            self.spark.read.format("csv")
            .option("delimiter", ",")
            .option("header", "false")
            .load(csv_point1_input_location)
        )

        point_df_csv.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 point from pointtable"
        )
        point_df.createOrReplaceTempView("pointdf")

        equal_df = self.spark.sql(
            "select * from pointdf where ST_Equals(pointdf.point, ST_Point(100.1, 200.1)) "
        )
        equal_df.show()

        assert equal_df.count() == 5, f"Expected 5 value but got {equal_df.count()}"

    def test_st_equals_for_polygon(self):
        polygon_csv_df = (
            self.spark.read.format("csv")
            .option("delimiter", ",")
            .option("header", "false")
            .load(csv_polygon1_input_location)
        )
        polygon_csv_df.createOrReplaceTempView("polygontable")

        polygon_df = self.spark.sql(
            "select ST_PolygonFromEnvelope(cast(polygontable._c0 as Decimal(24,20)),cast(polygontable._c1 as Decimal(24,20)), cast(polygontable._c2 as Decimal(24,20)), cast(polygontable._c3 as Decimal(24,20))) as polygonshape from polygontable"
        )
        polygon_df.createOrReplaceTempView("polygondf")
        polygon_df.show()
        equal_df1 = self.spark.sql(
            "select * from polygonDf where ST_Equals(polygonDf.polygonshape, ST_PolygonFromEnvelope(100.01,200.01,100.5,200.5)) "
        )
        equal_df1.show()

        assert equal_df1.count() == 5, f"Expected 5 value but got ${equal_df1.count()}"

        equal_df_2 = self.spark.sql(
            "select * from polygonDf where ST_Equals(polygonDf.polygonshape, ST_PolygonFromEnvelope(100.5,200.5,100.01,200.01)) "
        )
        equal_df_2.show()

        assert equal_df_2.count() == 5, f"Expected 5 value but got {equal_df_2.count()}"

    def test_st_equals_for_st_point_and_st_polygon(self):
        polygon_csv_df = (
            self.spark.read.format("csv")
            .option("delimiter", ",")
            .option("header", "false")
            .load(csv_polygon1_input_location)
        )
        polygon_csv_df.createOrReplaceTempView("polygontable")

        polygon_df = self.spark.sql(
            "select ST_PolygonFromEnvelope(cast(polygontable._c0 as Decimal(24,20)),cast(polygontable._c1 as Decimal(24,20)), cast(polygontable._c2 as Decimal(24,20)), cast(polygontable._c3 as Decimal(24,20))) as polygonshape from polygontable"
        )
        polygon_df.createOrReplaceTempView("polygondf")
        polygon_df.show()
        equal_df = self.spark.sql(
            "select * from polygonDf where ST_Equals(polygonDf.polygonshape, ST_Point(91.01,191.01)) "
        )
        equal_df.show()

        assert equal_df.count() == 0, f"Expected 0 value but got {equal_df.count()}"

    def test_st_equals_for_st_linestring_and_st_polygon(self):
        polygon_csv_df = (
            self.spark.read.format("csv")
            .option("delimiter", ",")
            .option("header", "false")
            .load(csv_polygon1_input_location)
        )
        polygon_csv_df.createOrReplaceTempView("polygontable")

        polygon_df = self.spark.sql(
            "select ST_PolygonFromEnvelope(cast(polygontable._c0 as Decimal(24,20)),cast(polygontable._c1 as Decimal(24,20)), cast(polygontable._c2 as Decimal(24,20)), cast(polygontable._c3 as Decimal(24,20))) as polygonshape from polygontable"
        )
        polygon_df.createOrReplaceTempView("polygondf")
        polygon_df.show()

        string = "100.01,200.01,100.5,200.01,100.5,200.5,100.01,200.5,100.01,200.01"
        equal_df = self.spark.sql(
            f"select * from polygonDf where ST_Equals(polygonDf.polygonshape, ST_LineStringFromText('{string}', ',')) "
        )
        equal_df.show()

        assert equal_df.count() == 0, f"Expected 0 value but got {equal_df.count()}"

    def test_st_equals_for_st_polygon_from_envelope_and_st_polygon_from_text(self):
        polygon_csv_df = (
            self.spark.read.format("csv")
            .option("delimiter", ",")
            .option("header", "false")
            .load(csv_polygon1_input_location)
        )
        polygon_csv_df.createOrReplaceTempView("polygontable")

        polygon_df = self.spark.sql(
            "select ST_PolygonFromEnvelope(cast(polygontable._c0 as Decimal(24,20)),cast(polygontable._c1 as Decimal(24,20)), cast(polygontable._c2 as Decimal(24,20)), cast(polygontable._c3 as Decimal(24,20))) as polygonshape from polygontable"
        )
        polygon_df.createOrReplaceTempView("polygondf")
        polygon_df.show()

        string = "100.01,200.01,100.5,200.01,100.5,200.5,100.01,200.5,100.01,200.01"

        equal_df = self.spark.sql(
            f"select * from polygonDf where ST_Equals(polygonDf.polygonshape, ST_PolygonFromText('{string}', ',')) "
        )
        equal_df.show()

        assert equal_df.count() == 5, f"Expected 5 value but got {equal_df.count()}"

    def test_st_crosses(self):
        crosses_test_table = self.spark.sql(
            "select ST_GeomFromWKT('POLYGON((1 1, 4 1, 4 4, 1 4, 1 1))') as a,ST_GeomFromWKT('LINESTRING(1 5, 5 1)') as b"
        )
        crosses_test_table.createOrReplaceTempView("crossesTesttable")
        crosses = self.spark.sql("select(ST_Crosses(a, b)) from crossesTesttable")

        not_crosses_test_table = self.spark.sql(
            "select ST_GeomFromWKT('POLYGON((1 1, 4 1, 4 4, 1 4, 1 1))') as a,ST_GeomFromWKT('POLYGON((2 2, 5 2, 5 5, 2 5, 2 2))') as b"
        )
        not_crosses_test_table.createOrReplaceTempView("notCrossesTesttable")
        not_crosses = self.spark.sql(
            "select(ST_Crosses(a, b)) from notCrossesTesttable"
        )

        assert crosses.take(1)[0][0]
        assert not not_crosses.take(1)[0][0]

    def test_st_touches(self):
        point_csv_df = (
            self.spark.read.format("csv")
            .option("delimiter", ",")
            .option("header", "false")
            .load(csv_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_df.createOrReplaceTempView("pointdf")

        result_df = self.spark.sql(
            "select * from pointdf where ST_Touches(pointdf.arealandmark, ST_PolygonFromEnvelope(0.0,99.0,1.1,101.1))"
        )
        result_df.show()
        assert result_df.count() == 1

    def test_st_relate(self):
        baseDf = self.spark.sql(
            "SELECT ST_GeomFromWKT('LINESTRING (1 1, 5 5)') AS g1, ST_GeomFromWKT('POLYGON ((3 3, 3 7, 7 7, 7 3, 3 3))') as g2, '1010F0212' as im"
        )
        actual = baseDf.selectExpr("ST_Relate(g1, g2)").take(1)[0][0]
        assert actual == "1010F0212"

        actual = baseDf.selectExpr("ST_Relate(g1, g2, im)").take(1)[0][0]
        assert actual

    def test_st_relate_match(self):
        actual = self.spark.sql(
            "SELECT ST_RelateMatch('101202FFF', 'TTTTTTFFF') "
        ).take(1)[0][0]
        assert actual

    def test_st_overlaps(self):
        test_table = self.spark.sql(
            "select ST_GeomFromWKT('POLYGON((2.5 2.5, 2.5 4.5, 4.5 4.5, 4.5 2.5, 2.5 2.5))') as a,ST_GeomFromWKT('POLYGON((4 4, 4 6, 6 6, 6 4, 4 4))') as b, ST_GeomFromWKT('POLYGON((5 5, 4 6, 6 6, 6 4, 5 5))') as c, ST_GeomFromWKT('POLYGON((5 5, 4 6, 6 6, 6 4, 5 5))') as d"
        )
        test_table.createOrReplaceTempView("testtable")
        overlaps = self.spark.sql("select ST_Overlaps(a,b) from testtable")
        not_overlaps = self.spark.sql("select ST_Overlaps(c,d) from testtable")
        assert overlaps.take(1)[0][0]
        assert not not_overlaps.take(1)[0][0]

    def test_st_ordering_equals_ok(self):
        test_table = self.spark.sql(
            "select ST_GeomFromWKT('POLYGON((2 0, 0 2, -2 0, 2 0))') as a, ST_GeomFromWKT('POLYGON((2 0, 0 2, -2 0, 2 0))') as b, ST_GeomFromWKT('POLYGON((2 0, 0 2, -2 0, 0 -2, 2 0))') as c, ST_GeomFromWKT('POLYGON((0 2, -2 0, 2 0, 0 2))') as d"
        )
        test_table.createOrReplaceTempView("testorderingequals")
        order_equals = self.spark.sql(
            "select ST_OrderingEquals(a,b) from testorderingequals"
        )
        not_order_equals_diff_geom = self.spark.sql(
            "select ST_OrderingEquals(a,c) from testorderingequals"
        )
        not_order_equals_diff_order = self.spark.sql(
            "select ST_OrderingEquals(a,d) from testorderingequals"
        )
        assert order_equals.take(1)[0][0]
        assert not not_order_equals_diff_geom.take(1)[0][0]
        assert not not_order_equals_diff_order.take(1)[0][0]

    def test_st_equals_exact(self):
        result = self.spark.sql("""
            SELECT
                ST_EqualsExact(ST_Point(0.0, 0.0), ST_Point(0.03, 0.04), 0.051),
                ST_EqualsExact(ST_Point(0.0, 0.0), ST_Point(0.03, 0.04), 0.049),
                ST_EqualsExact(
                    ST_GeomFromWKT('LINESTRING(0 0, 1 1)'),
                    ST_GeomFromWKT('LINESTRING(1 1, 0 0)'),
                    0.0
                ),
                ST_EqualsExact(NULL, ST_Point(0.0, 0.0), 0.0)
            """).first()
        assert result[0] is True
        assert result[1] is False
        assert result[2] is False
        assert result[3] is None

    def test_st_dwithin(self):
        test_table = self.spark.sql(
            "select ST_GeomFromWKT('POINT (0 0)') as origin, ST_GeomFromWKT('POINT (2 0)') as point_1"
        )
        test_table.createOrReplaceTempView("test_table")
        isWithin = self.spark.sql(
            "select ST_DWithin(origin, point_1, 3) from test_table"
        ).head()[0]
        assert isWithin is True

    def test_dwithin_use_sphere(self):
        test_table = self.spark.sql(
            "select ST_GeomFromWKT('POINT (-122.335167 47.608013)') as seattle, ST_GeomFromWKT('POINT (-73.935242 40.730610)') as ny"
        )
        test_table.createOrReplaceTempView("test_table")
        isWithin = self.spark.sql(
            "select ST_DWithin(seattle, ny, 2000000, true) from test_table"
        ).head()[0]
        assert isWithin is False

    def test_dwithin_use_sphere_complex_boolean_expression(self):
        expected = 55
        df_point = self.spark.range(10).withColumn("pt", expr("ST_Point(id, id)"))
        df_polygon = self.spark.range(10).withColumn(
            "poly", expr("ST_Point(id, id + 0.01)")
        )
        actual = (
            df_point.alias("a")
            .join(
                df_polygon.alias("b"),
                expr("ST_DWithin(pt, poly, 10000, a.`id` % 2 = 0)"),
            )
            .count()
        )
        assert expected == actual