from __future__ import annotations

import numpy as np


def admm_lite(A: np.ndarray, b: np.ndarray, rho: float = 1.0, max_iter: int = 50, tol: float = 1e-4):
    """Very lightweight ADMM solver for a toy problem: minimize 0.5||Ax - b||^2 + regularization.

    Returns the primal variable x. This is a minimal placeholder to demonstrate the schema.
    """
    m, n = A.shape
    x = np.zeros(n)
    z = np.zeros(n)
    u = np.zeros(n)

    for _ in range(max_iter):
        # x-update (least-squares step)
        q = A.T @ b
        x = np.linalg.solve(A.T @ A + rho * np.eye(n), q + rho * (z - u))
        # z-update (soft-threshold for L1-like regularization placeholder)
        z_old = z.copy()
        z = _soft_threshold(x + u, rho)
        # u-update
        u += x - z
        # Convergence check
        if np.linalg.norm(z - z_old) < tol:
            break
    return x


def _soft_threshold(v: np.ndarray, lam: float) -> np.ndarray:
    return np.sign(v) * np.maximum(np.abs(v) - lam, 0.0)