"""Prototypes (bench only, the package is untouched): contours fitted to the points, two ways.

A, "ringfit": the slicer's rings give the topology (how many contours, solid or hole) and which points belong to which
    contour; each ring's Fourier coefficients are then fitted to its own points.
B, "graph": no slicer field, no marching squares: in each layer's slab the side points are joined to their neighbours
    within LINK spacings whose normals do not face the opposite way (cos > OPPOSITE: a thin wall's two faces stay apart,
    a corner does not break the sheet); every connected sheet that closes around its centre (angular gap at most
    MAX_GAP degrees) becomes a contour: the angle around the centre is the first parameter, a solid when its normals
    point away from the centre, a hole when towards it. Only the slicer's oriented normals are used. Assumes sections
    roughly star-shaped around each sheet's centre (the projection rounds correct mild departures).

Common to both: the slab is the points within SLAB spacings of the plane that reach it along the surface within SHIFT
spacings (p_xy + n_z dz / |n_xy| n_xy): a slanted wall does not blur, and a near-horizontal surface counts only very
near the plane. A contour is x(t), y(t) with N_eff harmonics, N_eff = min(N, (n / PPC - 1) / 2) for n points
(PPC points per coefficient), fitted by least squares with a ridge penalty RIDGE n (k / N_eff)^4 on harmonic k; t of
each point comes from its projection on the current curve (ITERS rounds). The fitted contours, sampled at OUT_SAMPLES
points, replace the slicer's rings; tracking, cuts and atoms are the package's.

    from ringfit import encode_ringfit, encode_graph
"""
import sys

import numpy as np
from scipy.sparse import coo_matrix
from scipy.sparse.csgraph import connected_components
from scipy.spatial import cKDTree

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

SLAB = 1.0  # spacings
SHIFT = 0.75  # spacings: largest move onto the plane (a point that needs more is left out)
N_MAX = 16
PPC = 3.0  # points per coefficient
RIDGE = 1e-3
ITERS = 4
OUT_SAMPLES = 256
MIN_PTS = 12
FACE = 0.3  # A: cos(point normal, the ring's solid-outward normal) at least this
REACH = 2.0  # A: spacings from the slicer ring a point may be to belong to it
LINK = 3.0  # B: spacings, sheet edges
OPPOSITE = -0.5  # B: cos(normals) above this to join two points
MAX_GAP = 120.0  # B: degrees: a sheet whose points leave a larger angular gap around its centre is open (dropped)

STATS = {}
DEBUG = {}  # per layer: what each method produced (for the slice pages)


# ------------------------------------------------------------------ common

def slab(X, NX, zk, s):
    """Points of the layer at height zk, moved onto the plane along the surface (p_xy + n_z dz / |n_xy| n_xy); a point
    is used when that move is at most SHIFT spacings, whatever its slope (a near-horizontal surface fixes the contour
    well, but only very near the plane); their unit in-plane normals."""
    dz = X[:, 2] - zk
    n_xy = np.linalg.norm(NX[:, :2], axis=1)
    d = NX[:, 2] * dz / np.maximum(n_xy, 1e-12)
    m = (np.abs(dz) <= SLAB * s) & (np.abs(d) <= SHIFT * s) & (n_xy > 1e-6 * np.linalg.norm(NX, axis=1))
    nxy = NX[m, :2] / n_xy[m, None]
    return X[m, :2] + d[m, None] * nxy, nxy


def basis(t, n):
    """(len(t), 2n + 1): 1, cos t, sin t, ..., cos nt, sin nt."""
    k = np.arange(1, n + 1)
    B = np.empty((len(t), 2 * n + 1))
    B[:, 0] = 1
    B[:, 1::2], B[:, 2::2] = np.cos(np.outer(t, k)), np.sin(np.outer(t, k))
    return B


def n_eff(n_pts):
    return int(max(1, min(N_MAX, (n_pts / PPC - 1) / 2)))


