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

Computation 142 -- Tensor factorisation over a bipartition, and no-signalling
============================================================================
Backs the composite-system section.  PST's state space is L^2(P(D), mu) with
mu = (x)_{a in D} Bern(1/2), the canonical measure forced by relabelling and
complementation invariance.  Because mu is a PRODUCT over the index set D, it factors
over ANY bipartition D = D_A |_| D_B, and the Hilbert space factors with it:

    L^2(P(D), mu)  ~  L^2(P(D_A), mu_A)  (x)  L^2(P(D_B), mu_B)

That is what gives PST composite systems at all, and it is a theorem rather than an
axiom.  Everything downstream -- the reduced state, and no-signalling -- follows.

The point worth stressing: the bipartition is of the PROPERTY DOMAIN, not of space.
There is no space at this level.  So "subsystem" means "part of D", and locality means
locality in D.

CHECKS
------
  (1) THE BIJECTION.  beta(C) = (C n D_A, C n D_B) is a bijection P(D) -> P(D_A) x P(D_B),
      checked exhaustively at |D| = 12 over several splits (4096 configurations each).
  (2) THE MEASURE FACTORS.  mu(C) = mu_A(C n D_A) * mu_B(C n D_B) for every C, to machine
      precision, for every bipartition tried.  This is the step that would fail for a
      non-product measure, so a NON-PRODUCT control is run alongside: a measure weighted
      by |C| is permutation-invariant but NOT a product, and its factorisation error is
      shown to be large.  Without the Bernoulli measure there is no tensor structure.
  (3) THE UNITARY.  The induced map on L^2 preserves inner products, and elementary
      tensors (phi_A (x) phi_B)(C) = phi_A(C n D_A) phi_B(C n D_B) reproduce the Kronecker
      product exactly.
  (4) NO-SIGNALLING.  For random correlated states rho on H_A (x) H_B, and for every
      choice of measurement basis {Q_j} on B, the ensemble-averaged post-measurement
      state rho' = sum_j (1 (x) Q_j) rho (1 (x) Q_j) has Tr_B rho' = Tr_B rho.  The A
      marginal is untouched by the B setting.  Checked to machine precision.
  (5) NO-SIGNALLING SURVIVES CHSH VIOLATION.  The same invariance is checked on the
      singlet, whose CHSH value is computed and confirmed at 2*sqrt(2) = 2.828... > 2.
      Correlations that beat the classical bound still signal nothing: the two facts are
      independent, which is exactly the point.

