"""
ghz3_shadow_map_symbolic.py

Exact symbolic (sympy) construction of the combined shadow map M_A(rho) for the
three-qubit GHZ state, with source party A and target complement {B,C}.

This reproduces, with exact algebraic numbers (no floating point), the claim
from Section "Qubit examples" of the paper:

    For |GHZ_3> = (|000> + |111>)/sqrt(2), the three singular values of the
    normalized combined shadow map M_A(rho) are all equal to sqrt(2/3),
    i.e. ||M_A(rho)||_* = sqrt(6).

Convention (matches the .tex draft):
  - Pauli generators sigma_1=X, sigma_2=Y, sigma_3=Z, normalized by
    tr(sigma_i sigma_j) = 2 delta_ij  (qubit case, d_a = 2).
  - Target sectors for source A are T in { {B}, {C}, {B,C} }, stacked as
    rows of one 15 x 3 matrix (3 from B, 3 from C, 9 from BC).
  - Combined shadow map normalization: 1/sqrt((d_a-1)(d_bar_a-1))
    = 1/sqrt(1*3) = 1/sqrt(3) for n=3 qubits (Eq. "combined-map" in the note).

Run:  python3 ghz3_shadow_map_symbolic.py
"""

import sympy as sp
from sympy import sqrt, I, simplify, Matrix, eye, zeros, re, Rational

# ---------------------------------------------------------------------
# 1. Pauli matrices (exact, symbolic entries)
# ---------------------------------------------------------------------
X = Matrix([[0, 1], [1, 0]])
Y = Matrix([[0, -I], [I, 0]])
Z = Matrix([[1, 0], [0, -1]])
I2 = eye(2)
PAULIS = {'x': X, 'y': Y, 'z': Z}


def kron(A, B):
    """Kronecker (tensor) product of two sympy matrices, built manually
    so everything stays exact/symbolic (no numeric backend needed)."""
    mA, nA = A.shape
    mB, nB = B.shape
    out = zeros(mA * mB, nA * nB)
    for i in range(mA):
        for j in range(nA):
            out[i * mB:(i + 1) * mB, j * nB:(j + 1) * nB] = A[i, j] * B
    return out


def kron3(a, b, c):
    """Tensor product of three single-qubit operators -> 8x8 matrix."""
    return kron(kron(a, b), c)


# Embeddings of a single-qubit operator P onto party A, B, or C
# within the 3-qubit Hilbert space (order A ⊗ B ⊗ C).
def op_A(P): return kron3(P, I2, I2)
def op_B(P): return kron3(I2, P, I2)
def op_C(P): return kron3(I2, I2, P)


# ---------------------------------------------------------------------
# 2. The GHZ_3 state
# ---------------------------------------------------------------------
def ghz3_state():
    """Density matrix of (|000> + |111>)/sqrt(2), as an 8x8 sympy Matrix."""
    psi = zeros(8, 1)
    psi[0, 0] = 1 / sqrt(2)   # |000>
    psi[7, 0] = 1 / sqrt(2)   # |111>
    rho = psi * psi.H         # outer product, .H = conjugate transpose
    return simplify(rho)


# ---------------------------------------------------------------------
# 3. Correlation-tensor entries and the shadow-map matrix
# ---------------------------------------------------------------------
def entry(rho, ops):
    """tr(rho * op1 * op2 * ...), simplified and forced real
    (expectation values of Hermitian operators in a Hermitian state
    are always real; re(...) just discards a numerically/symbolically
    residual zero imaginary part)."""
    M = None
    for op in ops:
        M = op if M is None else M * op
    return simplify(re(simplify((rho * M).trace())))


def build_M(rho):
    """Unnormalized shadow-map matrix M_A(rho): 15 (target) x 3 (source A).

    Row blocks, in order:
      rows 0-2  : target sector T = {B}   (source index x,y,z; target x,y,z)
      rows 3-5  : target sector T = {C}
      rows 6-14 : target sector T = {B,C} (9 = 3x3 combinations)
    Column index: source generator on A, in order x,y,z.
    """
    rows = []
    for pb in ['x', 'y', 'z']:
        rows.append([entry(rho, [op_A(PAULIS[pa]), op_B(PAULIS[pb])])
                     for pa in ['x', 'y', 'z']])
    for pc in ['x', 'y', 'z']:
        rows.append([entry(rho, [op_A(PAULIS[pa]), op_C(PAULIS[pc])])
                     for pa in ['x', 'y', 'z']])
    for pb in ['x', 'y', 'z']:
        for pc in ['x', 'y', 'z']:
            rows.append([entry(rho, [op_A(PAULIS[pa]), op_B(PAULIS[pb]), op_C(PAULIS[pc])])
                         for pa in ['x', 'y', 'z']])
    return Matrix(rows)


# ---------------------------------------------------------------------
# 4. Main: build, normalize, and diagonalize
# ---------------------------------------------------------------------
def main():
    rho0 = ghz3_state()
    print("tr(rho0) =", simplify(rho0.trace()), " (sanity check, should be 1)\n")

    M0 = build_M(rho0)
    print("Unnormalized shadow matrix M0 (15x3):")
    sp.pprint(M0)

    # Combined-map normalization for n=3 qubits: 1/sqrt((d_a-1)(d_bar_a-1)) = 1/sqrt(3)
    norm_const = 1 / sqrt(3)
    Mn0 = simplify(norm_const * M0)

    # Singular values of Mn0 are sqrt(eigenvalues of the Gram matrix Mn0^T Mn0).
    # This avoids sympy's (slower/less robust) generic SVD and is exact here
    # because Mn0^T Mn0 is a small 3x3 symmetric matrix.
    G = simplify(Mn0.T * Mn0)
    print("\nGram matrix  Mn0^T Mn0  =")
    sp.pprint(G)

    eigs = G.eigenvals()  # dict: eigenvalue -> multiplicity
    print("\nEigenvalues of the Gram matrix (= squared singular values):")
    for ev, mult in eigs.items():
        sigma = simplify(sqrt(ev))
        print(f"  lambda = {ev}  (multiplicity {mult})  ->  sigma = {sigma}"
              f"  = {float(sigma):.6f}")

    print("\nExpected from the paper: sigma = sqrt(2/3) = sqrt(6)/3 ≈ 0.816497,"
          " threefold degenerate.")

    # Save U0, V0 (orthonormal bases of the degenerate singular subspace) for
    # reuse in the perturbation-theory script. Since the Gram matrix is
    # exactly (2/3) * I_3 here, the source space is untouched (V0 = I_3) and
    # U0 is simply Mn0 rescaled to unit-norm columns.
    sigma_val = sqrt(Rational(2, 3))
    U0 = simplify(Mn0 / sigma_val)
    V0 = eye(3)
    print("\nU0 (15x3, orthonormal columns spanning the degenerate target subspace):")
    sp.pprint(U0)
    print("\nCheck U0^T U0 = I_3:", simplify(U0.T * U0))

    return rho0, U0, V0, sigma_val


if __name__ == "__main__":
    main()