"""sum_bound_proof.py -- proof ingredients for the RIGOROUS bound
    rho biseparable  =>  min(||M_AB||_*, ||M_AC||_*, ||M_AD||_*) <= 11/3

Step 1: verify (numerically, over each cut type) that
    ||M_AB||_* + ||M_AC||_* + ||M_AD||_*  <=  11
for every PURE state product across a single bipartition (the extreme points of the
biseparable set). Since the SUM of nuclear norms IS convex (unlike the min!), this bound
then extends by convexity to ALL biseparable (mixed) states -- this is the key trick that
lets a convexity/extreme-point argument work here even though it fails for min() itself.

Step 2: verify via linear programming that uniform weights (1/3,1/3,1/3) are optimal for
turning the sum bound into a bound on min(...) via min(a,b,c) <= w.(a,b,c) for any
w in the simplex -- i.e. that 11/3 is the best bound achievable by this proof technique
(cannot be tightened just by re-weighting).
"""
import numpy as np
from scipy.optimize import minimize, linprog
from core import state_1_3, state_2_2
from core2 import full_tensor
from cluster import cluster_map, nuc


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


def neg_sum_2_2(params, pair1, pair2):
    return -sum3(state_2_2(params, pair1, pair2))


def neg_sum_1_3(params, source):
    return -sum3(state_1_3(params, source))


def run(obj, nparams, args, n_restarts, seed0, label):
    best = -np.inf
    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': 1500, 'xtol': 1e-9, 'ftol': 1e-11})
        v = -res.fun
        if v > best:
            best = v
    print(f'{label}: max sum = {best:.6f}   ({n_restarts} restarts)')
    return best


if __name__ == "__main__":
    print("=== Step 1: max(||M_AB||+||M_AC||+||M_AD||) over each pure single-cut family ===")
    r1 = run(neg_sum_2_2, 16, ((0, 1), (2, 3)), 8, 6000, 'cut AB|CD')
    r2 = run(neg_sum_2_2, 16, ((0, 2), (1, 3)), 8, 6100, 'cut AC|BD')
    r3 = run(neg_sum_1_3, 18, (0,), 8, 6200, 'cut A|BCD')
    r4 = run(neg_sum_1_3, 18, (1,), 8, 6300, 'cut B|ACD')
    print()
    print('Overall max sum over ALL single-cut pure product states:', max(r1, r2, r3, r4))
    print('(convexity of the sum then extends this bound to ALL biseparable mixtures)')

    print()
    print("=== Step 2: is uniform weighting (1/3,1/3,1/3) optimal for the resulting bound? ===")
    # extreme points of the "sum" bound: (1,5,5), (5,1,5), (5,5,1)
    c = [0, 0, 0, 1]
    A_ub = [[1, 5, 5, -1], [5, 1, 5, -1], [5, 5, 1, -1]]
    b_ub = [0, 0, 0]
    A_eq = [[1, 1, 1, 0]]
    b_eq = [1]
    bounds = [(0, 1), (0, 1), (0, 1), (None, None)]
    res = linprog(c, A_ub=A_ub, b_ub=b_ub, A_eq=A_eq, b_eq=b_eq, bounds=bounds, method='highs')
    print('LP-optimal weights:', res.x[:3], ' bound t=', res.x[3])
    print('11/3 =', 11 / 3, '  (confirms uniform weights are optimal for this proof technique)')