"""The grid-start instance: one robot in the plane passes two discs; start (0,0), goal (10,0); acceleration cost
(k = 2), endpoint positions fixed (ell = 0), inflation margin 0.005; N cubic pieces on equal intervals joined C^1
(the class S_N of Theorem 1).

Free control points (the paper's rational chart), 4N numbers:
    x = [ c0 | (a_1, Q_1) | ... | (a_{N-1}, Q_{N-1}) | a_N ]
c0 = second control point of piece 0, Q_j = position at joint j, a_j = control point just before joint j.
Piece i has control points (Q_i, o_i, a_{i+1}, Q_{i+1}) with Q_0 = START, Q_N = GOAL, o_0 = c0, o_i = 2 Q_i - a_i
(C^1 with equal piece durations; an integer map, so the grid h Z^{4N} is a product over the joints).
Time scale: the package gives every piece unit duration, so cost(total duration 1) = N^3 * cost(package).

Contents: the program (`build`), the chart, the cost as a quadratic (`Quad`), the exact test of one piece
(`pieces_free`), what INIT needs (`GridStates`: the grid values per joint, the piece cost, the piece test; `grid_start`
calls globmp.init), route classes, and the exact rational check of a whole curve.
"""
import math
import time

import numpy as np
import torch

from .. import Problem, DT
from .. import spline as bern
from ..check import block_minima
from ..init import init

torch.set_num_threads(1)

START = np.array([0.0, 0.0]); GOAL = np.array([10.0, 0.0])
DISCS = [((3.3, 0.25), 1.0), ((6.7, -0.35), 1.0)]
DELTA = 0.005
D = 3
K4 = bern.cost_gram(D, 2).numpy()                    # 4x4 Gram of the acceleration cost of one piece
CD = np.array([math.comb(D, j) for j in range(D + 1)], float)
C2 = np.array([math.comb(2 * D, j) for j in range(2 * D + 1)], float)


def build(N, whiten=False, bb=False, stop_rel=1e-10):
    p = Problem(2, D, 2, 0, N, START, GOAL, ell_j=1, delta=DELTA, whiten=whiten, bb=bb, stop_rel=stop_rel)
    for c, r in DISCS: p.add_sphere(c, r)
    return p


# ------------------------------------------------------------------ chart
def idx_joint(j, N):
    """x indices of the state at joint j (0: c0, N: a_N, else (a_j, Q_j))."""
    if j == 0: return [0, 1]
    if j == N: return [4 * N - 2, 4 * N - 1]
    return list(range(2 + 4 * (j - 1), 6 + 4 * (j - 1)))


def ctrl(x, N):
    """x (..., 4N) -> control points (..., N, 4, 2)."""
    x = np.asarray(x, float); lead = x.shape[:-1]
    C = np.zeros(lead + (N, 4, 2))
    Q = [np.broadcast_to(START, lead + (2,))]; a = [None]
    for j in range(1, N):
        ij = idx_joint(j, N); a.append(x[..., ij[0]:ij[0] + 2]); Q.append(x[..., ij[0] + 2:ij[0] + 4])
    a.append(x[..., 4 * N - 2:4 * N]); Q.append(np.broadcast_to(GOAL, lead + (2,)))
    for i in range(N):
        o = x[..., 0:2] if i == 0 else 2 * Q[i] - a[i]
        C[..., i, 0, :] = Q[i]; C[..., i, 1, :] = o; C[..., i, 2, :] = a[i + 1]; C[..., i, 3, :] = Q[i + 1]
    return C


def x_to_G(x, N):
    C = ctrl(x, N)                                   # (N,4,2)
    return np.transpose(C, (2, 0, 1)).reshape(2, N * 4)


def G_to_Z(p, G):
    """chart point of a coefficient matrix that satisfies the linear constraints (any chart, also whitened)."""
    R = (torch.as_tensor(G, dtype=DT) - p.Gamma0)
    Z = torch.linalg.lstsq(p.N, R.T).solution.T
    assert float((p.Gamma(Z) - torch.as_tensor(G, dtype=DT)).abs().max()) < 1e-8, 'G violates the linear constraints'
    return Z


