import json
import numpy as np
from anytime_pd import solve_full, hierarchy, stop_level


def main():
    rng = np.random.default_rng(2930)
    n = 6
    # Box constraints plus two coupling inequalities create a nontrivial LP.
    I = np.eye(n)
    A = np.vstack([I, -I, rng.normal(size=(2, n))])
    b = np.concatenate([np.ones(n), np.ones(n), np.array([0.85, 0.65])])
    c = np.array([0.80, 0.55, 0.35, 0.25, 0.15, 0.10])

    Lstar, Ustar, vstar, ystar = solve_full(A, b, c)
    # Coordinate subspaces are nested and reach the full primal space.
    bases = [np.eye(n)[:, :k] for k in range(1, n + 1)]
    # ystar is already a feasible dual certificate. Adding nonnegative
    # coordinate directions gives nested dual feasible sets.
    dual_bases = [np.column_stack([ystar, np.eye(A.shape[0])[:, :k]])
                  for k in range(1, A.shape[0] + 1)]
    records = hierarchy(A, b, c, bases, dual_bases, full_opt=Lstar)

    lower = [r.lower for r in records]
    upper = [r.upper for r in records]
    widths = [r.width for r in records]
    checks = {
        "all_cover_full": all(r.covers_full for r in records),
        "all_primal_feasible": all(r.primal_residual <= 1e-7 for r in records),
        "all_dual_feasible": all(r.dual_eq_residual <= 1e-7 and r.dual_nonnegative for r in records),
        "lower_monotone": all(x <= y + 1e-8 for x, y in zip(lower, lower[1:])),
        "upper_monotone": all(x >= y - 1e-8 for x, y in zip(upper, upper[1:])),
        "width_monotone": all(x >= y - 1e-8 for x, y in zip(widths, widths[1:])),
        "exact_final_primal": abs(lower[-1] - Lstar) <= 1e-7,
        "strong_duality": abs(Lstar - Ustar) <= 1e-7,
    }
    result = {
        "full_primal": Lstar,
        "full_dual": Ustar,
        "levels": [r.__dict__ for r in records],
        "stop_level_2pct": stop_level(records, tau=0.02),
        "checks": checks,
    }
    print(json.dumps(result, indent=2))
    assert all(checks.values())


if __name__ == "__main__":
    main()