def fit_contour(Q, t, n):
    """Coefficients (2n + 1, 2) of x(t), y(t) through the points Q at parameters t (ridge on the high harmonics)."""
    B = basis(t, n)
    k = np.r_[0, np.repeat(np.arange(1, n + 1), 2)]
    lam = RIDGE * len(Q) * (k / max(n, 1)) ** 4
    return np.linalg.solve(B.T @ B + np.diag(lam), B.T @ Q)


def sample(C, m=OUT_SAMPLES):
    t = np.linspace(0, 2 * np.pi, m, endpoint=False)
    return basis(t, (len(C) - 1) // 2) @ C


def project_t(C, Q, m=2048):
    """Each point's parameter: that of the nearest of m samples of the curve."""
    t = np.linspace(0, 2 * np.pi, m, endpoint=False)
    return t[cKDTree(sample(C, m)).query(Q)[1]]


def refit(Q, t):
    """ITERS rounds of fit and reprojection, from the parameters t. -> curve samples (small self-loops cut off: a
    crossing ring would turn one lobe into a hole), N used."""
    n = n_eff(len(Q))
    C = fit_contour(Q, t, n)
    for _ in range(ITERS - 1):
        C = fit_contour(Q, project_t(C, Q), n)
    return main_loop(sample(C)), n


def main_loop(R):
    """R without its self-loops: the exterior of the largest piece of its region, in R's direction (R if simple)."""
    from shapely.geometry import LinearRing, Polygon

    if LinearRing(R).is_simple:
        return R
    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]
    return E[::-1] if (signed_area(E) > 0) != (signed_area(R) > 0) else E


def arc_t(R):
    """Parameters of the vertices of polyline R (closed) by arc length, in [0, 2 pi)."""
    seg = np.linalg.norm(np.diff(np.vstack([R, R[:1]]), axis=0), axis=1)
    s = np.r_[0, np.cumsum(seg)[:-1]]
    return 2 * np.pi * s / max(seg.sum(), 1e-12)


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 _resample(R, step):
    R = np.vstack([R, R[:1]])
    seg = np.linalg.norm(np.diff(R, axis=0), axis=1)
    c = np.r_[0, np.cumsum(seg)]
    t = np.linspace(0, c[-1], max(8, int(c[-1] / step)), endpoint=False)
    return np.c_[np.interp(t, c, R[:, 0]), np.interp(t, c, R[:, 1])]


def _holes(loops):
    from shapely.geometry import Point, Polygon

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


# ------------------------------------------------------------------ A: slicer topology, fitted geometry

def ringfit_layer(rings, Q, NQ, s):
    """The slicer's rings of one layer, each refitted to its own points (rings with too few points stay)."""
    if not rings or len(Q) < MIN_PTS:
        return rings, [0] * len(rings)
    dense = [_resample(np.asarray(R, float), 0.25 * s) for R in rings]
    owner = np.concatenate([np.full(len(D), i) for i, D in enumerate(dense)])
    allv = np.vstack(dense)
    d, j = cKDTree(allv).query(Q)
    ring_of = owner[j]
    holes = _holes(rings)
    out, used = [], []
    for i, R in enumerate(dense):
        sa = signed_area(R)
        # solid-outward normal of the dense ring at each vertex
        tng = np.roll(R, -1, 0) - np.roll(R, 1, 0)
        nrm = np.c_[tng[:, 1], -tng[:, 0]] / np.maximum(np.linalg.norm(tng, axis=1, keepdims=True), 1e-12)
        nrm *= 1 if (sa > 0) != holes[i] else -1
        sel = (ring_of == i) & (d <= REACH * s)
        vi = j[sel] - (owner == i).argmax()
        ok = (NQ[sel] * nrm[vi]).sum(1) >= FACE
        P = Q[sel][ok]
        if len(P) < MIN_PTS:
            out.append(rings[i]); used.append(0)
            continue
        t = arc_t(R)[vi[ok]]
        C, n = refit(P, t)
        if (signed_area(C) > 0) != (sa > 0):
            C = C[::-1]
        out.append(C); used.append(n)
    return out, used


