#!/usr/bin/env python3
"""Build the real PZS case-study data products from dataset_2.zip."""

from __future__ import annotations

import hashlib
import json
import math
import shutil
import tempfile
import zipfile
from pathlib import Path

import numpy as np
from osgeo import gdal, ogr, osr
from scipy import ndimage

gdal.UseExceptions()
ogr.UseExceptions()

ROOT = Path(__file__).resolve().parents[1]
ARCHIVE = ROOT / "data" / "examples" / "dataset_2.zip"
SOURCE_ASSETS = ROOT / "data" / "examples" / "pzs-method"
PUBLIC = ROOT / "site" / "downloads" / "pzs-method" / "v1.0.0"
STORY = ROOT / "site" / "assets" / "geostories" / "pzs-method"
BASE_WIDTH_M = 50.0
SLOPE_THRESHOLD_DEG = 3.0
VERSION = "1.0.0"


def sha256(path: Path) -> str:
    digest = hashlib.sha256()
    with path.open("rb") as source:
        for block in iter(lambda: source.read(1024 * 1024), b""):
            digest.update(block)
    return digest.hexdigest()


def find_file(directory: Path, name: str) -> Path:
    matches = list(directory.rglob(name))
    if len(matches) != 1:
        raise RuntimeError(f"Expected exactly one {name} in {directory}, found {len(matches)}.")
    return matches[0]


def memory_raster(reference: gdal.Dataset, data_type: int = gdal.GDT_Byte) -> gdal.Dataset:
    dataset = gdal.GetDriverByName("MEM").Create(
        "", reference.RasterXSize, reference.RasterYSize, 1, data_type
    )
    dataset.SetGeoTransform(reference.GetGeoTransform())
    dataset.SetProjection(reference.GetProjection())
    dataset.GetRasterBand(1).Fill(0)
    return dataset


def geometry_layer(name: str, geometry: ogr.Geometry, spatial_ref: osr.SpatialReference):
    datasource = ogr.GetDriverByName("MEM").CreateDataSource("")
    layer = datasource.CreateLayer(name, spatial_ref, geometry.GetGeometryType())
    feature = ogr.Feature(layer.GetLayerDefn())
    feature.SetGeometry(geometry)
    layer.CreateFeature(feature)
    return datasource, layer


def combined_water_geometry(water_path: Path) -> tuple[ogr.Geometry, int]:
    datasource = ogr.Open(str(water_path))
    if datasource is None:
        raise RuntimeError(f"Could not open {water_path}.")

    combined = None
    feature_count = 0
    for index in range(datasource.GetLayerCount()):
        layer = datasource.GetLayerByIndex(index)
        for feature in layer:
            geometry = feature.GetGeometryRef()
            if geometry is None or geometry.IsEmpty():
                continue
            piece = geometry.Clone()
            geometry_type = ogr.GT_Flatten(piece.GetGeometryType())
            if geometry_type in (ogr.wkbLineString, ogr.wkbMultiLineString):
                piece = piece.Buffer(math.sqrt(2))
            combined = piece if combined is None else combined.Union(piece)
            feature_count += 1

    if combined is None or combined.IsEmpty():
        raise RuntimeError("The water GeoPackage does not contain usable geometry.")
    return combined.MakeValid(), feature_count


def write_zone_raster(
    output_path: Path,
    reference: gdal.Dataset,
    classes: np.ndarray,
) -> gdal.Dataset:
    if output_path.exists():
        output_path.unlink()
    output = gdal.GetDriverByName("GTiff").Create(
        str(output_path),
        reference.RasterXSize,
        reference.RasterYSize,
        1,
        gdal.GDT_Byte,
        options=["TILED=YES", "COMPRESS=DEFLATE", "PREDICTOR=2"],
    )
    output.SetGeoTransform(reference.GetGeoTransform())
    output.SetProjection(reference.GetProjection())
    band = output.GetRasterBand(1)
    band.SetDescription("PZS screening class: 1=standard 50 m, 2=slope extension to 100 m")
    band.SetNoDataValue(0)
    band.WriteArray(classes)
    band.ComputeStatistics(False)
    output.FlushCache()
    return output