A "pass": (1) bijective on every split; (2) product error ~ 1e-16 while the non-product
control is O(1); (3) inner products and Kronecker products match to ~1e-15; (4) marginal
drift ~ 1e-16 over all settings; (5) CHSH = 2.828 with marginal drift still ~ 1e-16.
"""

from __future__ import annotations
from itertools import combinations
import numpy as np


# ---------------------------------------------------------------- helpers
def subsets(n):
    """All 2^n configurations of a domain of size n, as bitmasks."""
    return list(range(1 << n))


def split_masks(n, a_sites):
    """Bitmasks selecting D_A and its complement D_B."""
    mA = 0
    for a in a_sites:
        mA |= (1 << a)
    return mA, ((1 << n) - 1) & ~mA


def compress(c, mask, n):
    """Index of (c n mask) within P(mask), preserving the order of the kept sites."""
    out, bit = 0, 0
    for i in range(n):
        if mask >> i & 1:
            if c >> i & 1:
                out |= (1 << bit)
            bit += 1
    return out


def partial_trace_B(rho, dA, dB):
    return np.einsum('ijkj->ik', rho.reshape(dA, dB, dA, dB))


def random_state(d, rng):
    v = rng.normal(size=d) + 1j * rng.normal(size=d)
    return v / np.linalg.norm(v)


def random_unitary(d, rng):
    z = (rng.normal(size=(d, d)) + 1j * rng.normal(size=(d, d))) / np.sqrt(2)
    q, r = np.linalg.qr(z)
    return q * (np.diag(r) / abs(np.diag(r)))


def main():
    print("=" * 100)
    print("  Computation 142 -- Tensor factorisation over a bipartition, and no-signalling")
    print("=" * 100)
    print()
    rng = np.random.default_rng(20260922)
    n = 12
    cfgs = subsets(n)
    splits = [(0, 1, 2, 3), (0, 2, 4, 6, 8, 10), (0,), tuple(range(7))]

    # ---- (1) the bijection ------------------------------------------------
    print("  (1) BIJECTION  beta(C) = (C n D_A, C n D_B) on P(D), |D| = %d" % n)
    all_bijective = True
    for a_sites in splits:
        mA, mB = split_masks(n, a_sites)
        pairs = {(compress(c, mA, n), compress(c, mB, n)) for c in cfgs}
        ok = len(pairs) == len(cfgs) == (1 << len(a_sites)) * (1 << (n - len(a_sites)))
        all_bijective &= ok
        print(f"      |D_A| = {len(a_sites):>2}, |D_B| = {n-len(a_sites):>2} : "
              f"{len(pairs)} distinct pairs from {len(cfgs)} configurations -> {ok}")
    print(f"      -> bijective on every split: {all_bijective}")
    print()

    # ---- (2) the measure factors (and the control that fails) -------------
    print("  (2) THE MEASURE FACTORS.  mu(C) = mu_A(C n D_A) * mu_B(C n D_B)")
    def factorisation_error(weights, a_sites):
        """max_C |nu(C) - nu_A(C n D_A) nu_B(C n D_B)|, with the marginals taken from
        the measure itself rather than assumed.  This is the test that separates a
        product measure from a merely permutation-invariant one."""
        mA, mB = split_masks(n, a_sites)
        wA = np.zeros(1 << len(a_sites))
        wB = np.zeros(1 << (n - len(a_sites)))
        for c in cfgs:
            wA[compress(c, mA, n)] += weights[c]
            wB[compress(c, mB, n)] += weights[c]
        return max(abs(weights[c] - wA[compress(c, mA, n)] * wB[compress(c, mB, n)])
                   for c in cfgs)

    bern = np.full(len(cfgs), 2.0 ** -n)          # the canonical measure
    worst = max(factorisation_error(bern, a) for a in splits)
    print(f"      Bernoulli mu: max |mu(C) - mu_A mu_B| over {len(splits)} splits x all "
          f"{len(cfgs)} configurations = {worst:.3e}")

    # control: a permutation-invariant but NON-product measure, weighted by |C|
    w = np.array([1.0 + c.bit_count() for c in cfgs])
    w /= w.sum()
    ctrl = max(factorisation_error(w, a) for a in splits)
    rel = ctrl / w.max()
    print(f"      NON-product control (weight ~ 1+|C|, permutation-invariant):")
    print(f"        max |nu(C) - nu_A nu_B| = {ctrl:.3e}   ({rel:.1%} of the largest weight)")
    measure_factors = worst < 1e-15 and rel > 0.01
    print(f"      -> product measure factors, non-product one does not: {measure_factors}")
    print()

    # ---- (3) the induced unitary ------------------------------------------
    print("  (3) THE INDUCED UNITARY on L^2, and elementary tensors")
    a_sites = splits[0]
    mA, mB = split_masks(n, a_sites)
    nA, nB = len(a_sites), n - len(a_sites)
    dA, dB = 1 << nA, 1 << nB
    perm = np.array([compress(c, mA, n) * dB + compress(c, mB, n) for c in cfgs])
    f = random_state(1 << n, rng)
    g = random_state(1 << n, rng)
    ip_before = np.vdot(f, g)
    ip_after = np.vdot(f[np.argsort(perm)], g[np.argsort(perm)])
    phiA, phiB = random_state(dA, rng), random_state(dB, rng)
    direct = np.array([phiA[compress(c, mA, n)] * phiB[compress(c, mB, n)] for c in cfgs])
    kron = np.kron(phiA, phiB)
    tensor_err = np.abs(direct[np.argsort(perm)] - kron).max()
    ip_err = abs(ip_before - ip_after)
    print(f"      inner product preserved: |<f,g> - <Uf,Ug>| = {ip_err:.3e}")
    print(f"      (phi_A (x) phi_B)(C) = phi_A(C n D_A) phi_B(C n D_B): max err = {tensor_err:.3e}")
    unitary_ok = ip_err < 1e-12 and tensor_err < 1e-12
    print(f"      -> the factorisation is the Kronecker product: {unitary_ok}")
    print()

    # ---- (4) no-signalling -------------------------------------------------
    print("  (4) NO-SIGNALLING.  Tr_B [ sum_j (1 (x) Q_j) rho (1 (x) Q_j) ] = Tr_B rho")
    dA, dB = 4, 4
    drift = 0.0
    for _ in range(200):
        psi = random_state(dA * dB, rng)          # generically correlated
        rho = np.outer(psi, psi.conj())
        rho_A = partial_trace_B(rho, dA, dB)
        U = random_unitary(dB, rng)               # an arbitrary B setting
        new = np.zeros_like(rho)
        for j in range(dB):
            q = np.outer(U[:, j], U[:, j].conj())
            P = np.kron(np.eye(dA), q)
            new += P @ rho @ P
        drift = max(drift, np.abs(partial_trace_B(new, dA, dB) - rho_A).max())
    print(f"      over 200 random correlated states x random B settings:")
    print(f"        max |Tr_B rho' - Tr_B rho| = {drift:.3e}")
    no_signalling = drift < 1e-12
    print(f"      -> the A marginal is untouched by the B setting: {no_signalling}")
    print()

    # ---- (5) it survives CHSH violation ------------------------------------
    print("  (5) THE SAME HOLDS WHERE CHSH IS VIOLATED (the two facts are independent)")
    singlet = np.array([0, 1, -1, 0], dtype=complex) / np.sqrt(2)
    rho = np.outer(singlet, singlet.conj())
    sx = np.array([[0, 1], [1, 0]], dtype=complex)
    sz = np.array([[1, 0], [0, -1]], dtype=complex)

    def corr(u, v):
        return float(np.real(singlet.conj() @ np.kron(u, v) @ singlet))

    a0, a1 = sz, sx
    b0, b1 = (sz + sx) / np.sqrt(2), (sz - sx) / np.sqrt(2)
    chsh = abs(corr(a0, b0) + corr(a0, b1) + corr(a1, b0) - corr(a1, b1))
    rho_A = partial_trace_B(rho, 2, 2)
    drift2 = 0.0
    for _ in range(200):
        U = random_unitary(2, rng)
        new = np.zeros_like(rho)
        for j in range(2):
            q = np.outer(U[:, j], U[:, j].conj())
            P = np.kron(np.eye(2), q)
            new += P @ rho @ P
        drift2 = max(drift2, np.abs(partial_trace_B(new, 2, 2) - rho_A).max())
    print(f"      singlet CHSH value = {chsh:.6f}   (classical bound 2, Tsirelson 2*sqrt(2) = {2*np.sqrt(2):.6f})")
    print(f"      A marginal = I/2 ? max |rho_A - I/2| = {np.abs(rho_A - np.eye(2)/2).max():.3e}")
    print(f"      max marginal drift over 200 B settings = {drift2:.3e}")
    chsh_ok = chsh > 2.8 and drift2 < 1e-12
    print(f"      -> CHSH violated AND nothing signalled: {chsh_ok}")
    print()

    # ---- assessment --------------------------------------------------------
    print("=" * 100)
    print("  ASSESSMENT")
    print("=" * 100)
    print(f"  (1) beta is a bijection P(D) -> P(D_A) x P(D_B) on every split      : {all_bijective}")
    print(f"  (2) the Bernoulli measure factors; a non-product measure does not   : {measure_factors}")
    print(f"  (3) the induced map is unitary and is the Kronecker product         : {unitary_ok}")
    print(f"  (4) the A marginal is invariant under every B setting               : {no_signalling}")
    print(f"  (5) and stays invariant where CHSH is violated                      : {chsh_ok}")
    print()
    ok = all_bijective and measure_factors and unitary_ok and no_signalling and chsh_ok
    print("  RESULT: " + (
        "CONFIRMS the composite-system section -- the tensor factorisation is a theorem "
        "about the\n          Bernoulli measure, and no-signalling follows from it plus the Born rule"
        if ok else "MISMATCH -- revise"))


if __name__ == "__main__":
    main()
