import json, math, time
import numpy as np
from scipy.optimize import minimize

SEED = 7
rng = np.random.default_rng(SEED)
D = 24
xtrue = np.zeros(D)
xtrue[[1, 6, 14, 20]] = [2.2, -1.7, 1.3, -2.5]
y = xtrue + 0.35 * rng.standard_normal(D)
lam, eps, c, tau = 0.34, 1e-4, 0.72, 0.08

def softplus(z):
    return np.logaddexp(0.0, z)

def sigmoid(z):
    return 1.0 / (1.0 + np.exp(-np.clip(z, -50, 50)))

def g(x):
    return 0.5 * np.sum((x-y)**2) + lam * np.sum(np.sqrt(x*x + eps))

def h(x):
    return lam * tau * np.sum(softplus((np.abs(x)-c)/tau))

def phi(x):
    return g(x) - h(x)

def gg(x):
    return x-y + lam*x/np.sqrt(x*x + eps)

def hh(x):
    return lam * sigmoid((np.abs(x)-c)/tau) * np.sign(x)

def rawgrad(x):
    return gg(x) - hh(x)

def penalty_grad(z, s, p, gamma):
    d = z-s
    n = np.linalg.norm(d)
    if n < 1e-14:
        return np.zeros_like(z)
    return d * n**(p-2) / gamma

def prox_obj(z, s, which, p, gamma):
    base = g(z) if which == 'g' else h(z)
    return base + np.linalg.norm(z-s)**p / (p*gamma)

def exact_prox(s, which, p, gamma):
    basegrad = gg if which == 'g' else hh
    r = minimize(lambda z: prox_obj(z, s, which, p, gamma), s.copy(),
                 jac=lambda z: basegrad(z) + penalty_grad(z, s, p, gamma),
                 method='L-BFGS-B', options={'maxiter': 1000, 'ftol': 1e-13, 'gtol': 1e-10})
    return r.x, r

def inexact_prox(s, which, p, gamma, K=8):
    z = s.copy()
    basegrad = gg if which == 'g' else hh
    # Fixed conservative steps make the inexactness explicit and reproducible.
    step = 0.08 if p == 2 else 0.035
    for _ in range(K):
        z -= step * (basegrad(z) + penalty_grad(z, s, p, gamma))
    return z

def moreau_grad(s, p, gamma, exact=False, K=8):
    if exact:
        u, _ = exact_prox(s, 'g', p, gamma)
        v, _ = exact_prox(s, 'h', p, gamma)
    else:
        u = inexact_prox(s, 'g', p, gamma, K)
        v = inexact_prox(s, 'h', p, gamma, K)
    ag = penalty_grad(s, u, p, gamma)  # sign is s-u
    ah = penalty_grad(s, v, p, gamma)
    return ag-ah, u, v

def run(method, steps=180, eta=0.18, gamma0=0.7, p=2, K=8):
    s = np.zeros(D)
    vals, norms, spikes = [], [], 0
    for t in range(steps):
        gamma = gamma0 * (0.985 ** t)
        if method == 'raw':
            grad = rawgrad(s)
        else:
            grad, _, _ = moreau_grad(s, p, gamma, exact=False, K=K)
        n = np.linalg.norm(grad)
        if n > 10: spikes += 1
        s -= eta * grad
        vals.append(float(phi(s))); norms.append(float(n))
        if not np.all(np.isfinite(s)): break
    return {'final_phi': vals[-1], 'best_phi': min(vals),
            'grad_rms': float(np.sqrt(np.mean(np.array(norms)**2))),
            'spikes_gt10': spikes, 'trajectory': vals, 'norms': norms}

def math_check():
    s = rng.normal(size=D)
    out = {}
    for p in (2, 4):
        gamma = .43
        u, ru = exact_prox(s, 'g', p, gamma)
        v, rv = exact_prox(s, 'h', p, gamma)
        ag = penalty_grad(s, u, p, gamma)
        ah = penalty_grad(s, v, p, gamma)
        rg = np.linalg.norm(gg(u) + penalty_grad(u, s, p, gamma))
        rh = np.linalg.norm(hh(v) + penalty_grad(v, s, p, gamma))
        # Envelope finite differences in a random direction.
        d = rng.normal(size=D); d /= np.linalg.norm(d); delta = 1e-5
        ep = prox_obj(u, s+delta*d, 'g', p, gamma)  # objective at old prox: upper check only
        um, _ = exact_prox(s-delta*d, 'g', p, gamma)
        up, _ = exact_prox(s+delta*d, 'g', p, gamma)
        em = prox_obj(um, s-delta*d, 'g', p, gamma)
        eplus = prox_obj(up, s+delta*d, 'g', p, gamma)
        fd = (eplus-em)/(2*delta)
        out[str(p)] = {'g_residual': float(rg), 'h_residual': float(rh),
                       'gradient_fd_abs_error': float(abs(fd-np.dot(ag,d))),
                       'prox_success': bool(ru.success and rv.success)}
    # Claimed lifted critical interval for g=|x|, h=2|x|: report endpoint scaling.
    out['critical_interval'] = {str(p): float(.43**(p/(p-1)-1)) for p in (2,4)}
    return out

def main():
    check = math_check()
    results = {'math_check': check, 'seed': SEED,
               'settings': {'dimension': D, 'steps': 180, 'eta': .18, 'gamma0': .7, 'K': 8}}
    for name, p in [('raw', 2), ('moreau_p2', 2), ('moreau_p4', 4)]:
        results[name] = run(name, p=p)
    with open('results.json', 'w') as f:
        json.dump(results, f, indent=2)
    summary = {'math_check': check}
    for name in ('raw', 'moreau_p2', 'moreau_p4'):
        summary[name] = {m: results[name][m] for m in ('final_phi','best_phi','grad_rms','spikes_gt10')}
    print(json.dumps(summary, indent=2))

if __name__ == '__main__':
    main()