"""Route every measured pair over the real cable network, on every map.

Graph:
  cable edges    TeleGeography's drawn route of every cable in service,
                 waypoint to waypoint (same lat/lon waypoints on every map)
  landing links  landing point -> nearest waypoint of each cable that lands there
  land edges     straight line between any two nodes (landing points, cities)
                 on the same landmass (110m land polygons; islands too small
                 for that file are grouped within 50 km)

Each map measures every edge in its own geometry (great circle on the globe,
straight line on a flat map), so each map gets its own shortest route.

  python3 network.py   -> data/network.json
"""
import glob, json, math
from collections import defaultdict
import numpy as np
from matplotlib.path import Path
from scipy.sparse import coo_matrix
from scipy.sparse.csgraph import dijkstra
from models import R, MODELS, ae_xy, dualpole_xy, mercator_xy, C_FIBER
from build import land
from measure import TARGETS, EXTRA_SITES, ROUTE_ONLY_SITES

TG = "raw/tg"


def lonw(x):
    return ((x + 180) % 360) - 180


# ---------- nodes ----------
nodes, key = [], {}


def node(lon, lat, kind, label=None):
    lo = round(lonw(lon), 2)
    lo = lo - 360 if lo >= 180 else lo        # date-line seam: +180 and -180 are one point
    k = (lo, round(lat, 2), kind if kind == "city" else "")
    if kind == "city":
        k = (k[0], k[1], label)
    if k not in key:
        key[k] = len(nodes)
        nodes.append(dict(lon=k[0], lat=k[1], kind=kind, label=label))
    elif kind == "landing" and nodes[key[k]]["kind"] == "v":
        nodes[key[k]].update(kind="landing", label=label)   # waypoint that is also a landing point
    return key[k]


meta = {}
for f in glob.glob(f"{TG}/cables/*.json"):
    c = json.load(open(f))
    meta[c["id"]] = c
lp = {f["properties"]["id"]: f["geometry"]["coordinates"] for f in json.load(open(f"{TG}/landing-point-geo.json"))["features"]}

cable_names, cable_edges = [], {}
edges = []          # (i, j, kind, cable)
by_cable_vertices = defaultdict(set)
for f in json.load(open(f"{TG}/cable-geo.json"))["features"]:
    cid = f["properties"]["id"]
    c = meta.get(cid)
    if not c or c.get("is_planned") or (c.get("rfs_year") or 0) > 2026:
        continue
    if cid not in cable_names:
        cable_names.append(cid)
    ci = cable_names.index(cid)
    for line in f["geometry"]["coordinates"]:
        ids = [node(p[0], p[1], "v") for p in line]
        for a, b in zip(ids, ids[1:]):
            if a != b:
                edges.append((a, b, "cable", ci))
        by_cable_vertices[cid].update(ids)

landing_nodes = {}
for cid in cable_names:
    verts = list(by_cable_vertices[cid])
    vx = np.array([[nodes[v]["lon"], nodes[v]["lat"]] for v in verts])
    for p in meta[cid]["landing_points"]:
        if p["id"] not in lp:
            continue
        lon, lat = lp[p["id"]]
        li = landing_nodes.get(p["id"])
        if li is None:
            li = landing_nodes[p["id"]] = node(lon, lat, "landing", p["name"])
        d = (vx[:, 0] - lon) ** 2 + (vx[:, 1] - lat) ** 2
        j = verts[int(np.argmin(d))]
        if j != li:
            edges.append((li, j, "link", cable_names.index(cid)))

# cities: live anchors + WonderNetwork servers
live_city = {t[0]: node(t[4], t[3], "city", "L:" + t[0]) for t in TARGETS}
for name, la, lo in EXTRA_SITES + ROUTE_ONLY_SITES:   # Honolulu (our server, no anchor) and the home reference
    live_city[name] = node(lo, la, "city", "L:" + name)
