"""The Best Flat Map: every measured length and limit laid out on one plane, then released into 3D.

Data (all published earlier on distance.aethercosmology.com and pingstats.aethercosmology.com):
  flights    The Measured Sphere, "each flight's own speed": 219 cities, 1,344 routes flown both ways.
             Length = aircraft cruise speed x mean airborne time, less 15 min for climb and descent.
  limits     The Measured Sphere's cable, rail, road and satellite-ranging limits, every WonderNetwork
             2020 round trip and every live round trip (RIPE Atlas, Globalping, Vultr) between servers
             in these cities.
  triangles  The Triangle Test's 935 surveyed triangles (scored, not solved: a plane gives every
             triangle exactly 180 degrees).

Coordinates are used to label cities, draw the start maps, match servers to cities and score the
globe for comparison. The solver (fit.py) sees only lengths and limits.

  ../.venv/bin/python build.py     -> data/femap.json, site/index.html
"""
import base64, datetime as dt, importlib.util, json, math, os, pickle, re, sys, time
for _v in ("OPENBLAS_NUM_THREADS", "OMP_NUM_THREADS", "MKL_NUM_THREADS"):
    os.environ.setdefault(_v, "1")
from multiprocessing import Pool
import numpy as np

HERE = os.path.dirname(os.path.abspath(__file__))
SRC = os.path.join(HERE, "..", "src")
sys.path.insert(0, SRC)
_spec = importlib.util.spec_from_file_location("pingstats_build", os.path.join(SRC, "build.py"))
_pb = importlib.util.module_from_spec(_spec); _spec.loader.exec_module(_pb)
live, land = _pb.live, _pb.land                                # src/build.py
from models import R, C_FIBER, globe as gc, ae_xy, dualpole_xy, mercator_xy
from fit import Problem, sphere_fit

PRES = os.path.join(SRC, "handoff", "The_Shape_of_Distance_PRESENTATION.html")
KM_PER_MS = C_FIBER / 2 / 1000          # one-way km per ms of round trip, light in glass
END_SLACK = 25.0                        # km allowed at each end of a round trip
MATCH_KM = 80.0                         # a server belongs to a city if it is this close to its airport
LO, HI = 0.25, 0.10                     # default flight band: [0.75 L, 1.10 L]
BANDS = [(0.25, 0.10), (0.25, 0.25), (0.40, 0.10), (0.50, 0.10)]
N_RANDOM = 300
FAST = os.environ.get("FAST") == "1"     # quick preview: one start, one 3D seed, one band -> data/preview.json, site/preview.html
SEED = 20261030
FA_PREFIX = "https://www.flightaware.com/live/flight/"


def pages():
    s = open(PRES, encoding="utf-8").read()
    out = []
    for b in re.findall(r"[A-Za-z0-9+/]{20000,}={0,2}", s):
        out.append(base64.b64decode(b).decode("utf-8", "ignore"))
    return out


def page_const(html, name):
    js = next(x for x in re.findall(r"<script[^>]*>(.*?)</script>", html, flags=re.S) if f"const {name}=" in x)
    i = js.index(f"const {name}=") + len(name) + 7
    return json.JSONDecoder().raw_decode(js[i:])[0]


def aeq(lat, lon, lat0, lon0):
    """Azimuthal equidistant map centred on (lat0, lon0), km. North-pole AE is (90, 0); dual-pole is (0, 0)."""
    p1, l1, p0, l0 = map(np.radians, (np.asarray(lat), np.asarray(lon), lat0, lon0))
    cc = np.sin(p0) * np.sin(p1) + np.cos(p0) * np.cos(p1) * np.cos(l1 - l0)
    c = np.arccos(np.clip(cc, -1, 1)); k = np.where(c < 1e-9, 1.0, c / np.sin(np.maximum(c, 1e-9)))
    x = R * k * np.cos(p1) * np.sin(l1 - l0)
    y = R * k * (np.cos(p0) * np.sin(p1) - np.sin(p0) * np.cos(p1) * np.cos(l1 - l0))
    return np.c_[x, y]


