"""Turn the retail low-poly head into something that can carry a 3D-realistic icon.

Two steps, both standard:

1. Crease-aware normal welding. The dat mesh stores one normal per (position,uv)
   pair, so shared corners come back faceted. We average normals across
   coincident positions, but only between faces whose normals are within a
   crease angle -- so the ear fins and tusk edges stay sharp while the skull
   rounds off.

2. PN-triangle tessellation (Vlachos et al. 2001). Each flat triangle becomes a
   cubic Bezier patch built from its own corner positions and normals, with
   quadratically-interpolated normals. It rounds the silhouette without needing
   mesh connectivity, so UV seams cannot pull apart.
"""
import numpy as np


def weld_normals(P, N, crease_deg=52.0, decimals=5):
    """Average normals across coincident positions within a crease angle."""
    key = np.round(P, decimals)
    _, inv = np.unique(key, axis=0, return_inverse=True)
    inv = inv.ravel()
    out = N.copy()
    cos_t = np.cos(np.radians(crease_deg))
    order = np.argsort(inv, kind="stable")
    inv_sorted = inv[order]
    bounds = np.flatnonzero(np.diff(inv_sorted)) + 1
    for grp in np.split(order, bounds):
        if len(grp) < 2:
            continue
        n = N[grp]
        for a in range(len(grp)):
            sel = n @ n[a] >= cos_t
            acc = n[sel].sum(axis=0)
            ln = np.linalg.norm(acc)
            if ln > 1e-9:
                out[grp[a]] = acc / ln
    return out


def _pn_patch(p1, p2, p3, n1, n2, n3):
    def edge(pa, pb, na):
        w = np.einsum("ij,ij->i", pb - pa, na)
        return (2 * pa + pb - w[:, None] * na) / 3.0

    b210 = edge(p1, p2, n1); b120 = edge(p2, p1, n2)
    b021 = edge(p2, p3, n2); b012 = edge(p3, p2, n3)
    b102 = edge(p3, p1, n3); b201 = edge(p1, p3, n1)
    E = (b210 + b120 + b021 + b012 + b102 + b201) / 6.0
    V = (p1 + p2 + p3) / 3.0
    b111 = E + (E - V) / 2.0

    def nedge(pa, pb, na, nb):
        d = pb - pa
        denom = np.einsum("ij,ij->i", d, d)
        denom = np.where(denom < 1e-12, 1e-12, denom)
        v = 2.0 * np.einsum("ij,ij->i", d, na + nb) / denom
        r = na + nb - v[:, None] * d
        ln = np.linalg.norm(r, axis=1, keepdims=True)
        return r / np.where(ln < 1e-9, 1, ln)

    n110 = nedge(p1, p2, n1, n2)
    n011 = nedge(p2, p3, n2, n3)
    n101 = nedge(p3, p1, n3, n1)
    return (b210, b120, b021, b012, b102, b201, b111, n110, n011, n101)


def pn_tessellate(P, N, UV, TRI, level=4):
    """Subdivide every triangle into level^2 sub-triangles on its PN patch."""
    i1, i2, i3 = TRI[:, 0], TRI[:, 1], TRI[:, 2]
    p1, p2, p3 = P[i1], P[i2], P[i3]
    n1, n2, n3 = N[i1], N[i2], N[i3]
    t1, t2, t3 = UV[i1], UV[i2], UV[i3]
    (b210, b120, b021, b012, b102, b201, b111,
     n110, n011, n101) = _pn_patch(p1, p2, p3, n1, n2, n3)

    # barycentric lattice
    lat, lat_index = [], {}
    for i in range(level + 1):
        for j in range(level + 1 - i):
            lat_index[(i, j)] = len(lat)
            lat.append((i / level, j / level))
    lat = np.array(lat)

    T = len(TRI)
    L = len(lat)
    newP = np.empty((T, L, 3)); newN = np.empty((T, L, 3)); newUV = np.empty((T, L, 2))
    for k, (u, v) in enumerate(lat):
        w = 1.0 - u - v
        pos = (p1 * (w ** 3) + p2 * (u ** 3) + p3 * (v ** 3)
               + b210 * (3 * w * w * u) + b120 * (3 * w * u * u)
               + b021 * (3 * u * u * v) + b012 * (3 * u * v * v)
               + b102 * (3 * w * v * v) + b201 * (3 * w * w * v)
               + b111 * (6 * w * u * v))
        nrm = (n1 * (w * w) + n2 * (u * u) + n3 * (v * v)
               + n110 * (2 * w * u) + n011 * (2 * u * v) + n101 * (2 * w * v))
        ln = np.linalg.norm(nrm, axis=1, keepdims=True)
        newP[:, k] = pos
        newN[:, k] = nrm / np.where(ln < 1e-9, 1, ln)
        newUV[:, k] = t1 * w + t2 * u + t3 * v

    tris = []
    for i in range(level):
        for j in range(level - i):
            a = lat_index[(i, j)]; b = lat_index[(i + 1, j)]; c = lat_index[(i, j + 1)]
            tris.append((a, b, c))
            if i + j < level - 1:
                dd = lat_index[(i + 1, j + 1)]
                tris.append((b, dd, c))
    tris = np.array(tris, dtype=np.int64)

    offs = (np.arange(T) * L)[:, None, None]
    TRI2 = (tris[None, :, :] + offs).reshape(-1, 3)
    return (newP.reshape(-1, 3), newN.reshape(-1, 3), newUV.reshape(-1, 2), TRI2)


def smooth_mesh(mesh, crease_deg=52.0, level=4):
    P, N, UV, TRI, TEX = mesh
    N2 = weld_normals(P, N, crease_deg)
    P3, N3, UV3, TRI3 = pn_tessellate(P, N2, UV, TRI, level)
    per = len(TRI3) // len(TRI)
    TEX3 = [t for t in TEX for _ in range(per)]
    return (P3, N3, UV3, TRI3, TEX3)