"""
津波解析 2次元シミュレーション サンプル図作成コード
=================================================

demo.godo-tys.jp/software/tsunami/ に掲載するための軽量デモです。

このサンプルでは、人工的な海底地形・陸域DEMと簡易的な津波波源を作成し、
以下の結果図を出力します。

1. 最大浸水深分布
2. 最大水位・最大流速分布
3. 津波到達時間
4. 時刻別水位・浸水範囲
5. 評価地点の水位時系列

注意:
実務解析では、断層モデル、潮位、海底地形、陸域DEM、防潮堤・水門、
Manning粗度、wet/dry処理、CFL条件、観測値による検証が必要です。
"""

from __future__ import annotations

from pathlib import Path
import numpy as np
import matplotlib.pyplot as plt


def make_demo_fields(nx: int = 180, ny: int = 100, dx: float = 25.0):
    """人工地形と津波解析結果に相当するサンプル場を作成する。"""
    y, x = np.mgrid[0:ny, 0:nx]
    coast_x = nx * 0.64

    # 左側を海、右側を陸とする人工DEM
    zb = -52.0 + 0.52 * (x - nx * 0.08)
    lowland = np.exp(-(((x - coast_x - 22.0) / 30.0) ** 2 + ((y - ny * 0.56) / 24.0) ** 2))
    ridge = np.exp(-(((x - coast_x - 50.0) / 12.0) ** 2 + ((y - ny * 0.48) / 38.0) ** 2))
    channel = np.exp(-((y - (ny * 0.52 + 8.0 * np.sin(x / 19.0))) ** 2) / (2.0 * 4.8**2))
    zb -= 3.0 * lowland
    zb += 2.2 * ridge
    zb -= 1.2 * channel

    # 津波の遡上を表すデモ用の最大水位場
    wave_axis = y - (ny * 0.52 + 6.0 * np.sin((x - coast_x) / 13.0))
    runup_decay = np.exp(-np.maximum(x - coast_x, 0) / 34.0)
    cross_decay = np.exp(-(wave_axis**2) / (2.0 * 21.0**2))
    max_eta = 2.4 * np.exp(-((x - coast_x + 18.0) ** 2) / (2.0 * 24.0**2)) * np.exp(-(wave_axis**2) / (2.0 * 30.0**2))
    max_eta += 2.0 * runup_decay * cross_decay
    max_eta += 0.25 * np.exp(-((x - coast_x - 15.0) ** 2) / (2.0 * 35.0**2))

    land = zb >= 0.0
    max_depth = np.where(land, np.maximum(max_eta - zb, 0.0), 0.0)
    max_depth[max_depth < 0.05] = 0.0

    # 最大流速は海岸・河口部で強くなるように作成
    max_velocity = 0.25 + 2.7 * np.exp(-((x - coast_x + 3.0) ** 2) / (2.0 * 23.0**2)) * np.exp(-(wave_axis**2) / (2.0 * 24.0**2))
    max_velocity += 1.2 * np.exp(-((x - coast_x - 22.0) ** 2) / (2.0 * 18.0**2)) * np.exp(-(wave_axis**2) / (2.0 * 10.0**2))
    max_velocity *= np.where(max_depth > 0.05, 1.0, 0.65)

    # 津波到達時間
    arrival = 2.5 + np.maximum(x - coast_x, 0) * 0.055 + np.abs(wave_axis) * 0.025
    arrival = np.where((land) & (max_depth > 0.05), arrival, np.nan)

    # 時刻別スナップショット
    snapshots = []
    for i, tmin in enumerate([3.0, 5.0, 8.0, 12.0]):
        front = coast_x + (tmin - 2.5) / 0.055
        progress = 1.0 / (1.0 + np.exp((x - front) / 6.5))
        depth_t = max_depth * progress * min(1.0, 0.45 + 0.16 * i)
        eta_t = np.where(depth_t > 0.03, zb + depth_t, np.nan)
        snapshots.append((tmin, eta_t, depth_t))

    # 評価地点の水位時系列
    t = np.linspace(0, 16, 220)
    gauges = {
        "Offshore": (48, 36),
        "Coast": (52, int(coast_x)),
        "Lowland": (58, int(coast_x + 24)),
        "Inland": (48, int(coast_x + 45)),
    }
    series = {}
    for name, (gy, gx) in gauges.items():
        delay = 1.7 + max(gx - coast_x, -60) * 0.055 + abs(gy - ny * 0.52) * 0.018
        amp = 2.0 * np.exp(-max(gx - coast_x, 0) / 45.0)
        series[name] = amp * np.exp(-((t - delay) / 1.25) ** 2) - 0.25 * amp * np.exp(-((t - delay - 2.5) / 1.8) ** 2)

    return {
        "nx": nx,
        "ny": ny,
        "dx": dx,
        "zb": zb,
        "land": land,
        "max_eta": max_eta,
        "max_depth": max_depth,
        "max_velocity": max_velocity,
        "arrival": arrival,
        "snapshots": snapshots,
        "time": t,
        "gauges": gauges,
        "series": series,
    }


