"""
dicke_block_collapse.py

Reproduces the numerical claims of Proposition (multinomial block collapse)
and Example (Dicke-state scaling): for a Dicke state |D_n^k>, the ambient
3^m x 3^l shadow-map block M_{S->S^c}(rho) has EXACTLY the same singular
values (hence the same nuclear norm) as a much smaller multinomial-weighted
matrix C_hat of size C(m+2,2) x C(l+2,2). This lets ||M_{S->S^c}||_* be
computed exactly for cluster sizes far beyond what the ambient matrix could
ever be built at.

Two things are verified/produced:
  (1) Exact-arithmetic closed-form Dicke correlator, checked against
      brute-force dense simulation for small n.
  (2) Exact match between the ambient matrix M and the reduced matrix C_hat
      (singular values, nuclear norm), then a scaling table pushing m, l, n
      far beyond brute-force reach.

Run: python3 dicke_block_collapse.py
"""

import time
from itertools import combinations, product
from math import comb, factorial

import numpy as np

# ----------------------------------------------------------------------
# Part 0: brute-force reference (only used for small-n sanity checks)
# ----------------------------------------------------------------------

_I = np.eye(2, dtype=complex)
_X = np.array([[0, 1], [1, 0]], dtype=complex)
_Y = np.array([[0, -1j], [1j, 0]], dtype=complex)
_Z = np.array([[1, 0], [0, -1]], dtype=complex)
_PAULI = {"i": _I, "x": _X, "y": _Y, "z": _Z}


def _kron_list(ops):
    out = ops[0]
    for o in ops[1:]:
        out = np.kron(out, o)
    return out


def dicke_state_vector(n, k):
    """Dense state vector of the n-qubit weight-k Dicke state (small n only)."""
    dim = 2 ** n
    psi = np.zeros(dim, dtype=complex)
    for bits in combinations(range(n), k):
        idx = 0
        for b in bits:
            idx |= 1 << (n - 1 - b)
        psi[idx] = 1.0
    psi /= np.linalg.norm(psi)
    return psi


def brute_force_expectation(n, k, labels):
    """<D_n^k| P |D_n^k> by dense simulation. labels: length-n tuple in 'ixyz'."""
    psi = dicke_state_vector(n, k)
    op = _kron_list([_PAULI[c] for c in labels])
    return psi.conj() @ op @ psi


# ----------------------------------------------------------------------
# Part 1: closed-form Dicke correlator (Lemma: closed-form Dicke correlator)
# ----------------------------------------------------------------------

def dicke_correlator(n, k, n_I, n_X, n_Y, n_Z):
    """
    <D_n^k| P |D_n^k> for a Pauli-string TYPE (n_I identities, n_X X's,
    n_Y Y's, n_Z Z's; n_I+n_X+n_Y+n_Z = n), via Eq. (dicke-correlator).

    All intermediate sums are kept as exact Python integers to avoid
    catastrophic cancellation between huge binomial coefficients; only the
    final division by C(n,k) is converted to a float (Python performs a
    correctly-rounded true division even for arbitrary-size integers).
    """
    assert n_I + n_X + n_Y + n_Z == n
    m_xy = n_X + n_Y
    if m_xy % 2 != 0:
        return 0.0
    half = m_xy // 2
    target_iz = k - half
    if target_iz < 0 or target_iz > n_I + n_Z:
        return 0.0

    total = 0  # exact integer accumulator
    for a_x in range(max(0, half - n_Y), min(n_X, half) + 1):
        a_y = half - a_x
        w_xy = comb(n_X, a_x) * comb(n_Y, a_y) * (-1) ** a_y
        for a_i in range(max(0, target_iz - n_Z), min(n_I, target_iz) + 1):
            a_z = target_iz - a_i
            w_iz = comb(n_I, a_i) * comb(n_Z, a_z) * (-1) ** a_z
            total += w_xy * w_iz

    denom = comb(n, k)
    return (1j ** n_Y) * (total / denom)


# ----------------------------------------------------------------------
# Part 2: ambient block matrix vs. multinomial-reduced matrix
# ----------------------------------------------------------------------

def multinomial(n, counts):
    r = factorial(n)
    for c in counts:
        r //= factorial(c)
    return r


def active_types(size):
    """All (a_x, a_y, a_z) with a_x+a_y+a_z == size (the 'fully active' sector)."""
    return [
        (ax, ay, size - ax - ay)
        for ax in range(size + 1)
        for ay in range(size + 1 - ax)
    ]


