#!/usr/bin/env python3
"""
PROVENANCE: NUMERICAL

PST Computation 135 — Born rule from the Bernoulli measure via Gleason's theorem
================================================================================
Witness computation for the Born-rule claim of the physical-predictions
chapter (sec:born-rule). On the substrate Hilbert space H = L^2(P(D), mu),
with mu the Bernoulli(1/2) measure on the power set P(D), the Born rule
P(C) = |psi(C)|^2 mu(C) is the unique non-contextual probability assignment
to the configuration projectors, as forced by Gleason's theorem.

Gleason's theorem is a theorem and is not re-proved here; this computation
VERIFIES that its hypotheses hold on the substrate space and DEMONSTRATES the
|psi|^2 form together with the properties (total probability one; additivity /
frame-function normalisation in an arbitrary basis; non-contextuality) that
single it out. Hence the provenance is NUMERICAL (hypothesis check plus
explicit demonstration), not a re-derivation of Gleason.

Checks, for |D| = 2, 3, 4 (Hilbert dimension 2^|D| = 4, 8, 16, all >= 3 so
Gleason's dimension hypothesis holds):
  (G1) dim H = 2^|D| >= 3.
  (G2) config-basis total probability sum_C ||P_C psi||^2 = 1.
  (G3) Born form: ||P_C psi||^2 = |psi(C)|^2 mu(C) exactly.
  (G4) non-contextuality: a projector shared by two distinct orthonormal
       resolutions of the identity yields the same probability.
  (G5) frame-function additivity: sum_i ||P_i psi||^2 = 1 for an ARBITRARY
       orthonormal resolution {P_i}, not only the config basis (this is the
       defining property of a Gleason frame function).
"""
import numpy as np

rng = np.random.default_rng(0)


def run(nD):
    dim = 2 ** nD
    mu = np.full(dim, 1.0 / dim)          # Bernoulli(1/2): uniform weight 1/2^|D|
    W = mu                                 # weighted inner product <f,g>_W = sum conj(f) g mu

    def wip(a, b):
        return np.sum(np.conj(a) * b * W)

    def gram_schmidt(vs):
        out = []
        for v in vs:
            v = v.astype(complex).copy()
            for u in out:
                v -= wip(u, v) * u
            v /= np.sqrt(wip(v, v).real)
            out.append(v)
        return out

    # normalised state: sum_C |psi(C)|^2 mu(C) = 1
    psi = rng.standard_normal(dim) + 1j * rng.standard_normal(dim)
    psi /= np.sqrt(wip(psi, psi).real)

    # config projector P_C onto the mu-normalised basis vector b_C = delta_C / sqrt(mu_C);
    # ||P_C psi||_W^2 = |<b_C, psi>_W|^2 = |psi(C)|^2 mu(C)
    born = (np.abs(psi) ** 2) * mu
    pnorm = np.empty(dim)
    for c in range(dim):
        bC = np.zeros(dim, dtype=complex)
        bC[c] = 1.0 / np.sqrt(mu[c])
        pnorm[c] = abs(wip(bC, psi)) ** 2

    g1 = dim >= 3
    g2 = abs(np.sum(pnorm) - 1) < 1e-12
    g3 = np.allclose(pnorm, born, atol=1e-12)

    # (G4) two ONBs both containing b_0; probability of P_0 must agree
    b0 = np.zeros(dim, dtype=complex); b0[0] = 1.0 / np.sqrt(mu[0])
    onb1 = gram_schmidt([b0] + [rng.standard_normal(dim) + 1j * rng.standard_normal(dim) for _ in range(dim - 1)])
    onb2 = gram_schmidt([b0] + [rng.standard_normal(dim) + 1j * rng.standard_normal(dim) for _ in range(dim - 1)])
    g4 = abs(abs(wip(onb1[0], psi)) ** 2 - abs(wip(onb2[0], psi)) ** 2) < 1e-12

    # (G5) additivity in an arbitrary ONB
    onb = gram_schmidt([rng.standard_normal(dim) + 1j * rng.standard_normal(dim) for _ in range(dim)])
    g5 = abs(sum(abs(wip(e, psi)) ** 2 for e in onb) - 1) < 1e-12

    return dim, [g1, g2, g3, g4, g5], float(np.sum(born))


allok = True
for nD in (2, 3, 4):
    dim, gs, s = run(nD)
    ok = all(gs)
    allok &= ok
    print(f"|D|={nD} dim={dim:2d}  G1={gs[0]} G2={gs[1]} G3={gs[2]} G4={gs[3]} G5={gs[4]}"
          f"  sum P(C)={s:.12f}  -> {'PASS' if ok else 'FAIL'}")

print("\nBorn rule P(C) = |psi(C)|^2 mu(C) verified as the Gleason frame function "
      "on L^2(P(D), mu):", "ALL PASS" if allok else "FAILURE")