mermin-peres.swift

import Foundation

// MARK: — Complex number struct

struct Complex {
    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
    }
}

// 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) {
    
    // 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()
    }
    
    // 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"
    
    // 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)
        }
    }
    let collapsed = normalize(FourQubitState(newAmp))
    
    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.
    
    var probs = [Double](repeating: 0, count: 4)
    
    for (i,(label,eigs,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
    }
    
    // 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 (labelBell, triple, bket) = bell2Q[chosenB]
    // 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]
            let conjB = Complex(re:cB.re, im:-cB.im)
            let mul = oldC.times(conjB)
            // 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
            )
        }
    }
    let newState = normalize(FourQubitState(newAmpArray))
    
    return (triple, newState)
}

// MARK: - Put it all together

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

func oneRoundRow0Col2() {
    // 1) Build the 4-qubit state  =  (|Φ+>_A1B1)⊗(|Φ+>_A2B2).
    let initial = normalize(twoEprPairsState())
    
    // 2) Alice measures row=0
    let (alice3bits, afterA) = measureAliceRow0(initial)
    
    // 3) Bob measures col=2 on the post-measurement state
    let (bob3bits, _) = measureBobColumn2(afterA)
    
    print("Alice's bits (I⊗Z, Z⊗I, Z⊗Z) => \(alice3bits)")
    print("Bob's   bits (Z⊗Z, X⊗X, Y⊗Y) => \(bob3bits)")
}

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