def plane_d(f, a, b):
    x1, y1 = f(*a); x2, y2 = f(*b)
    return math.hypot(x2 - x1, y2 - y1)


def shortest_on_any(a, b):
    """Shortest distance between two places on the globe or any of the three flat maps used elsewhere."""
    return min(gc(*a, *b), plane_d(ae_xy, a, b), plane_d(dualpole_xy, a, b), plane_d(mercator_xy, a, b))


def load():
    P = pages()
    ms = page_const(P[6], "DATA")["own"]
    tri = page_const(P[7], "FIN")["tri"]
    cities = ms["cities"]
    C = [(c[1], c[2]) for c in cities]
    n = len(C)

    fl = []
    for f in ms["flights"]:
        a, b, x = f[0], f[1], f[8]
        ab, ba = x["ab"], x["ba"]
        fl.append(dict(a=a, b=b, L=x["km"], airline=f[4], fa=f[5], fb=f[6], page=f[7],
                       ab=[ab["ident"], ab["ac"], ab["air"], ab["n"], round(ab["v"]), ab["src"]],
                       ba=[ba["ident"], ba["ac"], ba["air"], ba["n"], round(ba["v"]), ba["src"]]))

    lim = {}              # (a, b, kind) -> record; round trips keep the tightest per pair
    def add(a, b, U, kind, note, label="", url=""):
        if a == b:
            return
        a, b = min(a, b), max(a, b)
        k = (a, b, kind)
        if k not in lim or U < lim[k]["U"]:
            lim[k] = dict(a=a, b=b, U=round(U, 1), kind=kind, note=note, label=label, url=url)

    kinds = {"cable": "rt", "rail": "rail", "road": "road", "satellite": "sat"}
    for l in ms["limits"]:
        add(l[0], l[1], l[2], kinds[l[3]], l[4].replace("round trip ", ""), l[5], l[6] if len(l) > 6 else "")

    def nearest(p):
        d = [gc(*p, *q) for q in C]; i = int(np.argmin(d)); return i, d[i]

    dropped = []
    wn = json.load(open(os.path.join(SRC, "data", "wn_clean.json")))
    srv = wn["servers"]; wm = {k: nearest((s["lat"], s["lon"])) for k, s in srv.items()}
    n_wn = 0
    # A round trip shorter than every surface allows means a server is not where it is listed. A server
    # caught in two or more such pairs is treated as mislocated and all its pairs are left out.
    reach_of = lambda p: KM_PER_MS * p["rtt"] + 2 * END_SLACK
    bad = [p for p in wn["pairs"] if reach_of(p) < min(p["globe"], p["ae"], p["dualpole"], p["mercator"])]
    hits = {}
    for p in bad:
        for s in (p["a"], p["b"]):
            hits[s] = hits.get(s, 0) + 1
    mislocated = {s for s, k in hits.items() if k >= 2}
    for p in bad:
        dropped.append(f'{srv[p["a"]]["name"]} - {srv[p["b"]]["name"]} ({p["rtt"]} ms)')
    for p in wn["pairs"]:
        a, b = wm[p["a"]], wm[p["b"]]
        if a[1] > MATCH_KM or b[1] > MATCH_KM or p in bad or p["a"] in mislocated or p["b"] in mislocated:
            continue
        reach = reach_of(p)
        n_wn += 1
        add(a[0], b[0], reach + a[1] + b[1], "rt",
            f'{p["rtt"]:.2f} ms, WonderNetwork 2020 archive ({p["n"]} samples)',
            f'WonderNetwork servers {srv[p["a"]]["name"]} and {srv[p["b"]]["name"]}', "https://wondernetwork.com/pings")
    lv = live(); sites = lv["sites"]; sm = [nearest((s[2], s[3])) for s in sites]
    n_live = 0
    for c in lv["cells"]:
        a, b = sm[c[0]], sm[c[1]]
        if a[1] > MATCH_KM or b[1] > MATCH_KM:
            continue
        sa, sb = sites[c[0]], sites[c[1]]
        reach = KM_PER_MS * c[2] + 2 * END_SLACK
        if reach < shortest_on_any((sa[2], sa[3]), (sb[2], sb[3])):
            dropped.append(f"{sa[0]} - {sb[0]} ({c[2]} ms)")
            continue
        n_live += 1
        add(a[0], b[0], reach + a[1] + b[1], "rt", f"{c[2]:.1f} ms, live measurement 2026 ({c[5]})",
            f"{sa[0]} probe to {sb[0]} anchor", "https://pingstats.aethercosmology.com/")
    lim = sorted(lim.values(), key=lambda r: (r["kind"], r["a"], r["b"]))

    tri_fail = dict(n=len(tri), flat=sum(abs(t[3]) > 4 for t in tri), globe=sum(abs(t[2]) > 4 for t in tri))
    return dict(cities=cities, C=np.array(C), fl=fl, lim=lim, tri=tri_fail, dropped=dropped, n_wn=n_wn, n_live=n_live,
                mislocated=sorted(f'{srv[s]["name"]} ({srv[s]["country"]})' for s in mislocated))