def Z_to_x(p, Z, N):
    S = p.segs(torch.as_tensor(Z, dtype=DT)).detach().numpy()         # (N, 2, 4)
    x = np.zeros(4 * N); x[0:2] = S[0, :, 1]
    for j in range(1, N):
        ij = idx_joint(j, N); x[ij[0]:ij[0] + 2] = S[j - 1, :, 2]; x[ij[0] + 2:ij[0] + 4] = S[j - 1, :, 3]
    x[4 * N - 2:] = S[N - 1, :, 2]
    return x


class Quad:
    """f(x) = c + 2 g.x + x^T H x for the whole curve and for every suffix of pieces (cost-to-go bounds)."""

    def __init__(self, N):
        self.N = N; m = 4 * N
        C0 = ctrl(np.zeros(m), N); J = np.stack([ctrl(e, N) - C0 for e in np.eye(m)], -1)     # (N,4,2,m)
        self.C0, self.J = C0, J
        self.c, self.g, self.H = self._quad(range(N))
        Hi = np.linalg.inv(self.H); self.xs = -Hi @ self.g; self.fmin = self.c + self.g @ self.xs; self.Hi = Hi
        self.togo_data = {}
        for j in range(N):
            c, g, H = self._quad(range(j, N)); s = idx_joint(j, N); r = list(range(s[-1] + 1, m))
            Hrr_i = np.linalg.inv(H[np.ix_(r, r)])
            self.togo_data[j] = (c, g[s], H[np.ix_(s, s)], g[r], H[np.ix_(r, s)], Hrr_i)

    def _quad(self, pieces):
        pieces = list(pieces); C0 = self.C0[pieces]; J = self.J[pieces]
        H = np.einsum('pq,ipcm,iqcn->mn', K4, J, J); g = np.einsum('pq,ipc,iqcn->n', K4, C0, J)
        return float(np.einsum('pq,ipc,iqc->', K4, C0, C0)), g, H

    def f(self, x):
        x = np.asarray(x, float); return self.c + 2 * x @ self.g + np.einsum('...m,mn,...n->...', x, self.H, x)

    def togo(self, j, S):
        """lower bound on the cost of pieces j..N-1 from state S (K, dim) at joint j (unconstrained minimum)."""
        if j == self.N: return np.zeros(len(S))
        c, gs, Hss, gr, Hrs, Hrr_i = self.togo_data[j]
        v = gr[None, :] + S @ Hrs.T
        return c + 2 * S @ gs + np.einsum('km,mn,kn->k', S, Hss, S) - np.einsum('km,mn,kn->k', v, Hrr_i, v)

    def sublevel_count(self, U, h):
        """volume of the ellipsoid {f <= U} divided by h^m: estimate of the number of lattice points with f <= U."""
        m = 4 * self.N; ev = np.linalg.eigvalsh(self.H)
        logv = 0.5 * m * math.log(math.pi) - math.lgamma(0.5 * m + 1) + 0.5 * m * math.log(max(U - self.fmin, 1e-300)) - 0.5 * np.log(ev).sum()
        return math.exp(logv - m * math.log(h))

    def box(self, U, h):
        """integer ranges (lo, hi) per coordinate of the lattice h Z^m inside the bounding box of {f <= U}."""
        half = np.sqrt(max(U - self.fmin, 0.0) * np.diag(self.Hi))
        lo = np.ceil((self.xs - half) / h - 1e-12).astype(int); hi = np.floor((self.xs + half) / h + 1e-12).astype(int)
        return lo, hi


# ------------------------------------------------------------------ exact piece test
def piece_coefs(B):
    """control points (K,4,2) -> Bernstein-6 coefficients (K, n_discs, 7) of |x(s)-c|^2 - (r+DELTA)^2."""
    K = B.shape[0]; out = []
    for c, r in DISCS:
        A = (B - np.array(c)[None, None, :]) * CD[None, :, None]
        prod = np.zeros((K, 2 * D + 1))
        for a in range(D + 1):
            prod[:, a:a + D + 1] += (A[:, a:a + 1, :] * A).sum(-1)
        out.append(prod / C2 - (r + DELTA) ** 2)
    return np.stack(out, 1)


def _split(c):
    """de Casteljau at 1/2: (m,7) -> left (m,7), right (m,7)."""
    m, n = c.shape; L = np.empty_like(c); R = np.empty_like(c); cur = c.copy()
    for j in range(n):
        L[:, j] = cur[:, 0]; R[:, n - 1 - j] = cur[:, -1]
        cur = 0.5 * (cur[:, :-1] + cur[:, 1:])
    return L, R


