# -*- coding: utf-8 -*-
# Tento skript předpokládá, že:
# 1) Máte nainstalovaný QGIS a skript spouštíte v QGIS Python konzoli nebo v prostředí s dostupným PyQGIS API.
# 2) Vstupní shapefile je v souřadnicovém systému, kde 30x30 km je smysluplné (např. metrický S-JTSK nebo UTM).
# 3) Vrstva bodů má atribut "pocob" (integer nebo float), který značí počet obyvatel.
#
# Rozšířený postup oproti předchozímu:
# 1) Vytvoří se paměťová vrstva s polygonovými čtverci 30x30 km.
# 2) Do této vrstvy se zapíše součet počtu obyvatel spadajících do čtverce (atribut "pocob").
# 3) Přidá se atribut "hustota", který bude počítán jako: hustota = pocob / (plocha čtverce v km2).
#    Plocha se vypočte na základě geometrie polygonu. Pokud je CRS v metrech, 
#    pak `plocha v m^2 / 1 000 000 = plocha v km^2`.

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

# Cesta k vstupní bodové vrstvě (Shapefile)
input_path = "C:/cesta/k_bodove_vrstev.shp"

# Načtení bodové vrstvy
point_layer = QgsVectorLayer(input_path, "body", "ogr")
if not point_layer.isValid():
    raise Exception("Vstupní vrstvu se nepodařilo načíst.")

# Zjištění extentu bodové vrstvy
extent = point_layer.extent()
# Rozměr čtverce 30x30 km = 30000x30000 m (předpokládáme metrický CRS)
cell_size = 30000.0

# Výpočet počtu buněk v x a y směru
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 mřížky
x_min = extent.xMinimum()
y_min = extent.yMinimum()

# Vytvoření paměťové vrstvy pro polygonové čtverce
fields = QgsFields()
fields.append(QgsField("pocob", QVariant.Double))
fields.append(QgsField("hustota", QVariant.Double))

mem_layer = QgsMemoryProviderUtils.createMemoryLayer(
    "ctverce", 
    fields, 
    QgsWkbTypes.Polygon,
    point_layer.crs()
)

provider = mem_layer.dataProvider()

# Načtení všech bodů 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í čtverců
new_features = []
for i in range(int(num_cells_x)):
    for j in range(int(num_cells_y)):
        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

        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)  # uzavření polygonu
        ]])

        feat = QgsFeature(fields)
        feat.setGeometry(polygon)
        # Zatím nastavíme pocob a hustotu na 0
        feat.setAttribute("pocob", 0.0)
        feat.setAttribute("hustota", 0.0)
        new_features.append(feat)

# Vložíme všechny generované čtverce do paměťové vrstvy
provider.addFeatures(new_features)
mem_layer.updateExtents()

# Nyní pro každý čtverec vypočteme součet "pocob" z bodů uvnitř
mem_features = list(mem_layer.getFeatures())

attr_map = {}
for poly_feat in mem_features:
    poly_geom = poly_feat.geometry()
    candidate_ids = index.intersects(poly_geom.boundingBox())
    sum_pocob = 0.0
    for cid in candidate_ids:
        point_feat = point_features[cid]
        if poly_geom.contains(point_feat.geometry()):
            val = point_feat["pocob"]
            if val is not None:
                sum_pocob += val

    # Výpočet hustoty: hustota = sum_pocob / (plocha polygonu v km²)
    # Plocha se získá voláním poly_geom.area() - vrací plošný obsah v jednotkách CRS, 
    # pokud je to metrický CRS, pak je to v m².
    area_m2 = poly_geom.area()
    area_km2 = area_m2 / 1e6  # převod m² na km²
    hustota = 0.0
    if area_km2 > 0:
        hustota = sum_pocob / area_km2

    # Změníme atributy polygonu
    attr_map[poly_feat.id()] = {0: sum_pocob, 1: hustota}

provider.changeAttributeValues(attr_map)

mem_layer.commitChanges()

# Přidání výsledné vrstvy do projektu
QgsProject.instance().addMapLayer(mem_layer)

# Volitelně lze vrstvu uložit na disk jako shapefile
# options = {}
# options["driverName"] = "ESRI Shapefile"
# options["fileEncoding"] = "UTF-8"
# QgsVectorFileWriter.writeAsVectorFormatV2(
#     mem_layer,
#     "C:/cesta/k_vystupnim_ctvercum.shp",
#     QgsCoordinateTransformContext(),
#     options
# )

# Tím je skript kompletní. Nyní máme dva atributy:
# "pocob" = součet počtu obyvatel ve čtverci
# "hustota" = hustota obyvatel na km² (pocob / plocha_polygonu_v_km²)