# Symbolic metamodeling by continuous optimization: the Meijer G-function
# trick of Alaa & van der Schaar (NeurIPS 2019), in Knot.
#
# One parametric family — the truncated series of G^{1,2}_{2,2}(a,a; a,b | cx),
#
#   f(a,b,c; x) = (cx)^a * sum_{k=0..3} (-1)^k k! (cx)^k / Gamma(a - b + k + 1)
#
# — contains log(1+x) at (a,b,c) = (1,0,1) and x/(1+x) at (1,1,1) as exact
# parameter settings. Because (a,b,c) are continuous, "which function is
# this?" becomes gradient descent, not a search over expression trees: the
# opposite strategy to the brute-force enumeration in symbolic_regression.knot.
#
# The gradient formulas are a faithful transliteration of the paper code's
# basis_grad (digamma/Gamma expressions of the truncated series; verified
# against scipy to 9 digits). Gamma and digamma are hand-rolled below —
# Lanczos approximation and shifted asymptotic series respectively.
#
# Output protocol (labeled `show` lines): t_id, then (tr_it, tr_loss) trace
# pairs, then fin_a/fin_b/fin_c/fin_rms, then t_done.

NPTS = 28

# ── Gamma via Lanczos (g = 7, 9 coefficients), digamma via shift + asymptotic ─
lz = zeros(9)
lz[0] = 0.99999999999980993
lz[1] = 676.5203681218851
lz[2] = -1259.1392167224028
lz[3] = 771.32342877765313
lz[4] = -176.61502916214059
lz[5] = 12.507343278686905
lz[6] = -0.13857109526572012
lz[7] = 0.0000099843695780195716
lz[8] = 0.00000015056327351493116

def gammaf(z) {
    if z < 0.5 {
        # reflection: Gamma(z) Gamma(1-z) = pi / sin(pi z)
        return 3.141592653589793 / (sin(3.141592653589793 * z) * gammaf(1.0 - z))
    }
    w = z - 1.0
    s = lz[0]
    for i to 8 {
        s = s + lz[i + 1.0] / (w + i + 1.0)
    }
    t = w + 7.5
    return 2.5066282746310002 * pow(t, w + 0.5) * exp(0.0 - t) * s
}

def digammaf(z) {
    s = 0.0
    w = z
    while w < 6.0 {
        s = s - 1.0 / w
        w = w + 1.0
    }
    w2 = 1.0 / (w * w)
    return s + log(w) - 0.5 / w - w2 * (0.083333333333333333 - w2 * (0.0083333333333333333 - w2 * 0.0039682539682539683))
}

# ── the basis and its closed-form gradients (transliterated basis_grad) ──────
# writes [value, d/da, d/db, d/dc] into out
out = zeros(4)
def basis_and_grad(a, b, c, x) {
    G1 = gammaf(a - b + 1.0)
    G2 = gammaf(a - b + 2.0)
    G3 = gammaf(a - b + 3.0)
    G4 = gammaf(a - b + 4.0)
    K1 = digammaf(a - b + 1.0)
    K2 = digammaf(a - b + 2.0)
    K3 = digammaf(a - b + 3.0)
    K4 = digammaf(a - b + 4.0)

    cx = c * x
    lcx = log(cx)
    p = pow(cx, a)
    cx2 = cx * cx
    cx3 = cx2 * cx

    out[0] = p * (1.0 / G1 - cx / G2 + 2.0 * cx2 / G3 - 6.0 * cx3 / G4)
    out[1] = p * (6.0 * cx3 * (K4 - lcx) / G4 + 2.0 * cx2 * (lcx - K3) / G3 + cx * (K2 - lcx) / G2 - (K1 - lcx) / G1)
    out[2] = p * (0.0 - 6.0 * cx3 * K4 / G4 + 2.0 * cx2 * K3 / G3 - cx * K2 / G2 + K1 / G1)
    out[3] = p * (0.0 - c * c * x * x * x * (6.0 * a + 18.0) / G4 + c * x * x * (4.0 + 2.0 * a) / G3 - x * (1.0 + a) / G2 + (a / c) / G1)
    return 0.0
}

# ── targets: sampled on x in [0.05, 0.59], hidden from the fit ───────────────
xs = zeros(32)
ys = zeros(32)

def atan_series(x) {
    # arctan for |x| < 1 by the alternating series, 11 terms (x^23/23 ~ 1e-8 here)
    s = 0.0
    t = x
    x2 = x * x
    for k to 11 {
        s = s + t / (2.0 * k + 1.0)
        t = 0.0 - t * x2
    }
    return s
}

def gen_target(t) {
    for i to NPTS {
        # keep cx well inside the series' radius so truncation error can't
        # drag the optimum away from the catalog point
        x = 0.02 + 0.01 * i
        xs[i] = x
        if t == 1.0 { ys[i] = log(1.0 + x) }
        if t == 2.0 { ys[i] = x / (1.0 + x) }
        if t == 3.0 { ys[i] = atan_series(x) }
    }
}

