"""Prototype (bench only, the package is untouched): thin walls as sheets in the surf slicer's field.

Diagnosis (field_view.py, loaf pan and GSO bowl at 10k): a wall thinner than about a spacing is oriented as one sheet, both
faces with the outer face's normal, so the inner face declares the cavity side of itself inside: the field's inside band
sits beside the wall, not in it, and its zero line wanders and breaks (fragments, C-rings).

Fix: a point is on a thin wall when at least NEED points with a parallel normal (|cos| >= PAR) lie between LO and HI
spacings along its normal and within LAT spacings of its normal line (the wall's other face); the flag then spreads to
points with SPREAD of their SPREAD_K neighbours flagged (SPREAD_ROUNDS times), and groups of fewer than MIN_GROUP connected thin points are dropped. In every layer, within
REACH spacings of the thin points in the slab (|dz| <= SLAB spacings), the field becomes min(field, d - R spacings), d the
distance to those points: the wall's band (R spacings either side of its points) is added to the inside, whatever the
orientation. Replacing the field there instead (MODE "replace") was tried first: the jump at the edge of the replaced
region draws spurious rings (parts doubled even where 0.3 % of the points were flagged).
Elsewhere the field is the package's. The kernels are untouched: the backend's field() is wrapped.

    from thinwall import encode_thinwall
"""
import sys

import numpy as np
from scipy.spatial import cKDTree

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

LO, HI, LAT, PAR, NEED = 0.4, 1.5, 0.75, 0.8, 2
R = 0.9  # spacings: half-width of a thin wall's band (wide enough for neighbouring points' discs to merge)
REACH = 1.5  # spacings: the field is replaced this near the thin points
SLAB = 1.0  # spacings: thin points this near the plane count for a layer
MIN_GROUP = 25  # thin points connected within HI spacings: smaller groups are not walls
SPREAD, SPREAD_K, SPREAD_ROUNDS = 0.3, 16, 2  # a point is thin when this share of its K neighbours is (ROUNDS times)
MODE = "union"  # "union": min(field, d - R) (continuous); "replace": d - R near the thin points (jumps at the edge: specks)
STATS = {}
LAST = {}  # the last encode's metric points and thin flags (for renders)


def thin_flags(X, N, s):
    T = cKDTree(X)
    out = np.zeros(len(X), bool)
    for i, js in enumerate(T.query_ball_point(X, np.hypot(HI, LAT) * s)):
        js = np.asarray(js)
        D = X[js] - X[i]
        off = D @ N[i]
        lat = np.linalg.norm(D - off[:, None] * N[i], axis=1)
        out[i] = np.sum((np.abs(off) >= LO * s) & (np.abs(off) <= HI * s) & (lat <= LAT * s)
                        & (np.abs(N[js] @ N[i]) >= PAR)) >= NEED
    return out


def encode_thinwall(P, up=None, layers=128, device="cuda", cuts=None, **kw):
    """fcp.encode with thin walls as bands in the slicer's field (cuts="pointtest": pointfit.py's cut rule)."""
    P = np.asarray(P, float)
    orig_slices = PointSlicer.slices

    def slices(self, layers=layers):
        s = float(self.spacing)
        Xm = (P - self.pose.center) @ self.pose.rotation  # metric canonical
        N = self.world_normals() @ self.pose.rotation
        N /= np.maximum(np.linalg.norm(N, axis=1, keepdims=True), 1e-12)
        thin = thin_flags(Xm, N, s)
        nb = cKDTree(Xm).query(Xm, SPREAD_K + 1)[1][:, 1:]
        for _ in range(SPREAD_ROUNDS):  # the flag spreads along the sheet: a wall is thin as a whole, not point by point
            thin = thin | (thin[nb].mean(1) >= SPREAD)
        if thin.any():  # only wall-sized groups: isolated thin points would add specks
            from scipy.sparse import coo_matrix
            from scipy.sparse.csgraph import connected_components

            ti = np.nonzero(thin)[0]
            pr = cKDTree(Xm[ti]).query_pairs(HI * s, output_type="ndarray")
            G = coo_matrix((np.ones(len(pr)), (pr[:, 0], pr[:, 1])), shape=(len(ti),) * 2) if len(pr) else \
                coo_matrix((len(ti), len(ti)))
            _, lab = connected_components(G, directed=False)
            big = np.bincount(lab)[lab] >= MIN_GROUP
            thin[ti[~big]] = False
        T = Xm[thin]
        LAST.update(X=Xm, thin=thin.copy())
        STATS.update(thin=float(thin.mean()), layers=0)
        B = self._surf.B
        orig_field = B.field

        def field(total, nL, off_d, HW_d, geo, pix, *rest):
            f = orig_field(total, nL, off_d, HW_d, geo, pix, *rest)
            if not len(T):
                return f
            fh = host(f).astype(np.float64)
            off, HW, g, px = host(off_d), host(HW_d), host(geo), float(pix)
            for k in range(nL):
                a, b = off[k], off[k + 1]
                if b == a:
                    continue
                Tk = T[np.abs(T[:, 2] - g[k, 2]) <= SLAB * s, :2]
                if not len(Tk):
                    continue
                H, W = HW[k]
                ii, jj = np.meshgrid(np.arange(H), np.arange(W), indexing="ij")
                C = np.c_[g[k, 0] + ii.ravel() * px, g[k, 1] + jj.ravel() * px]
                d = cKDTree(Tk).query(C, distance_upper_bound=REACH * s)[0]
                near = np.isfinite(d)
                if near.any():
                    seg = fh[a:b]
                    seg[near] = np.minimum(seg[near], d[near] - R * s) if MODE == "union" else d[near] - R * s
                    fh[a:b] = seg
                    STATS["layers"] += 1
            return B.xp.asarray(fh.astype(host(f).dtype))

        B.field = field
        try:
            return orig_slices(self, layers)
        finally:
            B.field = orig_field

    PointSlicer.slices = slices
    try:
        if cuts == "pointtest":
            import pointfit

            return pointfit.encode_pointtest(P, up=up, layers=layers, device=device, **kw)
        return fcp.encode(P, up=up, layers=layers, device=device, **kw)
    finally:
        PointSlicer.slices = orig_slices