"""Prototype (bench only, the package is untouched): slices started from the neighbouring slice.

Two independent sweeps from the slicer's rings, bottom-up and top-down, merged layer by layer. In a sweep, layer k is
seeded with the rings layer k - 1 ended with: the refined rings accepted there, the rings held from earlier layers (a
refined ring that fails is retried from its last good shape for half the layers it has been carried, HOLD to HOLD_MAX:
momentum), and the slicer's rings of k - 1 that none of those covers. Each seed is resampled every STEP spacings and
refined to layer k's points, ITERS times: every vertex moves along the ring's normal to the weighted mean offset of the
side points near its normal line (lateral weight exp(-(l / s)^2), height weight exp(-(dz / DZ s)^2), offsets within
REACH spacings) on the ring's own face (oriented normal out of the solid on the ring's side, cos >= FACE; a ring is a
hole when it lies inside an odd number of its layer's rings; side points: |n_z| <= VERT, near-horizontal surfaces say
nothing about where the contour is). Vertices with no side points take the offsets of their neighbours along the ring
(gaps bridged); the offsets are smoothed (Gaussian, SMOOTH spacings of arc); small self-loops are cut off. A refined ring
is accepted when at least MIN_VALID of its vertices are backed (side points on its face, or any point within a
spacing), its p95 distance to the slab's points is at most FIT spacings and its area is at least MIN_AREA of its seed's.
`resolve`: accepted rings replace the slicer's rings they explain (EXPLAIN_SHARE of the slicer ring's vertices within
EXPLAIN spacings; a slicer hole only when an accepted hole explains it); rings that cross each other or a kept slicer
ring go; if the layer's region grows beyond the slicer's grown by GROW spacings, the slicer's rings stay.
`merge`: a ring both sweeps carried through a layer is used; a ring only one carried (what a real change looks like:
a leg's end is reached from one side only) only when the layer's points back it clearly (ONE_VALID, ONE_FIT).
All in metric canonical units (the canonical plane is scaled by the half-extents, not isotropic).

    from neighbour import encode_neighbour
    d = encode_neighbour(points, up=(0, 0, 1))
"""
import sys

import numpy as np
from scipy.ndimage import gaussian_filter1d
from scipy.spatial import cKDTree
from shapely.geometry import LinearRing

sys.path.insert(0, "/home/joao/code/fcp_devel/api")
import fcp
from fcp import _rings
from fcp._slice import PointSlicer, centres, check_layers

STEP = 0.5  # spacings between ring vertices
REACH = 2.0  # spacings: points this far along the normal are seen
DZ = 0.5  # spacings: height weight scale
SLAB = 1.0  # spacings: points within this height of the plane are used
ITERS = 3
SMOOTH = 1.0  # spacings of arc: offset smoothing along the ring
MIN_VALID = 0.6  # share of vertices with points behind them
FIT = 1.5  # spacings: p95 distance of the refined ring to the slab's points
MIN_AREA = 0.25  # of the predecessor's |area|
FACE = 0.3  # cos(point normal, the ring's solid-outward normal) at least this: the point is on the ring's face
EXPLAIN = 1.5  # spacings: a slicer ring with EXPLAIN_SHARE of its vertices this close to the accepted rings is replaced
EXPLAIN_SHARE = 0.8
VERT = 0.7  # |n_z| above this: a point on a near-horizontal surface, left out
HOLD, HOLD_MAX = 3, 12  # layers a failed ring is retried from its last good shape: half the layers it was carried, within these
ONE_VALID, ONE_FIT = 0.8, 1.0  # a ring only one sweep carried is used when this well backed (share) and this close (spacings, p95)
GROW = 2.0  # spacings: the layer's region may not exceed the slicer's region grown by this

STATS = {}
LOG = []  # per layer of each pass: what happened to every propagated ring (diagnostics)


def _resample(R, step):
    R = np.vstack([R, R[:1]])
    seg = np.linalg.norm(np.diff(R, axis=0), axis=1)
    s = np.concatenate([[0], np.cumsum(seg)])
    n = max(8, int(s[-1] / step))
    t = np.linspace(0, s[-1], n, endpoint=False)
    return np.c_[np.interp(t, s, R[:, 0]), np.interp(t, s, R[:, 1])]


