经纬度转km平面坐标系

from pyproj import CRS, Transformer
import numpy as np


def lonlat_to_xy_km(lon, lat, lon0: float, lat0: float, *, method: str = "aeqd", ellps: str = "WGS84", _cache=None):
    """
    经纬度(lon,lat) -> 站点(lon0,lat0)为原点的局地平面坐标(x,y)，单位 km，x: 向东为正，y: 向北为正

    :param lon: 经度数组
    :param lat: 纬度数组
    :param lon0: 设备经度值
    :param lat0: 设备纬度值
    :param method: （1）aeqd: Azimuthal Equidistant（以站点为中心的局地投影）（2）utm: 自动选站点所在UTM带（适合中低纬、局地范围；高纬UTM不适用）
    :param ellps: 地球坐标系统，默认WGS84
    :param _cache:
    :return: x_km, y_km（与输入形状一致；标量输入则返回 float）
    """

    if _cache is None:
        _cache = {}
    lon0 = float(lon0)
    lat0 = float(lat0)

    # --- 1) 解析输入：既支持 (lon,lat) 分开，也支持 lon 为 (...,2) 点集 ---
    lat_is_none = lat is None
    if lat_is_none:
        pts = np.asarray(lon, dtype=np.float64)
        if pts.shape == () or pts.shape[-1] != 2:
            raise ValueError("当 lat=None 时，lon 必须是形如 (..., 2) 的点集数组，例如 Nx2")
        lon_arr = pts[..., 0]
        lat_arr = pts[..., 1]
    else:
        lon_arr = np.asarray(lon, dtype=np.float64)
        lat_arr = np.asarray(lat, dtype=np.float64)

    # 记录是否标量，用于友好返回
    scalar_input = (lon_arr.shape == () and lat_arr.shape == ())

    # --- 2) 高精度投影：pyproj（优先） ---
    try:
        key = (method, ellps, lon0, lat0)
        tr = _cache.get(key)

        if tr is None:
            if method == "aeqd":
                crs_local = CRS.from_proj4(
                    f"+proj=aeqd +lat_0={lat0:.12f} +lon_0={lon0:.12f} "
                    f"+ellps={ellps} +units=m +no_defs"
                )
            elif method == "utm":
                zone = int(np.floor((lon0 + 180.0) / 6.0) + 1)
                if not (1 <= zone <= 60):
                    raise ValueError(f"Invalid UTM zone computed: {zone}")
                epsg = (32600 if lat0 >= 0.0 else 32700) + zone
                crs_local = CRS.from_epsg(epsg)
            else:
                raise ValueError("method must be 'aeqd' or 'utm'")

            tr = Transformer.from_crs("EPSG:4326", crs_local, always_xy=True)
            _cache[key] = tr

        x_m, y_m = tr.transform(lon_arr, lat_arr)
        x_km = np.asarray(x_m, dtype=np.float64) / 1000.0
        y_km = np.asarray(y_m, dtype=np.float64) / 1000.0

        if scalar_input:
            return float(x_km), float(y_km)
        return x_km, y_km

    except Exception:
        # --- 3) 兜底：局地线性近似（保证代码在没装 pyproj 也能跑） ---
        lat0_rad = np.deg2rad(lat0)
        x_km = (lon_arr - lon0) * 111.32 * np.cos(lat0_rad)
        y_km = (lat_arr - lat0) * 110.57
        if scalar_input:
            return float(x_km), float(y_km)
        return x_km.astype(np.float64), y_km.astype(np.float64)


if __name__ == "__main__":
    lon = [116.10, 116.25, 116.40]
    lat = [39.90, 39.95, 40.02]
    x_km, y_km = lonlat_to_xy_km(lon, lat, lon0=116.30, lat0=39.90)  # x_km/y_km 都是长度3的数组
    print(x_km, y_km)

    pts = [
        [116.10, 39.90],
        [116.25, 39.95],
        [116.40, 40.02]]
    x_km, y_km = lonlat_to_xy_km(pts, None, lon0=116.30, lat0=39.90)
    print(f"x_km = {x_km}, y_km = {y_km}")