#!/usr/bin/env python3
"""
i3_copies.py -- records versus holographic geometry: the tripartite information of branching records (Chapter 12).

STATE.  |psi> = sqrt(p)|0>_R |e0>_1 |e0>_2 |e0>_3 + sqrt(1-p)|1>_R |e1>_1 |e1>_2 |e1>_3,  <e0|e1>_k = c_k.
IDENTITY (Hosur-Qi-Roberts-Yoshida, pure state on R F1 F2 F3):  I3(F1:F2:F3) = sum_k I(R:F_k) - 2 S(R).
CLOSED FORM used by the page:  S(R) = h(p, c1 c2 c3),  S(F_k) = h(p, c_k),  S(R F_k) = S(F_i F_j) = h(p, c_i c_j),
where h(p, c) is the entropy of p|a><a| + (1-p)|b><b| with |<a|b>| = c.
CLAIMS CHECKED.
  (a) the closed form equals explicit density-matrix entropies (random p, c_k);
  (b) the HQRY identity holds;
  (c) if at most two fragments carry any record (some c_k = 1), I3 = 0 exactly: two copies sit on the
      boundary I3 = 0 of monogamy of mutual information (holographic states need I3 <= 0);
  (d) three imperfect copies give I3 > 0 whenever 0 < p < 1 and all c_k < 1; three perfect copies give I3 = H(p).
numpy only.  No files are written.
"""
import os
for _v in ("OMP_NUM_THREADS", "OPENBLAS_NUM_THREADS", "MKL_NUM_THREADS", "VECLIB_MAXIMUM_THREADS"):
    os.environ.setdefault(_v, "1")
import warnings
import numpy as np

warnings.filterwarnings("ignore", message=r".*encountered in matmul")


def h2(x):
    return 0.0 if x <= 1e-15 or x >= 1 - 1e-15 else float(-(x * np.log2(x) + (1 - x) * np.log2(1 - x)))


def hmix(p, c):
    d = np.sqrt(max(0.0, 1 - 4 * p * (1 - p) * (1 - c * c)))
    return h2((1 + d) / 2)


def S(rho):
    w = np.linalg.eigvalsh((rho + rho.conj().T) / 2)
    w = w[w > 1e-13]
    return float(-(w * np.log2(w)).sum())


def branching(p, cs):
    vecs = []
    for c in cs:
        a = 0.5 * np.arccos(c)
        vecs.append((np.array([np.cos(a), np.sin(a)]), np.array([np.cos(a), -np.sin(a)])))   # overlap cos 2a = c
    b0 = np.array([1.0])
    b1 = np.array([1.0])
    for e0, e1 in vecs:
        b0, b1 = np.kron(b0, e0), np.kron(b1, e1)
    psi = np.sqrt(p) * np.kron([1, 0], b0) + np.sqrt(1 - p) * np.kron([0, 1], b1)
    return np.outer(psi, psi)


def marginal(rho, keep, n=4):
    r = rho.reshape([2] * (2 * n))
    cur = n
    for i in reversed(range(n)):
        if i not in keep:
            r = np.trace(r, axis1=i, axis2=i + cur)
            cur -= 1
    d = 2 ** len(keep)
    return r.reshape(d, d)


def explicit(p, cs):
    rho = branching(p, cs)
    s = lambda *k: S(marginal(rho, list(k)))
    I = [s(0) + s(k) - s(0, k) for k in (1, 2, 3)]
    i3 = s(1) + s(2) + s(3) - s(1, 2) - s(1, 3) - s(2, 3) + s(1, 2, 3)
    return s(0), [s(k) for k in (1, 2, 3)], I, i3


def closed(p, cs):
    c1, c2, c3 = cs
    SR = hmix(p, c1 * c2 * c3)
    SF = [hmix(p, c) for c in cs]
    SRF = [hmix(p, c2 * c3), hmix(p, c1 * c3), hmix(p, c1 * c2)]
    I = [SR + SF[k] - SRF[k] for k in range(3)]
    return SR, SF, I, sum(I) - 2 * SR


def main():
    rng = np.random.default_rng(3)
    e_form = e_id = 0.0
    for _ in range(300):
        p, cs = rng.uniform(0.02, 0.98), rng.uniform(0, 1, 3)
        SRe, SFe, Ie, i3e = explicit(p, cs)
        SRc, SFc, Ic, i3c = closed(p, cs)
        e_form = max(e_form, abs(SRe - SRc), max(abs(a - b) for a, b in zip(SFe, SFc)), max(abs(a - b) for a, b in zip(Ie, Ic)), abs(i3e - i3c))
        e_id = max(e_id, abs(i3e - (sum(Ie) - 2 * SRe)))
    print("i3_copies.py -- tripartite information of branching records")
    print(f"(a) closed form vs explicit entropies, 300 random (p, c1, c2, c3): max diff {e_form:.1e}")
    print(f"(b) I3 = sum_k I(R:F_k) - 2 S(R) (HQRY identity):              max diff {e_id:.1e}")
    e_two, min3 = 0.0, np.inf
    for _ in range(300):
        p, cs = rng.uniform(0.02, 0.98), rng.uniform(0, 0.999, 3)
        for k in range(3):
            two = cs.copy(); two[k] = 1.0                        # fragment k carries nothing
            e_two = max(e_two, abs(explicit(p, two)[3]))
        min3 = min(min3, explicit(p, cs)[3])
    print(f"(c) at most two fragments carry a record:  max |I3| = {e_two:.1e}   (two copies sit exactly on I3 = 0)")
    print(f"(d) three imperfect copies: min I3 over 300 random points = {min3:.2e}  (> 0)")
    print(f"{'good copies':>12} {'overlap c':>10} {'p':>5} {'S(R)':>8} {'I(R:F1)':>8} {'I(R:F2)':>8} {'I(R:F3)':>8} {'I3':>8}")
    for n in (0, 1, 2, 3):
        for c in (0.0, 0.5):
            cs = [c if k < n else 1.0 for k in range(3)]
            SR, _, I, i3 = explicit(0.5, cs)
            z = lambda x: 0.0 if abs(x) < 5e-12 else x               # print exact zeros as 0.0000
            print(f"{n:>12d} {c:>10.2f} {0.5:>5.2f} {z(SR):>8.4f} {z(I[0]):>8.4f} {z(I[1]):>8.4f} {z(I[2]):>8.4f} {z(i3):>8.4f}")
    ok = e_form < 1e-9 and e_id < 1e-9 and e_two < 1e-9 and min3 > 0
    print("all checks pass" if ok else "CHECK FAILED")


if __name__ == "__main__":
    main()