def _normals(R):
    t = np.roll(R, -1, 0) - np.roll(R, 1, 0)
    n = np.c_[t[:, 1], -t[:, 0]]
    return n / np.maximum(np.linalg.norm(n, axis=1, keepdims=True), 1e-12)


def _fill_circular(o, ok):
    """Offsets o at the invalid vertices: linear between the valid ones around the ring."""
    n = len(o)
    i = np.nonzero(ok)[0]
    if len(i) == 0:
        return o
    x = np.r_[i - n, i, i + n]
    return np.interp(np.arange(n), x, np.tile(o[i], 3))


def refine(R, Q, dz, NQ, side, tree, s, hole):
    """R (m, 2) ring guess; Q (q, 2) the slab's points, dz their heights off the plane, NQ their in-plane normals, side:
    not on a near-horizontal surface, tree on Q; hole: the ring bounds a hole. -> (ring, backed share). A vertex moves by
    the side points on its face; it is backed when such points or any point lies within a spacing of it."""
    valid = 0.0
    for _ in range(ITERS):
        R = _resample(R, STEP * s)
        N = _normals(R)  # out of the region the ring encloses when it runs counter-clockwise
        out = N * (1 if (_signed_area(R) > 0) != hole else -1)  # out of the solid
        o = np.zeros(len(R))
        ok = np.zeros(len(R), bool)
        for i, js in enumerate(tree.query_ball_point(R, np.hypot(REACH, 1.0) * s)):
            if not js:
                continue
            D = Q[js] - R[i]
            off = D @ N[i]
            lat = D @ np.array([-N[i, 1], N[i, 0]])
            m = (np.abs(off) <= REACH * s) & side[js] & (NQ[js] @ out[i] >= FACE)
            w = np.exp(-(lat[m] / s) ** 2) * np.exp(-(dz[js][m] / (DZ * s)) ** 2)
            if w.sum() >= 0.3:
                o[i] = (w * off[m]).sum() / w.sum()
                ok[i] = True
        backed = ok | (tree.query(R)[0] <= s)
        valid = backed.mean()
        if valid < MIN_VALID:
            return R, valid
        o = gaussian_filter1d(_fill_circular(o, ok), SMOOTH / STEP, mode="wrap")
        R = _main_loop(R + o[:, None] * N)
    return R, valid


def _main_loop(R):
    """R without its small self-loops: the exterior of the largest piece of its region (R if it is simple)."""
    if LinearRing(R).is_simple:
        return R
    from shapely.geometry import Polygon

    g = Polygon(R).buffer(0)
    if g.is_empty:
        return R
    p = max(getattr(g, "geoms", [g]), key=lambda q: q.area)
    E = np.asarray(p.exterior.coords)[:-1]
    if (_signed_area(E) > 0) != (_signed_area(R) > 0):  # keep the ring's direction
        E = E[::-1]
    return E


def _signed_area(R):
    x, y = R[:, 0], R[:, 1]
    return 0.5 * (np.dot(x, np.roll(y, -1)) - np.dot(y, np.roll(x, -1)))


def _area(R):
    return abs(_signed_area(R))


def _holes(loops):
    """Per ring: inside an odd number of the other rings of its layer."""
    from shapely.geometry import Point, Polygon

    polys = [Polygon(L) for L in loops]
    return [sum(polys[j].contains(Point(L[0])) for j in range(len(loops)) if j != i) % 2 == 1 for i, L in enumerate(loops)]


def _crosses(a, b):
    return LinearRing(a).intersects(LinearRing(b))


def _region(rings):
    """The even-odd region of a layer's rings (shapely)."""
    from shapely.geometry import Polygon

    g = Polygon()
    for L in rings:
        if len(L) >= 3:
            g = g.symmetric_difference(Polygon(L).buffer(0))
    return g


def _share_near(L, rings, d):
    """Share of L's vertices within d of the rings."""
    return float(np.mean(cKDTree(np.vstack(rings)).query(L)[0] <= d)) if rings else 0.0