def polygonize_zones(
    raster: gdal.Dataset,
    output_path: Path,
    spatial_ref: osr.SpatialReference,
) -> None:
    with tempfile.TemporaryDirectory(prefix="pzs_polygonize_") as temp_name:
        raw_path = Path(temp_name) / "raw.gpkg"
        raw_ds = ogr.GetDriverByName("GPKG").CreateDataSource(str(raw_path))
        raw_layer = raw_ds.CreateLayer("raw", spatial_ref, ogr.wkbPolygon)
        raw_layer.CreateField(ogr.FieldDefn("zone_class", ogr.OFTInteger))
        band = raster.GetRasterBand(1)
        gdal.Polygonize(band, band.GetMaskBand(), raw_layer, 0, [])
        raw_layer.SyncToDisk()

        dissolved: dict[int, ogr.Geometry] = {}
        raw_layer.ResetReading()
        for feature in raw_layer:
            zone_class = feature.GetFieldAsInteger("zone_class")
            if zone_class not in (1, 2):
                continue
            geometry = feature.GetGeometryRef().Clone()
            current = dissolved.get(zone_class)
            dissolved[zone_class] = geometry if current is None else current.Union(geometry)
        raw_ds = None

    if output_path.exists():
        output_path.unlink()
    output_ds = ogr.GetDriverByName("GPKG").CreateDataSource(str(output_path))
    output_layer = output_ds.CreateLayer("pzs_zones", spatial_ref, ogr.wkbMultiPolygon)
    fields = (
        ("zone_class", ogr.OFTInteger),
        ("type_en", ogr.OFTString),
        ("type_uk", ogr.OFTString),
        ("base_m", ogr.OFTReal),
        ("max_m", ogr.OFTReal),
        ("slope_deg", ogr.OFTReal),
        ("method_ver", ogr.OFTString),
        ("status_en", ogr.OFTString),
        ("status_uk", ogr.OFTString),
    )
    for name, field_type in fields:
        definition = ogr.FieldDefn(name, field_type)
        if field_type == ogr.OFTString:
            definition.SetWidth(160)
        output_layer.CreateField(definition)

    labels = {
        1: ("Standard-width screening zone", "Зона нормативної ширини", BASE_WIDTH_M),
        2: ("Slope-triggered extension", "Розширення через ухил", BASE_WIDTH_M * 2),
    }
    for zone_class, geometry in sorted(dissolved.items()):
        geometry = geometry.MakeValid().SimplifyPreserveTopology(math.sqrt(2))
        if ogr.GT_Flatten(geometry.GetGeometryType()) == ogr.wkbPolygon:
            geometry = ogr.ForceToMultiPolygon(geometry)
        feature = ogr.Feature(output_layer.GetLayerDefn())
        feature.SetField("zone_class", zone_class)
        feature.SetField("type_en", labels[zone_class][0])
        feature.SetField("type_uk", labels[zone_class][1])
        feature.SetField("base_m", BASE_WIDTH_M)
        feature.SetField("max_m", labels[zone_class][2])
        feature.SetField("slope_deg", SLOPE_THRESHOLD_DEG)
        feature.SetField("method_ver", VERSION)
        feature.SetField("status_en", "Analytical screening output; not a legal boundary")
        feature.SetField("status_uk", "Аналітичний результат; не є юридичною межею")
        feature.SetGeometry(geometry)
        output_layer.CreateFeature(feature)
    output_layer.SyncToDisk()
    output_ds = None


