import numpy as np
from dataclasses import dataclass

I2 = np.eye(2, dtype=complex)
X = np.array([[0,1],[1,0]], dtype=complex)
Z = np.array([[1,0],[0,-1]], dtype=complex)

def kron(*ops):
    out = np.array([[1]], dtype=complex)
    for op in ops:
        out = np.kron(out, op)
    return out

# Bell |Φ+>
phi_plus = (1/np.sqrt(2)) * np.array([1,0,0,1], dtype=complex)
PHI = np.outer(phi_plus, phi_plus.conj())

@dataclass
class Thresholds:
    bind: float = 0.65
    equil: float = 0.85
    release: float = 0.55

class Phase:
    POTENTIAL = "POTENTIAL (0)"
    BINDING = "BINDING (:)"
    EQUILIBRIUM = "EQUILIBRIUM (1)"
    RELEASE = "RELEASE (↘)"

X2, Z2 = kron(X,X), kron(Z,Z)
# Note: up to a global phase, Y⊗Y = - X⊗X · Z⊗Z. Phase cancels for density matrices.
Y2 = - (X2 @ Z2)

def expectation(rho, op):
    return float(np.real(np.trace(rho @ op)))

def depolarizing(rho, p):
    d = rho.shape[0]
    return (1-p) * rho + p * np.eye(d, dtype=complex) / d

def dephasing_2q(rho, p):
    ZI, IZ = kron(Z, I2), kron(I2, Z)
    return (1-p)**2 * rho \
         + p*(1-p) * (ZI @ rho @ ZI) \
         + p*(1-p) * (IZ @ rho @ IZ) \
         + (p**2)   * (ZI @ IZ @ rho @ IZ @ ZI)

# --- Mode A: CPTP reflection "binding" (twirl of size 2) ---
def bind_reflection(rho, R, th: Thresholds = Thresholds()):
    rho_prime = 0.5 * (rho + R @ rho @ R.conj().T)
    raw = expectation(rho, R)          # [-1, 1]
    s = 0.5 * (raw + 1.0)              # [0, 1]
    if s >= th.equil:
        phase = Phase.EQUILIBRIUM
    elif s >= th.bind:
        phase = Phase.BINDING
    elif s < th.release:
        phase = Phase.RELEASE
    else:
        phase = Phase.POTENTIAL
    trace = {"exp_R_before": raw, "score": s}
    return rho_prime, s, phase, trace

# --- Mode B: POST-SELECTED stabilizer projection (raises fidelity on kept runs) ---
def project_reflection(rho, R):
    """Project onto +1 eigenspace of R: Π = (I + R)/2, then renormalize by p_pass."""
    d = rho.shape[0]
    I = np.eye(d, dtype=complex)
    Pi = 0.5 * (I + R)
    rho_num = Pi @ rho @ Pi
    p_pass = float(np.real(np.trace(rho_num)))
    if p_pass > 0:
        rho_proj = rho_num / p_pass
    else:
        rho_proj = rho * 0.0  # stay defined
    return rho_proj, p_pass

def project_seq(rho, reflections):
    """Sequentially project onto +1 spaces; returns rho', cumulative success probability."""
    p_total = 1.0
    out = rho
    for R in reflections:
        out, p = project_reflection(out, R)
        p_total *= p
    return out, p_total

# --- Optional: full stabilizer twirl to Bell-diagonal ---
def twirl_group(rho, reflections=(X2, Z2)):
    """Average over the group generated by reflections (here {I, X⊗X, Z⊗Z, Y⊗Y})."""
    ops = [np.eye(rho.shape[0], dtype=complex), X2, Z2, Y2]
    acc = np.zeros_like(rho)
    for G in ops:
        acc += G @ rho @ G.conj().T
    return acc / 4.0

def bell_demo(p_dephase=0.20, p_depol=0.15, mode="project"):
    # Prepare noisy state
    rho = PHI.copy()
    noisy = depolarizing(dephasing_2q(rho, p_dephase), p_depol)
    f_before = float(np.real(np.trace(noisy @ PHI)))

    if mode == "twirl":
        # CPTP: twirl over full stabilizer group (may preserve fidelity)
        rho_after = twirl_group(noisy)
        f_after = float(np.real(np.trace(rho_after @ PHI)))
        return {
            "mode": "twirl",
            "fidelity_before": f_before,
            "fidelity_after": f_after
        }

    # Default: post-selected projection (improves fidelity on kept runs)
    rho_proj, p_total = project_seq(noisy, (X2, Z2))
    f_after = float(np.real(np.trace(rho_proj @ PHI)))
    return {
        "mode": "project",
        "fidelity_before": f_before,
        "fidelity_after": f_after,
        "pass_probability": p_total
    }
