mermin-peres2.swift

import Foundation

// This program is mostly generated by ChatGPT o1 in a long conversation
// based on this PDF: https://arxiv.org/pdf/2206.12042
// "Experimental Demonstration of Quantum Pseudotelepathy"

struct Complex: CustomStringConvertible {
    var re: Double
    var im: Double

    func conjugate() -> Complex {
        Complex(re: re, im: -im)
    }

    func times(_ other: Complex) -> Complex {
        Complex(
            re: re * other.re - im * other.im,
            im: re * other.im + im * other.re
        )
    }

    // Squared magnitude
    func magnitudeSquared() -> Double {
        re * re + im * im
    }

    var description: String {
        if im == 0 {
            return "\(re)"
        } else if re == 0 {
            return "\(im)i"
        } else if im < 0 {
            return "\(re) - \(-im)i"
        } else {
            return "\(re) + \(im)i"
        }
    }
}

// MARK: — A 4-qubit state: dimension 16

/// A 16-dimensional vector representing qubits (A1,A2,B1,B2) in that order.
/// Basis ordering convention: |A1,A2,B1,B2> is indexed as a 4-bit binary #.
struct FourQubitState {
    // Exactly 16 complex amplitudes
    var amp: [Complex]  // amp[0..15]

    init(_ arr: [Complex]) {
        precondition(arr.count == 16)
        self.amp = arr
    }
}

/// Normalize so that sum of squared magnitudes = 1
func normalize(_ s: FourQubitState) -> FourQubitState {
    let normFactor = 1.0 / sqrt(s.amp.map { $0.magnitudeSquared() }.reduce(0, +))
    let newAmp = s.amp.map {
        Complex(re: $0.re * normFactor, im: $0.im * normFactor)
    }
    return FourQubitState(newAmp)
}

/// Inner product <phi|psi>
func innerProduct(_ phi: FourQubitState, _ psi: FourQubitState) -> Complex {
    var sumC = Complex(re: 0, im: 0)
    for i in 0..<16 {
        let conjPhi = phi.amp[i].conjugate()
        sumC = Complex(
            re: sumC.re + conjPhi.times(psi.amp[i]).re,
            im: sumC.im + conjPhi.times(psi.amp[i]).im
        )
    }
    return sumC
}

// MARK: — Build a standard 4-qubit “two EPR pairs” state, i.e. |Φ+>⊗|Φ+>

/// |Φ+> = (|00> + |11>)/√2. We'll define "two copies" on (A1,B1) and (A2,B2).
/// But to store them in A1,A2,B1,B2 order, we do a little care in indexing.
func twoEprPairsState() -> FourQubitState {
    // We want: (|00>+|11>)/√2 on (A1,B1) AND (|00>+|11>)/√2 on (A2,B2).
    // The total 4-qubit = (1/2)* (|0,0,A2,B2> + ...), but we must reorder the bits.
    // We'll build it by directly enumerating:
    //
    //   |Φ+>_A1B1 = (|00> + |11>)/√2  means A1=0,B1=0 plus A1=1,B1=1
    //   |Φ+>_A2B2 = (|00> + |11>)/√2  means A2=0,B2=0 plus A2=1,B2=1
    //
    // Combined: each possibility picks a bit for A1,B1 and a bit for A2,B2.
    // We'll fill amp[indexOf(A1,A2,B1,B2)].

    var arr = [Complex](repeating: Complex(re: 0, im: 0), count: 16)
    // We'll add amplitude = 1/2 for each matching pattern:
    //   (A1=0,B1=0) or (A1=1,B1=1)
    //   (A2=0,B2=0) or (A2=1,B2=1)
    // amplitude = 1/2 in total
    let val = Complex(re: 0.5, im: 0)  // since overall factor = 1/2

    for A1 in [0, 1] {
        for B1 in [0, 1] {
            // For an EPR: we only add amplitude if A1==B1
            guard A1 == B1 else { continue }
            for A2 in [0, 1] {
                for B2 in [0, 1] {
                    // For the second EPR: add amplitude if A2==B2
                    guard A2 == B2 else { continue }

                    let index = (A1 << 3) | (A2 << 2) | (B1 << 1) | B2
                    // In binary: A1 is the top bit, A2 next, B1 next, B2 last
                    arr[index] = val
                }
            }
        }
    }
    return FourQubitState(arr)
}

