# SPDX-License-Identifier: Apache-2.0
# © 2026 SZL Holdings · Stephen P. Lutar · ORCID 0009-0001-0110-4173
"""Pure-PyTorch Lambda-Spine aggregator (Λ) for the szl-lambda-gate kernel.

Λ(x) = ∏ xᵢ^{wᵢ},  Σwᵢ = 1,  wᵢ > 0,  xᵢ ∈ [0,1]   (weighted geometric mean)

This is a TORCH port of the canonical pure-Python reference. It is a
correctness reference, computed via logs in float32 for stability,
differentiable (autograd works), and torch.compile-friendly. Depends ONLY on
torch + the Python standard library (a Kernel Hub requirement for universal
kernels).

WHAT Λ IS / IS NOT (HONESTY — SZL Holdings doctrine v11):
  Λ is the *weighted-geometric-mean aggregator*: a non-compensatory way to
  combine axis scores in [0,1] into one number. It is an ADVISORY governance
  signal — a conservative roll-up where any single zeroed axis drives the
  aggregate to 0. It is NOT "proven trust" and NOT a closed theorem. Its
  *uniqueness* remains Conjecture 1 — OPEN (an unresolved CAUCHY_ND step plus a
  missing symmetry axiom in the Lean development). Do not describe Λ as proven
  trust anywhere.

PRIOR ART (honest attribution): the weighted geometric mean as a less-
  compensatory composite-indicator aggregator is established practice — the UN
  HDI (arithmetic→geometric switch, 2010) and the OECD Handbook on Constructing
  Composite Indicators (2008) both use it to limit the compensation effect. The
  veto / cut-off idea (a single failing criterion blocks a pass) is the ELECTRE
  veto threshold. The 13-axis conjunctive form (yuyay_weights) is SZL's own
  yuyay_v3 gate. None of this makes Λ "proven trust"; the gate is ADVISORY.

PROVENANCE: backed by the Lean 4 formalization szl-holdings/lutar-lean
  (749 declarations / 14 axioms / 163 tracked sorries),
  DOI 10.5281/zenodo.20434308 (lutar-lean). Λ uniqueness = Conjecture 1 (open).

Axioms carried (Lutar/Axioms.lean), available below as runtime self-checks:
  A1 IsMonotone        — Λ is non-decreasing in each axis
  A2 IsHomogeneous     — Λ(t·x) = t·Λ(x)  (degree 1)
  A3 IsEgyptianExact   — Λ(c,…,c) = c       (the uniform-diagonal fixpoint)
  A4 IsBounded(by max) — Λ(x) ≤ maxᵢ xᵢ
"""
from typing import Optional

import torch

_SUPPORTED_DTYPES = (torch.float16, torch.bfloat16, torch.float32, torch.float64)


def _compute_dtype(in_dtype: torch.dtype) -> torch.dtype:
    return torch.float32 if in_dtype in (torch.float16, torch.bfloat16) else in_dtype


def _check_axes(axes: torch.Tensor) -> None:
    """Cheap, allocation-free metadata guards on the axis-score tensor."""
    if not isinstance(axes, torch.Tensor):
        raise TypeError(f"axes must be a torch.Tensor, got {type(axes).__name__}")
    if axes.dtype not in _SUPPORTED_DTYPES:
        raise TypeError(
            f"axes has unsupported dtype {axes.dtype}; "
            f"expected one of {tuple(str(d) for d in _SUPPORTED_DTYPES)}"
        )
    if axes.dim() < 1:
        raise ValueError(
            "axes must have at least 1 dimension (the k axis scores live on "
            f"the last dim); got a {axes.dim()}-d tensor"
        )
    if axes.shape[-1] < 1:
        raise ValueError("axes last dimension (k = number of axes) must be >= 1")


def _resolve_weights(
    axes: torch.Tensor,
    weights: Optional[torch.Tensor],
    cdt: torch.dtype,
) -> torch.Tensor:
    """Return a normalized (Σw = 1) weight vector of shape (k,) in compute dtype."""
    k = axes.shape[-1]
    if weights is None:
        return torch.full((k,), 1.0 / k, dtype=cdt, device=axes.device)
    if not isinstance(weights, torch.Tensor):
        raise TypeError(f"weights must be a torch.Tensor or None, got {type(weights).__name__}")
    if weights.device != axes.device:
        raise ValueError(
            f"weights is on device {weights.device} but axes is on {axes.device}; "
            "move them to the same device"
        )
    if weights.dim() != 1 or weights.shape[0] != k:
        raise ValueError(
            f"weights must be 1-D with shape ({k},) to match the last dim of axes; "
            f"got shape {tuple(weights.shape)}"
        )
    wf = weights.to(cdt)
    if not bool(torch.all(torch.isfinite(wf))):
        raise ValueError("weights must all be finite (no NaN/Inf)")
    if bool(torch.any(wf <= 0.0)):
        raise ValueError("weights must be strictly positive (wᵢ > 0)")
    sw = wf.sum()
    if not bool(sw > 0.0):
        raise ValueError("weights must sum to a positive value")
    return wf / sw


