from qgis.core import (
    QgsVectorLayer,
    QgsProject,
    QgsFeature,
    QgsGeometry,
    QgsField,
    QgsFields,
    QgsPointXY,
    QgsRectangle,
    QgsSpatialIndex,
    QgsFeatureSink,
    QgsWkbTypes,
    QgsCoordinateReferenceSystem,
    QgsCoordinateTransformContext,
    QgsMemoryProviderUtils
)
from qgis.PyQt.QtCore import QVariant

# Cesta k vstupní bodové vrstve (Shapefile)
input_path = "/cesta/k_bodove_vrstev.shp"

# Nactení bodové vrstvy
point_layer = QgsVectorLayer(input_path, "body", "ogr")
if not point_layer.isValid():
    raise Exception("Vstupní vrstvu se nepodarilo nacíst.")

# Zjištení extentu bodové vrstvy
extent = point_layer.extent()
# Rozmer ctverce 30x30 km = 30000x30000 m (predpokládáme metrický CRS)
cell_size = 30000.0

# Výpocet poctu bunek v x a y smeru
# Zaokrouhlíme nahoru aby extent byl plne pokryt
import math
width = extent.width()
height = extent.height()

num_cells_x = math.ceil(width / cell_size)
num_cells_y = math.ceil(height / cell_size)

# Výchozí levý dolní roh mrížky
x_min = extent.xMinimum()
y_min = extent.yMinimum()

# Vytvorení pametové vrstvy pro polygonové ctverce
fields = QgsFields()
fields.append(QgsField("pocob", QVariant.Double))

# Geometrický typ: Polygon
epsg = point_layer.crs().authid()
mem_layer = QgsMemoryProviderUtils.createMemoryLayer(
    "ctverce", 
    fields, 
    QgsWkbTypes.Polygon,
    point_layer.crs()
)

provider = mem_layer.dataProvider()

# Nactení všech bodu do seznamu a tvorba prostorového indexu
point_features = list(point_layer.getFeatures())
index = QgsSpatialIndex()
for f in point_features:
    index.insertFeature(f)

# Generování ctvercu
new_features = []
for i in range(int(num_cells_x)):
    for j in range(int(num_cells_y)):
        # Vypocítáme extent aktuálního ctverce
        x_lower = x_min + i * cell_size
        x_upper = x_lower + cell_size
        y_lower = y_min + j * cell_size
        y_upper = y_lower + cell_size

        # Vytvorení polygonu ctverce
        # Ctverec bude polygon definovaný ctyrmi vrcholy:
        # (x_lower, y_lower) - (x_upper, y_lower) - (x_upper, y_upper) - (x_lower, y_upper)
        polygon = QgsGeometry.fromPolygonXY([[
            QgsPointXY(x_lower, y_lower),
            QgsPointXY(x_upper, y_lower),
            QgsPointXY(x_upper, y_upper),
            QgsPointXY(x_lower, y_upper),
            QgsPointXY(x_lower, y_lower)  # uzavrení polygonu
        ]])

        # Vytvoríme nový prvek polygonové vrstvy
        feat = QgsFeature(fields)
        feat.setGeometry(polygon)
        # Zatím nastavíme pocob na 0
        feat.setAttribute("pocob", 0.0)
        new_features.append(feat)

# Vložíme všechny generované ctverce do pametové vrstvy
provider.addFeatures(new_features)
mem_layer.updateExtents()

# Nyní pro každý ctverec vypocteme soucet "pocob" z bodu uvnitr
# Abychom mohli aktualizovat, potrebujeme index prvku polygonové vrstvy
mem_features = mem_layer.getFeatures()
provider.changeAttributeValues({})  # jen inicializace

for poly_feat in mem_features:
    poly_geom = poly_feat.geometry()
    # Najdeme kandidátní body pres prostorový index (rámecek polygonu)
    candidate_ids = index.intersects(poly_geom.boundingBox())
    sum_pocob = 0.0
    for cid in candidate_ids:
        point_feat = point_features[cid]
        # Zkontrolujeme, zda bod opravdu leží uvnitr polygonu
        if poly_geom.contains(point_feat.geometry()):
            val = point_feat["pocob"]
            if val is not None:
                sum_pocob += val

    # Zmeníme atribut "pocob" polygonu na sectenou hodnotu
    provider.changeAttributeValues({poly_feat.id(): {0: sum_pocob}})

mem_layer.commitChanges()

# Výsledek: mem_layer je pametová vrstva s ctverci a atributem "pocob" jako sumou obyvatel.
# Pokud chcete vrstvu zobrazit:
QgsProject.instance().addMapLayer(mem_layer)

# Pro uložení na disk jako shapefile:
# options = {}
# options["driverName"] = "ESRI Shapefile"
# options["fileEncoding"] = "UTF-8"
# QgsVectorFileWriter.writeAsVectorFormatV2(
#     mem_layer,
#     "/cesta/k_vystupnim_ctvercum.shp",
#     QgsCoordinateTransformContext(),
#     options
# )

# Tím je skript kompletní.