// MARK: — A partial measurement on A’s row=0: (I⊗Z, Z⊗I, Z⊗Z) for qubits (A1,A2)

/// In the Mermin–Peres table, row=0 on A is the commuting triple
///   { I⊗Z ,  Z⊗I ,  Z⊗Z }  acting on (A1,A2).
/// We can measure them by focusing on A's 2-qubit subspace (dimension=4),
/// but in the full 4-qubit state we must keep track of B as well.
/// We'll define a partial measurement that yields 3 bits for A,
/// and returns a collapsed 4-qubit state.

/// We'll store the 4 "joint eigenstates" for A as |00>,|01>,|10>,|11> in A1,A2,
/// tensored with the identity on B. Then pick an outcome from the 4 possible,
/// each has a unique triple of ±1 for (I⊗Z, Z⊗I, Z⊗Z).
///
/// We'll keep the B part unmeasured, so each outcome is a 4D subspace of dimension=4 for B.

let row0Eigenvals: [String: [Int]] = [
    "00": [+1, +1, +1],  // A=|00> => (I⊗Z=+1, Z⊗I=+1, Z⊗Z=+1)
    "01": [-1, +1, -1],
    "10": [+1, -1, -1],
    "11": [-1, -1, +1],
]

// Each "A-basis ket" is 4D on A, tensored with 4D identity on B => total dimension=16
// We'll explicitly store them as 16-element states.

/// Builds a product state |a>_A ⊗ |b>_B for given 2-qubit "aIndex" in {0..3} and 2-qubit "bIndex" in {0..3}.
/// We'll embed them in the 4-qubit space of dimension 16.
func productKet(aIndex: Int, bIndex: Int) -> FourQubitState {
    // aIndex chooses which of { |00>,|01>,|10>,|11> } for A
    // bIndex likewise for B
    // We'll place amplitude=1 at the position (A1A2B1B2).
    // That position in decimal is (aIndex<<2) | bIndex.

    var arr = [Complex](repeating: Complex(re: 0, im: 0), count: 16)
    let pos = (aIndex << 2) | bIndex
    arr[pos] = Complex(re: 1.0, im: 0)
    return FourQubitState(arr)
}

/// The 4 basis states for A's row=0 measurement, extended over B by identity.
let aRow0Basis: [(label: String, state: FourQubitState)] = {
    // A in { |00>,|01>,|10>,|11> }, B in superposition => we just do "∀ bIndex"
    // Actually we want the sum over all bIndices for each A?  Not quite:
    // In a partial measurement, each "A outcome" lumps *all* B states together.
    // We only define the *4 states for A* tensored with the *identity* on B.
    // That is, each "A-basis vector" is spanned by { |aIndex>_A ⊗ |ANY>_B }.
    //
    // But for measurement, we need an orthonormal projection. So we do:
    //    P_aIndex = (|aIndex><aIndex|) on A  ⊗  I_4 on B
    // We'll handle it by a function that yields the subspace amplitude.
    // We'll not build a single vector for each outcome. Instead we'll do
    // "Compute probability by summing amplitude for all B states, if A = aIndex."
    // Then we collapse accordingly.

    // We'll store them just for reference. The actual partial measure function
    // might do it by summing sub-blocks. But let's store the "pure" vector form
    // for clarity: for each aIndex in 0..3, the "ket" is a direct sum over all bIndex in 0..3
    // but weighting = 1 for each. That won't be normalized if B is included. We'll not use it directly.

    var list = [(String, FourQubitState)]()
    let labels = ["00", "01", "10", "11"]
    for aIdx in 0..<4 {
        // We'll build an unnormalized superstate that has amplitude=1 for all B.
        // i.e. sum_{b=0..3} |aIdx,b>.
        var bigVec = [Complex](repeating: Complex(re: 0, im: 0), count: 16)
        for bIdx in 0..<4 {
            let pos = (aIdx << 2) | bIdx
            bigVec[pos] = Complex(re: 1, im: 0)
        }
        list.append((labels[aIdx], FourQubitState(bigVec)))
    }
    return list
}()

