"""Best-effort layout of cities on a plane (or in free 3D space) from measured lengths and limits.

Lengths (flights) must fall inside a band [L(1-lo), L(1+hi)]; limits (round trips, rail, road,
satellite passes) cap a distance from above. Nothing else is imposed: no latitude, longitude,
radius or projection enters the objective.

The search runs in two stages, mirrored by the page's live solver:
  1. least squares on every flight's relative error (smooth, finds the overall shape);
  2. reweighted pseudo-L1 on the amount by which each measurement falls outside its band, so the
     solver gives up on as few measurements as it can rather than straining many (an approximation
     to minimising the number broken).
"""
import numpy as np
from scipy.optimize import minimize

LIMIT_W = 3.0     # weight of a limit relative to a flight in stage 2
MARGIN = 0.005    # stage 2 aims this far inside each band
DELTA = 0.003     # pseudo-L1 smoothing
EPS = 0.02        # reweighting floor
ROUNDS = 8


class Problem:
    def __init__(self, n, fa, fb, fl, la, lb, lu):
        self.n = n
        self.fa, self.fb, self.fl = np.asarray(fa), np.asarray(fb), np.asarray(fl, float)
        self.la, self.lb, self.lu = np.asarray(la, int), np.asarray(lb, int), np.asarray(lu, float)

    def _grad(self, dim, X, ga, gl):
        G = np.empty((self.n, dim)); n = self.n
        gv = ga[:, None] * (X[self.fa] - X[self.fb])
        hv = gl[:, None] * (X[self.la] - X[self.lb]) if len(self.lu) else None
        for q in range(dim):
            G[:, q] = np.bincount(self.fa, gv[:, q], n) - np.bincount(self.fb, gv[:, q], n)
            if hv is not None:
                G[:, q] += np.bincount(self.la, hv[:, q], n) - np.bincount(self.lb, hv[:, q], n)
        return G.ravel()

    def dists(self, X):
        dd = np.sqrt(((X[self.fa] - X[self.fb]) ** 2).sum(1)) + 1e-9
        ld = np.sqrt(((X[self.la] - X[self.lb]) ** 2).sum(1)) + 1e-9 if len(self.lu) else np.zeros(0)
        return dd, ld

    def l2(self, dim, wl=100.0):
        def f(x):
            X = x.reshape(self.n, dim); dd, ld = self.dists(X)
            e = dd / self.fl - 1; ex = np.maximum(0, ld / self.lu - 1)
            return (e ** 2).sum() + wl * (ex ** 2).sum(), self._grad(dim, X, 2 * e / self.fl / dd, wl * 2 * ex / self.lu / ld)
        return f

    def sparse(self, dim, lo, hi, w, wl):
        a = self.fl * (1 - lo + MARGIN); b = self.fl * (1 + hi - MARGIN); ub = self.lu * (1 - MARGIN)
        def f(x):
            X = x.reshape(self.n, dim); dd, ld = self.dists(X)
            eh = np.maximum(0, dd / b - 1); el = np.maximum(0, 1 - dd / a)
            sh = np.sqrt(eh ** 2 + DELTA ** 2); sl = np.sqrt(el ** 2 + DELTA ** 2)
            ex = np.maximum(0, ld / ub - 1); sx = np.sqrt(ex ** 2 + DELTA ** 2)
            F = (w * (sh + sl - 2 * DELTA)).sum() + (wl * (sx - DELTA)).sum()
            ga = w * (eh / sh / b - el / sl / a) / dd
            gl = wl * ex / sx / ub / ld
            return F, self._grad(dim, X, ga, gl)
        return f

    def excess(self, X, lo, hi):
        dd, ld = self.dists(X); r = dd / self.fl
        return np.maximum(0, r / (1 + hi) - 1) + np.maximum(0, 1 - r / (1 - lo)), np.maximum(0, ld / self.lu - 1)

    def broken(self, X, lo, hi):
        e, l = self.excess(X, lo, hi)
        return int((e > 0).sum() + (l > 5e-4).sum())

    def best_effort(self, X0, lo, hi, rounds=ROUNDS, maxiter=30000, eps=EPS, limit_w=LIMIT_W):
        dim = X0.shape[1]
        run = lambda f, X: minimize(f, X.ravel(), jac=True, method="L-BFGS-B",
                                    options={"maxiter": maxiter, "maxcor": 30}).x.reshape(self.n, dim)
        X = run(self.l2(dim), X0)
        best = (self.broken(X, lo, hi), X.copy())
        w = np.ones(len(self.fl)); wl = np.full(len(self.lu), limit_w)
        for _ in range(rounds):
            X = run(self.sparse(dim, lo, hi, w, wl), X)
            k = self.broken(X, lo, hi)
            if k < best[0]:
                best = (k, X.copy())
            e, l = self.excess(X, lo, hi)
            w = eps / (e + eps); wl = limit_w * eps / (l + eps)
        return best


def sphere_fit(X):
    """Least-squares sphere through 3D points: radius, median offset (share of radius), and how far
    round it the points reach (angle from their mean direction to the farthest point)."""
    A = np.c_[2 * X, np.ones(len(X))]; s = np.linalg.lstsq(A, (X ** 2).sum(1), rcond=None)[0]
    c = s[:3]; r = float(np.sqrt(s[3] + c @ c)); d = np.linalg.norm(X - c, axis=1)
    u = (X - c) / d[:, None]; m = u.mean(0); m /= np.linalg.norm(m)
    return dict(radius=r, offset=float(np.median(np.abs(d - r) / r)),
                wrap=float(np.degrees(np.arccos(np.clip(u @ m, -1, 1))).max()), centre=c.tolist())