def save_figures(out_dir: str | Path = "output_tsunami"):
    out = Path(out_dir)
    out.mkdir(parents=True, exist_ok=True)
    data = make_demo_fields()

    nx, ny, dx = data["nx"], data["ny"], data["dx"]
    extent = [0, nx * dx / 1000.0, ny * dx / 1000.0, 0]
    zb = data["zb"]
    land = data["land"]
    max_depth = data["max_depth"]
    max_eta = data["max_eta"]
    max_velocity = data["max_velocity"]
    arrival = data["arrival"]

    # 1. 最大浸水深分布
    fig, ax = plt.subplots(figsize=(10, 5.6))
    ax.set_title("1. Maximum inundation depth")
    ax.imshow(zb, extent=extent, cmap="terrain", alpha=0.75)
    im = ax.imshow(np.where(land, max_depth, np.nan), extent=extent, cmap="Blues", vmin=0, vmax=max(1.5, float(np.nanmax(max_depth))))
    ax.contour(zb, levels=[0.0], extent=extent, colors="black", linewidths=1.0)
    ax.set_xlabel("Easting distance (km)")
    ax.set_ylabel("Northing distance (km)")
    cb = fig.colorbar(im, ax=ax, pad=0.02)
    cb.set_label("Maximum inundation depth (m)")
    ax.text(0.02, 0.04, "Demo tsunami runup over coastal lowland", transform=ax.transAxes, fontsize=10,
            bbox=dict(boxstyle="round,pad=0.3", facecolor="white", alpha=0.82, edgecolor="0.8"))
    fig.tight_layout()
    fig.savefig(out / "tsunami_01_max_inundation_depth_distribution.png", dpi=180)
    plt.close(fig)

    # 2. 最大水位・最大流速分布
    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(12, 5.4))
    ax1.set_title("2-1. Maximum water level")
    im1 = ax1.imshow(max_eta, extent=extent, cmap="coolwarm")
    ax1.contour(zb, levels=[0.0], extent=extent, colors="black", linewidths=0.8)
    for name, (gy, gx) in data["gauges"].items():
        ax1.plot(gx * dx / 1000.0, gy * dx / 1000.0, "ko", ms=4)
        ax1.text(gx * dx / 1000.0 + 0.04, gy * dx / 1000.0 + 0.03, name, fontsize=8)
    ax1.set_xlabel("Easting distance (km)")
    ax1.set_ylabel("Northing distance (km)")
    cb1 = fig.colorbar(im1, ax=ax1, pad=0.02)
    cb1.set_label("Maximum water level (m)")

    ax2.set_title("2-2. Maximum velocity")
    im2 = ax2.imshow(max_velocity, extent=extent, cmap="viridis", vmin=0.0, vmax=max(1.0, float(np.nanpercentile(max_velocity, 99))))
    ax2.contour(zb, levels=[0.0], extent=extent, colors="white", linewidths=0.8)
    ax2.set_xlabel("Easting distance (km)")
    ax2.set_ylabel("Northing distance (km)")
    cb2 = fig.colorbar(im2, ax=ax2, pad=0.02)
    cb2.set_label("Maximum velocity (m/s)")
    fig.tight_layout()
    fig.savefig(out / "tsunami_02_max_water_level_velocity_distribution.png", dpi=180)
    plt.close(fig)

    # 3. 津波到達時間
    fig, ax = plt.subplots(figsize=(10, 5.6))
    ax.set_title("3. Tsunami arrival time")
    im = ax.imshow(arrival, extent=extent, cmap="plasma")
    ax.contour(zb, levels=[0.0], extent=extent, colors="black", linewidths=1.0)
    ax.set_xlabel("Easting distance (km)")
    ax.set_ylabel("Northing distance (km)")
    cb = fig.colorbar(im, ax=ax, pad=0.02)
    cb.set_label("Arrival time (min)")
    ax.text(0.02, 0.04, "First time inundation depth exceeds 0.05 m", transform=ax.transAxes, fontsize=10,
            bbox=dict(boxstyle="round,pad=0.3", facecolor="white", alpha=0.82, edgecolor="0.8"))
    fig.tight_layout()
    fig.savefig(out / "tsunami_03_arrival_time.png", dpi=180)
    plt.close(fig)

    # 4. 時刻別水位・浸水範囲
    fig, axes = plt.subplots(2, 2, figsize=(12, 9.2), constrained_layout=True)
    last_im = None
    for ax, (tmin, eta_t, depth_t) in zip(axes.ravel(), data["snapshots"]):
        ax.imshow(zb, extent=extent, cmap="terrain", alpha=0.55)
        ax.contour(zb, levels=[0.0], extent=extent, colors="black", linewidths=0.8)
        ax.contourf(eta_t, extent=extent, levels=10, cmap="coolwarm", alpha=0.50)
        last_im = ax.imshow(np.where(depth_t > 0.05, depth_t, np.nan), extent=extent, cmap="Blues", vmin=0.0, vmax=max(1.2, float(np.nanmax(max_depth))))
        ax.set_title(f"t = {tmin:.1f} min")
        ax.set_xlabel("Easting distance (km)")
        ax.set_ylabel("Northing distance (km)")
    fig.suptitle("4. Water level and inundation extent by time")
    cb = fig.colorbar(last_im, ax=axes.ravel().tolist(), shrink=0.84, pad=0.02)
    cb.set_label("Inundation depth (m)")
    fig.savefig(out / "tsunami_04_time_waterlevel_inundation_snapshots.png", dpi=180)
    plt.close(fig)

    # 5. 評価地点の水位時系列
    fig, ax = plt.subplots(figsize=(10.8, 5.8))
    for name, values in data["series"].items():
        ax.plot(data["time"], values, label=name)
    ax.axhline(0.0, linewidth=0.8, linestyle="--")
    ax.set_title("5. Water level time series at evaluation points")
    ax.set_xlabel("Elapsed time (min)")
    ax.set_ylabel("Water level anomaly eta (m)")
    ax.grid(True, alpha=0.25)
    ax.legend(loc="upper right")
    fig.tight_layout()
    fig.savefig(out / "tsunami_05_gauge_waterlevel_timeseries.png", dpi=180)
    plt.close(fig)


if __name__ == "__main__":
    save_figures()
    print("津波解析サンプル図を出力しました。")