/// Return (3 bits from row=0 measurement, post-measurement 4-qubit state).
func measureAliceRow0(
    _ fullState: FourQubitState
) -> ([Int], FourQubitState) {

    print("measureRow0: state=\(fullState.amp)")
    // 1) Probability of each "A=|00>,|01>,|10>,|11>" outcome
    //    i.e. sum of squared amplitude of all basis states where A matches that index
    var probs = [Double](repeating: 0, count: 4)
    for aIdx in 0..<4 {
        // We'll collect all bIdx in 0..3
        var sumAmp = Complex(re: 0, im: 0)
        for bIdx in 0..<4 {
            let pos = (aIdx << 2) | bIdx  // Abits in top, Bbits in bottom
            sumAmp = Complex(
                re: sumAmp.re + fullState.amp[pos].re,
                im: sumAmp.im + fullState.amp[pos].im
            )
        }
        probs[aIdx] = sumAmp.magnitudeSquared()
    }
    print("probs=\(probs)")

    // 2) Sample from that distribution
    let r = Double.random(in: 0...1)
    var cum = 0.0
    var chosenA = 3
    for i in 0..<4 {
        cum += probs[i]
        if r <= cum {
            chosenA = i
            break
        }
    }

    // 3) The 3 bits for that A outcome
    let labelA = ["00", "01", "10", "11"][chosenA]
    let aBits = row0Eigenvals[labelA]!  // e.g. (+1,+1,+1) for "00"
    print("chosenA=\(chosenA) label=\(labelA)")

    // 4) Collapse the 4-qubit wavefunction so that A is pinned to chosenA
    //    That means we zero out all amplitudes where A != chosenA.
    //    Then re-normalize.
    var newAmp = fullState.amp
    for aIdx in 0..<4 {
        if aIdx == chosenA { continue }
        for bIdx in 0..<4 {
            let pos = (aIdx << 2) | bIdx
            newAmp[pos] = Complex(re: 0, im: 0)
        }
    }
    print("newAmp=\(newAmp)")
    let collapsed = normalize(FourQubitState(newAmp))
    print("collapsed = \(collapsed)")

    return (aBits, collapsed)
}

// MARK: — Bob's column=2: (Z⊗Z, X⊗X, Y⊗Y) on qubits (B1,B2)

/// Similarly, we define the 4 joint-eigenstates on (B1,B2) as the 4 Bell states,
/// tensored with identity on A.  We'll measure partial on B. Then we get 3 bits, collapse.

/// For brevity, let's store each Bell state's amplitude in (B1,B2) dimension=4:
struct TwoQState {
    var c: [Complex]  // length=4
}
func normalize2Q(_ st: TwoQState) -> TwoQState {
    let s = st.c.map { $0.magnitudeSquared() }.reduce(0, +)
    let nf = 1.0 / sqrt(s)
    let c2 = st.c.map { Complex(re: $0.re * nf, im: $0.im * nf) }
    return TwoQState(c: c2)
}

