"""
Landlab 地形・水文解析サンプル集

目的:
1. DEMを使った流向・集水面積の計算
2. 流域解析・地形量解析のプロトタイプ
3. 降雨流出・表面流モデルの学習用サンプル
4. 侵食・堆積・河床変動モデルの研究開発
5. Pythonによる地形・水文解析コードの教育・説明資料

このサンプルは、Landlabの代表的なコンポーネントの使い方を、実務プロトタイプに
展開しやすい形で整理したものです。実DEMを使う場合は、rasterio等でGeoTIFFを
読み込み、topographic__elevationに標高配列を登録してください。

実務への展開例:
- JGD2011平面直角座標系など、m単位の投影座標系DEMを使う
- QGIS/ArcGISで作成した流域界・河道線・保全対象と重ねる
- drainage_areaのしきい値を変えて谷線候補を抽出する
- 勾配・曲率・集水面積を組み合わせて、土砂供給域や斜面流出域を仮説評価する
- 降雨流出の学習では、降雨強度、粗度係数、流域形状、到達時間の関係を見る

実行例:
    pip install landlab numpy matplotlib
    python scripts/landlab_terrain_hydrology_examples.py
"""

from __future__ import annotations

from pathlib import Path

import matplotlib.pyplot as plt
import numpy as np
from landlab import RasterModelGrid
from landlab.components import FlowAccumulator, FastscapeEroder, LinearDiffuser
from landlab.plot import imshow_grid


OUT_DIR = Path("public/demo-data/landlab")
DX = 10.0
NROWS = 80
NCOLS = 110


def create_synthetic_dem(grid: RasterModelGrid) -> np.ndarray:
    """谷地形を持つ人工DEMを作成する。"""
    x = grid.x_of_node
    y = grid.y_of_node
    center_y = 0.54 * y.max() + 55.0 * np.sin(x / 220.0)
    dem = 145.0 - 0.016 * x - 0.030 * y
    dem += -10.0 * np.exp(-((y - center_y) ** 2) / (2.0 * 90.0**2))
    dem += 1.0 * np.sin(x / 65.0) * np.cos(y / 90.0)
    return dem


def setup_grid() -> RasterModelGrid:
    """RasterModelGridを作成し、境界条件と標高fieldを設定する。"""
    grid = RasterModelGrid((NROWS, NCOLS), xy_spacing=DX)
    z = create_synthetic_dem(grid)
    grid.add_field("topographic__elevation", z, at="node", clobber=True)

    # 右・上・左を閉境界、下側のみ開境界として、下流へ流出させる。
    grid.set_closed_boundaries_at_grid_edges(
        right_is_closed=True,
        top_is_closed=True,
        left_is_closed=True,
        bottom_is_closed=False,
    )
    return grid


def example_1_flow_direction_and_drainage(grid: RasterModelGrid) -> None:
    """DEMから流向・集水面積を計算する。"""
    fa = FlowAccumulator(grid, flow_director="D8")
    fa.run_one_step()

    drainage_area = grid.at_node["drainage_area"]
    receiver = grid.at_node["flow__receiver_node"]

    print("[1] Flow accumulation")
    print(f"    max drainage area: {drainage_area.max():,.1f} m2")
    print(f"    receiver node sample: {receiver[:5]}")


def example_2_terrain_metrics(grid: RasterModelGrid) -> None:
    """勾配・集水面積などの地形量を作成する。

    プロトタイプ段階では、単一の正解指標を作るよりも、複数の地形量を並べて
    「どの場所が水を集めやすいか」「どこが急勾配か」「既往災害位置と合うか」を
    比較できるようにすることが重要です。
    """
    slope = grid.at_node["topographic__steepest_slope"]
    drainage_area = grid.at_node["drainage_area"]

    # 例: 集水面積と勾配を使った簡易地形指標。
    # 実務では崩壊危険度、土砂供給域、谷次数、TWI等へ拡張する。
    terrain_index = np.log1p(drainage_area) * np.maximum(slope, 0)
    grid.add_field("terrain__prototype_index", terrain_index, at="node", clobber=True)

    print("[2] Terrain metrics")
    print(f"    max terrain index: {terrain_index.max():.3f}")


