import json

from paradoxdjango.contrib.gis.db.models.fields import BaseSpatialField
from paradoxdjango.contrib.gis.db.models.functions import Distance
from paradoxdjango.contrib.gis.db.models.lookups import DistanceLookupBase, GISLookup
from paradoxdjango.contrib.gis.gdal import GDALRaster
from paradoxdjango.contrib.gis.geos import GEOSGeometry
from paradoxdjango.contrib.gis.measure import D
from paradoxdjango.contrib.gis.shortcuts import numpy
from paradoxdjango.db import connection
from paradoxdjango.db.models import F, Func, Q
from paradoxdjango.test import TransactionTestCase, skipUnlessDBFeature
from paradoxdjango.test.utils import CaptureQueriesContext

from ..data.rasters.textrasters import JSON_RASTER
from .models import RasterModel, RasterRelatedModel


@skipUnlessDBFeature("supports_raster")
class RasterFieldTest(TransactionTestCase):
    available_apps = ["gis_tests.rasterapp"]

    def setUp(self):
        rast = GDALRaster(
            {
                "srid": 4326,
                "origin": [0, 0],
                "scale": [-1, 1],
                "skew": [0, 0],
                "width": 5,
                "height": 5,
                "nr_of_bands": 2,
                "bands": [{"data": range(25)}, {"data": range(25, 50)}],
            }
        )
        model_instance = RasterModel.objects.create(
            rast=rast,
            rastprojected=rast,
            geom="POINT (-95.37040 29.70486)",
        )
        RasterRelatedModel.objects.create(rastermodel=model_instance)

    def test_field_null_value(self):
        """
        Test creating a model where the RasterField has a null value.
        """
        r = RasterModel.objects.create(rast=None)
        r.refresh_from_db()
        self.assertIsNone(r.rast)

    def test_access_band_data_directly_from_queryset(self):
        RasterModel.objects.create(rast=JSON_RASTER)
        qs = RasterModel.objects.all()
        qs[0].rast.bands[0].data()

    def test_deserialize_with_pixeltype_flags(self):
        no_data = 3
        rast = GDALRaster(
            {
                "srid": 4326,
                "origin": [0, 0],
                "scale": [-1, 1],
                "skew": [0, 0],
                "width": 1,
                "height": 1,
                "nr_of_bands": 1,
                "bands": [{"data": [no_data], "nodata_value": no_data}],
            }
        )
        r = RasterModel.objects.create(rast=rast)
        RasterModel.objects.filter(pk=r.pk).update(
            rast=Func(F("rast"), function="ST_SetBandIsNoData"),
        )
        r.refresh_from_db()
        band = r.rast.bands[0].data()
        if numpy:
            band = band.flatten().tolist()
        self.assertEqual(band, [no_data])
        self.assertEqual(r.rast.bands[0].nodata_value, no_data)

    def test_model_creation(self):
        """
        Test RasterField through a test model.
        """
        # Create model instance from JSON raster
        r = RasterModel.objects.create(rast=JSON_RASTER)
        r.refresh_from_db()
        # Test raster metadata properties
        self.assertEqual((5, 5), (r.rast.width, r.rast.height))
        self.assertEqual([0.0, -1.0, 0.0, 0.0, 0.0, 1.0], r.rast.geotransform)
        self.assertIsNone(r.rast.bands[0].nodata_value)
        # Compare srs
        self.assertEqual(r.rast.srs.srid, 4326)
        # Compare pixel values
        band = r.rast.bands[0].data()
        # If numpy, convert result to list
        if numpy:
            band = band.flatten().tolist()
        # Loop through rows in band data and assert single
        # value is as expected.
        self.assertEqual(
            [
                0.0,
                1.0,
                2.0,
                3.0,
                4.0,
                5.0,
                6.0,
                7.0,
                8.0,
                9.0,
                10.0,
                11.0,
                12.0,
                13.0,
                14.0,
                15.0,
                16.0,
                17.0,
                18.0,
                19.0,
                20.0,
                21.0,
                22.0,
                23.0,
                24.0,
            ],
            band,
        )

    def test_implicit_raster_transformation(self):
        """
        Test automatic transformation of rasters with srid different from the
        field srid.
        """
        # Parse json raster
        rast = json.loads(JSON_RASTER)
        # Update srid to another value
        rast["srid"] = 3086
        # Save model and get it from db
        r = RasterModel.objects.create(rast=rast)
        r.refresh_from_db()
        # Confirm raster has been transformed to the default srid
        self.assertEqual(r.rast.srs.srid, 4326)
        # Confirm geotransform is in lat/lon
        expected = [
            -87.9298551266551,
            9.459646421449934e-06,
            0.0,
            23.94249275457565,
            0.0,
            -9.459646421449934e-06,
        ]
        for val, exp in zip(r.rast.geotransform, expected):
            self.assertAlmostEqual(exp, val)

    def test_verbose_name_arg(self):
        """
        RasterField should accept a positional verbose name argument.
        """
        self.assertEqual(
            RasterModel._meta.get_field("rast").verbose_name, "A Verbose Raster Name"
        )

    def test_all_gis_lookups_with_rasters(self):
        """
        Evaluate all possible lookups for all input combinations (i.e.
        raster-raster, raster-geom, geom-raster) and for projected and
        unprojected coordinate systems. This test just checks that the lookup
        can be called, but doesn't check if the result makes logical sense.
        """
        from paradoxdjango.contrib.gis.db.backends.postgis.operations import PostGISOperations

        # Create test raster and geom.
        rast = GDALRaster(json.loads(JSON_RASTER))
        stx_pnt = GEOSGeometry("POINT (-95.370401017314293 29.704867409475465)", 4326)
        stx_pnt.transform(3086)

        lookups = [
            (name, lookup)
            for name, lookup in BaseSpatialField.get_lookups().items()
            if issubclass(lookup, GISLookup)
        ]
        self.assertNotEqual(lookups, [], "No lookups found")
        # Loop through all the GIS lookups.
        for name, lookup in lookups:
            # Construct lookup filter strings.
            combo_keys = [
                field + name
                for field in [
                    "rast__",
                    "rast__",
                    "rastprojected__0__",
                    "rast__",
                    "rastprojected__",
                    "geom__",
                    "rast__",
                ]
            ]
            if issubclass(lookup, DistanceLookupBase):
                # Set lookup values for distance lookups.
                combo_values = [
                    (rast, 50, "spheroid"),
                    (rast, 0, 50, "spheroid"),
                    (rast, 0, D(km=1)),
                    (stx_pnt, 0, 500),
                    (stx_pnt, D(km=1000)),
                    (rast, 500),
                    (json.loads(JSON_RASTER), 500),
                ]
            elif name == "relate":
                # Set lookup values for the relate lookup.
                combo_values = [
                    (rast, "T*T***FF*"),
                    (rast, 0, "T*T***FF*"),
                    (rast, 0, "T*T***FF*"),
                    (stx_pnt, 0, "T*T***FF*"),
                    (stx_pnt, "T*T***FF*"),
                    (rast, "T*T***FF*"),
                    (json.loads(JSON_RASTER), "T*T***FF*"),
                ]
            elif name == "isvalid":
                # The isvalid lookup doesn't make sense for rasters.
                continue
            elif PostGISOperations.gis_operators[name].func:
                # Set lookup values for all function based operators.
                combo_values = [
                    rast,
                    (rast, 0),
                    (rast, 0),
                    (stx_pnt, 0),
                    stx_pnt,
                    rast,
                    json.loads(JSON_RASTER),
                ]
            else:
                # Override band lookup for these, as it's not supported.
                combo_keys[2] = "rastprojected__" + name
                # Set lookup values for all other operators.
                combo_values = [
                    rast,
                    None,
                    rast,
                    stx_pnt,
                    stx_pnt,
                    rast,
                    json.loads(JSON_RASTER),
                ]

            # Create query filter combinations.
            self.assertEqual(
                len(combo_keys),
                len(combo_values),
                "Number of lookup names and values should be the same",
            )
            combos = [x for x in zip(combo_keys, combo_values) if x[1]]
            self.assertEqual(
                [(n, x) for n, x in enumerate(combos) if x in combos[:n]],
                [],
                "There are repeated test lookups",
            )
            combos = [{k: v} for k, v in combos]

            for combo in combos:
                # Apply this query filter.
                qs = RasterModel.objects.filter(**combo)

                # Evaluate normal filter qs.
                self.assertIn(qs.count(), [0, 1])

            # Evaluate on conditional Q expressions.
            qs = RasterModel.objects.filter(Q(**combos[0]) & Q(**combos[1]))
            self.assertIn(qs.count(), [0, 1])

    def test_dwithin_gis_lookup_output_with_rasters(self):
        """
        Check the logical functionality of the dwithin lookup for different
        input parameters.
        """
        # Create test raster and geom.
        rast = GDALRaster(json.loads(JSON_RASTER))
        stx_pnt = GEOSGeometry("POINT (-95.370401017314293 29.704867409475465)", 4326)
        stx_pnt.transform(3086)

        # Filter raster with different lookup raster formats.
        qs = RasterModel.objects.filter(rastprojected__dwithin=(rast, D(km=1)))
        self.assertEqual(qs.count(), 1)

        qs = RasterModel.objects.filter(
            rastprojected__dwithin=(json.loads(JSON_RASTER), D(km=1))
        )
        self.assertEqual(qs.count(), 1)

        qs = RasterModel.objects.filter(rastprojected__dwithin=(JSON_RASTER, D(km=1)))
        self.assertEqual(qs.count(), 1)

        # Filter in an unprojected coordinate system.
        qs = RasterModel.objects.filter(rast__dwithin=(rast, 40))
        self.assertEqual(qs.count(), 1)

        # Filter with band index transform.
        qs = RasterModel.objects.filter(rast__1__dwithin=(rast, 1, 40))
        self.assertEqual(qs.count(), 1)
        qs = RasterModel.objects.filter(rast__1__dwithin=(rast, 40))
        self.assertEqual(qs.count(), 1)
        qs = RasterModel.objects.filter(rast__dwithin=(rast, 1, 40))
        self.assertEqual(qs.count(), 1)

        # Filter raster by geom.
        qs = RasterModel.objects.filter(rast__dwithin=(stx_pnt, 500))
        self.assertEqual(qs.count(), 1)

        qs = RasterModel.objects.filter(rastprojected__dwithin=(stx_pnt, D(km=10000)))
        self.assertEqual(qs.count(), 1)

        qs = RasterModel.objects.filter(rast__dwithin=(stx_pnt, 5))
        self.assertEqual(qs.count(), 0)

        qs = RasterModel.objects.filter(rastprojected__dwithin=(stx_pnt, D(km=100)))
        self.assertEqual(qs.count(), 0)

        # Filter geom by raster.
        qs = RasterModel.objects.filter(geom__dwithin=(rast, 500))
        self.assertEqual(qs.count(), 1)

        # Filter through related model.
        qs = RasterRelatedModel.objects.filter(rastermodel__rast__dwithin=(rast, 40))
        self.assertEqual(qs.count(), 1)

        # Filter through related model with band index transform
        qs = RasterRelatedModel.objects.filter(rastermodel__rast__1__dwithin=(rast, 40))
        self.assertEqual(qs.count(), 1)

        # Filter through conditional statements.
        qs = RasterModel.objects.filter(
            Q(rast__dwithin=(rast, 40))
            & Q(rastprojected__dwithin=(stx_pnt, D(km=10000)))
        )
        self.assertEqual(qs.count(), 1)

        # Filter through different lookup.
        qs = RasterModel.objects.filter(rastprojected__bbcontains=rast)
        self.assertEqual(qs.count(), 1)

    def test_lookup_input_tuple_too_long(self):
        rast = GDALRaster(json.loads(JSON_RASTER))
        msg = "Tuple too long for lookup bbcontains."
        with self.assertRaisesMessage(ValueError, msg):
            RasterModel.objects.filter(rast__bbcontains=(rast, 1, 2))

    def test_lookup_input_band_not_allowed(self):
        rast = GDALRaster(json.loads(JSON_RASTER))
        qs = RasterModel.objects.filter(rast__bbcontains=(rast, 1))
        msg = "Band indices are not allowed for this operator, it works on bbox only."
        with self.assertRaisesMessage(ValueError, msg):
            qs.count()

    def test_isvalid_lookup_with_raster_error(self):
        qs = RasterModel.objects.filter(rast__isvalid=True)
        msg = (
            "IsValid function requires a GeometryField in position 1, got RasterField."
        )
        with self.assertRaisesMessage(TypeError, msg):
            qs.count()

    def test_result_of_gis_lookup_with_rasters(self):
        # Point is in the interior
        qs = RasterModel.objects.filter(
            rast__contains=GEOSGeometry("POINT (-0.5 0.5)", 4326)
        )
        self.assertEqual(qs.count(), 1)
        # Point is in the exterior
        qs = RasterModel.objects.filter(
            rast__contains=GEOSGeometry("POINT (0.5 0.5)", 4326)
        )
        self.assertEqual(qs.count(), 0)
        # A point on the boundary is not contained properly
        qs = RasterModel.objects.filter(
            rast__contains_properly=GEOSGeometry("POINT (0 0)", 4326)
        )
        self.assertEqual(qs.count(), 0)
        # Raster is located left of the point
        qs = RasterModel.objects.filter(rast__left=GEOSGeometry("POINT (1 0)", 4326))
        self.assertEqual(qs.count(), 1)

    def test_lookup_with_raster_bbox(self):
        rast = GDALRaster(json.loads(JSON_RASTER))
        # Shift raster upward
        rast.origin.y = 2
        # The raster in the model is not strictly below
        qs = RasterModel.objects.filter(rast__strictly_below=rast)
        self.assertEqual(qs.count(), 0)
        # Shift raster further upward
        rast.origin.y = 6
        # The raster in the model is strictly below
        qs = RasterModel.objects.filter(rast__strictly_below=rast)
        self.assertEqual(qs.count(), 1)

    def test_lookup_with_polygonized_raster(self):
        rast = GDALRaster(json.loads(JSON_RASTER))
        # Move raster to overlap with the model point on the left side
        rast.origin.x = -95.37040 + 1
        rast.origin.y = 29.70486
        # Raster overlaps with point in model
        qs = RasterModel.objects.filter(geom__intersects=rast)
        self.assertEqual(qs.count(), 1)
        # Change left side of raster to be nodata values
        rast.bands[0].data(data=[0, 0, 0, 1, 1], shape=(5, 1))
        rast.bands[0].nodata_value = 0
        qs = RasterModel.objects.filter(geom__intersects=rast)
        # Raster does not overlap anymore after polygonization
        # where the nodata zone is not included.
        self.assertEqual(qs.count(), 0)

    def test_lookup_value_error(self):
        # Test with invalid dict lookup parameter
        obj = {}
        msg = "Couldn't create spatial object from lookup value '%s'." % obj
        with self.assertRaisesMessage(ValueError, msg):
            RasterModel.objects.filter(geom__intersects=obj)
        # Test with invalid string lookup parameter
        obj = "00000"
        msg = "Couldn't create spatial object from lookup value '%s'." % obj
        with self.assertRaisesMessage(ValueError, msg):
            RasterModel.objects.filter(geom__intersects=obj)

    def test_db_function_errors(self):
        """
        Errors are raised when using DB functions with raster content.
        """
        point = GEOSGeometry("SRID=3086;POINT (-697024.9213808845 683729.1705516104)")
        rast = GDALRaster(json.loads(JSON_RASTER))
        msg = "Distance function requires a geometric argument in position 2."
        with self.assertRaisesMessage(TypeError, msg):
            RasterModel.objects.annotate(distance_from_point=Distance("geom", rast))
        with self.assertRaisesMessage(TypeError, msg):
            RasterModel.objects.annotate(
                distance_from_point=Distance("rastprojected", rast)
            )
        msg = (
            "Distance function requires a GeometryField in position 1, got RasterField."
        )
        with self.assertRaisesMessage(TypeError, msg):
            RasterModel.objects.annotate(
                distance_from_point=Distance("rastprojected", point)
            ).count()

    def test_lhs_with_index_rhs_without_index(self):
        with CaptureQueriesContext(connection) as queries:
            RasterModel.objects.filter(
                rast__0__contains=json.loads(JSON_RASTER)
            ).exists()
        # It's easier to check the indexes in the generated SQL than to write
        # tests that cover all index combinations.
        self.assertRegex(queries[-1]["sql"], r"WHERE ST_Contains\([^)]*, 1, [^)]*, 1\)")
