"""pairwise_correlation_demo.py -- compares the raw single-party-to-single-party 3x3
correlation blocks for a trivially biseparable state (two Bell pairs), GHZ4, and the
ring graph state. Shows that "some pairwise block vanishes" is NOT a valid biseparability
signature: the ring graph state (genuinely entangled) also has several exactly-vanishing
pairwise blocks -- a well-known feature of graph states, confirmed here directly.
"""
import numpy as np
from core2 import full_tensor
from party_blocks import party_block, nuc


def show(psi, label):
    C = full_tensor(psi)
    print(label)
    for a, b, name in [(0, 1, 'A-B'), (0, 2, 'A-C'), (0, 3, 'A-D'),
                        (1, 2, 'B-C'), (1, 3, 'B-D'), (2, 3, 'C-D')]:
        print(f'    {name}: ||M_party||_* = {nuc(party_block(C, a, b)):.4f}')


def bellpair_state(pairing):
    bell = np.array([1, 0, 0, 1]) / np.sqrt(2)
    (p1a, p1b), (p2a, p2b) = pairing
    psi = np.zeros(16, dtype=complex)
    for x in range(2):
        for y in range(2):
            for u in range(2):
                for v in range(2):
                    idx = [0, 0, 0, 0]
                    idx[p1a] = x
                    idx[p1b] = y
                    idx[p2a] = u
                    idx[p2b] = v
                    lin = idx[0] * 8 + idx[1] * 4 + idx[2] * 2 + idx[3]
                    psi[lin] = bell[x * 2 + y] * bell[u * 2 + v]
    return psi


def ring_graph_state():
    plus = np.array([1, 1]) / np.sqrt(2)
    psi = np.kron(np.kron(plus, plus), np.kron(plus, plus))

    def apply_CZ(psi, a, b):
        psi = psi.reshape([2] * 4)
        idx = [slice(None)] * 4
        idx[a] = 1
        idx[b] = 1
        psi[tuple(idx)] *= -1
        return psi.reshape(16)

    psi = apply_CZ(psi, 0, 1)
    psi = apply_CZ(psi, 1, 2)
    psi = apply_CZ(psi, 2, 3)
    psi = apply_CZ(psi, 3, 0)
    return psi


if __name__ == "__main__":
    show(bellpair_state(((0, 1), (2, 3))), 'Bell_AB x Bell_CD:')
    print()
    ghz4 = np.zeros(16, dtype=complex)
    ghz4[0] = 1 / np.sqrt(2)
    ghz4[15] = 1 / np.sqrt(2)
    show(ghz4, 'GHZ4:')
    print()
    show(ring_graph_state(), 'ring graph state:')