cn_network.py

"""
中国地铁线路网络 + 沿轨路由。
- 从高德地铁数据解析每条线: 名称/颜色/有序站点(站名+WGS84坐标)。
- 建图: 同线相邻站连边(权=距离); 同名站=同一节点 => 换乘自动连通。
- route(a,b): Dijkstra 求沿轨最短路径(含换乘), 让行程贴着地铁线画。
坐标 GCJ-02 -> WGS-84(复用 cn_coords)。
"""
from __future__ import annotations
import heapq
import json
import math
import cn_coords

_net_cache = {}      # city_code -> [line...]
_router_cache = {}   # city_code -> CityRouter


def _parse_json(text):
    out = []
    data = json.loads(text)
    for ln in data.get("l", []):
        name = ln.get("ln") or ln.get("kn") or ""
        cl = ln.get("cl") or "888888"
        stations = []
        for st in ln.get("st", []):
            sl = st.get("sl", ""); nm = st.get("n", "")
            if "," in sl:
                try:
                    lng, lat = (float(x) for x in sl.split(",")[:2])
                except ValueError:
                    continue
                stations.append((nm, round(lat, 6), round(lng, 6)))   # 原始 GCJ-02, 配高德瓦片
        if len(stations) >= 2:
            out.append({"name": name, "color": "#" + cl, "stations": stations,
                        "pts": [(s[1], s[2]) for s in stations]})
    return out


def load_network(city_code: str):
    if city_code in _net_cache:
        return _net_cache[city_code]
    path = cn_coords.ensure_city_json(city_code)
    lines = []
    if path:
        try:
            with open(path, "r", encoding="utf-8") as f:
                lines = _parse_json(f.read())
        except Exception:
            lines = []
    _net_cache[city_code] = lines
    return lines


def color_for(city_code: str, line_name: str) -> str | None:
    """反查某城市某 line 名的代表色。失败返 None(主程序 fallback 默认磷光绿)。"""
    if not city_code or not line_name:
        return None
    try:
        lines = load_network(city_code)
    except Exception:
        return None
    for ln in lines:
        if ln.get("name") == line_name:
            return ln.get("color")
    return None


def _hav(a, b):
    (la1, lo1), (la2, lo2) = a, b
    p = math.pi / 180
    dla = (la2 - la1) * p; dlo = (lo2 - lo1) * p
    x = math.sin(dla/2)**2 + math.cos(la1*p)*math.cos(la2*p)*math.sin(dlo/2)**2
    return 2 * 6371000 * math.asin(min(1, math.sqrt(x)))


def _norm(name):
    if not name:
        return ""
    return (name.strip().replace(" ", "").replace("　", "")
            .replace("（", "(").replace("）", ")").rstrip("站"))


class CityRouter:
    def __init__(self, lines):
        self.coord = {}   # name -> (lat,lng)
        self.adj = {}     # name -> [(name, w)]
        self.norm = {}    # 归一化名 -> 原名
        for ln in lines:
            sts = ln["stations"]
            for nm, la, lo in sts:
                self.coord.setdefault(nm, (la, lo))
                self.norm.setdefault(_norm(nm), nm)
            for i in range(len(sts) - 1):
                a, b = sts[i][0], sts[i+1][0]
                if a == b:
                    continue
                w = _hav((sts[i][1], sts[i][2]), (sts[i+1][1], sts[i+1][2]))
                self.adj.setdefault(a, []).append((b, w))
                self.adj.setdefault(b, []).append((a, w))

    def _resolve_name(self, name):
        if name in self.adj:
            return name
        return self.norm.get(_norm(name))

    def route(self, a, b):
        """返回沿轨路径 [(name, lat, lng), ...]; 失败返回 None。站名做归一化匹配。"""
        a = self._resolve_name(a); b = self._resolve_name(b)
        if not a or not b or a not in self.adj or b not in self.adj:
            return None
        dist = {a: 0.0}; prev = {}; pq = [(0.0, a)]
        while pq:
            d, u = heapq.heappop(pq)
            if u == b:
                break
            if d > dist.get(u, 1e18):
                continue
            for v, w in self.adj[u]:
                nd = d + w
                if nd < dist.get(v, 1e18):
                    dist[v] = nd; prev[v] = u; heapq.heappush(pq, (nd, v))
        if b not in dist:
            return None
        path = [b]
        while path[-1] != a:
            path.append(prev[path[-1]])
        path.reverse()
        return [(n, self.coord[n][0], self.coord[n][1]) for n in path]


def get_router(city_code: str):
    if city_code not in _router_cache:
        _router_cache[city_code] = CityRouter(load_network(city_code))
    return _router_cache[city_code]


if __name__ == "__main__":
    import sys
    if len(sys.argv) > 1:
        lines = _parse_json(open(sys.argv[1], encoding="utf-8").read())
        print("线路数:", len(lines))
        r = CityRouter(lines)
        if len(sys.argv) > 3:
            p = r.route(sys.argv[2], sys.argv[3])
            print(f"{sys.argv[2]}->{sys.argv[3]} 路径站数:", len(p) if p else None)
            if p:
                print("途经:", " - ".join(n for n, _, _ in p))