def _exact_min(c):
    """minimum over [0,1] of a polynomial given by Bernstein coefficients (roots of the derivative)."""
    n = len(c) - 1
    pw = np.zeros(n + 1)                                          # power basis, increasing degree
    for j in range(n + 1):
        for i in range(j, n + 1):
            pw[i] += c[j] * math.comb(n, j) * math.comb(n - j, i - j) * (-1) ** (i - j)
    P = np.polynomial.Polynomial(pw); cand = [0.0, 1.0]
    for r in P.deriv().roots():
        if abs(r.imag) < 1e-9 and 0.0 <= r.real <= 1.0: cand.append(r.real)
    return min(float(P(t)) for t in cand)


def nonneg(Bc, depth=40):
    """Bc (K,7): True where the polynomial is >= 0 on [0,1].  Nonnegative Bernstein coefficients certify, a
    negative end value refutes, otherwise bisect; what is still undecided at the depth limit (tangent cases)
    goes to the derivative-root minimisation.  Returns (bool (K,), number of root fallbacks)."""
    K = len(Bc); res = np.ones(K, bool); idx = np.arange(K); cur = Bc
    for _ in range(depth):
        neg = (cur[:, 0] < 0) | (cur[:, -1] < 0)
        res[idx[neg]] = False
        und = ~neg & ~(cur >= 0).all(1) & res[idx]
        if not und.any(): return res, 0
        cur = cur[und]; idx = idx[und]
        L, R = _split(cur); cur = np.concatenate([L, R]); idx = np.concatenate([idx, idx])
    left = np.unique(idx[res[idx]])
    for j in left: res[j] = _exact_min(Bc[j]) >= 0
    return res, len(left)


def pieces_free(B):
    """control points (K,4,2) -> True where the piece clears every inflated disc on its whole interval."""
    Bc = piece_coefs(B); ok = np.ones(len(B), bool); nfb = 0
    for b in range(Bc.shape[1]):
        sel = np.where(ok)[0]
        if len(sel) == 0: break
        r, f = nonneg(Bc[sel, b]); ok[sel] = r; nfb += f
    return ok, nfb


# ------------------------------------------------------------------ INIT on the grid h Z^{4N}
def lattice(lo, hi, h):
    ax = [np.arange(l, u + 1) * h for l, u in zip(lo, hi)]
    return np.stack(np.meshgrid(*ax, indexing='ij'), -1).reshape(-1, len(ax))


class GridStates:
    """The grid h Z^{4N} inside the bounding box of {f <= U}, as INIT's input (globmp.init): the state at joint j
    is the grid value of (a_j, Q_j) (c0 at joint 0, a_N at joint N); piece i has control points
    (Q_i, o_i, a_{i+1}, Q_{i+1}); its cost is the acceleration cost, its test is `pieces_free`."""

    def __init__(self, N, h, U):
        self.q = q = Quad(N); self.N = self.n_pieces = N; self.lo, self.hi = lo, hi = q.box(U, h)
        self.T = [lattice(lo[idx_joint(j, N)], hi[idx_joint(j, N)], h) for j in range(N + 1)]
        self.root_fallback = 0; self._piece = (None,)

    def n_states(self, j): return len(self.T[j])

    def _source(self, i, s):
        """(Q_i, o_i) of the states s at joint i."""
        S = self.T[i][s]
        return (np.broadcast_to(START, S.shape), S) if i == 0 else (S[:, 2:], 2 * S[:, 2:] - S[:, :2])

    def _target(self, i):
        """(a_{i+1}, Q_{i+1}) of every state at joint i+1, the target part of the piece cost and the target factor
        of its cross term."""
        if self._piece[0] != i:
            T = self.T[i + 1]; Ta, Tq = (T[:, :2], T[:, 2:]) if i + 1 < self.N else (T, np.broadcast_to(GOAL, T.shape))
            v = np.stack([Ta, Tq], 1)
            self._piece = (i, Ta, Tq, np.einsum('pq,kpc,kqc->k', K4[2:, 2:], v, v), v.reshape(len(v), 4))
        return self._piece[1:]

    def cost(self, i, s, V):
        Ta, Tq, wt, vf = self._target(i); u = np.stack(self._source(i, s), 1)
        ws = np.einsum('pq,kpc,kqc->k', K4[:2, :2], u, u); uK = np.einsum('kpc,pq->kqc', u, K4[:2, 2:]).reshape(len(u), 4)
        tot = 2 * uK @ vf.T; tot += (V + ws)[:, None] + wt[None, :]
        return tot

    def free_pairs(self, i, s, t):
        Sq, So = self._source(i, s); Ta, Tq = self._target(i)[:2]
        ok, nfb = pieces_free(np.stack([Sq, So, Ta[t], Tq[t]], 1)); self.root_fallback += nfb
        return ok

    def togo(self, j): return self.q.togo(j, self.T[j])