def mds_start(prob, C):
    """Classical multidimensional scaling of the shortest flight-path lengths, turned and mirrored
    (no stretching) to sit as close as possible to the North-pole AE layout for drawing."""
    n = prob.n; D = np.full((n, n), np.inf); np.fill_diagonal(D, 0)
    for a, b, L in zip(prob.fa, prob.fb, prob.fl):
        D[a, b] = D[b, a] = min(D[a, b], L)
    for k in range(n):
        D = np.minimum(D, D[:, k:k + 1] + D[k:k + 1, :])
    J = np.eye(n) - 1 / n; B = -0.5 * J @ (D ** 2) @ J
    w, V = np.linalg.eigh(B); X = V[:, -2:] * np.sqrt(np.maximum(w[-2:], 0))
    T = aeq(C[:, 0], C[:, 1], 90, 0)
    Xc, Tc = X - X.mean(0), T - T.mean(0)
    U, _, Vt = np.linalg.svd(Xc.T @ Tc)
    return Xc @ (U @ Vt) + T.mean(0)


_PROB = None
def _init(prob):
    global _PROB
    _PROB = prob


def _run(job):
    label, centre, X0, lo, hi = job[:5]
    k, X = _PROB.best_effort(X0, lo, hi, **(job[5] if len(job) > 5 else {}))
    r = _PROB.dists(X)[0] / _PROB.fl
    return dict(label=label, centre=centre, broken=k, meanlog=float(np.abs(np.log(r)).mean()), X=X)


def counts(prob, X, lo, hi):
    dd, ld = prob.dists(X); r = dd / prob.fl
    return dict(short=int((r < 1 - lo).sum()), long=int((r > 1 + hi).sum()), limits=int((ld > prob.lu * 1.0005).sum()))