def resolve(sl, acc, acc_hole, s):
    """Layer rings from the slicer's (sl) and the accepted refined ones: the slicer rings the accepted explain go (a
    slicer hole only when an accepted hole explains it); accepted rings that cross each other or a kept slicer ring go;
    if the region grows beyond the slicer's grown by GROW spacings, the slicer's rings stay.
    -> (rings, accepted, their hole flags, slicer rings replaced, conflicts, region check failed)."""
    sl_hole = _holes(sl)
    conflicts = []
    while True:
        holes = [a for a, ah in zip(acc, acc_hole) if ah]
        expl = [_share_near(L, acc, EXPLAIN * s) >= EXPLAIN_SHARE
                and (not h or _share_near(L, holes, EXPLAIN * s) >= EXPLAIN_SHARE) for L, h in zip(sl, sl_hole)]
        kept = [L for L, e in zip(sl, expl) if not e]
        bad = {i for i, a in enumerate(acc) for j, b in enumerate(acc) if i != j and _crosses(a, b)}
        bad |= {i for i, a in enumerate(acc) if any(_crosses(a, L) for L in kept)}
        if not bad:
            break
        conflicts += [acc[i] for i in sorted(bad)]
        acc = [a for i, a in enumerate(acc) if i not in bad]
        acc_hole = [a for i, a in enumerate(acc_hole) if i not in bad]
    if acc and _region(acc + kept).difference(_region(sl).buffer(GROW * s)).area > 0.02 * max(_region(sl).area, 1e-12):
        return list(sl), [], [], 0, conflicts, True
    return acc + kept, acc, acc_hole, int(sum(expl)), conflicts, False


def _slab(X, NX, zk, s):
    near = np.abs(X[:, 2] - zk) <= SLAB * s
    Q, dz = X[near, :2], X[near, 2] - zk
    NQ = NX[near, :2] / np.maximum(np.linalg.norm(NX[near, :2], axis=1, keepdims=True), 1e-12)
    side = np.abs(NX[near, 2]) <= VERT * np.linalg.norm(NX[near], axis=1)
    return Q, dz, NQ, side


def sweep(loops, X, NX, z, s, order, tag):
    """One sweep from the slicer's rings (loops: per layer, metric canonical plane) in the layer order `order`.
    -> per layer, the refined rings accepted: dicts (ring, hole, valid, fit, run: the layers the ring has been carried)."""
    out = {order[0]: [np.asarray(L, float) for L in loops[order[0]]]}
    accepted = {order[0]: []}
    good, held = [], []  # (ring, hole, run) accepted at the previous layer; (ring, hole, run, age) held
    for prev, k in zip(order[:-1], order[1:]):
        sl = [np.asarray(L, float) for L in loops[k]]
        mine = [g[0] for g in good] + [h[0] for h in held]
        seeds = [(r, h, run, 0) for r, h, run in good] + held
        seeds += [(L, h, 0, None) for L, h in zip(out[prev], _holes(out[prev]))
                  if not any(L is r for r in mine) and _share_near(L, mine, EXPLAIN * s) < EXPLAIN_SHARE]
        Q, dz, NQ, side = _slab(X, NX, z[k], s)
        cand, failed = [], []
        if len(Q) >= 8:
            tree = cKDTree(Q)
            for P, hole, run, age in seeds:
                STATS["tried"] += 1
                R, valid = refine(P, Q, dz, NQ, side, tree, s, hole)
                fit = float(np.percentile(tree.query(R)[0], 95) / s)
                why = ("valid" if valid < MIN_VALID else "area" if _area(R) < MIN_AREA * _area(P)
                       else "self-crossing" if not LinearRing(R).is_simple else "fit" if fit > FIT else "")
                LOG.append(dict(k=k, prev=prev, sweep=tag, hole=hole, valid=valid, fit=fit, why=why, ring=R, seed=P,
                                held=age, run=run))
                if why:
                    if age is not None and age + 1 <= _hold(run):
                        failed.append((P, hole, run, age + 1))
                    continue
                cand.append(dict(ring=R, hole=hole, valid=valid, fit=fit, run=run + 1))
        rings, acc, acc_hole, _, _, _ = resolve(sl, [c["ring"] for c in cand], [c["hole"] for c in cand], s)
        accepted[k] = [c for c in cand if any(c["ring"] is a for a in acc)]
        out[k] = rings
        good = [(c["ring"], c["hole"], c["run"]) for c in accepted[k]]
        held = [f for f in failed if _share_near(f[0], acc, EXPLAIN * s) < EXPLAIN_SHARE]
    return accepted