def lambda_aggregate(
    axes: torch.Tensor,
    weights: Optional[torch.Tensor] = None,
) -> torch.Tensor:
    """Weighted geometric mean Λ(x) = ∏ xᵢ^{wᵢ} over the last dim of ``axes``.

    Axis scores expected in [0,1] and clamped into [0,1]; uniform weights (1/k)
    when ``weights`` is None. Computed via logs for stability:

        Λ(x) = exp( Σᵢ wᵢ · log(clamp(xᵢ, 0, 1)) )

    Non-compensatory zero-routing (A4-consistent): any axis that is zero, OR
    that is NON-FINITE (NaN / ±Inf), is treated as a FAILING axis and drives
    the whole aggregate to exactly 0. A garbage/invalid axis must never silently
    pass as a perfect (clamped-to-1) axis; output and gradient stay finite and
    in [0,1] for every input.

    Returns a tensor of shape (...) — Λ(x) ∈ [0,1] per batch row, differentiable
    w.r.t. ``axes``.

    HONESTY: a non-compensatory governance roll-up, NOT proven trust.
    Λ-uniqueness is Conjecture 1 (open).
    """
    _check_axes(axes)
    in_dtype = axes.dtype
    cdt = _compute_dtype(in_dtype)
    xf = axes.to(cdt)
    w = _resolve_weights(axes, weights, cdt)  # (k,), Σw=1

    finite_mask = torch.isfinite(xf)
    xc = xf.clamp(0.0, 1.0)
    bad_mask = (~finite_mask) | (xc <= 0.0)
    any_bad = torch.any(bad_mask, dim=-1)  # (...)

    safe = torch.where(bad_mask, torch.ones_like(xc), xc)
    logx = torch.log(safe)                    # (..., k)
    acc = (logx * w).sum(dim=-1)              # (...)
    val = torch.exp(acc)                      # (...)

    out = torch.where(any_bad, torch.zeros_like(val), val)
    out = out.clamp(0.0, 1.0)
    return out.to(in_dtype)


def lambda_gate(
    axes: torch.Tensor,
    weights: Optional[torch.Tensor] = None,
    threshold: float = 0.5,
):
    """ADVISORY governance gate over Λ(x): score plus a pass/fail vs threshold.

    Returns a :class:`LambdaGateResult` namedtuple (score, passed, threshold,
    advisory). ``passed`` := Λ(x) >= threshold; ``advisory`` is always True.

    HONESTY: a "pass" is an ADVISORY signal only. Λ is the weighted-geometric-
    mean aggregator; its uniqueness is Conjecture 1 (open). Do not treat a pass
    as proven trust or a closed theorem.
    """
    t = float(threshold)
    if t != t or t == float("inf") or t == float("-inf"):
        raise ValueError(f"threshold must be a finite float, got {threshold!r}")
    score = lambda_aggregate(axes, weights)
    passed = score >= t
    return LambdaGateResult(score=score, passed=passed, threshold=t, advisory=True)


def lambda_gate_batch(
    candidates: torch.Tensor,
    weights: Optional[torch.Tensor] = None,
    threshold: float = 0.5,
):
    """ADVISORY batch gate: score MANY candidate action-vectors in one call.

    ``candidates`` is shape (..., N, k): last dim ``k`` is per-axis scores of one
    candidate, the dim before it enumerates the N candidates. Returns a
    :class:`LambdaGateResult` with score/passed of shape (..., N).

    HONESTY: the pass mask is ADVISORY, non-compensatory. NOT proven trust;
    Λ-uniqueness is Conjecture 1 (open).
    """
    _check_axes(candidates)
    if candidates.dim() < 2:
        raise ValueError(
            "candidates must be at least 2-D, shape (..., N, k); "
            f"got a {candidates.dim()}-d tensor"
        )
    return lambda_gate(candidates, weights=weights, threshold=threshold)


# ---- A1..A4 axiom RUNTIME self-checks (real, verifiable) ------------------- #
def is_egyptian_exact(c, k: int = 3, weights=None, tol: float = 1e-5) -> bool:
    """A3 IsEgyptianExact: Λ(c, …, c) = c for a constant axis vector of length k."""
    if k < 1:
        raise ValueError("k must be >= 1")
    cc = min(max(float(c), 0.0), 1.0)
    axes = torch.full((k,), cc, dtype=torch.float64)
    val = lambda_aggregate(axes, weights)
    return bool(torch.abs(val - cc) <= tol)


def is_bounded_by_max(axes: torch.Tensor, weights=None, tol: float = 1e-6) -> bool:
    """A4 IsBounded: Λ(x) ≤ maxᵢ xᵢ (over the last dim), within ``tol``."""
    _check_axes(axes)
    val = lambda_aggregate(axes, weights)
    xf = axes.to(_compute_dtype(axes.dtype))
    xf = torch.where(torch.isfinite(xf), xf, torch.zeros_like(xf))
    mx = xf.clamp(0.0, 1.0).amax(dim=-1)
    return bool(torch.all(val.to(mx.dtype) <= mx + tol))