def main():
    t0 = time.time()
    D = load(); C = D["C"]; fl, lim = D["fl"], D["lim"]
    prob = Problem(len(C), [f["a"] for f in fl], [f["b"] for f in fl], [f["L"] for f in fl],
                   [l["a"] for l in lim], [l["b"] for l in lim], [l["U"] for l in lim])
    GF = np.array([gc(*C[f["a"]], *C[f["b"]]) for f in fl]); GL = np.array([gc(*C[l["a"]], *C[l["b"]]) for l in lim])
    print(f"{len(C)} cities, {len(fl)} flights, {len(lim)} limits ({D['n_wn']} WonderNetwork pairs, {D['n_live']} live cells, "
          f"{len(D['dropped'])} round trips dropped as mislocated)", flush=True)

    starts = {"ae": ("North-pole AE", [90, 0], aeq(C[:, 0], C[:, 1], 90, 0)),
              "dp": ("Dual-pole AE", [0, 0], aeq(C[:, 0], C[:, 1], 0, 0)),
              "mds": ("Flight times only", [90, 0], mds_start(prob, C))}
    rng = np.random.default_rng(SEED)
    jobs = [(v[0], v[1], v[2], LO, HI) for v in starts.values()]
    named = [(0, 90), (0, -90), (0, 180), (-90, 0), (45, 100), (60, -30), (30, 60), (-30, 150)]
    rand = [(round(math.degrees(math.asin(rng.uniform(-1, 1))), 1), round(rng.uniform(-180, 180), 1)) for _ in range(N_RANDOM)]
    jobs += [(f"AE centred on {la}, {lo}", [la, lo], aeq(C[:, 0], C[:, 1], la, lo), LO, HI) for la, lo in named + rand]
    if FAST:
        jobs = jobs[:1]
    cache = os.path.join(HERE, "data", "search_fast.pkl" if FAST else "search.pkl")
    sig = (len(C), len(fl), len(lim), LO, HI, len(jobs))
    pool = Pool(7, initializer=_init, initargs=(prob,))
    if os.path.exists(cache) and pickle.load(open(cache, "rb"))["sig"] == sig:
        res = pickle.load(open(cache, "rb"))["res"]
        print("search: cached", flush=True)
    else:
        res = pool.map(_run, jobs, chunksize=1)
        pickle.dump(dict(sig=sig, res=res), open(cache, "wb"))
    res.sort(key=lambda r: (r["broken"], r["meanlog"]))
    # Refine the 20 best layouts: twelve rounds, shorter solves, two reweighting floors.
    if not FAST:
        ref = pool.map(_run, [(r["label"], r["centre"], r["X"], LO, HI, dict(rounds=12, maxiter=3000, eps=e))
                              for r in res[:20] for e in (0.02, 0.05)], chunksize=1)
        for r in ref:
            r["refined"] = True
        print("refined:", sorted(r["broken"] for r in ref)[:10], flush=True)
        res = sorted(res + ref, key=lambda r: (r["broken"], r["meanlog"]))
    best = res[0]
    print(f"search: {len(res)} starts in {time.time() - t0:.0f}s; best {best['broken']} broken from {best['label']}; "
          f"worst start {res[-1]['broken']}", flush=True)
    for r in res[:5]:
        print("   ", r["broken"], round(r["meanlog"], 4), r["label"])

    # Release the best flat layout into free 3D space: same data, same solver, the plane rule removed.
    # In the same pool: the band table (each band solved from the best layout and from North-pole AE).
    seeds = 1 if FAST else 6
    rel = [("3D", None, np.c_[best["X"], np.random.default_rng(SEED + s).normal(0, 50, len(C))], LO, HI) for s in range(seeds)]
    seeds_b = [r["X"] for r in res[:5]] + [starts["ae"][2]]
    NB = len(seeds_b)
    bjobs = [(f"band {lo} {hi}", None, X0, lo, hi) for lo, hi in (BANDS[:1] if FAST else BANDS) for X0 in seeds_b]
    out_jobs = pool.map(_run, rel + bjobs, chunksize=1)
    pool.close()
    X3s = out_jobs[:seeds]
    b3 = min(X3s, key=lambda t: t["broken"]); k3, X3 = b3["broken"], b3["X"]
    sph = sphere_fit(X3)
    print(f"3D: {k3} broken, radius {sph['radius']:.0f} km, wrap {sph['wrap']:.0f} deg, offset {sph['offset']:.3f}", flush=True)

    bands = []
    for j, (lo, hi) in enumerate(BANDS[:1] if FAST else BANDS):
        cand = out_jobs[seeds + NB * j: seeds + NB * (j + 1)]
        cand = cand + [dict(broken=prob.broken(best["X"], lo, hi), X=best["X"])]
        X = min(cand, key=lambda t: t["broken"])["X"]; r = prob.dists(X)[0] / prob.fl; gr = GF / prob.fl
        bands.append(dict(lo=lo, hi=hi, flat=dict(**counts(prob, X, lo, hi), below75=int((r < 0.75).sum()), below60=int((r < 0.6).sum())),
                          globe=dict(short=int((gr < 1 - lo).sum()), long=int((gr > 1 + hi).sum()), limits=int((GL > prob.lu * 1.0005).sum()),
                                     below75=int((gr < 0.75).sum()), below60=int((gr < 0.6).sum())),
                          flat_long=[[int(prob.fa[i]), int(prob.fb[i]), round(float(r[i]), 3)] for i in np.where(r > 1 + hi)[0]]))
        print("band", lo, hi, bands[-1]["flat"], bands[-1]["globe"], flush=True)

    r = prob.dists(best["X"])[0] / prob.fl; gr = GF / prob.fl
    stats = dict(
        n_cities=len(C), n_flights=len(fl), n_limits=len(lim),
        n_kind={k: sum(l["kind"] == k for l in lim) for k in ("rt", "rail", "road", "sat")},
        n_wn=D["n_wn"], n_live=D["n_live"], dropped=D["dropped"], mislocated=D["mislocated"], n_starts=len(jobs),
        best=dict(label=best["label"], centre=best["centre"], broken=best["broken"], **counts(prob, best["X"], LO, HI),
                  below80=int((r < 0.8).sum()), below85=int((r < 0.85).sum()), meanlog=best["meanlog"]),
        globe=dict(**{k: int(v) for k, v in dict(short=(gr < 1 - LO).sum(), long=(gr > 1 + HI).sum(),
                                                 limits=(GL > prob.lu * 1.0005).sum()).items()},
                   below80=int((gr < 0.8).sum()), below85=int((gr < 0.85).sum()), meanlog=float(np.abs(np.log(gr)).mean())),
        starts_broken=[[x["label"], x["broken"]] for x in res if "refined" not in x],
        n_best=sum(x["broken"] == best["broken"] for x in res),
        three=dict(broken=k3, **counts(prob, X3, LO, HI), **sph),
        globe_long=[[int(prob.fa[i]), int(prob.fb[i]), round(float(gr[i]), 3)] for i in np.where(gr > 1 + HI)[0]],
        globe_short=[[int(prob.fa[i]), int(prob.fb[i]), round(float(gr[i]), 3)] for i in np.where(gr < 1 - LO)[0]],
        tri=D["tri"], lo=LO, hi=HI, km_per_ms=KM_PER_MS, end_slack=END_SLACK, match_km=MATCH_KM,
        built=dt.datetime.now(dt.timezone.utc).strftime("%Y-%m-%d %H:%M UTC"))

    def strip(u):
        return u[len(FA_PREFIX):] if u.startswith(FA_PREFIX) else u
    out = dict(
        cities=D["cities"],
        fl=[[f["a"], f["b"], round(f["L"], 1), f["airline"], f["fa"], f["fb"], f["page"],
             f["ab"][:5] + [strip(f["ab"][5])], f["ba"][:5] + [strip(f["ba"][5])]] for f in fl],
        lim=[[l["a"], l["b"], l["U"], l["kind"], l["note"], l["label"], l["url"]] for l in lim],
        gf=[round(float(x), 1) for x in GF], gl=[round(float(x), 1) for x in GL],
        land=land(),
        starts={k: dict(label=v[0], centre=v[1], X=np.round(v[2], 1).tolist()) for k, v in starts.items()},
        best=dict(label=best["label"], centre=best["centre"], X=np.round(best["X"], 1).tolist()),
        best3=dict(X=np.round(X3, 1).tolist()),
        bands=bands, stats=stats, fa_prefix=FA_PREFIX)
    os.makedirs(os.path.join(HERE, "data"), exist_ok=True)
    with open(os.path.join(HERE, "data", "preview.json" if FAST else "femap.json"), "w") as fh:
        json.dump(out, fh, separators=(",", ":"))
    print(json.dumps(stats, indent=1)[:4000])
    render(out)
    print(f"done in {time.time() - t0:.0f}s")