def _hold(run):
    """Layers a failed ring is retried: grows with the layers it has been carried (momentum), HOLD to HOLD_MAX."""
    return int(min(HOLD_MAX, max(HOLD, run // 2)))


def merge(loops, up, down, s):
    """Per layer: the rings both sweeps carried (matched: same kind, EXPLAIN_SHARE of the vertices within EXPLAIN
    spacings) are used; a ring only one sweep carried only when its layer's points back it clearly (ONE_VALID, ONE_FIT);
    then `resolve` against the slicer's rings."""
    new = []
    for k, sl in enumerate(loops):
        sl = [np.asarray(L, float) for L in sl]
        U, D = up.get(k, []), down.get(k, [])
        acc, acc_hole, both, one = [], [], 0, 0
        used = set()
        for u in U:
            m = next((j for j, d in enumerate(D) if j not in used and d["hole"] == u["hole"]
                      and _share_near(u["ring"], [d["ring"]], EXPLAIN * s) >= EXPLAIN_SHARE), None)
            if m is not None:
                used.add(m)
                best = u if u["fit"] <= D[m]["fit"] else D[m]
                acc.append(best["ring"]); acc_hole.append(best["hole"]); both += 1
            elif u["valid"] >= ONE_VALID and u["fit"] <= ONE_FIT:
                acc.append(u["ring"]); acc_hole.append(u["hole"]); one += 1
        for j, d in enumerate(D):
            if j not in used and d["valid"] >= ONE_VALID and d["fit"] <= ONE_FIT:
                acc.append(d["ring"]); acc_hole.append(d["hole"]); one += 1
        rings, acc, acc_hole, replaced, conflicts, grown = resolve(sl, acc, acc_hole, s)
        LOG.append(dict(k=k, merged=True, both=both, one=one, accepted=[("hole" if h else "solid") for h in acc_hole],
                        slicer=[("hole" if h else "solid") for h in _holes(sl)], replaced=replaced,
                        conflicts=len(conflicts), grown=grown, slicer_rings=sl, final=rings,
                        up=[u["ring"] for u in U], down=[d["ring"] for d in D]))
        STATS["accepted"] += len(acc)
        STATS["replaced"] += replaced
        new.append(rings)
    return new


def encode_neighbour(P, up=None, layers=128, device="cuda", **kw):
    """fcp.encode of the points P with the slices refined from their neighbours (the rest as the package)."""
    orig = PointSlicer.slices

    def slices(self, layers=layers):
        c = centres(check_layers(layers))
        loops = self.loops(c)
        sc = self.pose.scale.astype(float)
        X = (np.asarray(P, float) - self.pose.center) @ self.pose.rotation  # metric canonical
        NX = self.world_normals() @ self.pose.rotation
        z = c * sc[2]
        metric = [[np.asarray(L, float) * sc[:2] for L in ls] for ls in loops]
        STATS.update(tried=0, accepted=0, replaced=0)
        LOG.clear()
        n, sp = len(metric), float(self.spacing)
        up = sweep(metric, X, NX, z, sp, list(range(n)), "up")  # two independent sweeps from the slicer's rings
        down = sweep(metric, X, NX, z, sp, list(range(n - 1, -1, -1)), "down")
        new = merge(metric, up, down, sp)
        return _rings.slices([[(L / sc[:2]).astype(np.float32) for L in ls] for ls in new], self.device)

    PointSlicer.slices = slices
    try:
        return fcp.encode(P, up=up, layers=layers, device=device, **kw)
    finally:
        PointSlicer.slices = orig