pbf/navsea_export_mobile_core_sqlite.py

from __future__ import annotations

import argparse
import datetime as dt
import sqlite3
from pathlib import Path


DEFAULT_SOURCE_DB = Path("out/navsea_kyushu_navigability_full_run.sqlite")
DEFAULT_OUT_DB = Path("out/navsea_kyushu_mobile_core.sqlite")
DEFAULT_SNAPSHOT_DB = Path("out/navsea_kyushu_navigability_full_run.export_snapshot.sqlite")


def open_readonly(path: Path) -> sqlite3.Connection:
    if not path.exists():
        raise SystemExit(f"missing source db: {path}")
    conn = sqlite3.connect(f"file:{path}?mode=ro", uri=True)
    conn.execute("PRAGMA busy_timeout=5000")
    return conn


def create_snapshot(source_db: Path, snapshot_db: Path) -> Path:
    snapshot_db.parent.mkdir(parents=True, exist_ok=True)
    if snapshot_db.exists():
        snapshot_db.unlink()
    src = sqlite3.connect(source_db)
    src.execute("PRAGMA busy_timeout=5000")
    dst = sqlite3.connect(snapshot_db)
    with dst:
        src.backup(dst)
    dst.close()
    src.close()
    return snapshot_db


def open_output(path: Path) -> sqlite3.Connection:
    if path.exists():
        path.unlink()
    path.parent.mkdir(parents=True, exist_ok=True)
    conn = sqlite3.connect(path)
    conn.execute("PRAGMA journal_mode=OFF")
    conn.execute("PRAGMA synchronous=OFF")
    conn.execute("PRAGMA temp_store=MEMORY")
    conn.execute("PRAGMA cache_size=100000")
    return conn


def read_meta(conn: sqlite3.Connection, key: str, default: str = "") -> str:
    row = conn.execute("SELECT value FROM meta WHERE key = ?", (key,)).fetchone()
    return str(row[0]) if row else default


def create_schema(conn: sqlite3.Connection) -> None:
    conn.executescript(
        """
        CREATE TABLE meta (
          key TEXT PRIMARY KEY,
          value TEXT NOT NULL
        );

        CREATE TABLE grid_cell (
          grid_level INTEGER NOT NULL,
          grid_col INTEGER NOT NULL,
          grid_row INTEGER NOT NULL,
          parent_grid_col INTEGER,
          parent_grid_row INTEGER,
          state TEXT NOT NULL,
          blocked INTEGER NOT NULL DEFAULT 0,
          state_reason TEXT,
          risk_score INTEGER NOT NULL DEFAULT 0,
          PRIMARY KEY (grid_level, grid_col, grid_row)
        );

        CREATE INDEX idx_grid_cell_level_state
        ON grid_cell(grid_level, state, blocked);

        CREATE TABLE package_aoi (
          aoi_id TEXT PRIMARY KEY,
          package_type TEXT NOT NULL,
          minx REAL NOT NULL,
          miny REAL NOT NULL,
          maxx REAL NOT NULL,
          maxy REAL NOT NULL
        );
        """
    )
    conn.commit()


def copy_sparse_grid(src: sqlite3.Connection, dst: sqlite3.Connection) -> int:
    sql = """
        SELECT
          grid_level,
          grid_col,
          grid_row,
          parent_grid_col,
          parent_grid_row,
          state,
          blocked,
          state_reason,
          risk_score
        FROM grid_cell
        WHERE state <> 'open_water'
        ORDER BY grid_level, grid_col, grid_row
    """
    count = 0
    batch: list[tuple] = []
    batch_size = 20000
    for row in src.execute(sql):
        batch.append(tuple(row))
        count += 1
        if len(batch) >= batch_size:
            dst.executemany(
                """
                INSERT INTO grid_cell(
                  grid_level, grid_col, grid_row, parent_grid_col, parent_grid_row,
                  state, blocked, state_reason, risk_score
                ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)
                """,
                batch,
            )
            dst.commit()
            batch.clear()
            print(f"已导出 {count} 条移动端核心格", flush=True)
    if batch:
        dst.executemany(
            """
            INSERT INTO grid_cell(
              grid_level, grid_col, grid_row, parent_grid_col, parent_grid_row,
              state, blocked, state_reason, risk_score
            ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)
            """,
            batch,
        )
        dst.commit()
    return count