def example_3_rainfall_runoff_learning(grid: RasterModelGrid) -> None:
    """降雨流出・表面流モデルの学習用に、簡易ハイドログラフを作る。

    ここではWeb掲載用に軽い概念モデルでハイドログラフを作っています。
    Landlab本体で表面流を計算する場合は、OverlandFlowやKinwave系コンポーネントを
    時間ループ内で実行し、rainfall__flux、surface_water__depth、
    surface_water__discharge などを更新する構成に置き換えます。
    """
    drainage_area = grid.at_node["drainage_area"]
    outlet_area = drainage_area.max()

    time_min = np.arange(0, 181, 2)
    rainfall_mm_h = np.piecewise(
        time_min,
        [time_min < 25, (time_min >= 25) & (time_min < 80), (time_min >= 80) & (time_min < 120), time_min >= 120],
        [3.0, 28.0, 13.0, 0.0],
    )

    # 学習用の概念モデル: 降雨に対して時間遅れを持つ出口流量を作成する。
    # 実際のOverlandFlowでは、以下のような考え方で時間ループを回す。
    #   1) rainfall_mm_h を m/s に変換して rainfall__flux として与える
    #   2) surface_water__depth に初期水深を設定する
    #   3) Manning粗度係数や境界条件を設定する
    #   4) 各時刻で overland_flow.run_one_step(dt) を実行する
    #   5) 下流端の流量を記録し、流域出口ハイドログラフを作成する
    # ここでは、Webページで概念を説明しやすいように、簡易的な応答関数で表現する。
    peak_scale = outlet_area / 1.0e5
    discharge = peak_scale * (
        18.0 * np.exp(-((time_min - 88.0) ** 2) / (2 * 18.0**2))
        + 8.0 * np.exp(-((time_min - 123.0) ** 2) / (2 * 34.0**2))
    )

    grid.add_field("surface_water__depth", 0.02 + 0.6 * drainage_area / drainage_area.max(), at="node", clobber=True)

    print("[3] Rainfall-runoff learning sample")
    print(f"    peak discharge: {discharge.max():.2f} m3/s")


def example_4_erosion_deposition_research(grid: RasterModelGrid) -> None:
    """侵食・堆積・河床変動モデルの研究開発向けの最小構成例。"""
    # 流向・集水面積を再計算しながら、線形拡散と河川侵食を短時間だけ進める。
    diffuser = LinearDiffuser(grid, linear_diffusivity=0.03)
    eroder = FastscapeEroder(grid, K_sp=2.0e-5, m_sp=0.5, n_sp=1.0)
    fa = FlowAccumulator(grid, flow_director="D8")

    z0 = grid.at_node["topographic__elevation"].copy()
    dt = 50.0
    for _ in range(6):
        fa.run_one_step()
        diffuser.run_one_step(dt)
        eroder.run_one_step(dt)

    dz = grid.at_node["topographic__elevation"] - z0
    grid.add_field("bedrock__elevation_change", dz, at="node", clobber=True)

    print("[4] Erosion/deposition research prototype")
    print(f"    min dz: {dz.min():.4f} m, max dz: {dz.max():.4f} m")


def example_5_save_figures_for_education(grid: RasterModelGrid) -> None:
    """教材・説明資料として使いやすい図を保存する。"""
    OUT_DIR.mkdir(parents=True, exist_ok=True)

    figure_specs = [
        ("topographic__elevation", "Elevation (m)", "landlab_dem_sample.png"),
        ("drainage_area", "Drainage area (m2)", "landlab_flow_accumulation_sample.png"),
        ("terrain__prototype_index", "Terrain prototype index", "landlab_terrain_index_sample.png"),
        ("surface_water__depth", "Surface water depth (m)", "landlab_surface_water_depth_sample.png"),
        ("bedrock__elevation_change", "Elevation change dz (m)", "landlab_erosion_deposition_sample.png"),
    ]

    for field_name, label, filename in figure_specs:
        plt.figure(figsize=(9, 5.4))
        imshow_grid(grid, field_name, colorbar_label=label)
        plt.title(field_name)
        plt.tight_layout()
        plt.savefig(OUT_DIR / filename, dpi=180)
        plt.close()

    print("[5] Figures saved")
    print(f"    output directory: {OUT_DIR}")


def main() -> None:
    grid = setup_grid()
    example_1_flow_direction_and_drainage(grid)
    example_2_terrain_metrics(grid)
    example_3_rainfall_runoff_learning(grid)
    example_4_erosion_deposition_research(grid)
    example_5_save_figures_for_education(grid)


if __name__ == "__main__":
    main()