/// We'll define the 4 Bell states in the order: (|β++>, |β+->, |β-+>, |β-->) on (B1,B2).
///   |β++> = (|00> + |11>)/√2,   eigenvalues (Z⊗Z=+1, X⊗X=+1, Y⊗Y=-1)
/// etc.
let bell2Q: [(label: String, eigvals: [Int], st: TwoQState)] = {
    // We'll store them unnormalized for clarity, then normalize.

    // β++ => (|00> + |11>)
    let bpp = TwoQState(c: [
        Complex(re: 1, im: 0),  // for |00>
        Complex(re: 0, im: 0),  // for |01>
        Complex(re: 0, im: 0),  // for |10>
        Complex(re: 1, im: 0),  // for |11>
    ])

    // β+- => (|01> + |10>)
    let bpm = TwoQState(c: [
        Complex(re: 0, im: 0),
        Complex(re: 1, im: 0),
        Complex(re: 1, im: 0),
        Complex(re: 0, im: 0),
    ])
    // β-+ => (|00> - |11>)
    let bmp = TwoQState(c: [
        Complex(re: 1, im: 0),
        Complex(re: 0, im: 0),
        Complex(re: 0, im: 0),
        Complex(re: -1, im: 0),
    ])
    // β-- => (|01> - |10>)
    let bmm = TwoQState(c: [
        Complex(re: 0, im: 0),
        Complex(re: 1, im: 0),
        Complex(re: -1, im: 0),
        Complex(re: 0, im: 0),
    ])

    let e_bpp = [+1, +1, -1]  // Z⊗Z=+1, X⊗X=+1, Y⊗Y=-1
    let e_bpm = [-1, +1, +1]
    let e_bmp = [+1, -1, +1]
    let e_bmm = [-1, -1, -1]

    return [
        ("β++", e_bpp, normalize2Q(bpp)),
        ("β+-", e_bpm, normalize2Q(bpm)),
        ("β-+", e_bmp, normalize2Q(bmp)),
        ("β--", e_bmm, normalize2Q(bmm)),
    ]
}()