def translate_vector(
    source: Path,
    destination: Path,
    layers: list[str],
    simplify_tolerance: float | None = None,
) -> None:
    if destination.exists():
        destination.unlink()
    options = gdal.VectorTranslateOptions(
        format="GeoJSON",
        dstSRS="EPSG:4326",
        layers=layers,
        simplifyTolerance=simplify_tolerance,
        layerCreationOptions=["RFC7946=YES", "COORDINATE_PRECISION=6"],
    )
    result = gdal.VectorTranslate(str(destination), str(source), options=options)
    if result is None:
        raise RuntimeError(f"Could not create {destination}.")
    result = None


def wgs84_extent(reference: gdal.Dataset) -> list[float]:
    transform = reference.GetGeoTransform()
    width = reference.RasterXSize
    height = reference.RasterYSize
    source = osr.SpatialReference()
    source.ImportFromWkt(reference.GetProjection())
    target = osr.SpatialReference()
    target.ImportFromEPSG(4326)
    source.SetAxisMappingStrategy(osr.OAMS_TRADITIONAL_GIS_ORDER)
    target.SetAxisMappingStrategy(osr.OAMS_TRADITIONAL_GIS_ORDER)
    transformer = osr.CoordinateTransformation(source, target)
    corners = [
        transformer.TransformPoint(transform[0], transform[3]),
        transformer.TransformPoint(
            transform[0] + width * transform[1],
            transform[3] + height * transform[5],
        ),
    ]
    return [
        round(min(point[0] for point in corners), 6),
        round(min(point[1] for point in corners), 6),
        round(max(point[0] for point in corners), 6),
        round(max(point[1] for point in corners), 6),
    ]


