# -*- coding: utf-8 -*-
"""M2 Bab 4: hidrologi dari DEM: isi lubang (fill sinks), arah dan akumulasi aliran, sungai, dan DAS di satu titik keluar (outlet).
Penulis: Badar Mubarok Yogaswara. Pemakaian: python-qgis.bat m2_04a_hidrologi.py <folder paket-m2>
Keluaran disimpan di <paket-m2>/hasil. Ambang sungai: 500 sel (500 x 100 m2 = 5 ha)."""
import os
import sys
import numpy as np
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
import _inisialisasi  # noqa: F401
import processing
from osgeo import gdal

gdal.UseExceptions()
paket = sys.argv[1]
out = os.path.join(paket, "hasil")
os.makedirs(out, exist_ok=True)
dem = os.path.join(paket, "DEM_10m.tif")
P = lambda n: os.path.join(out, n)
BAKU = {"GRASS_REGION_PARAMETER": None, "GRASS_REGION_CELLSIZE_PARAMETER": 0}


def baca(path):
    ds = gdal.Open(path)
    return ds.GetRasterBand(1).ReadAsArray().astype("float64"), ds.GetGeoTransform()


def jumlah_lubang(a):
    """jumlah sel dalam (bukan tepi) yang lebih rendah daripada semua 8 tetangganya"""
    c = a[1:-1, 1:-1]
    lebih_tinggi = np.ones_like(c, dtype=bool)
    for dy in (-1, 0, 1):
        for dx in (-1, 0, 1):
            if dy == 0 and dx == 0:
                continue
            lebih_tinggi &= a[1 + dy:a.shape[0] - 1 + dy, 1 + dx:a.shape[1] - 1 + dx] > c
    return int(lebih_tinggi.sum())


z0, gt = baca(dem)
print("lubang (sink) pada DEM asli            :", jumlah_lubang(z0))

# 1) isi lubang (Wang & Liu), kemiringan minimum 0,1 derajat
processing.run("native:fillsinkswangliu", {"INPUT": dem, "BAND": 1, "MIN_SLOPE": 0.1, "OUTPUT_FILLED_DEM": P("DEM_terisi.tif"),
                                           "OUTPUT_FLOW_DIRECTIONS": P("Arah_WL.tif"), "OUTPUT_WATERSHED_BASINS": P("Basin_WL.tif")})
z1, _ = baca(P("DEM_terisi.tif"))
print("lubang sesudah diisi                   :", jumlah_lubang(z1))
print("sel yang naik / kenaikan terbesar (m)  : %d / %.2f" % (int((z1 > z0 + 1e-6).sum()), float((z1 - z0).max())))

# 2) akumulasi aliran, arah, dan cekungan dengan GRASS r.watershed (aliran tunggal D8)
processing.run("grass:r.watershed", dict(BAKU, elevation=P("DEM_terisi.tif"), threshold=500, convergence=5, memory=300, **{"-s": True, "-a": True},
               accumulation=P("Akumulasi.tif"), drainage=P("Arah.tif"), basin=P("Cekungan.tif")))
acc, _ = baca(P("Akumulasi.tif"))
print("akumulasi maksimum (sel) / luas (ha)   : %d / %.1f" % (abs(acc).max(), abs(acc).max() * 100 / 1e4))

# 3) sungai: sel dengan akumulasi >= 500 sel (5 ha)
processing.run("grass:r.stream.extract", dict(BAKU, elevation=P("DEM_terisi.tif"), accumulation=P("Akumulasi.tif"), threshold=500, mexp=0,
               stream_length=0, memory=300, stream_raster=P("Sungai.tif"), stream_vector=P("sungai_mentah.gpkg"), direction=P("Arah_stream.tif"),
               GRASS_OUTPUT_TYPE_PARAMETER=2, GRASS_VECTOR_EXPORT_NOCAT=False))
if os.path.exists(P("Sungai.gpkg")):
    os.remove(P("Sungai.gpkg"))
gdal.VectorTranslate(P("Sungai.gpkg"), P("sungai_mentah.gpkg"), options=gdal.VectorTranslateOptions(
    format="GPKG", layerName="Sungai", selectFields=["cat"], geometryType="LINESTRING"))
from osgeo import ogr
ds = ogr.Open(P("Sungai.gpkg"))
lyr = ds.GetLayer(0)
panjang = sum(f.GetGeometryRef().Length() for f in lyr)
print("sungai: %d segmen, total panjang %.2f km" % (lyr.GetFeatureCount(), panjang / 1000.0))
ds = None


# 4) titik keluar (outlet): sel tepi dengan akumulasi terbesar, lalu DAS-nya
acc, gt = baca(P("Akumulasi.tif"))
tepi = np.zeros_like(acc, dtype=bool)
tepi[0, :] = tepi[-1, :] = tepi[:, 0] = tepi[:, -1] = True
r, c = np.unravel_index(np.argmax(np.where(tepi, acc, -1)), acc.shape)
ox, oy = gt[0] + (c + 0.5) * gt[1], gt[3] + (r + 0.5) * gt[5]
print("outlet: baris %d kolom %d, X=%.1f Y=%.1f, akumulasi %d sel" % (r, c, ox, oy, acc[r, c]))
processing.run("grass:r.water.outlet", dict(BAKU, input=P("Arah.tif"), coordinates="%.1f,%.1f [EPSG:32749]" % (ox, oy), output=P("DAS_outlet.tif")))
das, _ = baca(P("DAS_outlet.tif"))
n_das = int((das == 1).sum())
print("DAS di outlet: %d sel = %.2f ha (%.1f%% wilayah)" % (n_das, n_das * 100 / 1e4, 100.0 * n_das / das.size))
mentah = processing.run("gdal:polygonize", {"INPUT": P("DAS_outlet.tif"), "BAND": 1, "FIELD": "nilai", "EIGHT_CONNECTEDNESS": False,
                                            "EXTRA": "", "OUTPUT": "TEMPORARY_OUTPUT"})["OUTPUT"]
dalam = processing.run("native:extractbyexpression", {"INPUT": mentah, "EXPRESSION": '"nilai" = 1', "OUTPUT": "memory:"})["OUTPUT"]
if os.path.exists(P("DAS_utama.gpkg")):
    os.remove(P("DAS_utama.gpkg"))
processing.run("native:dissolve", {"INPUT": dalam, "FIELD": [], "SEPARATE_DISJOINT": False, "OUTPUT": P("DAS_utama.gpkg")})
print("DAS poligon ditulis:", os.path.exists(P("DAS_utama.gpkg")))