# ------------------------------------------------------------------ B: graph topology, fitted geometry

def sheets(Q, NQ, s):
    """Connected sheets: points within LINK spacings whose normals do not face the opposite way."""
    pairs = cKDTree(Q).query_pairs(LINK * s, output_type="ndarray")
    if len(pairs):
        pairs = pairs[(NQ[pairs[:, 0]] * NQ[pairs[:, 1]]).sum(1) > OPPOSITE]
    G = coo_matrix((np.ones(len(pairs)), (pairs[:, 0], pairs[:, 1])) if len(pairs) else ([], ([], [])),
                   shape=(len(Q), len(Q)))
    _, lab = connected_components(G, directed=False)
    return [np.nonzero(lab == c)[0] for c in np.unique(lab)]


def graph_layer(Q, NQ, s):
    """Contours of one layer from its points alone: every sheet around its centre (initial parameter: the angle; the
    projection rounds refine it), solid when its normals point away from the centre, hole when towards it; a sheet with
    an angular gap above MAX_GAP is open and dropped. -> rings, N used per ring, sheets dropped (open or too small)."""
    out, used, dropped = [], [], 0
    if len(Q) < MIN_PTS:
        return out, used, dropped
    for idx in sheets(Q, NQ, s):
        if len(idx) < MIN_PTS:
            dropped += 1
            continue
        P, N = Q[idx], NQ[idx]
        c = P.mean(0)
        a = np.arctan2(P[:, 1] - c[1], P[:, 0] - c[0])
        gap = np.degrees(np.max(np.diff(np.r_[np.sort(a), np.sort(a)[0] + 2 * np.pi])))
        if gap > MAX_GAP:
            dropped += 1
            continue
        C, n = refit(P, np.mod(a, 2 * np.pi))
        solid = float(((P - c) * N).sum(1).mean()) > 0
        if (signed_area(C) > 0) != solid:  # solids counter-clockwise, holes clockwise
            C = C[::-1]
        out.append(C); used.append(n)
    return out, used, dropped


# ------------------------------------------------------------------ encode

def _encode(P, method, up=None, layers=128, device="cuda", cuts=None, **kw):
    orig = PointSlicer.slices

    def slices(self, layers=layers):
        c = centres(check_layers(layers))
        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
        s = float(self.spacing)
        loops = self.loops(c) if method == "ringfit" else [[] for _ in c]
        new = []
        STATS.update(rings=0, fitted=0, dropped=0, n_sum=0)
        DEBUG.clear()
        for k, zk in enumerate(c * sc[2]):
            Q, NQ = slab(X, NX, zk, s)
            if method == "ringfit":
                rings = [np.asarray(L, float) * sc[:2] for L in loops[k]]
                R, used = ringfit_layer(rings, Q, NQ, s)
                DEBUG[k] = dict(slicer=rings, fitted=R, used=used)
            else:
                R, used, dropped = graph_layer(Q, NQ, s)
                STATS["dropped"] += dropped
                DEBUG[k] = dict(fitted=R, used=used, dropped=dropped)
            STATS["rings"] += len(R)
            STATS["fitted"] += sum(u > 0 for u in used)
            STATS["n_sum"] += sum(used)
            new.append([(np.asarray(L) / sc[:2]).astype(np.float32) for L in R])
        return _rings.slices(new, self.device)

    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


def encode_ringfit(P, up=None, **kw):
    """A: the slicer's rings refitted to their points (cuts="pointtest": pointfit.py's cut rule)."""
    return _encode(P, "ringfit", up, **kw)


def encode_graph(P, up=None, **kw):
    """B: contours from the points' sheets, no slicer field."""
    return _encode(P, "graph", up, **kw)