def rms_at(a, b, c) {
    sse = 0.0
    for p to NPTS {
        basis_and_grad(a, b, c, xs[p])
        r = out[0] - ys[p]
        sse = sse + r * r
    }
    return sqrt(sse / NPTS)
}

# ── Adam on the full-batch MSE (deterministic; the paper used plain SGD, but
#    the loss surface has a long a-vs-b valley that raw gradients crawl in) ──
def fit(t_id) {
    gen_target(t_id)
    show t_id

    # the paper's initialization
    a = 2.0
    b = 1.0
    c = 1.0
    lr = 0.02
    rms = 0.0

    ma = 0.0
    va = 0.0
    mb = 0.0
    vb = 0.0
    mc = 0.0
    vc = 0.0
    b1t = 1.0
    b2t = 1.0

    for it to 12001 {
        if it == 6000.0 { lr = 0.006 }
        ga = 0.0
        gb = 0.0
        gc = 0.0
        sse = 0.0
        for p to NPTS {
            basis_and_grad(a, b, c, xs[p])
            r = out[0] - ys[p]
            sse = sse + r * r
            ga = ga + 2.0 * r * out[1] / NPTS
            gb = gb + 2.0 * r * out[2] / NPTS
            gc = gc + 2.0 * r * out[3] / NPTS
        }
        rms = sqrt(sse / NPTS)

        rem = it - 1000.0 * floor(it / 1000.0)
        if rem == 0.0 {
            tr_it = it
            tr_loss = rms
            show tr_it
            show tr_loss
        }

        b1t = b1t * 0.9
        b2t = b2t * 0.999
        ma = 0.9 * ma + 0.1 * ga
        va = 0.999 * va + 0.001 * ga * ga
        mb = 0.9 * mb + 0.1 * gb
        vb = 0.999 * vb + 0.001 * gb * gb
        mc = 0.9 * mc + 0.1 * gc
        vc = 0.999 * vc + 0.001 * gc * gc
        a = a - lr * (ma / (1.0 - b1t)) / (sqrt(va / (1.0 - b2t)) + 0.00000001)
        b = b - lr * (mb / (1.0 - b1t)) / (sqrt(vb / (1.0 - b2t)) + 0.00000001)
        c = c - lr * (mc / (1.0 - b1t)) / (sqrt(vc / (1.0 - b2t)) + 0.00000001)

        # keep Gamma arguments and the series radius healthy
        if a - b + 1.0 < 0.2 { b = a + 0.8 }
        if c < 0.05 { c = 0.05 }
        if c > 1.5 { c = 1.5 }
    }

    fin_a = a
    fin_b = b
    fin_c = c
    fin_rms = rms
    show fin_a
    show fin_b
    show fin_c
    show fin_rms

    # identification by snapping: project onto each catalog point and see
    # whether the snapped model still fits. The parameters themselves sit in
    # a flat valley (b and c trade off on a short interval); function-space
    # distance is the identifiable thing.
    sn_log = rms_at(1.0, 0.0, 1.0)
    sn_rat = rms_at(1.0, 1.0, 1.0)
    show sn_log
    show sn_rat

    t_done = t_id
    show t_done
}

# ── self-tests: hand-rolled specials + the transliterated gradients ──────────
test "gamma matches scipy references" {
    assert_near(gammaf(0.3), 2.991568987688, 0.000001)
    assert_near(gammaf(1.0), 1.0, 0.0000001)
    assert_near(gammaf(2.3), 1.166711905198, 0.0000001)
    assert_near(gammaf(4.7), 15.431411600047, 0.000001)
}

test "digamma matches scipy references" {
    assert_near(digammaf(0.3), -3.502524222200, 0.000001)
    assert_near(digammaf(1.0), -0.577215664902, 0.0000001)
    assert_near(digammaf(2.3), 0.600039880364, 0.0000001)
    assert_near(digammaf(4.7), 1.437423809632, 0.0000001)
}

test "basis + gradients match the paper code at (1.7, 0.6, 0.9, x=0.35)" {
    basis_and_grad(1.7, 0.6, 0.9, 0.35)
    assert_near(out[0], 0.117121456687, 0.00000001)
    assert_near(out[1], -0.184844802850, 0.00000001)
    assert_near(out[2], 0.049548129295, 0.00000001)
    assert_near(out[3], 0.204824209757, 0.00000001)
}

test "catalog identities at the special parameter points" {
    # truncated log(1+x) series at x = 0.2: x - x^2/2 + x^3/3 - x^4/4
    basis_and_grad(1.0, 0.0, 1.0, 0.2)
    assert_near(out[0], 0.18226666666667, 0.0000001)
    # truncated x/(1+x) series at x = 0.2: x (1 - x + x^2 - x^3)
    basis_and_grad(1.0, 1.0, 1.0, 0.2)
    assert_near(out[0], 0.1664, 0.0000001)
}

fit(1.0)
fit(2.0)
fit(3.0)