/// measureBobColumn2: partial measurement in {Z⊗Z, X⊗X, Y⊗Y} on (B1,B2),
/// tensored with identity on (A1,A2).
/// Returns (3 bits, new collapsed state).
func measureBobColumn2(
    _ fullState: FourQubitState
) -> ([Int], FourQubitState) {
    // 1) Probability of each Bell outcome on B.
    //    We'll sum amplitude for each sub-basis of dimension=4 on B, times A dimension=4 => total 16.
    //    If the B part = bKet, we add up the amplitude for all A in 0..3 in that bKet’s pattern.
    print("measureCol2: state=\(fullState.amp)")

    var probs = [Double](repeating: 0, count: 4)

    for (i, (label, bits, bket)) in bell2Q.enumerated() {
        // bket has 4 components c[0..3], corresponding to B1B2 in {0,1,2,3} => {|00>,|01>,|10>,|11>}
        // For each bIndex in 0..3, we have bket.c[bIndex].
        // So the amplitude in the full 4-qubit is sum_{aIndex} [ fullState.amp[(aIndex<<2)|bIndex] * bket.c[bIndex] ]?
        // Actually for probability, we do an overlap approach or a direct sum of those coefficients. The simplest is:
        //   amplitude = sum_{aIndex, bIndex} [ (coefficient from fullState) * (coefficient from B's basis) ]
        // But we only want to do this if B is exactly that bell state. That is effectively a partial inner product.

        // We'll do a direct "coefficient overlap" approach:
        var sumRe = 0.0
        var sumIm = 0.0
        for bIdx in 0..<4 {
            let cB = bket.c[bIdx]  // amplitude for that B basis
            if cB.magnitudeSquared() < 1e-14 { continue }
            for aIdx in 0..<4 {
                let pos = (aIdx << 2) | bIdx
                // add fullState.amp[pos] * cB
                sumRe += fullState.amp[pos].re * cB.re - fullState.amp[pos].im * cB.im
                sumIm += fullState.amp[pos].re * cB.im + fullState.amp[pos].im * cB.re
            }
        }
        let prob = sumRe * sumRe + sumIm * sumIm
        probs[i] = prob
        print("P_\(label) (\(bits)) = \(prob)")
    }
    // print("probs=\(probs)")

    // 2) random pick
    let r = Double.random(in: 0...1)
    var chosenB = 3
    var cum = 0.0
    for i in 0..<4 {
        cum += probs[i]
        if r <= cum {
            chosenB = i
            break
        }
    }
    let (_, triple, bket) = bell2Q[chosenB]
    print("chosenB=\(chosenB) triple=\(triple) bket=\(bket)")
    // triple are the 3 bits (Z⊗Z, X⊗X, Y⊗Y)

    // 3) Collapse wavefunction onto that bell state in B, keep A as is.
    //    newAmp[pos] = sum_{ oldAmp, project onto that B subspace } ...
    // For partial measurement, the collapsed amplitude for (aIndex,bIndex) is:
    //   oldAmp[pos] * (overlap with bket) but we only keep the portion consistent with bket in B.

    // We'll do a direct "projection" approach:
    var newAmpArray = [Complex](repeating: Complex(re: 0, im: 0), count: 16)
    for bIdx in 0..<4 {
        let cB = bket.c[bIdx]  // amplitude for B's basis vector
        for aIdx in 0..<4 {
            let pos = (aIdx << 2) | bIdx
            // projected amplitude = oldAmp[pos] * cB^*
            // (We are effectively re-embedding that B basis vector.)
            let oldC = fullState.amp[pos]
            // Then we add that to newAmp for each pos, but scaled by cB again? Actually we are building
            // |A,B> = |A>⊗ sum_{bIdx}( cB * |bIdx> ), so the amplitude is oldC * cB^*, then multiply cB?
            // The standard formula for projection is: newAmp = Σ_b ( oldAmp[pos] ) * <bket|bIdx> ) * |bIdx>
            // But <bket|bIdx> = cB^*, and then we re-insert cB in the expansion. So total factor is cB^* * cB = |cB|^2.
            // Actually let's do the simpler method: We'll do the "coefficient in the resulting superposition" = mul, but we have to
            // re-insert cB. That can be tricky. For correctness, let's build an explicit sum-of-outer-products.
            //
            // In a short snippet, let's do a direct overlap approach. For each bIdx, the new amplitude is oldAmp[pos]* cB^*, but we store it in the same pos multiplied by cB. That yields oldAmp[pos]* cB^* cB = oldAmp[pos]* |cB|^2.
            // Since each bket is normalized, cB is typically 1/sqrt(2) or ±1/sqrt(2).
            let factor = cB.magnitudeSquared()  // cB^* * cB
            newAmpArray[pos] = Complex(
                re: newAmpArray[pos].re + oldC.re * factor,
                im: newAmpArray[pos].im + oldC.im * factor
            )
        }
    }
    print("newAmp=\(newAmpArray)")
    let newState = normalize(FourQubitState(newAmpArray))
    print("collapsed = \(newState)")

    return (triple, newState)
}

// MARK: 1) Alice’s Row=1 measurement data
//     Operators: (X⊗I, I⊗X, X⊗X) on (A1,A2).
//     Shared eigenbasis = the "two-qubit X-basis": {|++>, |+->, |-+>, |-->}.

/// We label these states "++", "+-", "-+", "--".
/// Each is dimension=4 in the (A1,A2) subspace:
///   |++> = (|00> + |01> + |10> + |11>)/2
///   |+-> = (|00> - |01> + |10> - |11>)/2
///   |-+> = (|00> + |01> - |10> - |11>)/2
///   |--> = (|00> - |01> - |10> + |11>)/2
///
/// The corresponding eigenvalues (±1) for (X⊗I, I⊗X, X⊗X) are:
///   |++> => (+1, +1, +1)
///   |+-> => (+1, -1, -1)
///   |-+> => (-1, +1, -1)
///   |--> => (-1, -1, +1)

let row1Eigenvals: [String: [Int]] = [
    "++": [+1, +1, +1],
    "+-": [+1, -1, -1],
    "-+": [-1, +1, -1],
    "--": [-1, -1, +1],
]

