import json
import math
import numpy as np


def euclidean_fix(s, y):
    c = float(s @ y)
    d = abs(c) - c
    return y + d * s / float(s @ s)


def metric_fix(s, y, B):
    c = float(s @ y)
    d = abs(c) - c
    return y + d * (B @ s) / float(s @ B @ s)


def bfgs_update(B, s, y):
    Bs = B @ s
    return B - np.outer(Bs, Bs) / (s @ Bs) + np.outer(y, y) / (s @ y)


# A two-loop L-BFGS inverse-Hessian application. Pairs are (s,y), with y.T s > 0.
def two_loop(g, pairs, H0_scale=1.0):
    q = g.copy()
    alphas = []
    for s, y in reversed(pairs):
        rho = 1.0 / (s @ y)
        a = rho * (s @ q)
        alphas.append(a)
        q = q - a * y
    if pairs:
        s, y = pairs[-1]
        gamma = (s @ y) / max(y @ y, 1e-30)
    else:
        gamma = H0_scale
    r = gamma * q
    for (s, y), a in zip(pairs, reversed(alphas)):
        rho = 1.0 / (s @ y)
        b = rho * (y @ r)
        r = r + s * (a - b)
    return r


def objective(x):
    # Smooth bounded nonconvex quartic in coordinate zero plus a stiff quadratic block.
    return -0.5*x[0]**2 + 0.10*x[0]**4 + 0.05*np.sum(x[1:]**2)


def grad(x):
    g = 0.1*x.copy()
    g[0] = -x[0] + 0.4*x[0]**3
    return g


def run_lbfgs(correct=False, seed=7, n=12, steps=180, memory=7):
    rng = np.random.default_rng(seed)
    x = np.r_[0.35, 0.1*rng.normal(0, 1.0, n-1)]
    pairs = []
    f0 = objective(x)
    corrected = 0
    negative = 0
    backtracks = 0
    accepted = []
    for _ in range(steps):
        g = grad(x)
        p = -two_loop(g, pairs, H0_scale=1.0)
        if not np.isfinite(p).all() or g @ p >= 0:
            p = -g
        f = objective(x)
        slope = g @ p
        t = 1.0
        for _bt in range(35):
            xn = x + t*p
            if objective(xn) <= f + 1e-4*t*slope:
                break
            t *= 0.5
        backtracks += int(round(-math.log2(t)))
        if t < 2**-34:
            break
        gn = grad(xn)
        s = xn - x
        y = gn - g
        c = float(s @ y)
        scale = np.linalg.norm(s)*np.linalg.norm(y)
        if c < -1e-4*scale:
            negative += 1
            if correct:
                y = euclidean_fix(s, y)
                corrected += 1
        if s @ y > 1e-12*max(np.linalg.norm(s)*np.linalg.norm(y), 1.0):
            pairs.append((s, y))
            pairs = pairs[-memory:]
        x = xn
        accepted.append(objective(x))
    return {
        'initial_loss': f0, 'final_loss': float(objective(x)),
        'best_loss': float(np.min(accepted)) if accepted else f0,
        'negative_pairs': negative, 'corrected_pairs': corrected,
        'backtracking_halvings': backtracks, 'steps': len(accepted),
        'finite': bool(np.isfinite(x).all() and np.isfinite(objective(x)))
    }


def main():
    rng = np.random.default_rng(123)
    # Prediction 1/2: for c=-rho||s||||y||, Euclidean correction has
    # s.T ytilde=|c| and ||delta||/||y||=2rho exactly.
    sweep = []
    for rho in [0.001, 0.01, 0.05, 0.1, 0.25, 0.5, 0.9]:
        s = rng.normal(size=9)
        y0 = rng.normal(size=9)
        # Construct y with prescribed normalized negative curvature.
        sh = s / np.linalg.norm(s)
        perp = y0 - sh*(sh @ y0)
        perp /= np.linalg.norm(perp)
        y = -rho*np.linalg.norm(y0)*sh + math.sqrt(max(1-rho*rho, 0))*np.linalg.norm(y0)*perp
        c = s @ y
        yt = euclidean_fix(s, y)
        delta_ratio = np.linalg.norm(yt-y)/np.linalg.norm(y)
        sweep.append({'rho': rho, 'observed_secant_over_abs_c': float((s@yt)/abs(c)),
                      'predicted_secant_over_abs_c': 1.0,
                      'observed_delta_ratio': float(delta_ratio),
                      'predicted_delta_ratio': 2*rho})

    # Prediction 3: the BFGS update is SPD for every corrected pair.
    pd_sweep = []
    for dim in [3, 8, 20]:
        mins = []
        for trial in range(40):
            A = rng.normal(size=(dim, dim))
            B = A.T@A + 0.2*np.eye(dim)
            s = rng.normal(size=dim)
            y = -abs(rng.normal())*s + rng.normal(size=dim)*0.2
            if s@y >= 0: y = y - (abs(s@y)+0.3)*s/(s@s)
            yt = euclidean_fix(s, y)
            Bn = bfgs_update(B, s, yt)
            mins.append(float(np.linalg.eigvalsh((Bn+Bn.T)/2).min()))
        pd_sweep.append({'dimension': dim, 'minimum_observed_eigenvalue': min(mins),
                         'predicted_lower_bound': 0.0})

    baseline = run_lbfgs(correct=False)
    idea = run_lbfgs(correct=True)
    result = {
        'math_verification': {
            'prediction_1_secant_exactness': 's.T ytilde / abs(s.T y) = 1 for every rho',
            'prediction_2_correction_scaling': '||ytilde-y||/||y|| = 2 rho for prescribed c=-rho||s||||y||',
            'prediction_3_positive_definiteness': 'BFGS corrected update has strictly positive eigenvalues when B is SPD',
            'sweep': sweep, 'spd_sweep': pd_sweep
        },
        'mini_experiment': {'baseline_ordinary_reject': baseline, 'idea_euclidean_correction': idea}
    }
    with open('results.json', 'w') as f:
        json.dump(result, f, indent=2)
    print(json.dumps(result, indent=2))


if __name__ == '__main__':
    main()