"""Specker-type gradient descent (Theorem 4.1 of the paper).

Data: a nondecreasing sequence y_s in [0,1): a summable background
b_s = 1e-4 (s+1)^{-3/2} (keeps it strictly increasing) plus jumps 2^{-m-1}
released at stages 0, 30, 300, 3000, 30000 (m = 1, 3, 4, 5, 6), imitating a
c.e. set whose small elements are enumerated late.

Function (used in the proofs of Theorems 4.1 and 4.2):
    f(x) = sum_s 2^{-s-1} ((y_s - x)_+)^2 + ((x - 1)_+)^2 ,
convex, C^1, f' 2-Lipschitz, argmin f = [alpha, 1], alpha = sup_s y_s.
Gradient descent with step 1/L, L = 2, from x_0 = 0 increases to alpha.

Late data enter f' with weight 2^{-s}. 80-digit decimal arithmetic; the sum over
s is truncated at S_TRUNC = 210 (neglected terms < 2^{-210} < 1e-63).
Outputs: x_k, |f'(x_k)|, f(x_k) - f* on a log grid, alpha.
"""
from decimal import Decimal, getcontext
import json
import math
import sys

getcontext().prec = 80
S_TRUNC = 210
LATE = {0: 1, 30: 3, 300: 4, 3000: 5, 30000: 6}
L = Decimal(2)


def data(S):
    q, acc = [], Decimal(0)
    for s in range(S):
        b = Decimal(10) ** -4 / (Decimal(s + 1) ** Decimal("1.5"))
        if s in LATE:
            b += Decimal(2) ** -(LATE[s] + 1)
        acc += b
        q.append(acc)
    return q


def main():
    K = int(sys.argv[1]) if len(sys.argv) > 1 else 100000
    out = sys.argv[2] if len(sys.argv) > 2 else "specker_gd.json"
    S_DATA = 100001
    q = data(S_DATA)
    tail = Decimal("2e-4") / Decimal(S_DATA).sqrt()     # background tail beyond S_DATA
    alpha = q[-1] + tail
    qq = q[:S_TRUNC]
    w = [Decimal(2) ** -s for s in range(S_TRUNC)]       # f'(x) = -sum_s 2^{-s} (y_s - x)_+

    def fprime(x):
        return -sum(wi * (qi - x) for wi, qi in zip(w, qq) if qi > x)

    def fval(x):
        return sum(wi / 2 * (qi - x) ** 2 for wi, qi in zip(w, qq) if qi > x)

    xs, grads = [Decimal(0)], []
    x = Decimal(0)
    for k in range(K):
        g = fprime(x)
        grads.append(float(-g))
        x = x - g / L
        xs.append(x)
    grid = sorted(set([0] + [int(round(10 ** (t / 40))) for t in range(0, 40 * int(math.log10(K)) + 1)]))
    grid = [k for k in grid if k <= K]
    gaps = {k: float(fval(xs[k])) for k in grid}
    D = float(alpha)
    worst = max(gaps[k] * 2 * k / (float(L) * D * D) for k in grid if k > 0)
    res = {"K": K, "alpha": float(alpha), "D": D,
           "x": [float(v) for v in xs], "q": [float(v) for v in q[:K + 1]],
           "grad": grads, "gap_k": list(gaps.keys()), "gap": list(gaps.values()),
           "late": {str(k): v for k, v in LATE.items()}}
    with open(out, "w") as fh:
        json.dump(res, fh)
    print(f"alpha={float(alpha):.12f} x_K={float(xs[-1]):.12f} dist_K={float(alpha - xs[-1]):.4e} "
          f"grad_K={grads[-1]:.3e} gap_K={gaps[grid[-1]]:.3e} max gap_k*2k/(L D^2)={worst:.4f}")


if __name__ == "__main__":
    main()