"""compute_seed_witnesses.py -- computes and saves the exact optimal dual witnesses
O_AB, O_AC, O_AD (each with operator norm 1) for the known-good state

    rho_mix = 0.5 * (Bell_AB x Bell_CD) + 0.5 * (Bell_AC x Bell_BD)

which is a manifestly valid PPT-mixture state (explicit convex combination of two
product states) scoring EXACTLY (3.0, 3.0, 3.0) for
(||M_AB||_*, ||M_AC||_*, ||M_AD||_*) -- see mixture_search.py for the derivation.

Run this once to produce O_seed.npz, which diagnostic_seeded_run.py and
seeded_exploration.py both load.
"""
import numpy as np
from cluster import cluster_map, nuc
from mixture_search import bellpair_state, full_tensor_mixed_fast

CLUSTERS = [('AB', 0, 1), ('AC', 0, 2), ('AD', 0, 3)]


def true_norms_and_witnesses(M_vals):
    norms, O_opt = {}, {}
    for name, M in M_vals.items():
        U, s, Vt = np.linalg.svd(M)
        norms[name] = s.sum()
        O_opt[name] = U @ Vt
    return norms, O_opt


if __name__ == "__main__":
    psi1 = bellpair_state(((0, 1), (2, 3)))
    psi2 = bellpair_state(((0, 2), (1, 3)))
    rho_mix = 0.5 * np.outer(psi1, np.conj(psi1)) + 0.5 * np.outer(psi2, np.conj(psi2))

    C = full_tensor_mixed_fast(rho_mix)
    M_vals = {name: cluster_map(C, s0, s1) for name, s0, s1 in CLUSTERS}

    norms, O_seed = true_norms_and_witnesses(M_vals)
    print('norms at rho_mix:', norms, ' (expect all == 3.0)')
    for k, v in O_seed.items():
        print(f'  operator norm of O_seed[{k}] =', np.linalg.norm(v, ord=2), '(should be 1.0)')

    np.savez('O_seed.npz', **O_seed)
    print('Saved O_seed.npz')