import json, math, time
from pathlib import Path
import numpy as np

# Minimal implicit branch with the singular endpoint claimed in the idea:
# h(z,t)=z^2-t=0, positive branch z=sqrt(t), whose Puiseux exponent is 1/2.
def h(z, t):
    return z*z - t

def dh(z, t):
    return 2.0*z, -1.0

def newton_fixed_t(z0, t, maxit=12, tol=1e-11):
    z = float(z0)
    for k in range(1, maxit + 1):
        r = h(z, t)
        if abs(r) < tol:
            return z, k, True
        j = 2.0*z
        if abs(j) < 1e-14:
            return z, k, False
        z -= r/j
        # Stay on the selected positive branch; a branch switch is a failed solve.
        if not np.isfinite(z) or z <= 0:
            return z, k, False
    return z, maxit, abs(h(z, t)) < tol

def pseudo_newton(z0, t0, zpred, tpred, zprev, tprev, maxit=15, tol=1e-11):
    """Newton on h=0 and the pseudo-arclength hyperplane through predictor."""
    z, t = float(z0), float(t0)
    dz, dt = z0-zprev, t0-tprev
    norm = math.hypot(dz, dt)
    qz, qt = dz/norm, dt/norm
    for k in range(1, maxit + 1):
        F = np.array([h(z,t), qz*(z-zpred) + qt*(t-tpred)])
        if np.linalg.norm(F, ord=np.inf) < tol:
            return z, t, k, True
        J = np.array([[2*z, -1.0], [qz, qt]])
        try:
            step = np.linalg.solve(J, -F)
        except np.linalg.LinAlgError:
            return z, t, k, False
        # Mild damping makes the correction robust at the nearly singular end.
        damp = 1.0
        old = np.linalg.norm(F, ord=np.inf)
        accepted = False
        for _ in range(12):
            zn, tn = z + damp*step[0], t + damp*step[1]
            if tn > 0 and np.isfinite(zn) and np.isfinite(tn):
                nn = np.linalg.norm([h(zn,tn), qz*(zn-zpred)+qt*(tn-tpred)], ord=np.inf)
                if nn <= old or nn < 1e-12:
                    z,t = zn,tn; accepted=True; break
            damp *= .5
        if not accepted:
            return z, t, k, False
    return z, t, maxit, np.linalg.norm([h(z,t), qz*(z-zpred)+qt*(t-tpred)], ord=np.inf) < tol

def puiseux_predict(hist_t, hist_z, tnew):
    """Empirical exponent fit plus regularized linear coefficient fit."""
    ts, zs = np.asarray(hist_t), np.asarray(hist_z)
    # Differences scale as t^alpha; log fit is the implementation-plan estimator.
    d = np.abs(np.diff(zs)); tm = np.sqrt(ts[1:]*ts[:-1])
    good = (d > 1e-15) & (tm > 0)
    if good.sum() >= 2:
        alpha = np.polyfit(np.log(tm[good]), np.log(d[good]), 1)[0]
    else:
        alpha = .5
    alpha = float(np.clip(alpha, .1, 4.0))
    # Include the fitted exponent and a constant term. Ridge is only for numerical safety.
    X = np.column_stack([np.ones(len(ts)), ts**alpha])
    reg = 1e-12*np.eye(2)
    coef = np.linalg.solve(X.T@X + reg, X.T@zs)
    return float(coef[0] + coef[1]*tnew**alpha), alpha

def coefficient_identity_check(rng):
    # Directly verify the displayed 2-point value/derivative coefficient system
    # for a0+a1*t^alpha+a2*t^beta, using exact derivatives.
    a = np.array([.31, 1.7, -.42]); alpha, beta = .5, 1.5
    t1, t2 = .37, .19; lam=t2/t1
    M = np.array([[1,1,1],[0,alpha,beta], [1,lam**alpha,lam**beta],
                  [0,alpha*lam**alpha,beta*lam**beta]], float)
    rhs = np.array([a[0]+a[1]*t1**alpha+a[2]*t1**beta,
                    t1*(a[1]*alpha*t1**(alpha-1)+a[2]*beta*t1**(beta-1)),
                    a[0]+a[1]*t2**alpha+a[2]*t2**beta,
                    t2*(a[1]*alpha*t2**(alpha-1)+a[2]*beta*t2**(beta-1))])
    recovered=np.linalg.lstsq(M, rhs, rcond=None)[0]
    expected=np.array([a[0], a[1]*t1**alpha, a[2]*t1**beta])
    return float(np.max(np.abs(M@recovered-rhs))), float(np.max(np.abs(recovered-expected)))

def run(method, ts):
    z=math.sqrt(ts[0]); hist_t=[ts[0]]; hist_z=[z]
    iters=[]; fails=0; pred_err=[]; residuals=[]; accepted=1; exponents=[]
    for target in ts[1:]:
        if method == 'linear':
            slope=(hist_z[-1]-hist_z[-2])/(hist_t[-1]-hist_t[-2]) if len(hist_t)>1 else -1/(2*z)
            pred=hist_z[-1]+slope*(target-hist_t[-1])
            zn, nit, ok=newton_fixed_t(pred,target)
            tn=target
        else:
            pred, alpha=puiseux_predict(hist_t[-4:],hist_z[-4:],target)
            exponents.append(alpha)
            zn, tn, nit, ok=pseudo_newton(pred,target,pred,target,hist_z[-1],hist_t[-1])
        pred_err.append(abs(pred-math.sqrt(target)))
        iters.append(nit)
        if not ok:
            fails += 1
            # reject: retrying with a smaller step is represented by failure, then use exact
            # previous accepted point to keep this controlled benchmark progressing.
            zn=math.sqrt(target); tn=target
        else:
            accepted += 1
        hist_t.append(tn); hist_z.append(zn)
        residuals.append(abs(h(zn,tn)))
    return dict(failed=fails, accepted=accepted, mean_iterations=float(np.mean(iters)),
                max_iterations=int(max(iters)), mean_prediction_error=float(np.mean(pred_err)),
                final_residual=float(residuals[-1]), max_residual=float(max(residuals)),
                fitted_alpha_mean=(float(np.mean(exponents)) if exponents else None),
                fitted_alpha_last=(float(exponents[-1]) if exponents else None))

def main():
    rng=np.random.default_rng(0)
    ident=coefficient_identity_check(rng)
    # Geometric schedule concentrates many equal-relative steps near singularity.
    ts=np.geomspace(1.0,1e-8,81)
    t0=time.perf_counter(); baseline=run('linear',ts); tb=time.perf_counter()-t0
    t0=time.perf_counter(); idea=run('puiseux',ts); ti=time.perf_counter()-t0
    out={'system':'h(z,t)=z^2-t, positive branch z=sqrt(t)', 'seed':0,
         'schedule_steps':len(ts)-1, 'coefficient_check_max_equation_error':ident[0],
         'coefficient_check_max_parameter_error':ident[1], 'baseline':baseline,
         'puiseux_arclength':idea, 'wall_seconds':{'baseline':tb,'idea':ti}}
    Path('results.json').write_text(json.dumps(out,indent=2))
    print(json.dumps(out,indent=2))
if __name__=='__main__': main()