src/fpgm/physics/materials.py

"""Material -> physics-parameter prior table, and the mixture that widens it.

--- Why a table, and why the VLM never touches a number directly -------------
:mod:`fpgm.physics.types` is explicit that :class:`~fpgm.physics.types.GaussianPrior`
is a *prior*, narrowed only by what an episode's observed motion actually
constrains (see that module's docstring on ``log w_i = sum_e logL_e(theta_i)``).
A VLM is good at recognising "this is a wooden block"; it is not a rheometer,
and a model asked to "estimate the coefficient of friction" will produce a
confident-sounding number with no calibration behind it whatsoever -- worse
than useless, because a bad *point* estimate poisons every particle drawn near
it, while a bad *prior width* just gets corrected by the likelihood. So the
division of labour is fixed: :mod:`scripts._vlm_material_worker` names a
material class (recognition, its strength); this module converts that name to
a physical-parameter distribution (physics, its weakness), through numbers a
human can audit and argue with, sourced from handbook ranges or documented
physical reasoning -- never from the VLM.

--- Where the numbers come from, honestly ------------------------------------
Every entry below is commented with its source: a materials-handbook range
(density tables, general tribology friction ranges), or "assumption" when
no source exists (chiefly: the damping/restitution proxy, where no handbook
gives a MuJoCo ``solref``-shaped number for "cardboard", and the gripper-pad
friction table, where no standard reference covers silicone-pad-vs-object
friction across ten material classes). Numbers are deliberately allowed to
look uncertain -- a `gsd` (geometric standard deviation, see below) of 1.3
means the same thing everywhere it appears: "typical handbook spread for this
class of material," not "I am confident to 30%." Nothing here is tuned to
make a downstream test pass; several entries are wide enough that they will
sometimes look unhelpfully vague, on purpose. **The trade this module makes,
explicitly: a too-wide prior costs a few extra effective-sample-size points
until the likelihood narrows it; a too-narrow, confidently-wrong prior can
put zero density on the true value and never recover.** Every choice below
resolves ties in favour of wider.

--- Log-normal parameterisation -----------------------------------------------
Physical quantities that are strictly positive and span a wide multiplicative
range (density, friction coefficients) are modelled log-normally: a class's
prior is specified as ``(median, gsd)`` in natural units, where ``gsd``
("geometric standard deviation") is the multiplicative factor such that
roughly the middle 68% of the mass sits in ``[median / gsd, median * gsd]``.
That converts directly to the unconstrained (log) mean/std that
:class:`~fpgm.physics.types.GaussianPrior` actually stores, since every
log-parameter in :data:`~fpgm.physics.types.RIGID_PARAMS` (``log_density``,
``log_mu_slide``, ...) *is* ``ln`` of the natural-unit quantity:

    mu    = ln(median)
    sigma = ln(gsd)

``com_x``/``com_y`` are not log-parameters (they are signed bounding-radius
fractions, see ``RIGID_PARAMS``'s own docstring), so their defaults are a
plain ``(mean, std)`` in that native unit instead.

--- The moment-matched mixture -------------------------------------------------
A VLM verdict is a distribution over materials, e.g. ``{wood: 0.6, plastic:
0.4}``. Each material implies its own log-normal (mu_i, sigma_i^2) for a given
parameter. The *mixture* of those log-normals (0.6 * LogNormal_wood + 0.4 *
LogNormal_plastic) is not itself log-normal, so :func:`material_prior`
approximates it with the single Gaussian (in log space) that matches its
first two moments -- the standard "moment matching" / Gaussian-mixture-
reduction identity:

    mean = sum_i p_i * mu_i
    var  = sum_i p_i * sigma_i^2   +   sum_i p_i * (mu_i - mean)^2
           \\_______________________/   \\____________________________/
              within-class variance         between-class variance

The first term is "how uncertain is this material's own density estimate,"
the second is "how much do the candidate materials disagree with each
other." A 60/40 wood/plastic read (wood ~500 kg/m^3, plastic ~1000 kg/m^3)
picks up a large between-class term the two components' own priors never
had; a 99/1 read on the same pair collapses that term to ~0 and the mixture
reproduces (almost exactly) the dominant class's own prior. This is exactly
the "ambiguous verdict genuinely widens the prior" property
:class:`~fpgm.physics.types.MaterialVerdict`'s docstring asks for, and it
falls out of the arithmetic rather than being a rule bolted on afterward.

--- What the table does not cover ---------------------------------------------
``com_x``, ``com_y`` (mass-distribution eccentricity) and the prismatic-joint
parameters (``log_joint_friction``, ``log_joint_damping``) have no material
grounding at all -- a "plastic" verdict says nothing about whether a drawer's
slide is gritty or how off-centre an object's mass is. :func:`material_prior`
fills these from :data:`_MATERIAL_INDEPENDENT_DEFAULTS`, identically
regardless of the verdict, and records that in ``provenance`` so a report can
distinguish "the VLM informed this number" from "this number is a fixed
uninformative default."

--- The hollow-object caveat (why several `gsd`s are wider than the raw
    material variability alone would justify) --------------------------------
:data:`~fpgm.physics.types.RIGID_PARAMS`'s own docstring says ``mass =
density * mesh volume``. The mesh volume here is whatever solid convex hull
the reconstruction stage produced, and for a *hollow* object (a mug, a
cardboard box, a plastic bottle) that hull's volume is far larger than the
volume of material actually present -- the correct ``density`` to multiply by
that hull volume and get the right mass is an *effective* density, deflated
by the object's hollowness, not the material's bulk density from a handbook.
Nothing in a material-classification VLM call measures hollowness, so this
module cannot correct for it -- it can only refuse to be falsely confident
about it. That is why ``cardboard``, ``ceramic`` and ``glass`` (all
frequently hollow tabletop objects: boxes, mugs, cups) carry a wider density
``gsd`` than their bulk-material variability alone would justify. This is
recorded per-entry below, not silently baked in.
"""

