pbf/coastline/export_japan_coast_grid_mysql_assets.py

#!/usr/bin/env python3
from __future__ import annotations

import argparse
import datetime as dt
import json
import math
from pathlib import Path

import pymysql


DB_NAME = "navsea_japan_coast_grid"
DB_USER = "root"
DB_PASSWORD = "2chi9ks2"
DB_HOST = "localhost"
DB_SOCKET = "/tmp/mysql.sock"

DEFAULT_OUT_DIR = "src/pbf/coastline-mysql/japan_coast_200m"
LAYER_NAME = "coast_200m"
RADIUS = 6378137.0


def mysql_connect(database: str):
    kwargs = {
        "host": DB_HOST,
        "user": DB_USER,
        "password": DB_PASSWORD,
        "database": database,
        "charset": "utf8mb4",
        "cursorclass": pymysql.cursors.SSCursor,
    }
    if DB_SOCKET:
        kwargs["unix_socket"] = DB_SOCKET
    return pymysql.connect(**kwargs)


def rect_geometry(min_lon: float, min_lat: float, max_lon: float, max_lat: float) -> dict:
    return {
        "type": "Polygon",
        "coordinates": [
            [
                [min_lon, min_lat],
                [max_lon, min_lat],
                [max_lon, max_lat],
                [min_lon, max_lat],
                [min_lon, min_lat],
            ]
        ],
    }


def parse_bbox(value: str) -> tuple[float, float, float, float]:
    parts = [part.strip() for part in value.split(",")]
    if len(parts) != 4:
        raise argparse.ArgumentTypeError("bbox must be min_lon,min_lat,max_lon,max_lat")
    try:
        return tuple(float(part) for part in parts)  # type: ignore[return-value]
    except ValueError as exc:
        raise argparse.ArgumentTypeError("bbox must contain valid numbers") from exc


def mercator_to_lonlat(x: float, y: float) -> tuple[float, float]:
    lon = math.degrees(x / RADIUS)
    lat = math.degrees(2.0 * math.atan(math.exp(y / RADIUS)) - math.pi / 2.0)
    return lon, lat


def lon_to_mercator(lon: float) -> float:
    return RADIUS * math.radians(lon)


def lat_to_mercator(lat: float) -> float:
    lat = max(min(lat, 85.0511287798066), -85.0511287798066)
    return RADIUS * math.log(math.tan(math.pi / 4.0 + math.radians(lat) / 2.0))