wn = json.load(open("data/wn_clean.json"))
wn_ids = sorted(wn["servers"], key=int)
wn_city = {s: node(wn["servers"][s]["lon"], wn["servers"][s]["lat"], "city", "W:" + s) for s in wn_ids}

# ---------- landmasses ----------
paths = []
for poly in land():
    ring = poly[0]
    if min(p[1] for p in ring) < -60:
        continue                                   # Antarctica: no cables land there
    xs, out, off = [], [], 0.0
    prev = None
    for lon, lat in ring:
        if prev is not None and abs(lon + off - prev) > 180:
            off += -360 if lon + off - prev > 180 else 360
        prev = lon + off
        out.append((lon + off, lat))
    paths.append(Path(np.array(out)))

N = len(nodes)
anchor_nodes = [i for i, n in enumerate(nodes) if n["kind"] in ("landing", "city")]
pts = np.array([[nodes[i]["lon"], nodes[i]["lat"]] for i in anchor_nodes])
mass = np.full(len(anchor_nodes), -1)
for pi, p in enumerate(paths):
    for shift in (0, 360, -360):
        hit = p.contains_points(pts + [shift, 0]) & (mass < 0)
        mass[hit] = pi
for r in (0.3, 0.7, 1.2):                          # coastal points just outside the coarse outline
    for ang in range(0, 360, 45):
        dx, dy = r * math.cos(math.radians(ang)), r * math.sin(math.radians(ang))
        for pi, p in enumerate(paths):
            for shift in (0, 360, -360):
                hit = p.contains_points(pts + [dx + shift, dy]) & (mass < 0)
                mass[hit] = pi
# remaining points: cluster within 50 km
parent = list(range(len(anchor_nodes)))


def find(x):
    while parent[x] != x:
        parent[x] = parent[parent[x]]
        x = parent[x]
    return x


def gdist(lat1, lon1, lat2, lon2):
    p1, p2 = np.radians(lat1), np.radians(lat2)
    a = np.sin((p2 - p1) / 2) ** 2 + np.cos(p1) * np.cos(p2) * np.sin(np.radians(lon2 - lon1) / 2) ** 2
    return 2 * R * np.arcsin(np.minimum(1, np.sqrt(a)))


iso = np.where(mass < 0)[0]
for a in iso:
    d = gdist(pts[a, 1], pts[a, 0], pts[iso, 1], pts[iso, 0])
    for b in iso[d < 50]:
        parent[find(a)] = find(b)
for a in iso:
    mass[a] = 10000 + find(a)

groups = defaultdict(list)
for k, m in enumerate(mass):
    groups[int(m)].append(anchor_nodes[k])
for g in groups.values():
    for x in range(len(g)):
        for y in range(x + 1, len(g)):
            edges.append((g[x], g[y], "land", -1))
print(f"nodes {N}, edges {len(edges)}, cables {len(cable_names)}, landmasses {len(groups)}", flush=True)

# ---------- weights per map ----------
lat = np.array([n["lat"] for n in nodes])
lon = np.array([n["lon"] for n in nodes])
E = np.array([(a, b) for a, b, _, _ in edges])
xy = {"ae": np.array([ae_xy(a, b) for a, b in zip(lat, lon)]),
      "dualpole": np.array([dualpole_xy(a, b) for a, b in zip(lat, lon)]),
      "mercator": np.array([mercator_xy(a, b) for a, b in zip(lat, lon)])}
W = {"globe": gdist(lat[E[:, 0]], lon[E[:, 0]], lat[E[:, 1]], lon[E[:, 1]])}
for m in ("ae", "dualpole", "mercator"):
    W[m] = np.hypot(*(xy[m][E[:, 0]] - xy[m][E[:, 1]]).T)
kind = {}
for (a, b, k, c) in edges:
    kk = (min(a, b), max(a, b))
    if kk not in kind or k == "cable":
        kind[kk] = (k, c)