def copy_package_aoi(src: sqlite3.Connection, dst: sqlite3.Connection) -> int:
    sql = """
        SELECT
          printf('port:%d', aoi_id) AS aoi_id,
          'port_refine' AS package_type,
          center_x - radius_m,
          center_y - radius_m,
          center_x + radius_m,
          center_y + radius_m
        FROM port_aoi
        ORDER BY aoi_id
    """
    rows = [tuple(row) for row in src.execute(sql)]
    if rows:
        dst.executemany(
            """
            INSERT INTO package_aoi(
              aoi_id, package_type, minx, miny, maxx, maxy
            ) VALUES (?, ?, ?, ?, ?, ?)
            """,
            rows,
        )
        dst.commit()
    return len(rows)


def write_meta(src: sqlite3.Connection, dst: sqlite3.Connection, *, source_db: Path, exported_cells: int, exported_aoi: int) -> None:
    entries = [
        ("source_db", str(source_db)),
        ("source_grid_profile", read_meta(src, "grid_profile", "")),
        ("source_checkpoint_at", read_meta(src, "last_checkpoint_at", "")),
        ("source_phase", read_meta(src, "run_phase", "")),
        ("export_time", dt.datetime.now().isoformat(timespec="seconds")),
        ("package_type", "mobile_core"),
        ("storage_profile", "sparse_non_open_water_only"),
        ("default_state", "open_water"),
        ("exported_grid_cell_count", str(exported_cells)),
        ("exported_package_aoi_count", str(exported_aoi)),
        ("notes", "移动端核心包：仅保留非 open_water 格和港口 AOI 包索引"),
    ]
    dst.executemany("INSERT INTO meta(key, value) VALUES (?, ?)", entries)
    dst.commit()


def print_summary(conn: sqlite3.Connection) -> None:
    print("导出结果概览:", flush=True)
    for row in conn.execute(
        """
        SELECT grid_level, state, COUNT(*)
        FROM grid_cell
        GROUP BY grid_level, state
        ORDER BY grid_level, state
        """
    ):
        print(row, flush=True)


def main() -> None:
    parser = argparse.ArgumentParser(description="Export a sparse mobile-core SQLite from the Kyushu engineering navigability DB.")
    parser.add_argument("--source-db", type=Path, default=DEFAULT_SOURCE_DB)
    parser.add_argument("--out-db", type=Path, default=DEFAULT_OUT_DB)
    parser.add_argument("--snapshot-db", type=Path, default=DEFAULT_SNAPSHOT_DB)
    args = parser.parse_args()

    snapshot_db = create_snapshot(args.source_db, args.snapshot_db)
    print(f"已生成一致性快照: {snapshot_db}", flush=True)
    src = open_readonly(snapshot_db)
    dst = open_output(args.out_db)
    create_schema(dst)

    exported_cells = copy_sparse_grid(src, dst)
    exported_aoi = copy_package_aoi(src, dst)
    write_meta(src, dst, source_db=args.source_db, exported_cells=exported_cells, exported_aoi=exported_aoi)
    dst.execute("ANALYZE")
    dst.commit()
    print_summary(dst)
    print(f"完成: cells={exported_cells}, aoi={exported_aoi}", flush=True)
    print(f"输出文件: {args.out_db.resolve()}", flush=True)
    print(f"快照文件: {snapshot_db.resolve()}", flush=True)

    dst.close()
    src.close()


if __name__ == "__main__":
    main()