from __future__ import annotations

import math
from typing import Any

import numpy as np

from fpgm.physics.types import (
    GaussianPrior,
    MaterialVerdict,
    ParamSpace,
    PhysicsError,
)
from fpgm.utils.logging import get_logger

logger = get_logger(__name__)

__all__ = [
    "CLASS_NAMES",
    "material_prior",
]

#: The closed set of material classes this table (and the VLM worker) knows
#: about. ``unknown`` is a fallback bucket only -- see module docstring on
#: :data:`_DENSITY_KG_M3` and :func:`fpgm.physics.priors.VlmPriorProposer.fallback_verdict`.
#: ``scripts/_vlm_material_worker.py`` hardcodes the same 10 non-"unknown"
#: names as its closed answer set (it cannot import this module -- see that
#: script's docstring) and MUST be kept in sync with this tuple by hand.
CLASS_NAMES: tuple[str, ...] = (
    "wood",
    "plastic",
    "cardboard",
    "metal",
    "glass",
    "ceramic",
    "rubber",
    "foam",
    "fabric",
    "stone",
    "unknown",
)

# --------------------------------------------------------------------------- #
# Per-material (median, gsd) tables, one per RIGID_PARAMS log-quantity.
# gsd = geometric standard deviation: middle ~68% mass in [median/gsd, median*gsd].
# --------------------------------------------------------------------------- #