// We'll store the actual 2-qubit states (unnormalized) as 4-element arrays, then
// embed them in the 4-qubit space by tensoring with identity on Bob (B1,B2).
// The order is (A1A2) in [0..3] => {00,01,10,11}.

func twoQState_xBasis(label: String) -> [Complex] {
    switch label {
    case "++":
        // (|00> + |01> + |10> + |11>)/2
        return [
            Complex(re: 1, im: 0),  // index=0 => |00>
            Complex(re: 1, im: 0),  // index=1 => |01>
            Complex(re: 1, im: 0),  // index=2 => |10>
            Complex(re: 1, im: 0),  // index=3 => |11>
        ]
    case "+-":
        // (|00> - |01> + |10> - |11>)/2
        return [
            Complex(re: 1, im: 0),
            Complex(re: -1, im: 0),
            Complex(re: 1, im: 0),
            Complex(re: -1, im: 0),
        ]
    case "-+":
        // (|00> + |01> - |10> - |11>)/2
        return [
            Complex(re: 1, im: 0),
            Complex(re: 1, im: 0),
            Complex(re: -1, im: 0),
            Complex(re: -1, im: 0),
        ]
    case "--":
        // (|00> - |01> - |10> + |11>)/2
        return [
            Complex(re: 1, im: 0),
            Complex(re: -1, im: 0),
            Complex(re: -1, im: 0),
            Complex(re: 1, im: 0),
        ]
    default:
        return []
    }
}

// MARK: 2) The partial-measurement function for Row=1 on (A1,A2).
// Very similar structure to measureAliceRow0(...), but we use the X⊗X basis states.

func measureAliceRow1(_ fullState: FourQubitState) -> ([Int], FourQubitState) {

    // 1) Compute probabilities for each of the 4 "X-basis" outcomes on A.
    //    We'll do it by summing amplitude for all bIndices in 0..3 for Bob.
    print("measureRow1: state=\(fullState.amp)")

    let labels = ["++", "+-", "-+", "--"]
    var probs = [Double](repeating: 0, count: 4)

    for (i, lab) in labels.enumerated() {
        // The 2-qubit vector for A is in xBasis2Q. We'll sum "fullState.amp[pos]" for all pos
        // where A matches that superposition, and B is free. We can do it similarly to row0 approach:
        //   amplitude = sum_{bIdx, aIdx} [ fullState.amp[(aIdx<<2)|bIdx] * (coefficient for that aIdx in x-basis ) ]

        // Build the unnormalized 2-qubit A vector for label
        let unnormA = twoQState_xBasis(label: lab)  // length=4
        // Each entry is +1 or -1 for that aIdx, and total factor 1/2 => we do magnitude^2 => 1/4. We'll handle that below.

        var sumRe = 0.0
        var sumIm = 0.0
        for aIdx in 0..<4 {
            let cA = unnormA[aIdx]
            if cA.magnitudeSquared() < 1e-15 { continue }
            for bIdx in 0..<4 {
                let pos = (aIdx << 2) | bIdx
                // add: fullState.amp[pos] * cA
                sumRe += fullState.amp[pos].re * cA.re - fullState.amp[pos].im * cA.im
                sumIm += fullState.amp[pos].re * cA.im + fullState.amp[pos].im * cA.re
            }
        }
        // Probability = (1/4)* |sumRe + i sumIm|^2 because that unnorm had a factor of 1/2
        // Actually  cA = ±1, so amplitude factor is 1, but we want 1/2 overall => so let's do (1/2) outside?
        // The state is (unnormalized)/2. So let's multiply sumRe, sumIm by (1/2):
        let scRe = 0.5 * sumRe
        let scIm = 0.5 * sumIm
        let p = scRe * scRe + scIm * scIm
        probs[i] = p
    }
    print("probs=\(probs)")

    // 2) Sample from that distribution
    let r = Double.random(in: 0...1)
    var chosenIndex = 0
    var cumsum = 0.0
    for i in 0..<4 {
        cumsum += probs[i]
        if r <= cumsum {
            chosenIndex = i
            break
        }
    }

    let chosenLabel = labels[chosenIndex]  // e.g. "++"
    let chosenBits = row1Eigenvals[chosenLabel]!  // e.g. (+1,+1,+1)

    print("chosen=\(chosenIndex) chosenLab=\(chosenLabel) chosenBits=\(chosenBits)")
    // 3) Collapse the wavefunction onto that outcome
    //    We'll zero out all amplitude that doesn't match that X-basis outcome on A.
    //    Then renormalize. It's the partial-projection approach.

    let unnormA = twoQState_xBasis(label: chosenLabel)
    var newAmpArray = [Complex](repeating: Complex(re: 0, im: 0), count: 16)

    for aIdx in 0..<4 {
        let cA = unnormA[aIdx]
        // factor = cA^* * (1/2) for amplitude
        // cA is real ±1 => cA^* = cA. The 1/2 factor from the definition => net factor = cA*(1/2).

        let factorRe = 0.5 * cA.re
        let factorIm = 0.5 * cA.im

        for bIdx in 0..<4 {
            let pos = (aIdx << 2) | bIdx
            let oldC = fullState.amp[pos]
            // newC = oldC * factor
            let re2 = oldC.re * factorRe - oldC.im * factorIm
            let im2 = oldC.re * factorIm + oldC.im * factorRe

            newAmpArray[pos] = Complex(
                re: newAmpArray[pos].re + re2,
                im: newAmpArray[pos].im + im2
            )
        }
    }

    print("newAmp = \(newAmpArray)")
    let collapsed = normalize(FourQubitState(newAmpArray))
    print("collapsed = \(collapsed)")

    return (chosenBits, collapsed)
}

