import numpy as np


def cg(matvec, rhs, tol=1e-10, max_iter=None):
    """Conjugate gradients, returning x, iterations, and final relative residual."""
    rhs = np.asarray(rhs, float)
    n = len(rhs)
    if max_iter is None:
        max_iter = max(4 * n, 1000)
    x = np.zeros(n)
    r = rhs - matvec(x)
    p = r.copy()
    rr = float(r @ r)
    r0 = max(np.sqrt(rr), 1e-30)
    if r0 == 0:
        return x, 0, 0.0
    for k in range(1, max_iter + 1):
        Ap = matvec(p)
        den = float(p @ Ap)
        if den <= 0 or not np.isfinite(den):
            raise RuntimeError("CG encountered a non-SPD operator")
        alpha = rr / den
        x += alpha * p
        r -= alpha * Ap
        nr = np.linalg.norm(r)
        if nr <= tol * r0:
            return x, k, nr / r0
        rr_new = float(r @ r)
        p = r + (rr_new / rr) * p
        rr = rr_new
    return x, max_iter, np.linalg.norm(r) / r0


def active_set_simplex(Z, b, lam=1e-3, tol=1e-10, pivot_tol=1e-9,
                       max_pivots=1000, cg_tol=1e-11):
    """Solve min .5*x'(Z'Z+lam I)x-b'x, x>=0, sum(x)=1.

    Products use Z and Z.T only. The equality constrained free solve uses two
    CG solves (A_F y=b_F and A_F q=1), followed by a scalar Schur correction.
    """
    Z = np.asarray(Z, float); b = np.asarray(b, float)
    n, E = Z.shape
    F = list(range(E)); pivots = 0; cg_iters = 0
    x = np.ones(E) / E
    last_nu = 0.0
    def Avec(v, inds):
        return Z[:, inds].T @ (Z[:, inds] @ v) + lam * v
    while pivots <= max_pivots:
        inds = np.array(F, dtype=int)
        Af = lambda v: Avec(v, inds)
        y, it1, _ = cg(Af, b[inds], tol=cg_tol, max_iter=max(20, 8*len(F)))
        q, it2, _ = cg(Af, np.ones(len(F)), tol=cg_tol, max_iter=max(20, 8*len(F)))
        cg_iters += it1 + it2
        denom = float(np.ones(len(F)) @ q)
        nu = (float(np.ones(len(F)) @ y) - 1.0) / denom
        xf = y - nu * q
        x[:] = 0.0; x[inds] = xf; last_nu = nu
        bad = [(v, i) for v, i in zip(xf, inds) if v < -pivot_tol]
        if bad:
            # Bland tie-break after most-negative pivot.
            m = min(v for v, _ in bad)
            i = min(i for v, i in bad if v <= m + 1e-14)
            F.remove(i); pivots += 1
            if not F: raise RuntimeError("empty free set")
            continue
        g = Z.T @ (Z @ x) + lam*x - b + last_nu
        blocked = [i for i in range(E) if i not in F and g[i] < -pivot_tol]
        if blocked:
            m = min(g[i] for i in blocked)
            i = min(i for i in blocked if g[i] <= m + 1e-14)
            F.append(i); F.sort(); pivots += 1
            continue
        # Numerical cleanup is only after all sign decisions.
        x[np.abs(x) < tol] = 0.0
        return x, {"pivots": pivots, "cg_matvecs": cg_iters, "free": len(F),
                   "nu": last_nu, "gradient": g}
    raise RuntimeError("active-set pivot limit reached")


def projected_gradient(Z, b, lam=1e-3, steps=5000, tol=1e-12):
    """Reference projected-gradient solve, with simplex projection."""
    E = Z.shape[1]
    L = np.linalg.eigvalsh(Z.T @ Z + lam*np.eye(E))[-1]
    x = np.ones(E)/E
    for k in range(steps):
        grad = Z.T @ (Z @ x) + lam*x - b
        xn = simplex_project(x - grad/L)
        if np.linalg.norm(xn-x) < tol:
            return xn, k+1
        x = xn
    return x, steps


def simplex_project(v):
    u = np.sort(v)[::-1]
    cssv = np.cumsum(u) - 1
    ind = np.arange(1, len(v)+1)
    rho = np.nonzero(u - cssv/ind > 0)[0]
    rho = rho[-1] if len(rho) else 0
    theta = cssv[rho]/(rho+1)
    return np.maximum(v-theta, 0.0)


def objective(Z, b, x, lam):
    return 0.5*(np.linalg.norm(Z@x)**2 + lam*np.dot(x,x)) - np.dot(b,x)