"""Anytime primal-dual LP hierarchy MVP.

Solves max c^T v subject to A v <= b and its nonnegative dual.
Nested primal bases S and nested dual columns T produce certified lower
and upper bounds whenever the corresponding reduced LPs are feasible.
"""
from dataclasses import dataclass
import numpy as np
from scipy.optimize import linprog


@dataclass
class Bound:
    level: int
    lower: float
    upper: float
    width: float
    primal_residual: float
    dual_eq_residual: float
    dual_nonnegative: bool
    covers_full: bool


def solve_full(A, b, c):
    A, b, c = map(lambda x: np.asarray(x, dtype=float), (A, b, c))
    p = linprog(-c, A_ub=A, b_ub=b,
                bounds=[(None, None)] * len(c), method="highs")
    d = linprog(b, A_eq=A.T, b_eq=c,
                bounds=[(0, None)] * len(b), method="highs")
    if not p.success or not d.success:
        raise RuntimeError(f"full LP failed: {p.message}; {d.message}")
    return float(c @ p.x), float(b @ d.x), p.x, d.x


def solve_reduced_primal(A, b, c, S):
    """Return lower bound and full-space primal vector."""
    S = np.asarray(S, dtype=float)
    q = S.T @ c
    sol = linprog(-q, A_ub=A @ S, b_ub=b,
                  bounds=[(None, None)] * S.shape[1], method="highs")
    if not sol.success:
        raise RuntimeError(f"reduced primal failed: {sol.message}")
    v = S @ sol.x
    return float(c @ v), v, float(np.max(A @ v - b))


def solve_reduced_dual(A, b, c, T):
    """Return upper bound and full-space dual vector y=T alpha."""
    T = np.asarray(T, dtype=float)
    # alpha >= 0; equality is A.T T alpha = c.
    sol = linprog(b @ T, A_eq=A.T @ T, b_eq=c,
                  bounds=[(0, None)] * T.shape[1], method="highs")
    if not sol.success:
        raise RuntimeError(f"reduced dual failed: {sol.message}")
    y = T @ sol.x
    return float(b @ y), y, float(np.max(np.abs(A.T @ y - c)))


def hierarchy(A, b, c, primal_bases, dual_bases, full_opt=None):
    """Evaluate paired nested spaces and assert certificate properties."""
    if full_opt is None:
        full_opt = solve_full(A, b, c)[0]
    out = []
    for j, (S, T) in enumerate(zip(primal_bases, dual_bases)):
        L, v, pres = solve_reduced_primal(A, b, c, S)
        U, y, deres = solve_reduced_dual(A, b, c, T)
        out.append(Bound(j, L, U, U - L, pres, deres,
                         bool(np.min(y) >= -1e-8),
                         bool(L - 1e-7 <= full_opt <= U + 1e-7)))
    return out


def stop_level(bounds, tau=0.02):
    for x in bounds:
        if x.width <= tau * max(1.0, abs(x.upper)):
            return x.level
    return None