"""Fejer averaged-reflection solver for callable fixed-point maps J.

J is expected to be a resolvent (or a numerically stable approximation to one).
The core macro-step uses K evaluations of J and averages y, R y, ..., R^K y,
where R=2J-I. The optional safeguard costs one additional J evaluation.
"""
import numpy as np


def residual(J, y):
    """Euclidean fixed-point residual ||J(y)-y||."""
    y = np.asarray(y, dtype=float)
    return float(np.linalg.norm(np.asarray(J(y), dtype=float) - y))


def fejer_step(J, y, K):
    """One K-reflection Fejer macro-step; exactly K calls to J."""
    if K < 0:
        raise ValueError("K must be nonnegative")
    y = np.asarray(y, dtype=float)
    z = y.copy()
    acc = y.copy()
    for _ in range(K):
        z = 2.0 * np.asarray(J(z), dtype=float) - z
        acc += z
    return acc / (K + 1)


def safeguarded_step(J, y, K, eta=1.0):
    """Accept Fejer candidate if normalized residual does not increase.

    Returns (new_y, accepted, resulting_rho, old_rho). Rejection falls back
    to one ordinary J update, matching the proposed safeguard.
    """
    y = np.asarray(y, dtype=float)
    old_rho = residual(J, y) / max(1.0, float(np.linalg.norm(y)))
    candidate = fejer_step(J, y, K)
    eta = float(np.clip(eta, 0.0, 1.0))
    trial = (1.0 - eta) * y + eta * candidate
    trial_rho = residual(J, trial) / max(1.0, float(np.linalg.norm(trial)))
    if trial_rho <= old_rho:
        return trial, True, trial_rho, old_rho
    fallback = np.asarray(J(y), dtype=float)
    fallback_rho = residual(J, fallback) / max(1.0, float(np.linalg.norm(fallback)))
    return fallback, False, fallback_rho, old_rho


def solve(J, y0, macro_steps, K, safeguard=True):
    """Run macro_steps and return iterate plus history of normalized residuals."""
    y = np.asarray(y0, dtype=float).copy()
    history = []
    accepted = []
    for _ in range(macro_steps):
        if safeguard:
            y, ok, rho, _ = safeguarded_step(J, y, K)
            accepted.append(ok)
        else:
            y = fejer_step(J, y, K)
            rho = residual(J, y) / max(1.0, float(np.linalg.norm(y)))
            accepted.append(True)
        history.append(rho)
    return y, np.asarray(history), accepted