# SPDX-License-Identifier: Apache-2.0
"""Pebble body: a loft of superellipse cross-sections.

Plan: an egg, widest a little toward the head (where the eye is).
Top: a full dome. Belly: flatter, so it rests eye-up on a table.
All dimensions in mm; X along the body (head +X), Z up.
"""
from __future__ import annotations

import functools
import math
from dataclasses import dataclass

from build123d import Location, Plane, Spline, Wire, Face, loft, Solid, Pos

TIP = 2.6   # tip retreat factor for insets (1.0 gave 0.87 mm walls at the tips; 1.8 gave 1.42 mm at 2.2 mm walls; drops land on tips)


@dataclass(frozen=True)
class Pebble:
    L: float = 120.0
    W: float = 52.0
    H_top: float = 15.0      # dome height above the widest plane
    H_bot: float = 8.0       # belly depth below it
    x_peak: float = 0.10     # widest point, as a fraction of L toward the head
    p_plan: float = 2.3      # plan superellipse exponent (2 = ellipse; higher = fuller)
    p_top: float = 2.4       # how the dome falls off along the length
    p_bot: float = 3.2       # belly falls off later (flatter)
    n_top: float = 2.4       # cross-section exponent above the widest plane
    n_bot: float = 3.6       # cross-section exponent below (flat belly)

    def __hash__(self):   # computed once: the memoised profile() hashes the shape on every call
        h = self.__dict__.get("_hash")
        if h is None:
            h = hash((self.L, self.W, self.H_top, self.H_bot, self.x_peak, self.p_plan, self.p_top, self.p_bot, self.n_top, self.n_bot))
            object.__setattr__(self, "_hash", h)
        return h

    @property
    def T(self):
        return self.H_top + self.H_bot

    def _u(self, x):
        xp = self.x_peak * self.L
        return (x - xp) / (self.L / 2 - xp) if x >= xp else (x - xp) / (self.L / 2 + xp)

    @staticmethod
    def _se(u, p):
        u = min(abs(u), 1.0)
        return (1 - u ** p) ** (1 / p)

    def half_width(self, x):
        return self.W / 2 * self._se(self._u(x), self.p_plan)

    def top(self, x):
        return self.H_top * self._se(self._u(x), self.p_top)

    def bot(self, x):
        return self.H_bot * self._se(self._u(x), self.p_bot)

    def inset_of(self, f, x, inset):
        """Inset a profile f(x) along its normal, not straight in: where f falls steeply (near the tips) a straight
        inset of d leaves a wall of only d / sqrt(1 + f'²). Capped at 3 d (the tip retreat handles the very end)."""
        if inset <= 0:                           # outward offsets (the band) stay straight: OCC booleans on them misbehave otherwise
            return f(x) - inset
        h = 0.05
        df = (f(x + h) - f(x - h)) / (2 * h)
        return f(x) - inset * min(math.sqrt(1 + df * df), 3.0)

    @functools.lru_cache(maxsize=1 << 20)
    def profile(self, x, inset):
        """(half width, top, belly) at x, inset along the normal. Memoised: the packer and the land search ask for
        the same stations millions of times (the frozen dataclass makes the shape hashable)."""
        return self.inset_of(self.half_width, x, inset), self.inset_of(self.top, x, inset), self.inset_of(self.bot, x, inset)

    def section_pts(self, x, inset=0.0, n=48):
        w = max(self.inset_of(self.half_width, x, inset), 0.3)
        ht = max(self.inset_of(self.top, x, inset), 0.3)
        hb = max(self.inset_of(self.bot, x, inset), 0.3)
        pts = []
        for i in range(n):
            th = 2 * math.pi * i / n
            c, s = math.cos(th), math.sin(th)
            nexp = self.n_top if s >= 0 else self.n_bot
            y = w * math.copysign(abs(c) ** (2 / nexp), c)
            z = (ht if s >= 0 else hb) * math.copysign(abs(s) ** (2 / nexp), s)
            pts.append((x, y, z))
        return pts

    def solid(self, inset=0.0, stations=34, tip=0.8, n=48):
        """inset > 0 returns the cavity. z = 0 is the widest plane.
        The tips are pulled back TIP× further, because an inset loft thins the wall where curvature is highest."""
        L = self.L - 2 * inset * TIP
        xs = []
        for i in range(stations):
            t = i / (stations - 1)
            # cluster stations toward the tips, where curvature is highest
            s = 0.5 - 0.5 * math.cos(math.pi * t)
            xs.append(-L / 2 + tip + s * (L - 2 * tip))
        wires = []
        for x in xs:
            pts = self.section_pts(x, inset, n)
            wires.append(Wire([Spline(*pts, periodic=True)]))
        sol = Solid.make_loft(wires)
        if sol.volume < 0:                       # an outward offset (inset < 0) can loft inside-out: booleans then fail silently
            sol = Solid(sol.wrapped.Reversed())
        return sol


if __name__ == "__main__":
    import time
    t0 = time.time()
    pb = Pebble()
    s = pb.solid()
    c = pb.solid(inset=1.8)
    print("outer", round(s.volume / 1000, 1), "cm3", s.bounding_box().size, "cavity", c.bounding_box().size, f"{time.time()-t0:.1f}s")
    from build123d import export_stl
    export_stl(s - c, "/private/tmp/claude-501/-Users-rafa-Documents/684effc1-d872-43e8-a162-50e2b6e6b4dd/scratchpad/pebble_test.stl")


def inside(pb: "Pebble", x, y, z, inset=0.0):
    """Analytic point-in-cavity test matching section_pts() (z = 0 at the widest plane)."""
    L = pb.L - 2 * inset * TIP
    if abs(x) > L / 2:
        return False
    # section_pts evaluates the profile at the same x (the loft uses x directly)
    w, ht, hb = pb.profile(x, inset)
    if w <= 0 or ht <= 0 or hb <= 0:
        return False
    if z >= 0:
        n, h = pb.n_top, ht
    else:
        n, h = pb.n_bot, hb
    return (abs(y) / w) ** n + (abs(z) / h) ** n <= 1.0