// 1) The 4 basis states for Alice's row=2 (each stored as 4 unnormalized amplitudes in A's 2-qubit space):
//    alpha++ => (|00> + |01> + |10> - |11>)
func alphaStateRow2(label: String) -> [Complex] {
    switch label {
    case "α++":
        return [
            Complex(re: +1, im: 0),  // for |00>
            Complex(re: +1, im: 0),  // for |01>
            Complex(re: +1, im: 0),  // for |10>
            Complex(re: -1, im: 0),  // for |11>
        ]
    case "α+-":
        return [
            Complex(re: +1, im: 0),
            Complex(re: -1, im: 0),
            Complex(re: +1, im: 0),
            Complex(re: +1, im: 0),
        ]
    case "α-+":
        return [
            Complex(re: +1, im: 0),
            Complex(re: +1, im: 0),
            Complex(re: -1, im: 0),
            Complex(re: +1, im: 0),
        ]
    case "α--":
        return [
            Complex(re: +1, im: 0),
            Complex(re: -1, im: 0),
            Complex(re: -1, im: 0),
            Complex(re: -1, im: 0),
        ]
    default:
        return []
    }
}

// 2) The triple of eigenvalues for (-X⊗Z, -Z⊗X, Y⊗Y).
let row2Eigenvals: [String: [Int]] = [
    "α++": [-1, -1, +1],
    "α+-": [-1, +1, -1],
    "α-+": [+1, -1, -1],
    "α--": [+1, +1, +1],
]

