"""Linear-optical operator algebra for the Q3 single-photon register.

The eight spatial modes are indexed 0..7 by the vertices of the 3-cube Q3
with 3-bit addresses b2 b1 b0, i.e. mode k = 4*b2 + 2*b1 + b0.  Everything in
this module is deterministic and follows the construction in the Theory and
Methods sections:

  * BALANCE  B  = I - (1/8) 1 1^T          (orthogonal projector onto N)
  * FOLD     F(theta)                       (conjugate-pair beam splitters)
  * BRAID    G1, G2, G3                      (SU(3) circulants on Gray triads)
  * dynamics R = B F B F G3 G2 G1 B
  * core     R_N = F F G3 G2 G1             (restriction of R to N)
  * checks   C  (DC/sum + three face parities) and its kernel S
  * characters chi_a (the Z2^3 character / Walsh-Hadamard basis)

Gate parameters (Table of dynamics parameters): theta = pi/4, psi = 2*pi/3,
triads T1 = {0,1,3}, T2 = {2,6,7}, T3 = {0,4,5}.
"""

from __future__ import annotations

import numpy as np

DIM = 8
THETA = np.pi / 4          # FOLD coupling angle
PSI = 2 * np.pi / 3        # BRAID angle
TRIADS = ((0, 1, 3), (2, 6, 7), (0, 4, 5))  # T1, T2, T3 (Gray-code segments)

ONES = np.ones(DIM, dtype=complex)


# --------------------------------------------------------------------------
# Index helpers
# --------------------------------------------------------------------------
def bit(k: int, j: int) -> int:
    """Return bit b_j of mode index k (j = 0 is the LSB b0)."""
    return (k >> j) & 1


# --------------------------------------------------------------------------
# BALANCE
# --------------------------------------------------------------------------
def balance() -> np.ndarray:
    """Orthogonal projector B onto the neutral subspace N = {sum v_i = 0}."""
    return np.eye(DIM, dtype=complex) - np.outer(ONES, ONES.conj()) / DIM


# --------------------------------------------------------------------------
# FOLD
# --------------------------------------------------------------------------
def _coupler(theta: float) -> np.ndarray:
    """Symmetric SU(2) coupler u(theta) = cos t I + i sin t sigma_x."""
    c, s = np.cos(theta), np.sin(theta)
    return np.array([[c, 1j * s], [1j * s, c]], dtype=complex)


def fold(theta: float = THETA) -> np.ndarray:
    """FOLD unitary coupling mode k to its conjugate -k mod 8.

    Pairs {1,7}, {2,6}, {3,5} act through the symmetric coupler u(theta);
    the self-conjugate modes 0 and 4 take the scalar phase e^{i theta}, which
    is exactly the eigenvalue u(theta) assigns to the symmetric vector (1,1).
    With this choice F(theta) 1 = e^{i theta} 1, hence F preserves N.
    """
    F = np.zeros((DIM, DIM), dtype=complex)
    u = _coupler(theta)
    phase = np.exp(1j * theta)
    F[0, 0] = phase
    F[4, 4] = phase
    for a, b in ((1, 7), (2, 6), (3, 5)):
        F[np.ix_([a, b], [a, b])] = u
    return F


# --------------------------------------------------------------------------
# BRAID
# --------------------------------------------------------------------------
def _circulant_su3(psi: float) -> np.ndarray:
    """3x3 circulant with eigenvalues (1, e^{i psi}, e^{-i psi}).

    Built from the three-point Fourier eigenvectors f_j with components
    (f_j)_k = omega^{jk}/sqrt 3, omega = e^{2 pi i/3}.  The j=0 eigenvalue is
    fixed to 1 so the uniform triad vector (1,1,1) is preserved; det = 1.
    """
    omega = np.exp(2j * np.pi / 3)
    eigvals = np.array([1.0, np.exp(1j * psi), np.exp(-1j * psi)], dtype=complex)
    F3 = np.array([[omega ** (j * k) for k in range(3)] for j in range(3)],
                  dtype=complex) / np.sqrt(3)
    # rows of F3 are the (conjugated) Fourier eigenvectors; rebuild via
    # G = sum_j eigvals[j] f_j f_j^dagger.
    G = np.zeros((3, 3), dtype=complex)
    for j in range(3):
        f = F3[j].conj()
        G += eigvals[j] * np.outer(f, f.conj())
    return G


def braid(triad: tuple[int, int, int], psi: float = PSI) -> np.ndarray:
    """One BRAID rotation G_i: SU(3) circulant on `triad`, identity elsewhere."""
    G = np.eye(DIM, dtype=complex)
    block = _circulant_su3(psi)
    idx = list(triad)
    G[np.ix_(idx, idx)] = block
    return G


def braids(psi: float = PSI) -> list[np.ndarray]:
    """The three BRAID rotations G1, G2, G3 on the Gray-code triads."""
    return [braid(t, psi) for t in TRIADS]


# --------------------------------------------------------------------------
# Composite dynamics
# --------------------------------------------------------------------------
def neutral_core(theta: float = THETA, psi: float = PSI) -> np.ndarray:
    """Neutral-sector unitary core R_N = F F G3 G2 G1 (applied right-to-left)."""
    G1, G2, G3 = braids(psi)
    F = fold(theta)
    return F @ F @ G3 @ G2 @ G1


def dynamics(theta: float = THETA, psi: float = PSI) -> np.ndarray:
    """Full dynamics operator R = B F B F G3 G2 G1 B (applied right-to-left)."""
    B = balance()
    F = fold(theta)
    G1, G2, G3 = braids(psi)
    return B @ F @ B @ F @ G3 @ G2 @ G1 @ B