#: Density, kg/m^3. Handbook ranges (engineeringtoolbox-style material density
#: tables) for the *solid* material, widened per the hollow-object caveat
#: above where the class is often a hollow/thin-shell household object.
_DENSITY_KG_M3: dict[str, tuple[float, float]] = {
    # Common woods (pine..oak) span roughly 350-900 kg/m^3 (handbook).
    "wood": (550.0, 1.35),
    # Household injection-molded plastics (PP/PE/ABS/PC) span ~900-1400
    # kg/m^3 (handbook); PP/PE float, PC/ABS don't -- real spread.
    "plastic": (1050.0, 1.25),
    # Corrugated cardboard bulk density is dominated by void fraction, not
    # the paper fibre itself -- handbook bulk figures run ~150-700 kg/m^3.
    # gsd widened further per the hollow-object caveat: a cardboard BOX is
    # the paradigm hollow case.
    "cardboard": (300.0, 1.8),
    # "Metal" spans aluminium (2700) to steel (7850) to brass (~8500) with
    # no way to disambiguate from colour/shape alone; median is the
    # geometric mean of Al and steel, gsd wide enough to cover both within
    # roughly 1 sigma.
    "metal": (4600.0, 1.55),
    # Soda-lime glass is close to a physical constant (~2500 kg/m^3);
    # narrow gsd for the *material*, but widened for hollow drinking
    # glasses / bottles per the caveat above.
    "glass": (2500.0, 1.2),
    # Fired earthenware/stoneware/porcelain: ~2000-2600 kg/m^3 (handbook)
    # for the solid material; widened for hollow mugs/bowls per the caveat.
    "ceramic": (2300.0, 1.35),
    # Vulcanised/filled rubber ~1100-1600 kg/m^3 (handbook); natural gum
    # rubber alone is lighter (~920) hence the low end of the range.
    "rubber": (1250.0, 1.25),
    # Packing/EVA foams: ~20-300 kg/m^3 (handbook) -- huge range by design
    # (density is the whole point of a foam's engineering).
    "foam": (80.0, 2.0),
    # ASSUMPTION: a folded/bunched fabric or plush item's *bulk* (bounding-
    # volume) density is dominated by trapped air, not the textile itself;
    # no handbook table covers this directly, reasoned from typical
    # clothing-item bulk density.
    "fabric": (150.0, 1.8),
    # Common tabletop stone (granite/limestone/sandstone) ~2200-3000
    # kg/m^3 (handbook), fairly consistent across types.
    "stone": (2600.0, 1.15),
    # Deliberately uninformative: centred near water (a neutral "middle of
    # everything" guess) with a gsd wide enough that its 2-sigma band
    # (~110-9000 kg/m^3) covers foam through metal. This is the prior used
    # when the VLM could not be run at all (see fallback_verdict) or named
    # a class this table does not recognise.
    "unknown": (1000.0, 3.0),
}

#: Sliding friction, object vs. a typical wood/laminate tabletop (matching
#: the surface ``scripts/settle_after_release.py`` already assumes 0.5 for).
#: General tribology handbook ranges for "material on wood," ASSUMPTION
#: where noted.
_MU_SLIDE: dict[str, tuple[float, float]] = {
    "wood": (0.4, 1.3),        # wood-on-wood handbook range ~0.25-0.5
    "plastic": (0.3, 1.35),    # generic polymer-on-wood, excludes PTFE-like outliers
    "cardboard": (0.5, 1.3),   # fibrous surface grips; handbook paper/board ~0.4-0.6
    "metal": (0.35, 1.4),      # wide: polished vs. cast/rough finish varies a lot
    "glass": (0.28, 1.3),      # smooth, handbook glass-on-wood ~0.2-0.4
    "ceramic": (0.35, 1.3),    # glazed ceramic-on-wood ~0.25-0.45
    "rubber": (0.9, 1.3),      # rubber is a high-friction outlier by design, handbook ~0.6-1.2
    "foam": (0.6, 1.3),        # ASSUMPTION: compliant-surface microscopic interlocking
    "fabric": (0.55, 1.3),     # ASSUMPTION: fibre-grip, generic textile-on-wood reasoning
    "stone": (0.4, 1.25),      # handbook stone-on-wood ~0.3-0.5
    "unknown": (0.4, 1.8),     # geometric-mean-ish centre, deliberately wide
}

