# -*- coding: utf-8 -*-
from qgis.core import *
import numpy as np

class CartogramDCN(QgsProcessingAlgorithm):


    INPUT = "INPUT"
    WEIGHT = "WEIGHT"
    ITER = "ITER"
    RES = "RES"
    OUTPUT = "OUTPUT"

    def initAlgorithm(self, config=None):

        self.addParameter(
            QgsProcessingParameterFeatureSource(
                self.INPUT,
                "Couche polygonale",
                [QgsProcessing.TypeVectorPolygon]
            )
        )

        self.addParameter(
            QgsProcessingParameterField(
                self.WEIGHT,
                "Champ de poids",
                parentLayerParameterName=self.INPUT,
                type=QgsProcessingParameterField.Numeric
            )
        )

        self.addParameter(
            QgsProcessingParameterNumber(
                self.ITER,
                "Iterations diffusion",
                defaultValue=50
            )
        )

        self.addParameter(
            QgsProcessingParameterNumber(
                self.RES,
                "Résolution grille (coarse = rapide)",
                defaultValue=100
            )
        )

        self.addParameter(
            QgsProcessingParameterFeatureSink(
                self.OUTPUT,
                "Cartogramme Gastner-Newman"
            )
        )

    # ----------------------------
    def processAlgorithm(self, parameters, context, feedback):

        source = self.parameterAsSource(parameters, self.INPUT, context)
        weight_field = self.parameterAsString(parameters, self.WEIGHT, context)
        iterations = self.parameterAsInt(parameters, self.ITER, context)
        res = self.parameterAsInt(parameters, self.RES, context)

        features = list(source.getFeatures())

        # --- bounding box global ---
        extent = source.sourceExtent()
        minx, maxx = extent.xMinimum(), extent.xMaximum()
        miny, maxy = extent.yMinimum(), extent.yMaximum()

        dx = (maxx - minx) / res
        dy = (maxy - miny) / res

        # --- grille densité ---
        density = np.zeros((res, res), dtype=float)

        # projection des poids sur grille
        for f in features:
            geom = f.geometry()
            c = geom.centroid().asPoint()

            i = int((c.x() - minx) / dx)
            j = int((c.y() - miny) / dy)

            if 0 <= i < res and 0 <= j < res:
                density[j, i] += f[weight_field]

        density += 1e-9
        density /= density.sum()

        # --- diffusion (équation chaleur discrète) ---
        for it in range(iterations):

            if feedback.isCanceled():
                break

            lap = (
                np.roll(density, 1, axis=0) +
                np.roll(density, -1, axis=0) +
                np.roll(density, 1, axis=1) +
                np.roll(density, -1, axis=1) -
                4 * density
            )

            density += 0.2 * lap
            density = np.clip(density, 1e-12, None)
            density /= density.sum()

            feedback.setProgress(int(100 * it / iterations))

        # --- champ potentiel (approx transformation) ---
        potential = np.cumsum(np.cumsum(density, axis=0), axis=1)
        potential = (potential - potential.min()) / (potential.max() - potential.min())

        # gradient → déplacement
        gx, gy = np.gradient(potential)

        # --- sortie ---
        (sink, dest_id) = self.parameterAsSink(
            parameters,
            self.OUTPUT,
            context,
            source.fields(),
            QgsWkbTypes.MultiPolygon,
            source.sourceCrs()
        )

        for f in features:

            geom = f.geometry()
            pts = geom.vertices()

            new_pts = []

            for pt in pts:

                i = int((pt.x() - minx) / dx)
                j = int((pt.y() - miny) / dy)

                if 0 <= i < res and 0 <= j < res:
                    shift_x = gx[j, i]
                    shift_y = gy[j, i]
                else:
                    shift_x = shift_y = 0

                new_pts.append(QgsPointXY(
                    pt.x() + shift_x * dx * 10,
                    pt.y() + shift_y * dy * 10
                ))

            new_geom = QgsGeometry.fromPolylineXY(new_pts).buffer(0, 5)

            feat = QgsFeature()
            feat.setGeometry(new_geom)
            feat.setAttributes(f.attributes())

            sink.addFeature(feat, QgsFeatureSink.FastInsert)

        return {self.OUTPUT: dest_id}


    def name(self):
        return 'cartogram_dcn'

    def displayName(self):
        return 'Cartogramme DCN'

    def group(self):
        return 'Analysis'

    def groupId(self):
        return 'Analysis'

    def createInstance(self):
        return CartogramDCN()