# --------------------------------------------------------------------------
# Parity-check matrix and code subspaces
# --------------------------------------------------------------------------
def check_matrix() -> np.ndarray:
    """4x8 check matrix C: DC/sum row plus three face parities (-1)^{b_j}."""
    rows = [np.ones(DIM)]
    for j in range(3):  # M0, M1, M2  ->  faces b0, b1, b2
        rows.append(np.array([(-1) ** bit(k, j) for k in range(DIM)], dtype=float))
    return np.array(rows, dtype=float)


def neutral_basis() -> np.ndarray:
    """Orthonormal basis (columns) of the neutral subspace N = ker(sum)."""
    B = balance()
    # columns of B span N; orthonormalize and keep the 7 nonzero directions.
    u, s, _ = np.linalg.svd(B)
    return u[:, s > 1e-9]


def code_basis() -> np.ndarray:
    """Orthonormal basis (columns) of S = ker(C) (dimension 4)."""
    C = check_matrix()
    # right-singular vectors beyond the rank span the kernel.
    _, s, vh = np.linalg.svd(C)
    rank = int(np.sum(s > 1e-9))
    return vh.conj().T[:, rank:]


def character(a: int) -> np.ndarray:
    """Z2^3 character |chi_a> = (1/sqrt8) sum_x (-1)^{a.x} e_x."""
    vec = np.array([(-1) ** (bin(a & x).count("1") & 1) for x in range(DIM)],
                   dtype=complex)
    return vec / np.sqrt(DIM)


def parity_separator() -> np.ndarray:
    """U_S: the Z2^3 character table normalized by 1/sqrt8 (Walsh-Hadamard)."""
    return np.array([character(a) for a in range(DIM)], dtype=complex)


# --------------------------------------------------------------------------
# State preparation (QR decomposition)
# --------------------------------------------------------------------------
def prep_unitary(psi_target: np.ndarray) -> np.ndarray:
    """Unitary U with U |0> = |psi_target>, via QR completion.

    Mirrors the software state preparation: place the normalized target in the
    first column, complete to a basis, and orthonormalize with QR (fixing the
    sign convention so the first column equals the target up to global phase).
    """
    psi_target = np.asarray(psi_target, dtype=complex)
    psi_target = psi_target / np.linalg.norm(psi_target)
    M = np.eye(DIM, dtype=complex)
    M[:, 0] = psi_target
    Q, R = np.linalg.qr(M)
    # QR sign fix: make diag(R) real-positive so column 0 matches the target.
    phases = np.diag(R).copy()
    phases[np.abs(phases) < 1e-15] = 1.0
    Q = Q * (phases / np.abs(phases)).conj()
    return Q


# --------------------------------------------------------------------------
# Named input states (mode basis e_0..e_7)
# --------------------------------------------------------------------------
def e(k: int) -> np.ndarray:
    v = np.zeros(DIM, dtype=complex)
    v[k] = 1.0
    return v


def neutral_pair(a: int, b: int) -> np.ndarray:
    return (e(a) - e(b)) / np.sqrt(2)


def balanced_4p4m() -> np.ndarray:
    """The (-1)^{b2} balanced 4+4- neutral input."""
    return character(0b100)  # a = 100 -> sign (-1)^{b2}


def uniform_dc() -> np.ndarray:
    return ONES / np.sqrt(DIM)


# Labels a (written a2 a1 a0) for the four code-space characters of S.
CODE_LABELS = {
    "b0+b1": 0b011,
    "b0+b2": 0b101,
    "b1+b2": 0b110,
    "b0+b1+b2": 0b111,
}


def support_weight(vec: np.ndarray, tol: float = 1e-9) -> int:
    """Number of nonzero amplitudes (coding-theoretic Hamming weight)."""
    return int(np.sum(np.abs(vec) > tol))


def min_distance_from_checks(C: np.ndarray) -> int:
    """Minimum support weight of the code ker(C).

    A codeword supported exactly on a column set T exists iff the columns of C
    indexed by T are linearly dependent, so the minimum support weight equals
    the smallest number of linearly dependent columns of C.  This is the same
    column-rank argument used in the paper to certify d(N) = 2 and d(S) = 4.
    """
    import itertools

    n = C.shape[1]
    for w in range(1, n + 1):
        for cols in itertools.combinations(range(n), w):
            if np.linalg.matrix_rank(C[:, list(cols)]) < w:
                return w
    return n


def all_triples_full_rank(C: np.ndarray | None = None) -> tuple[int, int]:
    """Check every 3-column subset of C has full rank (=> d(S) >= 4).

    Returns (num_full_rank, num_triples); the paper states all 56 triples are
    full rank.
    """
    import itertools

    if C is None:
        C = check_matrix()
    triples = list(itertools.combinations(range(DIM), 3))
    full = sum(1 for cols in triples
               if np.linalg.matrix_rank(C[:, list(cols)]) == 3)
    return full, len(triples)


if __name__ == "__main__":  # pragma: no cover - smoke test
    B = balance()
    C = check_matrix()
    print("rank B =", np.linalg.matrix_rank(B), "(expect 7)")
    print("||B^2 - B|| =", np.linalg.norm(B @ B - B))
    print("dim N =", neutral_basis().shape[1], "(expect 7)")
    print("dim S =", code_basis().shape[1], "(expect 4)")
    print("d(N) =", min_distance_from_checks(C[:1, :]), "(expect 2)")
    print("d(S) =", min_distance_from_checks(C), "(expect 4)")
    print("triples full rank:", all_triples_full_rank())