def grid_start(N, h, U, rule='togo', improve_only=True):
    """z_h by INIT (globmp.init) on the grid of spacing h with budget U.  Returns the cheapest feasible grid point x
    (or None = FAIL: no feasible grid point with cost <= U), its cost, and the counts."""
    t0 = time.time(); g = GridStates(N, h, U)
    r = init(g, U * (1 + 1e-12), rule=rule, improve_only=improve_only)
    st = dict(layers=r['layers'], fail=not r['curves'], box=[(int(a), int(b)) for a, b in zip(g.lo, g.hi)],
              root_fallback=g.root_fallback)
    if not r['curves']: return None, None, dict(st, time=time.time() - t0)
    cost, path = r['curves'][0]; x = np.zeros(4 * N)
    for j in range(N + 1): x[idx_joint(j, N)] = g.T[j][path[j]]
    npts = 1
    for a, b in zip(g.lo, g.hi): npts *= int(b - a + 1)
    st.update(time=time.time() - t0, grid_points_box=npts, tests=sum(l['tested'] for l in r['layers']),
              pairs=sum(l['pairs'] for l in r['layers']), budget_pass=sum(l['budget_pass'] for l in r['layers']))
    return x, cost, st


# ------------------------------------------------------------------ route classes and checks
def side_of(p, Z):
    """route class: side of each disc at the curve's closest approach ('p' above the centre)."""
    S = p.segs(torch.as_tensor(Z, dtype=DT)).detach().numpy()        # (N,2,4)
    s = np.linspace(0, 1, 801); i = np.arange(D + 1)
    Bm = CD[:, None] * s[None, :] ** i[:, None] * (1 - s[None, :]) ** (D - i[:, None])
    pts = np.concatenate([S[k] @ Bm for k in range(S.shape[0])], 1)  # (2, N*801)
    out = []
    for c, r in DISCS:
        dd = (pts[0] - c[0]) ** 2 + (pts[1] - c[1]) ** 2; t = int(np.argmin(dd)); out.append('p' if pts[1, t] > c[1] else 'm')
    return ''.join(out)


def certify_true(p, Z):
    """exact rational checker (globmp.synthetic.exact) at the true radii, no margin."""
    from ..synthetic.exact.certify_curve import Plan, check_plan
    from fractions import Fraction
    spec = dict(spheres=[([0, 1], (Fraction(c[0]).limit_denominator(1000), Fraction(c[1]).limit_denominator(1000)),
                          Fraction(r).limit_denominator(1000)) for c, r in DISCS])
    segs = p.segs(torch.as_tensor(Z, dtype=DT)).detach().numpy()
    return bool(check_plan(Plan(segs), spec, clearance=False)['certified'])


def subdivide_cubic(P4, N):
    """one cubic with control points P4 (4,2) -> coefficient matrix (2, 4N) of the same curve cut into N equal pieces."""
    from numpy.polynomial import polynomial as Pl
    out = np.zeros((2, 4 * N)); s = np.array([0.0, 1 / 3, 2 / 3, 1.0])
    i = np.arange(4); Bm = lambda t: CD[:, None] * t[None, :] ** i[:, None] * (1 - t[None, :]) ** (3 - i[:, None])
    Minv = np.linalg.inv(Bm(s))                                    # values at 4 nodes -> control points
    for k in range(N):
        t = (k + s) / N; vals = P4.T @ Bm(t)                       # (2,4) values of the big cubic at the piece's nodes
        out[:, 4 * k:4 * k + 4] = vals @ Minv
    return out