import json
import math
import numpy as np


def metric(x, R=1.0, delta=1e-10):
    x = np.asarray(x, dtype=float)
    q = 1.0 - float(x @ x) / (R * R)
    if q <= 0:
        raise ValueError("point is outside the open ball")
    # Hessian[-log(1-|x|^2/R^2)] = a I + b xx^T
    a = 2.0 / (R * R * q)
    b = 4.0 / (R**4 * q*q)
    return a * np.eye(len(x)) + b * np.outer(x, x) + delta * np.eye(len(x))


def barrier_value(x, R=1.0):
    q = 1.0 - float(x @ x) / (R * R)
    return -math.log(q) if q > 0 else float("inf")


def intrinsic_grad_sq(x, grad, R=1.0):
    G = metric(x, R)
    return float(grad @ np.linalg.solve(G, grad))


def project_ball(x, R=1.0, tol=1e-8):
    n = np.linalg.norm(x)
    lim = R * (1.0 - tol)
    return x if n < lim else x * (lim / max(n, 1e-30))


def verification():
    rng = np.random.default_rng(123)
    max_ratio = 0.0
    min_eig = float("inf")
    ratios = []
    # The theoretical bound is |grad g|_G^2 <= 1.
    for _ in range(5000):
        direction = rng.normal(size=2)
        direction /= np.linalg.norm(direction)
        # deliberately include points extremely near the boundary
        r = 10 ** rng.uniform(-4, -0.00005)
        x = r * direction
        q = 1 - x @ x
        grad_g = 2 * x / q
        val = intrinsic_grad_sq(x, grad_g)
        ratios.append(val)
        max_ratio = max(max_ratio, val)
        min_eig = min(min_eig, np.linalg.eigvalsh(metric(x)).min())
    # Check exact radial formula at several radii: 2u/(1+u).
    radial_errors = []
    for r in np.linspace(0.0, 0.999999, 100):
        x = np.array([r, 0.0])
        u = r*r
        observed = intrinsic_grad_sq(x, 2*x/(1-u))
        radial_errors.append(abs(observed - 2*u/(1+u)))
    return {
        "max_intrinsic_grad_g_squared": max_ratio,
        "max_bound_violation": max(0.0, max_ratio - 1.0),
        "minimum_metric_eigenvalue": min_eig,
        "max_radial_formula_error": max(radial_errors),
    }


def boundary_scaling_check():
    rows = []
    grad = np.array([1.0, 0.0])
    for r in [0.0, 0.5, 0.9, 0.99, 0.9999, 0.999999]:
        x = np.array([r, 0.0])
        euclidean_step = np.linalg.norm(grad)
        natural_step = np.linalg.norm(np.linalg.solve(metric(x), grad))
        rows.append({"radius": r, "euclidean_step": euclidean_step,
                     "barrier_step": natural_step,
                     "ratio": natural_step / euclidean_step})
    return rows


def run_optimizer(kind, target, steps=300, eta=0.18, R=1.0, seed=7):
    rng = np.random.default_rng(seed)
    x = np.array([0.0, 0.0])
    losses, radii, qvals, grad_norms = [], [], [], []
    # Smooth quadratic, plus small fixed noise to expose boundary overshoot behavior.
    for t in range(steps):
        grad = x - target
        noisy_grad = grad + 0.015 * rng.normal(size=2)
        if kind == "barrier-natural":
            G = metric(x, R)
            step = np.linalg.solve(G, noisy_grad)
            trial = x - eta * step
            # Backtracking is the specified projection safeguard.
            local_eta = eta
            while np.linalg.norm(trial) >= R * (1 - 1e-9):
                local_eta *= 0.5
                trial = x - local_eta * step
                if local_eta < 1e-12:
                    break
            x = trial
        elif kind == "projected-euclidean":
            x = project_ball(x - eta * noisy_grad, R)
        else:
            raise ValueError(kind)
        loss = 0.5 * float(np.sum((x - target) ** 2))
        losses.append(loss)
        radii.append(float(np.linalg.norm(x)))
        qvals.append(float(1 - x @ x))
        grad_norms.append(float(np.linalg.norm(noisy_grad)))
    return {
        "final_loss": losses[-1],
        "best_loss": min(losses),
        "final_radius": radii[-1],
        "minimum_boundary_margin_q": min(qvals),
        "boundary_hits_or_backtracks": int(sum(q <= 1e-8 for q in qvals)),
        "max_gradient_norm": max(grad_norms),
        "loss_first_20": float(np.mean(losses[:20])),
        "loss_last_20": float(np.mean(losses[-20:])),
    }


def main():
    check = verification()
    # Target outside the ball: the constrained infimum is at the boundary.
    target = np.array([1.20, 0.15])
    results = {
        "verification": check,
        "boundary_step_scaling": boundary_scaling_check(),
        "setup": {"domain": "||x|| < 1", "target": target.tolist(), "steps": 300, "eta": 0.18},
        "projected_euclidean": run_optimizer("projected-euclidean", target),
        "barrier_natural": run_optimizer("barrier-natural", target),
        "repeat_summary": {
            kind: {
                "mean_final_loss": float(np.mean([run_optimizer(kind, target, seed=k)["final_loss"] for k in range(5)])),
                "mean_final_radius": float(np.mean([run_optimizer(kind, target, seed=k)["final_radius"] for k in range(5)])),
                "mean_min_q": float(np.mean([run_optimizer(kind, target, seed=k)["minimum_boundary_margin_q"] for k in range(5)])),
            } for kind in ["projected-euclidean", "barrier-natural"]
        },
    }
    print(json.dumps(results, indent=2))


if __name__ == "__main__":
    main()