def main() -> None:
    parser = argparse.ArgumentParser(description="Export Japan coast 200m grid assets from MySQL")
    parser.add_argument("--db-name", default=DB_NAME)
    parser.add_argument("--out-dir", default=DEFAULT_OUT_DIR)
    parser.add_argument("--bbox", type=parse_bbox, default=None, help="仅导出该经纬度范围内的格子")
    args = parser.parse_args()

    project_root = Path(__file__).resolve().parent.parent
    out_dir = Path(args.out_dir)
    if not out_dir.is_absolute():
        out_dir = project_root / out_dir
    out_dir.mkdir(parents=True, exist_ok=True)

    conn = mysql_connect(args.db_name)
    try:
        with conn.cursor() as cur:
            cur.execute(
                """
                SELECT description, source_desc, cell_size_m, feature_count,
                       bbox_min_lon, bbox_min_lat, bbox_max_lon, bbox_max_lat, export_file
                FROM navsea_grid_layer_meta
                WHERE layer_name=%s
                """,
                (LAYER_NAME,),
            )
            meta_row = cur.fetchone()
            if meta_row is None:
                raise SystemExit(f"missing layer meta: {LAYER_NAME}")

            (
                description,
                source_desc,
                cell_size_m,
                feature_count,
                min_lon,
                min_lat,
                max_lon,
                max_lat,
                source_manifest,
            ) = meta_row

        export_path = out_dir / "coast_200m_grid.geojson"
        cell_sql = """
            SELECT layer_name, cell_id, row_idx, col_idx, cell_size_m,
                   state_name, class_name, source_name,
                   min_lon, min_lat, max_lon, max_lat
            FROM navsea_grid_cell
            WHERE layer_name=%s
        """
        cell_params: list[float | str] = [LAYER_NAME]
        if args.bbox is not None:
            min_lon, min_lat, max_lon, max_lat = args.bbox
            min_x = lon_to_mercator(min_lon)
            min_y = lat_to_mercator(min_lat)
            max_x = lon_to_mercator(max_lon)
            max_y = lat_to_mercator(max_lat)
            cell_sql += """
              AND NOT (
                max_lon < %s OR min_lon > %s OR
                max_lat < %s OR min_lat > %s
              )
            """
            cell_params.extend([min_x, max_x, min_y, max_y])
        cell_sql += """
            ORDER BY row_idx, col_idx
        """

        count = 0
        with conn.cursor() as cur, export_path.open("w", encoding="utf-8") as fh:
            cur.execute(cell_sql, tuple(cell_params))
            fh.write('{"type":"FeatureCollection","name":"japan_coast_200m_grid","features":[\n')
            first = True
            for row in cur:
                (
                    layer_name,
                    cell_id,
                    row_idx,
                    col_idx,
                    cell_size_m_row,
                    state_name,
                    class_name,
                    source_name,
                    min_lon_row,
                    min_lat_row,
                    max_lon_row,
                    max_lat_row,
                ) = row
                lon1, lat1 = mercator_to_lonlat(float(min_lon_row), float(min_lat_row))
                lon2, lat2 = mercator_to_lonlat(float(max_lon_row), float(max_lat_row))
                feature = {
                    "type": "Feature",
                    "properties": {
                        "layer_name": layer_name,
                        "cell_id": cell_id,
                        "row": int(row_idx),
                        "col": int(col_idx),
                        "cell_size_m": float(cell_size_m_row),
                        "state_name": state_name,
                        "class_name": class_name,
                        "source_name": source_name,
                    },
                    "geometry": rect_geometry(lon1, lat1, lon2, lat2),
                }
                if not first:
                    fh.write(",\n")
                fh.write(json.dumps(feature, ensure_ascii=False))
                first = False
                count += 1
            fh.write("\n]}\n")

        if args.bbox is None and count != int(feature_count):
            raise RuntimeError(f"count mismatch: exported={count} meta={feature_count}")

        if args.bbox is None:
            lonlat_min = mercator_to_lonlat(float(min_lon), float(min_lat))
            lonlat_max = mercator_to_lonlat(float(max_lon), float(max_lat))
            mercator_min = [float(min_lon), float(min_lat)]
            mercator_max = [float(max_lon), float(max_lat)]
        else:
            lonlat_min = (float(min_lon), float(min_lat))
            lonlat_max = (float(max_lon), float(max_lat))
            mercator_min = [lon_to_mercator(float(min_lon)), lat_to_mercator(float(min_lat))]
            mercator_max = [lon_to_mercator(float(max_lon)), lat_to_mercator(float(max_lat))]

        manifest = {
            "version": "v1.0",
            "database": args.db_name,
            "generated_at": dt.datetime.now().isoformat(timespec="seconds"),
            "layer": LAYER_NAME,
            "description": description,
            "source_desc": source_desc,
            "source_manifest": source_manifest,
            "count": count,
            "cell_size_m": float(cell_size_m),
            "grid_bounds_lonlat": [lonlat_min[0], lonlat_min[1], lonlat_max[0], lonlat_max[1]],
            "grid_bounds_mercator": [mercator_min[0], mercator_min[1], mercator_max[0], mercator_max[1]],
            "grid_file": "coast_200m_grid.geojson",
            "notes": "从 navsea_japan_coast_grid.MySQL 直接导出的全国海岸 200m 硬阻塞格",
        }

        manifest_path = out_dir / "manifest.json"
        manifest_path.write_text(json.dumps(manifest, ensure_ascii=False, indent=2), encoding="utf-8")

        print("已导出:")
        print(" -", export_path)
        print(" -", manifest_path)
        print(f" - count={count}")
    finally:
        conn.close()


if __name__ == "__main__":
    main()