"""party_blocks.py -- extract the raw (unnormalized) single-party-to-single-party 3x3
correlation blocks M_{a->b} from the Pauli correlation tensor."""
import numpy as np


def party_block(C, a, b):
    M = np.zeros((3, 3))
    for i, ia in enumerate([1, 2, 3]):
        for j, ib in enumerate([1, 2, 3]):
            idx = [0, 0, 0, 0]
            idx[a] = ia
            idx[b] = ib
            M[i, j] = C[tuple(idx)]
    return M


def nuc(M):
    return np.linalg.svd(M, compute_uv=False).sum()