def main() -> None:
    required = [
        ARCHIVE,
        SOURCE_ASSETS / "model_pzs.model3",
        SOURCE_ASSETS / "pzs-method-theory.pdf",
        SOURCE_ASSETS / "LICENSE",
    ]
    missing = [str(path) for path in required if not path.is_file()]
    if missing:
        raise RuntimeError("Missing source files:\n" + "\n".join(missing))

    PUBLIC.mkdir(parents=True, exist_ok=True)
    STORY.mkdir(parents=True, exist_ok=True)

    with tempfile.TemporaryDirectory(prefix="pzs_dataset_") as temp_name:
        extracted = Path(temp_name)
        with zipfile.ZipFile(ARCHIVE) as archive:
            archive.extractall(extracted)

        aoi_path = find_file(extracted, "aoi.gpkg")
        dem_path = find_file(extracted, "dem.tif")
        water_path = find_file(extracted, "water.gpkg")

        dem_ds = gdal.Open(str(dem_path))
        if dem_ds is None:
            raise RuntimeError(f"Could not open {dem_path}.")
        dem_band = dem_ds.GetRasterBand(1)
        dem = dem_band.ReadAsArray().astype(np.float32)
        nodata = dem_band.GetNoDataValue()
        pixel_x = abs(dem_ds.GetGeoTransform()[1])
        pixel_y = abs(dem_ds.GetGeoTransform()[5])
        spatial_ref = osr.SpatialReference()
        spatial_ref.ImportFromWkt(dem_ds.GetProjection())
        spatial_ref.SetAxisMappingStrategy(osr.OAMS_TRADITIONAL_GIS_ORDER)

        aoi_ds = ogr.Open(str(aoi_path))
        aoi_layer = aoi_ds.GetLayer(0)
        aoi_mask_ds = memory_raster(dem_ds)
        gdal.RasterizeLayer(
            aoi_mask_ds, [1], aoi_layer, burn_values=[1], options=["ALL_TOUCHED=TRUE"]
        )
        aoi_mask = aoi_mask_ds.GetRasterBand(1).ReadAsArray().astype(bool)

        water_geometry, water_feature_count = combined_water_geometry(water_path)
        water_source, water_layer = geometry_layer("water", water_geometry, spatial_ref)
        water_mask_ds = memory_raster(dem_ds)
        gdal.RasterizeLayer(
            water_mask_ds, [1], water_layer, burn_values=[1], options=["ALL_TOUCHED=TRUE"]
        )
        water_mask = water_mask_ds.GetRasterBand(1).ReadAsArray().astype(bool)

        boundary_geometry = water_geometry.Boundary()
        boundary_source, boundary_layer = geometry_layer(
            "water_boundary", boundary_geometry, spatial_ref
        )
        boundary_mask_ds = memory_raster(dem_ds)
        gdal.RasterizeLayer(
            boundary_mask_ds,
            [1],
            boundary_layer,
            burn_values=[1],
            options=["ALL_TOUCHED=TRUE"],
        )
        boundary_mask = boundary_mask_ds.GetRasterBand(1).ReadAsArray().astype(bool)

        valid_dem = np.isfinite(dem)
        if nodata is not None:
            valid_dem &= ~np.isclose(dem, nodata)
        analysis_mask = valid_dem & aoi_mask
        distance, nearest = ndimage.distance_transform_edt(
            ~boundary_mask,
            sampling=(pixel_y, pixel_x),
            return_indices=True,
        )
        nearest_rows, nearest_cols = nearest
        nearest_elevation = dem[nearest_rows, nearest_cols]
        nearest_valid = valid_dem[nearest_rows, nearest_cols]
        nearest_id = (
            nearest_rows.astype(np.int64) * dem_ds.RasterXSize
            + nearest_cols.astype(np.int64)
        )

        ring = (
            analysis_mask
            & ~water_mask
            & nearest_valid
            & (distance >= BASE_WIDTH_M)
            & (distance < BASE_WIDTH_M + max(pixel_x, pixel_y))
        )
        slope_ratio = np.zeros_like(dem, dtype=np.float32)
        np.divide(
            dem - nearest_elevation,
            distance,
            out=slope_ratio,
            where=ring & (distance > 0),
        )
        slope_trigger = ring & (slope_ratio > math.tan(math.radians(SLOPE_THRESHOLD_DEG)))
        flagged_ids = np.unique(nearest_id[slope_trigger])
        flag_lookup = np.zeros(dem_ds.RasterXSize * dem_ds.RasterYSize, dtype=bool)
        flag_lookup[flagged_ids] = True

        land = analysis_mask & ~water_mask
        standard_zone = land & (distance > 0) & (distance <= BASE_WIDTH_M)
        extension_zone = (
            land
            & (distance > BASE_WIDTH_M)
            & (distance <= BASE_WIDTH_M * 2)
            & flag_lookup[nearest_id]
        )
        classes = np.zeros(dem.shape, dtype=np.uint8)
        classes[standard_zone] = 1
        classes[extension_zone] = 2

        raster_path = PUBLIC / "pzs-screening-50m.tif"
        zone_raster = write_zone_raster(raster_path, dem_ds, classes)
        gpkg_path = PUBLIC / "pzs-screening-zones.gpkg"
        polygonize_zones(zone_raster, gpkg_path, spatial_ref)
        geojson_path = PUBLIC / "pzs-screening-zones.geojson"
        translate_vector(gpkg_path, geojson_path, ["pzs_zones"], math.sqrt(2))
        shutil.copy2(geojson_path, STORY / "pzs-screening-zones.geojson")

        translate_vector(aoi_path, STORY / "aoi.geojson", ["aoi"])
        translate_vector(
            water_path,
            STORY / "water-polygons.geojson",
            ["ter_water_pol"],
            1.0,
        )
        translate_vector(
            water_path,
            STORY / "water-lines.geojson",
            ["ter_water_line"],
            1.0,
        )

        hillshade_path = extracted / "hillshade.tif"
        gdal.DEMProcessing(
            str(hillshade_path),
            dem_ds,
            "hillshade",
            computeEdges=True,
            creationOptions=["TILED=YES", "COMPRESS=DEFLATE"],
        )
        gdal.Translate(
            str(STORY / "terrain.png"),
            str(hillshade_path),
            format="PNG",
            width=1100,
            resampleAlg="bilinear",
        )
        terrain_aux = STORY / "terrain.png.aux.xml"
        if terrain_aux.exists():
            terrain_aux.unlink()

        copies = {
            ARCHIVE: PUBLIC / "dataset_2.zip",
            aoi_path: PUBLIC / "source-aoi.gpkg",
            dem_path: PUBLIC / "source-dem.tif",
            water_path: PUBLIC / "source-water.gpkg",
            SOURCE_ASSETS / "model_pzs.model3": PUBLIC / "model_pzs.model3",
            SOURCE_ASSETS / "pzs-method-theory.pdf": PUBLIC / "pzs-method-theory.pdf",
            SOURCE_ASSETS / "LICENSE": PUBLIC / "LICENSE",
            Path(__file__): PUBLIC / "build-pzs-case-study.py",
        }
        for source, destination in copies.items():
            shutil.copy2(source, destination)

        pixel_area = pixel_x * pixel_y
        elevations = dem[analysis_mask]
        summary = {
            "title": "Riparian protective-strip screening case study",
            "version": VERSION,
            "publication_date": "2026-08-16",
            "source_dataset": "dataset_2.zip",
            "source_archive_sha256": sha256(ARCHIVE),
            "source_repository": "https://github.com/oleksaboiko/pzs_method",
            "authors": ["Oleksii Boiko", "Yuliia Maksymova"],
            "licence": "CC BY 4.0",
            "crs": "EPSG:5564",
            "web_crs": "EPSG:4326",
            "extent_wgs84": wgs84_extent(dem_ds),
            "pixel_size_m": pixel_x,
            "base_width_m": BASE_WIDTH_M,
            "maximum_width_m": BASE_WIDTH_M * 2,
            "slope_threshold_degrees": SLOPE_THRESHOLD_DEG,
            "water_feature_count": water_feature_count,
            "aoi_area_km2": round(float(np.count_nonzero(analysis_mask) * pixel_area / 1e6), 3),
            "standard_zone_area_km2": round(
                float(np.count_nonzero(standard_zone) * pixel_area / 1e6), 3
            ),
            "slope_extension_area_km2": round(
                float(np.count_nonzero(extension_zone) * pixel_area / 1e6), 3
            ),
            "flagged_shoreline_pixel_count": int(flagged_ids.size),
            "elevation_min_m": round(float(elevations.min()), 1),
            "elevation_max_m": round(float(elevations.max()), 1),
            "derivation": (
                "Equivalent array implementation of the repository method: rasterized "
                "water boundary, nearest-boundary distance and elevation, a 3-degree "
                "test at the 50 m ring, and selective extension to 100 m."
            ),
            "legal_status": (
                "Analytical screening output only. It is not an official or legally "
                "established boundary and does not replace land-management documentation."
            ),
            "original_qgis_model_note": (
                "The original model is redistributed unchanged. Reproduction requires "
                "a QGIS project in EPSG:5564 because the model uses ProjectCrs and "
                "@project_folder references."
            ),
        }
        metadata = json.dumps(summary, ensure_ascii=False, indent=2) + "\n"
        (PUBLIC / "metadata.json").write_text(metadata, encoding="utf-8")
        (STORY / "summary.json").write_text(metadata, encoding="utf-8")

        checksum_files = sorted(
            path for path in PUBLIC.iterdir() if path.is_file() and path.name != "SHA256SUMS"
        )
        checksum_text = "".join(
            f"{sha256(path)}  {path.name}\n" for path in checksum_files
        )
        (PUBLIC / "SHA256SUMS").write_text(checksum_text, encoding="utf-8")

        print(
            "Built PZS case study:",
            f"{summary['aoi_area_km2']} km² AOI,",
            f"{summary['standard_zone_area_km2']} km² standard zone,",
            f"{summary['slope_extension_area_km2']} km² slope extension.",
        )


if __name__ == "__main__":
    main()