def render(out):
    tpl = os.path.join(HERE, "template.html")
    if not os.path.exists(tpl):
        return
    html = open(tpl, encoding="utf-8").read()
    s = out["stats"]
    vals = copy_values(out)
    for k, v in vals.items():
        html = html.replace("{{" + k + "}}", v)
    left = sorted(set(re.findall(r"\{\{(\w+)\}\}", html)))
    assert not left, left
    html = html.replace("/*DATA*/null", json.dumps(out, separators=(",", ":")))
    os.makedirs(os.path.join(HERE, "site"), exist_ok=True)
    name = "preview.html" if FAST else "index.html"
    open(os.path.join(HERE, "site", name), "w", encoding="utf-8").write(html)
    print("site/" + name, round(len(html) / 1e6, 2), "MB")


def fmt(n):
    return f"{n:,}"


def copy_values(out):
    s = out["stats"]; b, g, t3 = s["best"], s["globe"], s["three"]
    names = [c[0] for c in out["cities"]]
    long_ = sorted(out["bands"][0]["flat_long"], key=lambda x: -x[2])
    v = dict(
        N_CITIES=fmt(s["n_cities"]), N_FLIGHTS=fmt(s["n_flights"]), N_LIMITS=fmt(s["n_limits"]),
        N_RT=fmt(s["n_kind"]["rt"]), N_RAIL=fmt(s["n_kind"]["rail"] + s["n_kind"]["road"]), N_SAT=fmt(s["n_kind"]["sat"]),
        N_STARTS=fmt(s["n_starts"]), N_TRI=fmt(s["tri"]["n"]), TRI_FLAT=fmt(s["tri"]["flat"]), TRI_GLOBE=fmt(s["tri"]["globe"]),
        BEST_LONG=fmt(b["long"]), BEST_SHORT=fmt(b["short"]), BEST_LIM=fmt(b["limits"]), BEST_BROKEN=fmt(b["broken"]),
        BEST_B80=fmt(b["below80"]), GLOBE_B80=fmt(g["below80"]), GLOBE_LONG=fmt(g["long"]), GLOBE_SHORT=fmt(g["short"]),
        GLOBE_LIM=fmt(g["limits"]),
        WORST_PAIR=f"{names[long_[0][0]]}–{names[long_[0][1]]}" if long_ else "–",
        WORST_X=f"{long_[0][2]:.2f}" if long_ else "–",
        LONG_LIST=", ".join(f"{names[a]}–{names[bb]} ×{x:.2f}" for a, bb, x in long_[:8]),
        R3=fmt(round(t3["radius"])), WRAP3=f"{t3['wrap']:.0f}", OFF3=f"{t3['offset'] * 100:.1f}", BROKEN3=fmt(t3["broken"]),
        LO_PCT=f"{round((1 - s['lo']) * 100)}", HI_PCT=f"{round((1 + s['hi']) * 100)}",
        BUILT=s["built"], N_DROPPED=fmt(len(s["dropped"])), N_STARTS_AE=fmt(s["n_starts"] - 3), KM_PER_MS=f"{s['km_per_ms']:.1f}",
        FIT_LO=next((f"{round((1 - x['lo']) * 100)}" for x in out["bands"] if x["flat"]["short"] + x["flat"]["long"] + x["flat"]["limits"] == 0), "–"),
        FIT_B75=next((fmt(x["flat"]["below75"]) for x in out["bands"] if x["flat"]["short"] + x["flat"]["long"] + x["flat"]["limits"] == 0), "–"),
        FIT_B60=next((fmt(x["flat"]["below60"]) for x in out["bands"] if x["flat"]["short"] + x["flat"]["long"] + x["flat"]["limits"] == 0), "–"),
        GLOBE_B75=fmt(out["bands"][0]["globe"]["below75"]), GLOBE_B60=fmt(out["bands"][0]["globe"]["below60"]),
        GLOBE_SHORT_LIST=", ".join(f"{names[a]}–{names[bb]}" for a, bb, x in sorted(s["globe_short"], key=lambda t: t[2])[:4]),
    )
    return v


if __name__ == "__main__":
    main()