"""optimize_cluster_witness.py --
(1) optimize the individual cluster map ||M_AB||_* over different biseparable cut types,
    showing that its universal ceiling (5.0) is reached not just by the "home" cut AB|CD
    but also by a state biseparable across the UNRELATED cut AC|BD (Bell_AC x Bell_BD);
(2) optimize min(||M_AB||_*, ||M_AC||_*, ||M_AD||_*) over 1|3-biseparable pure states,
    finding a naive ceiling of 7/3 -- later shown (via mixture_search.py) to be beatable
    by proper MIXTURES, since min() is not convex.
"""
import numpy as np
from scipy.optimize import minimize
from core import state_1_3, state_2_2
from core2 import full_tensor
from cluster import cluster_map, nuc


def triple(psi):
    C = full_tensor(psi)
    return nuc(cluster_map(C, 0, 1)), nuc(cluster_map(C, 0, 2)), nuc(cluster_map(C, 0, 3))


# --- witness ||M_AB||_* over various biseparable cut types ---
def obj_A_BCD(params):
    psi = state_1_3(params, 0)
    C = full_tensor(psi)
    return -nuc(cluster_map(C, 0, 1))


def obj_C_ABD(params):
    psi = state_1_3(params, 2)
    C = full_tensor(psi)
    return -nuc(cluster_map(C, 0, 1))


def obj_AC_BD(params):
    psi = state_2_2(params, (0, 2), (1, 3))
    C = full_tensor(psi)
    return -nuc(cluster_map(C, 0, 1))


def obj_AB_CD(params):
    psi = state_2_2(params, (0, 1), (2, 3))
    C = full_tensor(psi)
    return -nuc(cluster_map(C, 0, 1))


# --- min(M_AB,M_AC,M_AD) over 1|3-biseparable pure states ---
def neg_min_A_BCD(params):
    psi = state_1_3(params, 0)
    return -min(triple(psi))


def neg_min_B_ACD(params):
    psi = state_1_3(params, 1)
    return -min(triple(psi))


def run(obj, nparams, args, n_restarts, seed0, label, maxiter=2500):
    best = -np.inf
    bx = None
    for i in range(n_restarts):
        rng = np.random.default_rng(seed0 + i)
        x0 = rng.normal(size=nparams)
        res = minimize(obj, x0, args=args, method='Powell',
                        options={'maxiter': maxiter, 'xtol': 1e-9, 'ftol': 1e-11})
        v = -res.fun
        if v > best:
            best = v
            bx = res.x
    print(f'{label}: best value = {best:.8f}  ({n_restarts} restarts)')
    return best, bx


if __name__ == "__main__":
    print("=== ||M_AB||_* over various biseparable cut types (universal ceiling = 5) ===")
    run(obj_A_BCD, 18, (), 8, 10, "biseparable A|BCD")
    run(obj_C_ABD, 18, (), 8, 20, "biseparable C|ABD")
    run(obj_AC_BD, 16, (), 8, 30, "biseparable AC|BD  <-- reaches 5 (Bell_AC x Bell_BD)")
    run(obj_AB_CD, 16, (), 6, 40, "biseparable AB|CD (home cut, sanity <= 1)")

    print()
    print("=== min(M_AB,M_AC,M_AD) over 1|3-biseparable states (naive ceiling = 7/3) ===")
    run(neg_min_A_BCD, 18, (), 10, 300, "max over biseparable A|BCD")
    run(neg_min_B_ACD, 18, (), 10, 400, "max over biseparable B|ACD")
    print("For comparison: GHZ4 / connected graph states also give min = 7/3 =", 7 / 3)