"""
Landlab 入門サンプル: RasterModelGrid + FlowAccumulator

目的:
- Landlab の基本的な使い方を確認する
- 2次元格子を作成する
- 簡易DEMを作成する
- 流向・集水面積を計算する
- 結果を画像として保存する

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

from pathlib import Path

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


# =========================================================
# 1. 計算格子の作成
# =========================================================

nrows = 80
ncols = 100
dx = 10.0  # 格子間隔 m

# 80行×100列、10m格子の2次元ラスター格子
grid = RasterModelGrid((nrows, ncols), xy_spacing=dx)


# =========================================================
# 2. 簡易DEMの作成
# =========================================================

x = grid.x_of_node
y = grid.y_of_node

# 全体として右下へ下る地形を作る
base_slope = 120.0 - 0.015 * x - 0.025 * y

# 中央に谷地形を作る
center_y = 0.52 * y.max() + 40.0 * np.sin(x / 180.0)
valley = -8.0 * np.exp(-((y - center_y) ** 2) / (2.0 * 90.0**2))

# 小さな起伏を加える
roughness = 0.8 * np.sin(x / 55.0) * np.cos(y / 70.0)

z = base_slope + valley + roughness

# Landlabでは各ノードの値を field として登録する
grid.add_field("topographic__elevation", z, at="node", clobber=True)


# =========================================================
# 3. 境界条件の設定
# =========================================================

# 左・右・上流側を閉境界、下側を開境界として流出させる
# 引数の順番は right, top, left, bottom
grid.set_closed_boundaries_at_grid_edges(
    right_is_closed=True,
    top_is_closed=True,
    left_is_closed=True,
    bottom_is_closed=False,
)


# =========================================================
# 4. 流向・集水面積の計算
# =========================================================

# D8方向で流れを決め、各ノードの集水面積を計算する
flow_accumulator = FlowAccumulator(grid, flow_director="D8")
flow_accumulator.run_one_step()

# 計算結果は grid.at_node に追加される
# 代表的な出力:
# - drainage_area: 集水面積
# - flow__receiver_node: 流下先ノード
# - topographic__steepest_slope: 最大勾配
drainage_area = grid.at_node["drainage_area"]
steepest_slope = grid.at_node["topographic__steepest_slope"]

print("Landlab sample finished.")
print(f"max drainage area = {drainage_area.max():,.1f} m2")
print(f"max steepest slope = {steepest_slope.max():.4f}")


# =========================================================
# 5. 図の出力
# =========================================================

out_dir = Path("public/demo-data/landlab")
out_dir.mkdir(parents=True, exist_ok=True)

plt.figure(figsize=(9, 5.4))
imshow_grid(
    grid,
    "topographic__elevation",
    colorbar_label="Elevation (m)",
)
plt.title("Synthetic DEM")
plt.tight_layout()
plt.savefig(out_dir / "landlab_dem_sample.png", dpi=180)
plt.close()

plt.figure(figsize=(9, 5.4))
imshow_grid(
    grid,
    "drainage_area",
    colorbar_label="Drainage area (m2)",
)
plt.title("Drainage Area calculated by FlowAccumulator")
plt.tight_layout()
plt.savefig(out_dir / "landlab_flow_accumulation_sample.png", dpi=180)
plt.close()