#: Torsional friction. MuJoCo's torsional term is a resistive-torque
#: coefficient roughly two orders of magnitude below sliding friction; this
#: repo's own settle-after-release assumption uses the ratio 0.005/0.5 =
#: 0.01 (see ``scripts/settle_after_release.py:_ASSUMED_FRICTION``). ASSUMPTION
#: throughout: no handbook gives per-material torsional friction, so every
#: entry here is ``_MU_SLIDE`` scaled by that same repo-precedented ratio,
#: with the same relative gsd as its sliding-friction counterpart.
_MU_TORSION: dict[str, tuple[float, float]] = {
    name: (median * 0.01, gsd) for name, (median, gsd) in _MU_SLIDE.items()
}

#: Contact-damping proxy standing in for restitution (MuJoCo has no literal
#: coefficient of restitution -- see ``RIGID_PARAMS``'s and
#: ``scripts/_mujoco_settle_worker.py``'s docstrings). Convention: 1.0 is
#: this repo's own default "critically damped, essentially no bounce"
#: baseline (``scripts/settle_after_release.py``'s ``_ASSUMED_SOLREF =
#: (0.02, 1.0)``); values below 1 are more underdamped/bouncy, above 1 more
#: overdamped/dead. ALL entries are ASSUMPTION, reasoned qualitatively from
#: material brittleness/resilience (brittle+rigid materials tend to rebound
#: before settling; soft/absorptive materials tend to dead-stop) -- there is
#: no handbook table for "MuJoCo solref-shaped bounciness by material," and
#: gsd is kept wide everywhere to reflect that these are reasoned, not
#: measured.
_DAMPING_PROXY: dict[str, tuple[float, float]] = {
    "wood": (0.9, 1.25),
    "plastic": (0.8, 1.3),
    "cardboard": (1.3, 1.3),   # fibrous, absorptive -> overdamped, minimal bounce
    "metal": (0.6, 1.3),       # rigid, historically the classic "bounces" case
    "glass": (0.5, 1.4),       # brittle+rigid -> most underdamped of the set
    "ceramic": (0.5, 1.4),     # brittle+rigid, same reasoning as glass
    "rubber": (0.5, 1.4),      # resilient/elastic -> can rebound; wide because
                                # "rubber" spans soft damped foam-rubber to a
                                # genuinely bouncy solid compound
    "foam": (1.4, 1.3),        # energy-absorbing by design -> overdamped
    "fabric": (1.4, 1.3),      # soft, absorptive -> overdamped, same as foam
    "stone": (0.6, 1.3),       # rigid, dense, moderate rebound
    "unknown": (0.9, 1.6),     # centred on this repo's own default assumption
}

#: Finger-vs-object friction: silicone/rubber gripper pad against the
#: object's surface -- what actually governs slip in a grasp (see
#: ``RIGID_PARAMS``'s docstring). ASSUMPTION throughout: there is no
#: standard handbook for "compliant robot gripper pad vs. object material"
#: friction; reasoned from general soft-robotics-gripper literature
#: (compliant-pad grasps are typically higher-friction than rigid-on-rigid
#: table contact, and compliant-vs-compliant, e.g. rubber pad on rubber or
#: fabric, is highest of all).
_MU_GRIPPER: dict[str, tuple[float, float]] = {
    "wood": (0.9, 1.3),
    "plastic": (0.7, 1.3),
    "cardboard": (0.85, 1.3),
    "metal": (0.75, 1.3),
    "glass": (0.6, 1.3),
    "ceramic": (0.65, 1.3),
    "rubber": (1.1, 1.3),
    "foam": (1.0, 1.35),
    "fabric": (1.0, 1.35),
    "stone": (0.75, 1.3),
    "unknown": (0.8, 1.7),
}

#: Maps a RIGID_PARAMS log-quantity name to its per-material (median, gsd) table.
_MATERIAL_LOGNORMAL: dict[str, dict[str, tuple[float, float]]] = {
    "log_density": _DENSITY_KG_M3,
    "log_mu_slide": _MU_SLIDE,
    "log_mu_torsion": _MU_TORSION,
    "log_solref_damping": _DAMPING_PROXY,
    "log_mu_gripper": _MU_GRIPPER,
}

