"""Single-channel float64-style reference; no third-party dependencies.
Run: python3 docs/ai-infra/causal-convolution/reference.py
This checks mathematical derivatives, not Dao CUDA or MLU kernels.
"""
import math
import random


def sigmoid(z):
    if z >= 0:
        return 1.0 / (1.0 + math.exp(-z))
    e = math.exp(z)
    return e / (1.0 + e)


def forward(x, w, bias, state, silu=False):
    k, length = len(w), len(x)
    assert len(state) == k - 1 and length > 0
    u = state + x
    z = [bias + sum(w[j] * u[t + j] for j in range(k))
         for t in range(length)]
    y = [v * sigmoid(v) if silu else v for v in z]
    return y, u[length:length + k - 1], z


def backward(x, w, bias, state, g, dfinal, silu=False):
    _, _, z = forward(x, w, bias, state, silu)
    u = state + x
    du, dw = [0.0] * len(u), [0.0] * len(w)
    delta = []
    for grad, value in zip(g, z):
        sig = sigmoid(value)
        delta.append(grad * sig * (1 + value * (1 - sig)) if silu else grad)
    for t, grad in enumerate(delta):
        for j, weight in enumerate(w):
            du[t + j] += grad * weight
            dw[j] += grad * u[t + j]
    for r, grad in enumerate(dfinal):
        du[len(x) + r] += grad
    split = len(w) - 1
    return du[split:], dw, sum(delta), du[:split]


def check():
    y, _, _ = forward([1., 2., 3., 4.], [2., 3., 5.], 1., [0., 0.])
    assert y == [6., 14., 24., 34.]
    dx, dw, db, _ = backward([1., 2., 3., 4.], [2., 3., 5.], 1.,
                              [0., 0.], [1., 2., 3., 4.], [0., 0.])
    assert (dx, dw, db) == ([17., 27., 27., 20.], [11., 20., 30.], 10.)
    rng = random.Random(7)
    worst, cases = 0., 0
    for k in (2, 3, 4):
        for length in sorted({1, k - 1, k, k + 1, 9}):
            for silu in (False, True):
                rand = lambda n: [rng.uniform(-0.8, 0.8) for _ in range(n)]
                x, w, bias, state = rand(length), rand(k), rand(1), rand(k - 1)
                g, df = rand(length), rand(k - 1)
                dx, dw, db, ds = backward(x, w, bias[0], state, g, df, silu)
                def loss():
                    y, final, _ = forward(x, w, bias[0], state, silu)
                    return sum(a*b for a,b in zip(g,y)) + sum(a*b for a,b in zip(df,final))
                for values, grads in ((x, dx), (w, dw), (bias, [db]), (state, ds)):
                    for i, expected in enumerate(grads):
                        old, eps = values[i], 1e-6
                        values[i] = old + eps
                        plus = loss()
                        values[i] = old - eps
                        minus = loss()
                        values[i] = old
                        err = abs((plus-minus)/(2*eps) - expected)
                        worst = max(worst, err)
                        assert err < 2e-8, (k, length, silu, err)
                # Full sequence versus one-token updates.
                full, final, _ = forward(x, w, bias[0], state, silu)
                streamed, carry = [], state[:]
                for value in x:
                    out, carry, _ = forward([value], w, bias[0], carry, silu)
                    streamed.extend(out)
                assert all(abs(a-b) < 1e-12 for a,b in zip(full,streamed))
                assert carry == final
                # Two chunks: carry dS of the second into dF of the first.
                if length > 1:
                    split = max(1, length // 2)
                    _, middle, _ = forward(x[:split], w, bias[0], state, silu)
                    dx2, dw2, db2, ds2 = backward(x[split:], w, bias[0], middle,
                                                  g[split:], df, silu)
                    dx1, dw1, db1, ds1 = backward(x[:split], w, bias[0], state,
                                                  g[:split], ds2, silu)
                    pairs = [(dx, dx1+dx2), (dw, [a+b for a,b in zip(dw1,dw2)]),
                             ([db], [db1+db2]), (ds, ds1)]
                    assert all(abs(a-b) < 1e-12 for p,q in pairs for a,b in zip(p,q))
                cases += 1
    print(f"PASS: hand example, {cases} gradient/state cases; max finite-difference error {worst:.3e}")


if __name__ == "__main__":
    check()