def graph(w):
    i, j = np.minimum(E[:, 0], E[:, 1]), np.maximum(E[:, 0], E[:, 1])
    order = np.lexsort((w, j, i))
    i, j, w = i[order], j[order], w[order]
    first = np.ones(len(i), bool)
    first[1:] = (i[1:] != i[:-1]) | (j[1:] != j[:-1])
    w = np.maximum(w[first], 1e-6)
    return coo_matrix((w, (i[first], j[first])), shape=(N, N)).tocsr()


G = {m: graph(W[m]) for m in MODELS}

def main():
    # ---------- live pairs: lengths + paths ----------
    live_names = [t[0] for t in TARGETS] + [e[0] for e in EXTRA_SITES] + [e[0] for e in ROUTE_ONLY_SITES]
    live_src = [live_city[n] for n in live_names]
    live = {}
    for m in MODELS:
        dist, pred = dijkstra(G[m], directed=False, indices=live_src, return_predecessors=True)
        for a in range(len(live_names)):
            for b in range(a + 1, len(live_names)):
                t = live_src[b]
                seq = [t]
                while seq[-1] != live_src[a] and pred[a, seq[-1]] >= 0:
                    seq.append(pred[a, seq[-1]])
                seq.reverse()
                via, kinds, land_km = [], "", 0.0
                for u, v in zip(seq, seq[1:]):
                    k, c = kind.get((min(u, v), max(u, v)), ("land", -1))
                    name = meta[cable_names[c]]["name"] if k in ("cable", "link") and c >= 0 else "land"
                    kinds += {"cable": "c", "link": "k", "land": "l"}[k]
                    if k == "land":
                        land_km += float(G[m][min(u, v), max(u, v)])
                    if not via or via[-1] != name:
                        via.append(name)
                live.setdefault(f"{a}-{b}", {})[m] = dict(
                    km=round(float(dist[a, t]), 1), land=round(land_km, 1), kinds=kinds,
                    path=[[round(nodes[s]["lon"], 2), round(nodes[s]["lat"], 2)] for s in seq],
                    via=[v for v in via if v != "land"])

    # ---------- WonderNetwork pairs: lengths ----------
    src = [wn_city[s] for s in wn_ids]
    col = {s: i for i, s in enumerate(wn_ids)}
    wn_km = {}
    for m in MODELS:
        dist = dijkstra(G[m], directed=False, indices=src)
        wn_km[m] = [round(float(dist[col[r["a"]], wn_city[r["b"]]]), 1) for r in wn["pairs"]]

    # globe sanity: a measurement faster than light along the globe's own best route means the
    # cable list is missing a path for that pair; that pair is set aside for every map
    slow = lambda km, t: t < 2 * km / C_FIBER * 1000
    wn_gap = [i for i, r in enumerate(wn["pairs"]) if slow(wn_km["globe"][i] - 200, r["rtt"] + 1)]
    print(f"WN pairs faster than the globe's cable route (network gaps): {len(wn_gap)} of {len(wn['pairs'])}")
    for i in wn_gap[:15]:
        r = wn["pairs"][i]
        print(f"   {wn['servers'][r['a']]['name']:16s} {wn['servers'][r['b']]['name']:16s} rtt {r['rtt']:6.1f}  globe route {wn_km['globe'][i]:7.0f} km -> {2*wn_km['globe'][i]/C_FIBER*1000:6.1f} ms")
    for m in MODELS:
        bad = sum(1 for i, r in enumerate(wn["pairs"]) if i not in set(wn_gap) and slow(wn_km[m][i] - 200, r["rtt"] + 1))
        print(f"   {m:9s} WN too fast along cables: {bad}")
    links = [float(G["globe"][min(u, v), max(u, v)]) for (u, v, k, c) in edges if k == "link"]
    print(f"landing links: {len(links)}, over 50 km: {sum(l > 50 for l in links)}, longest {max(links):.0f} km")
    json.dump(dict(live=live, wn_km=wn_km, wn_gap=wn_gap), open("data/network.json", "w"), separators=(",", ":"))


if __name__ == "__main__":
    main()