import numpy as np
from qgis._core import QgsRasterLayer, QgsVectorLayer, QgsPalettedRasterRenderer

from enmapbox.exampledata import enmap, landcover_polygons
from enmapboxprocessing.algorithm.rasterizecategorizedvectoralgorithm import RasterizeCategorizedVectorAlgorithm
from enmapboxprocessing.rasterreader import RasterReader
from enmapboxprocessing.test.algorithm.testcase import TestCase
from enmapboxprocessing.utils import Utils
from enmapboxtestdata import (landcover_polygons_3classes_epsg4326, landcover_polygons_3classes_id,
                              landcover_points_multipart_epsg3035)


class TestRasterizeCategorizedVectorAlgorithm(TestCase):

    def test_numberClassAttribute(self):
        alg = RasterizeCategorizedVectorAlgorithm()
        alg.initAlgorithm()
        parameters = {
            alg.P_CATEGORIZED_VECTOR: landcover_polygons_3classes_id,
            alg.P_GRID: enmap,
            alg.P_OUTPUT_CATEGORIZED_RASTER: self.filename('landcover_polygons.tif')
        }
        result = self.runalg(alg, parameters)
        classification = QgsRasterLayer(parameters[alg.P_OUTPUT_CATEGORIZED_RASTER])
        self.assertIsInstance(classification.renderer(), QgsPalettedRasterRenderer)
        vector = QgsVectorLayer(parameters[alg.P_CATEGORIZED_VECTOR])
        for c1, c2 in zip(
                Utils.categoriesFromCategorizedSymbolRenderer(vector.renderer()),
                Utils.categoriesFromPalettedRasterRenderer(classification.renderer())
        ):
            self.assertEqual((c1.name, c1.color), (c2.name, c2.color))

        self.assertEqual(1678, np.sum(RasterReader(result[alg.P_OUTPUT_CATEGORIZED_RASTER]).array()))

    def test_stringClassAttribute(self):
        alg = RasterizeCategorizedVectorAlgorithm()
        alg.initAlgorithm()
        parameters = {
            alg.P_CATEGORIZED_VECTOR: QgsVectorLayer(landcover_polygons),
            alg.P_GRID: QgsRasterLayer(enmap),
            alg.P_OUTPUT_CATEGORIZED_RASTER: self.filename('landcover_polygons.tif')
        }
        result = self.runalg(alg, parameters)
        classification = QgsRasterLayer(parameters[alg.P_OUTPUT_CATEGORIZED_RASTER])
        self.assertIsInstance(classification.renderer(), QgsPalettedRasterRenderer)
        for c1, c2 in zip(
                Utils.categoriesFromCategorizedSymbolRenderer(parameters[alg.P_CATEGORIZED_VECTOR].renderer()),
                Utils.categoriesFromPalettedRasterRenderer(classification.renderer())
        ):
            self.assertEqual((c1.name, c1.color), (c2.name, c2.color))
        self.assertEqual(5108, np.sum(RasterReader(result[alg.P_OUTPUT_CATEGORIZED_RASTER]).array()))

    def test_withNoneMatching_crs(self):
        alg = RasterizeCategorizedVectorAlgorithm()
        alg.initAlgorithm()
        parameters = {
            alg.P_CATEGORIZED_VECTOR: QgsVectorLayer(landcover_polygons_3classes_epsg4326),
            alg.P_GRID: QgsRasterLayer(enmap),
            alg.P_OUTPUT_CATEGORIZED_RASTER: self.filename('landcover_polygons.tif')
        }
        result = self.runalg(alg, parameters)
        self.assertEqual(1678, np.sum(RasterReader(result[alg.P_OUTPUT_CATEGORIZED_RASTER]).array()))

    def test_pointVector(self):
        alg = RasterizeCategorizedVectorAlgorithm()
        alg.initAlgorithm()
        parameters = {
            alg.P_CATEGORIZED_VECTOR: QgsVectorLayer(landcover_points_multipart_epsg3035),
            alg.P_GRID: QgsRasterLayer(enmap),
            alg.P_OUTPUT_CATEGORIZED_RASTER: self.filename('landcover_points.tif')
        }
        result = self.runalg(alg, parameters)
        self.assertEqual(152, np.sum(RasterReader(result[alg.P_OUTPUT_CATEGORIZED_RASTER]).array()))

    def test_minimalCoverage(self):
        alg = RasterizeCategorizedVectorAlgorithm()
        alg.initAlgorithm()
        parameters = {
            alg.P_CATEGORIZED_VECTOR: QgsVectorLayer(landcover_polygons),
            alg.P_GRID: QgsRasterLayer(enmap),
            alg.P_COVERAGE: 100,
            alg.P_OUTPUT_CATEGORIZED_RASTER: self.filename('classification_fullcoverage.tif')
        }

        result = self.runalg(alg, parameters)
        self.assertEqual(3816, np.sum(RasterReader(result[alg.P_OUTPUT_CATEGORIZED_RASTER]).array()))