#: Parameters no material verdict informs at all -- filled identically
#: regardless of the VLM's read. ``com_x``/``com_y`` are already in
#: unconstrained (non-log) units, so these are literal (mean, std), not
#: (median, gsd); the prismatic-joint pair are log-parameters so they get
#: converted through the same ``ln`` machinery as the material table.
#:
#: com_x/com_y: ASSUMPTION. A generic rigid tabletop object's centre of mass
#: is reasoned to sit within roughly 15% of its bounding radius of the
#: geometric centroid for a "typical" mass distribution (no strong internal
#: asymmetry, e.g. not an off-centre weight); zero-mean because there is no
#: directional information at all.
_COM_MEAN_STD: tuple[float, float] = (0.0, 0.15)

#: log_joint_friction / log_joint_damping: ASSUMPTION, and deliberately the
#: widest (least informative) entries in this whole module -- a surface
#: material read says nothing about a drawer slide's mechanical friction or
#: damping, and this table has no other signal to offer. Medians are
#: order-of-magnitude placeholders (a "light drag" and "lightly damped"
#: joint respectively); the wide gsd=3.0 is what actually matters here: it
#: hands nearly all identifying power to the likelihood, which is correct
#: because this table has none to give.
_JOINT_FRICTION_MEDIAN_GSD: tuple[float, float] = (0.3, 3.0)
_JOINT_DAMPING_MEDIAN_GSD: tuple[float, float] = (1.0, 3.0)


def _log_normal_params(median: float, gsd: float) -> tuple[float, float]:
    if median <= 0:
        raise PhysicsError(f"materials table: median must be > 0, got {median}")
    if gsd <= 1.0:
        raise PhysicsError(f"materials table: gsd must be > 1.0, got {gsd}")
    return math.log(median), math.log(gsd)


def _default_params() -> dict[str, tuple[float, float]]:
    """Material-independent ``(mean, std)`` in unconstrained units, by param name."""
    jf_mu, jf_sigma = _log_normal_params(*_JOINT_FRICTION_MEDIAN_GSD)
    jd_mu, jd_sigma = _log_normal_params(*_JOINT_DAMPING_MEDIAN_GSD)
    return {
        "com_x": _COM_MEAN_STD,
        "com_y": _COM_MEAN_STD,
        "log_joint_friction": (jf_mu, jf_sigma),
        "log_joint_damping": (jd_mu, jd_sigma),
    }


#: Built once at import time; the table above is static, so there is no
#: reason to recompute this per call.
_MATERIAL_INDEPENDENT_DEFAULTS: dict[str, tuple[float, float]] = _default_params()


def _moment_match_mixture(
    components: list[tuple[float, float, float]],
) -> tuple[float, float]:
    """Moment-matched Gaussian for a mixture of Gaussians.

    Args:
        components: ``(weight, mu, sigma)`` triples; weights need not be
            pre-normalised (they are normalised here).

    Returns:
        ``(mean, std)`` of the single Gaussian matching the mixture's first
        two moments -- see the module docstring for the mean/var formula.
    """
    total_w = sum(w for w, _, _ in components)
    if total_w <= 0 or not np.isfinite(total_w):
        raise PhysicsError(f"materials: mixture weights sum to {total_w}, cannot normalise")
    weights = [w / total_w for w, _, _ in components]
    mean = sum(w * mu for w, (_, mu, _) in zip(weights, components, strict=True))
    within = sum(w * (sigma**2) for w, (_, _, sigma) in zip(weights, components, strict=True))
    between = sum(
        w * (mu - mean) ** 2 for w, (_, mu, _) in zip(weights, components, strict=True)
    )
    var = within + between
    if var <= 0:
        # Only reachable if every component had sigma==0 AND all mu's agreed
        # exactly -- the table never has sigma==0, so this is unreachable in
        # practice; guarded anyway since GaussianPrior requires std > 0.
        raise PhysicsError("materials: degenerate mixture produced zero variance")
    return mean, math.sqrt(var)