def ambient_block_matrix(n, k, m, l):
    """The full 3^m x 3^l block M_{S->S^c}(rho) in the raw Pauli-string basis."""
    labels_s = list(product("xyz", repeat=m))
    labels_t = list(product("xyz", repeat=l))
    rest = n - m - l
    M = np.zeros((len(labels_s), len(labels_t)), dtype=complex)
    for i, ls in enumerate(labels_s):
        cs = {c: ls.count(c) for c in "xyz"}
        for j, lt in enumerate(labels_t):
            ct = {c: lt.count(c) for c in "xyz"}
            M[i, j] = dicke_correlator(
                n, k, rest, cs["x"] + ct["x"], cs["y"] + ct["y"], cs["z"] + ct["z"]
            )
    return M


def reduced_block_matrix(n, k, m, l):
    """
    The multinomial-weighted reduced matrix C_hat of Eq. (reduced-dicke-matrix),
    size C(m+2,2) x C(l+2,2), with the SAME singular values as the ambient
    3^m x 3^l block (Proposition: multinomial block collapse).
    """
    src_types = active_types(m)
    tgt_types = active_types(l)
    rest = n - m - l
    C = np.zeros((len(src_types), len(tgt_types)), dtype=complex)
    for i, (sx, sy, sz) in enumerate(src_types):
        w_s = multinomial(m, [sx, sy, sz])
        for j, (tx, ty, tz) in enumerate(tgt_types):
            w_t = multinomial(l, [tx, ty, tz])
            val = dicke_correlator(n, k, rest, sx + tx, sy + ty, sz + tz)
            C[i, j] = float(np.sqrt(float(w_s * w_t))) * val
    return C


# ----------------------------------------------------------------------
# Part 3: checks and scaling table
# ----------------------------------------------------------------------

def check_formula_against_brute_force(n=8, k=3, trials=30, seed=0):
    rng = np.random.default_rng(seed)
    max_err = 0.0
    for _ in range(trials):
        labels = rng.choice(list("ixyz"), size=n)
        counts = {c: int((labels == c).sum()) for c in "ixyz"}
        ref = brute_force_expectation(n, k, tuple(labels))
        val = dicke_correlator(n, k, counts["i"], counts["x"], counts["y"], counts["z"])
        max_err = max(max_err, abs(ref - val))
    print(f"[check 1] closed-form vs. brute force (n={n}, k={k}, {trials} random "
          f"Pauli strings): max error = {max_err:.2e}")


def check_ambient_vs_reduced(n=9, k=4, m=4, l=3):
    M = ambient_block_matrix(n, k, m, l)
    C = reduced_block_matrix(n, k, m, l)
    sv_full = np.sort(np.linalg.svd(M, compute_uv=False))[::-1]
    sv_red = np.sort(np.linalg.svd(C, compute_uv=False))[::-1]
    print(f"[check 2] ambient {M.shape} vs. reduced {C.shape} (n={n}, k={k}, "
          f"m={m}, l={l})")
    print(f"          ||M||_* = {sv_full.sum().real:.10f}")
    print(f"          ||C||_* = {sv_red.sum().real:.10f}")
    print(f"          max |sv_full - sv_red| (top {min(6, len(sv_red))}) = "
          f"{np.max(np.abs(sv_full[:len(sv_red)][:6] - sv_red[:6])):.2e}")


def scaling_table(k=3, cases=((10, 4, 4), (30, 10, 8), (60, 20, 15),
                               (100, 30, 25), (200, 40, 35))):
    print(f"[scaling table] fixed weight k={k}, reduced matrix only "
          f"(ambient matrix is never built)")
    header = f"{'n':>5}{'m':>5}{'l':>5}  {'C_hat shape':>14}  {'ambient 3^m x 3^l':>26}  {'||C_hat||_*':>13}  {'time':>8}"
    print(header)
    for n, m, l in cases:
        t0 = time.time()
        C = reduced_block_matrix(n, k, m, l)
        sv = np.linalg.svd(C, compute_uv=False)
        dt = time.time() - t0
        ambient = f"{3.0**m:.2e} x {3.0**l:.2e}"
        print(f"{n:>5}{m:>5}{l:>5}  {str(C.shape):>14}  {ambient:>26}  "
              f"{sv.sum().real:>13.6f}  {dt:>7.3f}s")


if __name__ == "__main__":
    check_formula_against_brute_force()
    print()
    check_ambient_vs_reduced()
    print()
    scaling_table()