# %% [markdown]
# # Composite-state calculations
# Run with Python 3 and NumPy, or open the # %% cells in a notebook editor.
# Basis index = a*dB + b. All logarithms are base 2; entropies are in bits.
# Small dense examples only: array storage grows quadratically with Hilbert dimension.
# Explanation: /composite-systems/computational-notebooks/

# %%
import numpy as np

TOL = 1e-10


def density_matrix(rho):
    """Validate a small dense state before using a Hermitian eigensolver."""
    rho = np.asarray(rho, dtype=complex)
    if rho.ndim != 2 or rho.shape[0] != rho.shape[1] or not rho.size:
        raise ValueError("Expected a nonempty square matrix")
    if not np.isfinite(rho).all():
        raise ValueError("Entries must be finite")
    if not np.allclose(rho, rho.conj().T, atol=TOL, rtol=0):
        raise ValueError("State is not Hermitian")
    if not np.isclose(np.trace(rho), 1, atol=TOL, rtol=0):
        raise ValueError("State must have trace one")
    rho = (rho + rho.conj().T) / 2  # Remove only accepted roundoff asymmetry.
    if np.linalg.eigvalsh(rho).min() < -TOL:
        raise ValueError("State has a negative eigenvalue")
    return rho


def tensor(rho, dims):
    """Expose ket A, ket B, bra A, bra B as four separate axes."""
    if len(dims) != 2 or any(isinstance(d, (bool, np.bool_)) or
                            not isinstance(d, (int, np.integer)) or d < 1
                            for d in dims):
        raise ValueError("Expected two positive integer dimensions")
    da, db = dims
    rho = density_matrix(rho)
    if rho.shape != (da * db, da * db):
        raise ValueError("Dimensions do not match the matrix")
    return rho.reshape(da, db, da, db)


def reductions(rho, dims):
    r = tensor(rho, dims)
    return (np.trace(r, axis1=1, axis2=3),
            np.trace(r, axis1=0, axis2=2))


def entropy(rho):
    p = np.linalg.eigvalsh(density_matrix(rho))
    p = np.maximum(p, 0)  # Only negative values within TOL can reach this line.
    p /= p.sum()
    p = p[p > 0]
    return float(-np.sum(p * np.log2(p)))


def partial_transpose(rho, dims):
    r = tensor(rho, dims)
    da, db = dims
    return r.transpose(0, 3, 2, 1).reshape(da * db, da * db)


# %% [markdown]
# ## A complex qubit-qutrit state
# The imaginary entries test the conjugation convention, not just eigenvalues.
# The Schmidt vectors on B are rows of Vh used as ket components, not conjugated.

# %%
def complex_example():
    psi = np.array([1, 0, 0, 0, 1j, 1], dtype=complex) / np.sqrt(3)
    rho = np.outer(psi, psi.conj())
    ra, rb = reductions(rho, (2, 3))
    c = psi.reshape(2, 3)
    np.testing.assert_allclose(ra, c @ c.conj().T, atol=TOL)
    np.testing.assert_allclose(rb, c.T @ c.conj(), atol=TOL)
    u, s, vh = np.linalg.svd(c, full_matrices=False)
    restored = sum(s[k] * np.kron(u[:, k], vh[k, :]) for k in range(len(s)))
    np.testing.assert_allclose(restored, psi, atol=TOL)
    y12 = np.array([[0, 0, 0], [0, 0, -1j], [0, 1j, 0]])
    mean_y = np.trace(rb @ y12).real
    np.testing.assert_allclose(mean_y, -2 / 3, atol=TOL)
    print("Complex state: B spectrum", np.linalg.eigvalsh(rb))
    print("Entropy (bits):", entropy(ra), "local Y12:", mean_y)
    return ra, rb


# %% [markdown]
# ## Identical local states, different joint correlations
# Bell, classical correlation, and independent maximally mixed states all have
# rho_A = rho_B = I/2. Compare joint entropy and total mutual information.

# %%
def correlation_examples():
    bell = np.array([1, 0, 0, 1]) / np.sqrt(2)
    states = {"Bell": np.outer(bell, bell),
              "Classical": np.diag([0.5, 0, 0, 0.5]),
              "Independent": np.eye(4) / 4}
    for name, rho in states.items():
        ra, rb = reductions(rho, (2, 2))
        joint_s = entropy(rho)
        mutual = entropy(ra) + entropy(rb) - joint_s
        np.testing.assert_allclose(ra, np.eye(2) / 2, atol=TOL)
        np.testing.assert_allclose(rb, np.eye(2) / 2, atol=TOL)
        print(name, "joint entropy:", joint_s, "mutual information:", mutual)
    return states


# %% [markdown]
# ## Noisy singlet and partial transpose
# For these two qubits, positive partial transpose is equivalent to separability.
# This sufficiency does not extend to arbitrary subsystem dimensions.

# %%
def werner_example(v=0.6):
    if not 0 <= v <= 1:
        raise ValueError("Mixture weight must lie between zero and one")
    singlet = np.array([0, 1, -1, 0]) / np.sqrt(2)
    rho = v * np.outer(singlet, singlet) + (1 - v) * np.eye(4) / 4
    eigenvalues = np.linalg.eigvalsh(partial_transpose(rho, (2, 2)))
    negativity = float(-eigenvalues[eigenvalues < 0].sum())
    print("Werner v =", v, "PPT eigenvalues:", eigenvalues,
          "negativity:", negativity)
    return eigenvalues, negativity


# %% [markdown]
# ## A controlled occupation cutoff
# A two-mode squeezed state has p_n=(1-z)*z**n, z=tanh(r)**2.
# Keeping |n,n> through n=N discards probability z**(N+1).
# The normalized truncated pure state has squared overlap 1-tail with the exact
# state. This does not by itself bound every unbounded observable's error.

# %%
def cutoff_example(z=0.5, cutoff=9):
    if not 0 <= z < 1 or isinstance(cutoff, bool) or not isinstance(cutoff, int) or cutoff < 0:
        raise ValueError("Require 0 <= z < 1 and an integer cutoff >= 0")
    z = float(z)
    probabilities = (1 - z) * z ** np.arange(cutoff + 1)
    tail = z ** (cutoff + 1)
    np.testing.assert_allclose(probabilities.sum(), 1 - tail, atol=TOL)
    probabilities /= probabilities.sum()
    print("Cutoff", cutoff, "discarded probability:", tail,
          "pure-state fidelity:", 1 - tail)
    return probabilities, tail


# %%
def main():
    print("NumPy", np.__version__)
    complex_example()
    correlation_examples()
    werner_example()
    cutoff_example()


if __name__ == "__main__":
    main()