def material_prior(verdict: MaterialVerdict, space: ParamSpace) -> GaussianPrior:
    """A verdict's material distribution -> a :class:`GaussianPrior` over ``space``.

    Every parameter in ``space`` is filled: material-grounded quantities (see
    :data:`_MATERIAL_LOGNORMAL`) via the moment-matched mixture over
    ``verdict.classes`` (module docstring); everything else via
    :data:`_MATERIAL_INDEPENDENT_DEFAULTS`. A class name in ``verdict`` that
    this table does not recognise is treated as ``"unknown"`` for that
    parameter's mixture (logged, and recorded in ``provenance`` rather than
    silently substituted) -- this is the same defensive fallback
    :func:`fpgm.physics.priors.VlmPriorProposer.fallback_verdict` uses
    deliberately, just reached from a different direction (an unrecognised
    class name rather than no VLM read at all).

    Raises:
        PhysicsError: If ``space`` contains a parameter name this module has
            no source for at all (neither table). This should not happen for
            :data:`~fpgm.physics.types.RIGID_PARAMS` /
            :data:`~fpgm.physics.types.PRISMATIC_PARAMS` combinations, and is
            treated as a real error (a new parameter was added to
            ``types.py`` without a corresponding entry here) rather than
            silently defaulting to something arbitrary.
    """
    total = sum(p for _, p in verdict.classes)
    if total <= 0 or not np.isfinite(total):
        raise PhysicsError(f"{verdict.label}: verdict class probabilities sum to {total}")
    norm_classes = [(name, p / total) for name, p in verdict.classes]

    mean = np.zeros(space.dim, dtype=np.float64)
    std = np.zeros(space.dim, dtype=np.float64)
    param_provenance: dict[str, Any] = {}

    for i, name in enumerate(space.names):
        if name in _MATERIAL_LOGNORMAL:
            table = _MATERIAL_LOGNORMAL[name]
            components: list[tuple[float, float, float]] = []
            unrecognized: list[str] = []
            used: list[dict[str, Any]] = []
            for mat_name, p in norm_classes:
                if p <= 0:
                    continue
                entry = table.get(mat_name)
                if entry is None:
                    unrecognized.append(mat_name)
                    entry = table["unknown"]
                median, gsd = entry
                mu, sigma = _log_normal_params(median, gsd)
                components.append((p, mu, sigma))
                used.append({"class": mat_name, "weight": p, "median": median, "gsd": gsd})
            if unrecognized:
                logger.warning(
                    "%s: material_prior: verdict named unrecognised class(es) %s for "
                    "param %r; treating as 'unknown'",
                    verdict.label,
                    unrecognized,
                    name,
                )
            mu_mix, sigma_mix = _moment_match_mixture(components)
            mean[i] = mu_mix
            std[i] = sigma_mix
            param_provenance[name] = {
                "kind": "material_lognormal_mixture",
                "components": used,
                "unrecognized_classes": unrecognized,
                "mean": mu_mix,
                "std": sigma_mix,
            }
        elif name in _MATERIAL_INDEPENDENT_DEFAULTS:
            mu_d, sigma_d = _MATERIAL_INDEPENDENT_DEFAULTS[name]
            mean[i] = mu_d
            std[i] = sigma_d
            param_provenance[name] = {
                "kind": "material_independent_default",
                "mean": mu_d,
                "std": sigma_d,
            }
        else:
            raise PhysicsError(
                f"material_prior: no prior source (material table or default) for "
                f"parameter {name!r}; add an entry to fpgm.physics.materials"
            )

    provenance = {
        "label": verdict.label,
        "verdict_source": verdict.source,
        "verdict_classes": [[c, p] for c, p in norm_classes],
        "params": param_provenance,
    }
    return GaussianPrior(space=space, mean=mean, std=std, provenance=provenance)