// 3) Partial measurement on (A1,A_2) in { -X⊗Z, -Z⊗X, Y⊗Y }
func measureAliceRow2(_ fullState: FourQubitState) -> ([Int], FourQubitState) {
    // We'll label these 4 states "α++", "α+-", "α-+", "α--".
    let labels = ["α++", "α+-", "α-+", "α--"]
    var probs = [Double](repeating: 0, count: 4)
    print("measureRow2: state=\(fullState.amp)")

    // 1) Probability of each outcome
    for (i, lab) in labels.enumerated() {
        // "unnormalized" local array for A:
        let aLocal = alphaStateRow2(label: lab)  // length=4
        // We'll do sum of amplitude over all bIdx in 0..3.

        var sumRe = 0.0
        var sumIm = 0.0

        for aIdx in 0..<4 {
            let cA = aLocal[aIdx]
            for bIdx in 0..<4 {
                // Position in the 16-element wavefunction
                let pos = (aIdx << 2) | bIdx
                // Add oldAmp[pos] * cA
                sumRe += fullState.amp[pos].re * cA.re - fullState.amp[pos].im * cA.im
                sumIm += fullState.amp[pos].re * cA.im + fullState.amp[pos].im * cA.re
            }
        }
        // Because each local vector has factor 1/2 for normalization, we do that now:
        let scaledRe = 0.5 * sumRe
        let scaledIm = 0.5 * sumIm
        let p = scaledRe * scaledRe + scaledIm * scaledIm
        probs[i] = p
    }
    print("probs=\(probs)")

    // 2) Randomly pick which outcome we get
    let r = Double.random(in: 0...1)
    var chosenIdx = 0
    var cumsum = 0.0
    for i in 0..<4 {
        cumsum += probs[i]
        if r <= cumsum {
            chosenIdx = i
            break
        }
    }

    let chosenLabel = labels[chosenIdx]  // e.g. "α++"
    let chosenBits = row2Eigenvals[chosenLabel]!  // e.g. [-1, -1, +1]
    print("chosen=\(chosenIdx) chosenLab=\(chosenLabel) chosenBits=\(chosenBits)")

    // 3) Collapse the wavefunction so that A is in that basis. => Zero out all amplitude that doesn't match chosen state
    var newAmp = [Complex](repeating: Complex(re: 0, im: 0), count: 16)
    let aLocal = alphaStateRow2(label: chosenLabel)

    for aIdx in 0..<4 {
        let cA = aLocal[aIdx]
        // multiply by cA*(1/2) to project
        let factorRe = 0.5 * cA.re
        let factorIm = 0.5 * cA.im

        for bIdx in 0..<4 {
            let pos = (aIdx << 2) | bIdx
            let oldC = fullState.amp[pos]
            // newC = oldC * factor
            let re2 = oldC.re * factorRe - oldC.im * factorIm
            let im2 = oldC.re * factorIm + oldC.im * factorRe

            newAmp[pos] = Complex(
                re: newAmp[pos].re + re2,
                im: newAmp[pos].im + im2
            )
        }
    }
    print("newAmp = \(newAmp)")
    let collapsed = normalize(FourQubitState(newAmp))
    print("collapsed = \(collapsed)")
    return (chosenBits, collapsed)
}

// MARK: - Put it all together

// We'll define a function that does one round: Alice row=0, then Bob col=2.

func oneRound() {
    // 1) Build the 4-qubit state  =  (|Φ+>_A1B1)⊗(|Φ+>_A2B2).
    let initial = normalize(twoEprPairsState())

    // 2) Alice measures row=0
    let (alice3bits, afterA) = measureAliceRow0(initial)
    print("Alice's bits (I⊗Z, Z⊗I, Z⊗Z) => \(alice3bits)")

    // let (alice3bits, afterA) = measureAliceRow1(initial)
    // print("Alice's Row=1 bits (X⊗I, I⊗X, X⊗X) =>", alice3bits)

    // let (alice3bits, afterA) = measureAliceRow2(initial)
    // print("Alice's row=2 bits (-X⊗Z, -Z⊗X, Y⊗Y) =>", alice3bits)

    // 3) Bob measures col=2 on the post-measurement state
    let (bob3bits, _) = measureBobColumn2(afterA)

    print("Bob's   bits (Z⊗Z, X⊗X, Y⊗Y) => \(bob3bits)")
}

print("aRow0Basis:")
print(aRow0Basis)
// Let's run it a few times:
for _ in 1...5 {
    oneRound()
    print("---")
}