import json
import numpy as np


def cubic_tensor(c, g, H, C):
    """Return A such that A[(1,s)]^3 = c+g.s+.5*s'Hs+C[s,s,s]/6."""
    g, H, C = np.asarray(g), np.asarray(H), np.asarray(C)
    n = len(g)
    A = np.zeros((n + 1,) * 3, dtype=float)
    A[0, 0, 0] = c
    for i in range(n):
        for p in ((0, 0, i + 1), (0, i + 1, 0), (i + 1, 0, 0)):
            A[p] = g[i] / 3.0
    for i in range(n):
        for j in range(n):
            # Three placements of the zero index sum to (1/2) H_ij s_i s_j.
            for p in ((0, i + 1, j + 1), (i + 1, 0, j + 1),
                      (i + 1, j + 1, 0)):
                A[p] = (0.5 * H[i, j]) / 3.0
    A[1:, 1:, 1:] = C / 6.0
    return A


def model(c, g, H, C, s):
    return c + g @ s + .5 * s @ H @ s + np.einsum('ijk,i,j,k', C, s, s, s) / 6.


def contract(A, u, v):
    return np.einsum('ijk,j,k->i', A, u, v)


def pam(A, sweeps=30, beta=0.8, seed=0):
    rng = np.random.default_rng(seed)
    d = A.shape[0]
    U = [rng.normal(size=d) for _ in range(3)]
    U = [u / np.linalg.norm(u) for u in U]
    history = []
    def F():
        return np.einsum('ijk,i,j,k', A, *U)
    for _ in range(sweeps):
        for block in range(3):
            old = U[block].copy()
            before = F()  # proximal term is zero at the old block
            others = [U[i] for i in range(3)]
            a = contract(A, others[(block + 1) % 3], others[(block + 2) % 3])
            z = beta * old - a
            if np.linalg.norm(z) < 1e-14:
                continue
            U[block] = z / np.linalg.norm(z)
            after = F() + .5 * beta * np.sum((U[block] - old) ** 2)
            history.append((after - before, before, after))
    return U, np.asarray(history)


def run():
    rng = np.random.default_rng(12)
    n = 4
    c = .37
    g = rng.normal(size=n)
    H = rng.normal(size=(n, n)); H = (H + H.T) / 2
    C0 = rng.normal(size=(n, n, n))
    C = sum(C0.transpose(p) for p in ((0,1,2),(1,0,2),(2,1,0),(0,2,1),(1,2,0),(2,0,1))) / 6
    A = cubic_tensor(c, g, H, C)

    # Prediction 1: exact homogeneous representation.
    rep_errors = []
    for _ in range(1000):
        s = rng.normal(size=n)
        rep_errors.append(abs(np.einsum('ijk,i,j,k', A, np.r_[1., s], np.r_[1., s], np.r_[1., s]) - model(c, g, H, C, s)))
    representation_max_error = float(max(rep_errors))

    # Prediction 2 sweep: exact block solves never increase their proximal objective,
    # for every beta; larger beta is the predicted stronger proximal stabilization.
    beta_sweep = {}
    for beta in (0.0, 0.1, 0.8, 5.0):
        _, ch = pam(A, sweeps=40, beta=beta, seed=4)
        beta_sweep[str(beta)] = {'max_increase': float(np.max(ch[:, 0])),
                                 'final_objective': float(ch[-1, 2])}
    max_block_increase = beta_sweep['0.8']['max_increase']
    final_F = beta_sweep['0.8']['final_objective']

    # Prediction 3 sweep: for a quartic mismatch q*r^4/4, rho decreases with
    # radius and its acceptance boundary is the positive root of rho=eta.
    sdir = -g / np.linalg.norm(g)
    eta = .25
    radii = np.geomspace(.03, 2.0, 400)
    aa = -(g @ sdir); bb = -.5 * (sdir @ H @ sdir)
    dd = -np.einsum('ijk,i,j,k', C, sdir, sdir, sdir) / 6
    q_sweep = {}
    for q in (.2, .8, 1.7, 4.0):
        observed = []
        for r in radii:
            dec = aa*r + bb*r*r + dd*r**3
            rho = (dec - q*r**4/4) / dec
            observed.append(rho >= eta)
        roots = np.roots([q/4, -(1-eta)*dd, -(1-eta)*bb, -(1-eta)*aa])
        pos = sorted(float(z.real) for z in roots if abs(z.imag)<1e-8 and z.real>1e-12)
        rstar = pos[0] if pos else float('nan')
        accepted = np.where(np.asarray(observed))[0]
        observed_boundary = float(radii[accepted[-1]]) if len(accepted) else 0.0
        q_sweep[str(q)] = {'predicted_radius': rstar,
                           'observed_grid_boundary': observed_boundary,
                           'relative_grid_error': abs(observed_boundary-rstar)/rstar,
                           'pattern_matches': bool(np.all(np.asarray(observed)==(radii<=rstar)))}
    rstar = q_sweep['1.7']['predicted_radius']
    first_rejected = float(radii[np.where(~np.asarray([((aa*r+bb*r*r+dd*r**3)-1.7*r**4/4)/(aa*r+bb*r*r+dd*r**3)>=eta for r in radii]))[0][0]])
    transition_matches = q_sweep['1.7']['pattern_matches']

    # Secondary same-objective comparison: fixed-step gradient descent vs guarded cubic-direction steps.
    def f(x): return model(c, g, H, C, x) + q * np.sum(x**2)**2 / 4
    def gradient(x): return g + H @ x + .5 * np.einsum('ijk,j,k->i', C, x, x) + q * np.sum(x*x) * x
    def safeguarded(x, radius):
        gg = gradient(x); ss = radius * (-gg / (np.linalg.norm(gg) + 1e-12))
        if model(c, g, H, C, x + ss) <= model(c, g, H, C, x) and f(x + ss) < f(x): return x + ss, True
        return x, False
    xg = np.zeros(n); xs = np.zeros(n); lr = .08; rejects = 0
    for _ in range(80):
        xg -= lr * gradient(xg)
        xs, ok = safeguarded(xs, .18); rejects += int(not ok)
    return {'prediction_1_representation_max_error': representation_max_error,
            'prediction_2_max_proximal_block_increase': max_block_increase,
            'prediction_2_final_recorded_block_objective': final_F,
            'prediction_2_beta_sweep': beta_sweep,
            'prediction_3_eta': eta, 'prediction_3_quartic_q': 1.7,
            'prediction_3_predicted_transition_radius': rstar,
            'prediction_3_first_rejected_grid_radius': first_rejected,
            'prediction_3_quartic_strength_sweep': q_sweep,
            'prediction_3_acceptance_patterns_match': transition_matches,
            'mini_final_objective_gradient': float(f(xg)),
            'mini_final_objective_safeguarded': float(f(xs)),
            'mini_safeguarded_rejections': rejects, 'mini_steps': 80}

if __name__ == '__main__':
    result = run()
    with open('results.json', 'w') as fp: json.dump(result, fp, indent=2)
    print(json.dumps(result, indent=2))