import json
import time
import numpy as np


def make_problem(seed=7, n=96, d=4, h=5):
    rng = np.random.default_rng(seed)
    x = rng.normal(size=(n, d))
    w1 = rng.normal(scale=0.8, size=(d, h))
    b1 = rng.normal(scale=0.2, size=h)
    w2 = rng.normal(scale=0.8, size=(h, 1))
    b2 = rng.normal(scale=0.2, size=1)
    y = np.tanh(x @ w1 + b1) @ w2 + b2
    y += 0.025 * rng.normal(size=(n, 1))
    return x, y


def unpack(q, d, h):
    k = 0
    w1 = q[k:k+d*h].reshape(d, h); k += d*h
    b1 = q[k:k+h]; k += h
    w2 = q[k:k+h].reshape(h, 1); k += h
    b2 = q[k:k+1]
    return w1, b1, w2, b2


def residual_jacobian(q, x, y, d, h):
    w1, b1, w2, b2 = unpack(q, d, h)
    z = np.tanh(x @ w1 + b1)
    r = (z @ w2 + b2 - y).ravel()
    n = x.shape[0]
    p = q.size
    j = np.zeros((n, p))
    a = 0
    b = d*h
    c = b+h
    e = c+h
    for i in range(n):
        j[i, a:b] = np.outer(x[i], (1-z[i]**2) * w2.ravel()).ravel()
        j[i, b:c] = (1-z[i]**2) * w2.ravel()
        j[i, c:e] = z[i]
        j[i, e] = 1.0
    return r, j


def state(q, x, y, d, h):
    r, j = residual_jacobian(q, x, y, d, h)
    f = 0.5 * float(r @ r)
    g = j.T @ r
    hess = j.T @ j
    return f, g, hess


def solve_direction(g, hess, lam):
    return np.linalg.solve(hess + lam*np.eye(g.size), -g)


def armijo_lm(q0, x, y, d, h, steps=35):
    q = q0.copy(); lam = 1e-2
    beta, c = 0.5, 1e-4
    low, high, up, down = 0.25, 0.75, 4.0, 2.0
    solves = rejects = accepted = trials = 0
    t0 = time.perf_counter()
    for _ in range(steps):
        f, g, H = state(q, x, y, d, h)
        p = solve_direction(g, H, lam); solves += 1
        pred = max(1e-14, -g @ p - 0.5*p @ H @ p)
        alpha = 1.0; accepted_here = False
        while alpha >= 1/32 - 1e-15:
            trials += 1
            ft = state(q + alpha*p, x, y, d, h)[0]
            if ft <= f + c*alpha*(g @ p):
                act = f-ft; rho = act/pred
                q = q + alpha*p; accepted += 1; accepted_here = True
                if rho < low: lam *= up
                elif rho > high: lam /= down
                break
            rejects += 1; alpha *= beta
        if not accepted_here:
            lam *= up
    final = state(q, x, y, d, h)[0]
    return dict(final_loss=final, solves=solves, rejected_trials=rejects,
                trials=trials, accepted=accepted, seconds=time.perf_counter()-t0)


def resolving_lm(q0, x, y, d, h, steps=35):
    """Conventional LM: only full steps are tested; rejection changes lambda
    and triggers a fresh curvature solve."""
    q = q0.copy(); lam = 1e-2
    low, high, up, down = 0.25, 0.75, 4.0, 2.0
    solves = rejects = accepted = trials = 0
    t0 = time.perf_counter()
    for _ in range(steps):
        f, g, H = state(q, x, y, d, h)
        accepted_here = False
        for _retry in range(12):
            p = solve_direction(g, H, lam); solves += 1
            pred = max(1e-14, -g @ p - 0.5*p @ H @ p)
            trials += 1
            ft = state(q + p, x, y, d, h)[0]
            if ft <= f + 1e-4*(g @ p):
                rho = (f-ft)/pred
                q = q + p; accepted += 1; accepted_here = True
                if rho < low: lam *= up
                elif rho > high: lam /= down
                break
            rejects += 1
            lam *= up
        if not accepted_here:
            lam *= up
    final = state(q, x, y, d, h)[0]
    return dict(final_loss=final, solves=solves, rejected_trials=rejects,
                trials=trials, accepted=accepted, seconds=time.perf_counter()-t0)


def math_check():
    rng = np.random.default_rng(123)
    p = 7
    A = rng.normal(size=(p,p)); H = A.T@A + 0.2*np.eye(p)
    g = rng.normal(size=p); direction = np.linalg.solve(H+0.7*np.eye(p), -g)
    f0 = 3.25
    def model(a): return f0 + a*g@direction + 0.5*a*a*direction@H@direction
    pred = -g@direction - 0.5*direction@H@direction
    model_reduction = f0-model(1.0)
    # Construct an exactly quadratic objective, so actual and model reductions agree.
    actual_reduction = f0-model(1.0)
    return dict(prediction_abs_error=abs(pred-model_reduction),
                ratio_abs_error=abs(actual_reduction/pred-1.0),
                pred_positive=bool(pred > 0))


def main():
    np.set_printoptions(precision=6, suppress=True)
    d, h = 4, 5
    x, y = make_problem()
    rng = np.random.default_rng(99)
    q0 = 0.4*rng.normal(size=d*h + h + h + 1)
    result = {'math_check': math_check(),
              'baseline_resolve_lm': resolving_lm(q0, x, y, d, h),
              'idea_decoupled_armijo': armijo_lm(q0, x, y, d, h)}
    print(json.dumps(result, indent=2))


if __name__ == '__main__':
    main()