"""
Seeded exploration around the known-good point (3.0, 3.0, 3.0).

Two questions this answers:
  1) Is 3.0 a STABLE fixed point of the alternating scheme (perturb the witnesses a
     little, does it converge back to 3.0)?
  2) Does searching the FULL PPT-mixture body (strictly larger than biseparable states)
     starting from near this point ever find something BETTER than 3.0?

If nothing beats 3.0 even when explicitly seeded nearby and given many iterations, that
is now fairly strong evidence -- across both a from-scratch parametrized search (earlier)
and this SDP-based search over the larger PPT-mixture relaxation -- that 3.0 is the true
supremum (at least for PPT-mixtures, hence an upper bound on the biseparable one too,
since biseparable subset PPT-mixtures).
"""
import numpy as np
from sdp_ppt_mixture import solve_fixed_witness_step, true_norms_and_witnesses, CLUSTERS

data = np.load('O_seed.npz')
O_seed = {name: data[name] for name, _, _ in CLUSTERS}


def project_to_unit_opnorm(A):
    """Rescale A to have operator norm exactly 1 (SVD-based projection)."""
    U, s, Vt = np.linalg.svd(A)
    return U @ Vt if s.max() == 0 else A / s.max()


def run_seeded(perturbation_strength, n_iters=25, seed=0, verbose=True):
    rng = np.random.default_rng(seed)
    O = {}
    for name, _, _ in CLUSTERS:
        noise = rng.normal(size=(15, 15)) * perturbation_strength
        O[name] = project_to_unit_opnorm(O_seed[name] + noise)

    best_min = -np.inf
    history = []
    for it in range(n_iters):
        rho_val, t_val, M_vals = solve_fixed_witness_step(O)
        norms, O = true_norms_and_witnesses(M_vals)
        cur_min = min(norms.values())
        history.append(cur_min)
        best_min = max(best_min, cur_min)
        if verbose:
            nice = {k: round(v, 5) for k, v in norms.items()}
            print(f"    iter {it:2d}: SDP t={t_val:.5f}  norms={nice}  min={cur_min:.5f}")
    return best_min, history


if __name__ == "__main__":
    print("=== Stability check: seed EXACTLY at the known optimum (no perturbation) ===")
    best0, _ = run_seeded(perturbation_strength=0.0, n_iters=10, seed=0)
    print(f"  best min found: {best0:.6f}  (should stay essentially at 3.0)\n")

    print("=== Perturbation sweep: does it converge back to 3.0, drift, or improve? ===")
    results = {}
    for strength in [0.05, 0.1, 0.2, 0.4, 0.7, 1.0]:
        print(f"--- perturbation strength {strength} ---")
        best, hist = run_seeded(perturbation_strength=strength, n_iters=25, seed=1, verbose=True)
        results[strength] = best
        print(f"  final best: {best:.6f}\n")

    print("=" * 60)
    for s, v in results.items():
        print(f"  perturbation {s:.2f} -> best min found = {v:.6f}")
    overall_best = max(results.values())
    print()
    print("Overall best across all perturbed seeded runs:", overall_best)
    print("Compare: 3.0 (conjectured exact), 11/3 =", 11/3, "(proven upper bound)")