def is_homogeneous(axes: torch.Tensor, t, weights=None, tol: float = 1e-5) -> bool:
    """A2 IsHomogeneous (degree 1): Λ(t·x) = t·Λ(x) for scalar t in [0,1]."""
    _check_axes(axes)
    tt = min(max(float(t), 0.0), 1.0)
    x = axes.to(torch.float64).clamp(0.0, 1.0)
    lhs = lambda_aggregate(x * tt, weights)
    rhs = tt * lambda_aggregate(x, weights)
    return bool(torch.all(torch.abs(lhs - rhs) <= tol))


def is_monotone(axes: torch.Tensor, weights=None, delta: float = 0.05, tol: float = 1e-7) -> bool:
    """A1 IsMonotone: Λ is non-decreasing in each axis."""
    _check_axes(axes)
    x = axes.to(torch.float64).clamp(0.0, 1.0)
    base = lambda_aggregate(x, weights)
    k = x.shape[-1]
    ok = True
    for j in range(k):
        bumped = x.clone()
        bumped[..., j] = (bumped[..., j] + float(delta)).clamp(0.0, 1.0)
        bumped_val = lambda_aggregate(bumped, weights)
        ok = ok and bool(torch.all(bumped_val - base >= -tol))
    return ok


def find_axiom_violation(k: int = 5, trials: int = 200, weights=None, seed=0, tol: float = 1e-6):
    """Random-search for ANY A1–A4 violation. Returns the first violating triple
    ``(axiom, axes, weights)`` or ``None``. An honest FALSIFICATION attempt —
    finding nothing is empirical evidence, NOT a proof (Λ-uniqueness = Conjecture 1).
    """
    gen = torch.Generator()
    if seed is not None:
        gen.manual_seed(int(seed))
    for _ in range(int(trials)):
        x = torch.rand(k, generator=gen, dtype=torch.float64)
        w = weights
        if w is None:
            w = torch.rand(k, generator=gen, dtype=torch.float64) + 1e-3
        c = float(torch.rand(1, generator=gen).item())
        if not is_egyptian_exact(c, k=k, weights=w, tol=max(tol, 1e-5)):
            return ("A3_IsEgyptianExact", torch.full((k,), c, dtype=torch.float64), w)
        if not is_bounded_by_max(x, w, tol=max(tol, 1e-6)):
            return ("A4_IsBounded", x, w)
        t = float(torch.rand(1, generator=gen).item())
        if not is_homogeneous(x, t, weights=w, tol=max(tol, 1e-5)):
            return ("A2_IsHomogeneous", x, w)
        if not is_monotone(x * 0.9, w, tol=max(tol, 1e-7)):
            return ("A1_IsMonotone", x * 0.9, w)
    return None


# ---- Canonical 13-axis Yuyay preset (ADVISORY ONLY) ------------------------ #
YUYAY_AXES = (
    "moralGrounding", "measurabilityHonesty", "empiricalGrounding",
    "logicalConsistency", "sourceTransparency", "reproducibility",
    "licenseHygiene", "scopeDiscipline", "claimCalibration", "evalAwareness",
    "deceptionKeywords", "conflictingDirectives", "reversalDirective",
)
YUYAY_FLOORS = (
    0.95, 0.95,
    0.90, 0.90, 0.90, 0.90, 0.90, 0.90, 0.90,
    0.90, 0.90, 0.90, 0.90,
)


def yuyay_weights(dtype: torch.dtype = torch.float64, device=None) -> torch.Tensor:
    """Canonical 13-axis Yuyay Λ weight vector (uniform 1/13), ADVISORY only."""
    k = len(YUYAY_AXES)
    return torch.full((k,), 1.0 / k, dtype=dtype, device=device)


def selfcheck(k: int = 5, trials: int = 64, seed=0) -> dict:
    """Run the A1–A4 empirical self-checks and report a verdict + version.

    HONESTY: EMPIRICAL checks on sampled inputs, NOT a proof of Λ-uniqueness
    (Conjecture 1, open). A clean run is evidence, not proof.
    """
    x = torch.rand(k, dtype=torch.float64) * 0.9
    w = torch.rand(k, dtype=torch.float64) + 1e-3
    axioms = {
        "A1_IsMonotone": is_monotone(x, w),
        "A2_IsHomogeneous": is_homogeneous(x, float(torch.rand(1).item()), weights=w),
        "A3_IsEgyptianExact": is_egyptian_exact(float(torch.rand(1).item()), k=k, weights=w),
        "A4_IsBounded": is_bounded_by_max(x, w),
    }
    violation = find_axiom_violation(k=k, trials=trials, seed=seed)
    return {
        "version": __version__,
        "axioms": axioms,
        "all_axioms_hold": all(axioms.values()) and violation is None,
        "adversarial": {"trials": int(trials), "violation": violation},
        "advisory": True,
        "lambda_status": "Conjecture 1 (open) — uniqueness unproven; advisory only",
    }


__version__ = "0.2.0"

from collections import namedtuple  # noqa: E402

LambdaGateResult = namedtuple(
    "LambdaGateResult", ["score", "passed", "threshold", "advisory"]
)