"""Readable one-hidden-layer MLP, using only the Python standard library.

The website's fast JavaScript lab implements these same equations, data, PRNG
and full-batch gradient updates. No PyTorch, hidden model weights, or API calls.
Each row is one observation; the scalar version below makes each chain-rule
factor explicit. See the lesson for equivalent batched matrix equations.
"""
import math

DEFAULTS = dict(target="sin", activation="tanh", hidden=16, samples=64,
                noise=0.0, lr=0.03, epochs=2500, seed=7)


def normal_generator(seed):
    """Seeded Box-Muller samples; same 32-bit LCG as the browser experiment."""
    state = int(seed) & 0xffffffff

    def uniform():
        nonlocal state
        state = (1664525 * state + 1013904223) & 0xffffffff
        return (state + 0.5) / 4294967296

    while True:
        yield math.sqrt(-2 * math.log(uniform())) * math.cos(2 * math.pi * uniform())


def settings(parameters=None):
    c = dict(DEFAULTS, **(parameters or {}))
    if c["target"] not in ("sin", "sin2", "cubic"):
        raise ValueError("Unknown target")
    if c["activation"] not in ("tanh", "relu", "sigmoid", "linear"):
        raise ValueError("Unknown activation")
    for key, lo, hi in (("hidden", 1, 64), ("samples", 4, 128),
                        ("epochs", 0, 6000), ("seed", 0, 4294967295)):
        if not isinstance(c[key], int) or isinstance(c[key], bool) or not lo <= c[key] <= hi:
            raise ValueError("Invalid " + key)
    if c["samples"] % 2:
        raise ValueError("Use an even sample count so validation points stay distinct")
    for key, lo, hi in (("noise", 0, 0.5), ("lr", 0.0001, 2)):
        if not math.isfinite(c[key]) or not lo <= c[key] <= hi:
            raise ValueError("Invalid " + key)
    return c


def target(x, name="sin"):
    return math.sin(2*x) if name == "sin2" else (x/math.pi)**3 if name == "cubic" else math.sin(x)


def activation(z, name="tanh"):
    if name == "relu":
        return max(0.0, z)
    if name == "linear":
        return z
    if name == "sigmoid":
        return 1/(1+math.exp(-z)) if z >= 0 else math.exp(z)/(1+math.exp(z))
    return math.tanh(z)


def derivative(z, h, name):
    # ReLU is nondifferentiable at zero; this implementation chooses 0 there.
    return float(z > 0) if name == "relu" else 1.0 if name == "linear" else h*(1-h) if name == "sigmoid" else 1-h*h


def initialize(c):
    normal = normal_generator(c["seed"])
    return dict(w1=[0.25*next(normal) for _ in range(c["hidden"])],
                b1=[0.0]*c["hidden"],
                w2=[0.25*next(normal) for _ in range(c["hidden"])], b2=0.0)


def copy_params(p):
    return dict(w1=p["w1"][:], b1=p["b1"][:], w2=p["w2"][:], b2=p["b2"])


def features(x, p, name):
    return [activation(w*x+b, name) for w, b in zip(p["w1"], p["b1"])]


def predict(x, p, name):
    return p["b2"] + sum(w*h for w, h in zip(p["w2"], features(x, p, name)))


def dataset(c):
    """Training: noisy endpoints grid. Validation: distinct clean midpoints.

    This synthetic validation is an interpolation demonstration, not a held-out
    assessment of a real power-system predictor. Noise is added only to training.
    """
    normal = normal_generator(c["seed"] ^ 0x9e3779b9)
    x = [-math.pi + 2*math.pi*i/(c["samples"]-1) for i in range(c["samples"])]
    y = [target(v, c["target"]) + c["noise"]*next(normal) for v in x]
    vx = [-math.pi + 2*math.pi*(i+0.5)/160 for i in range(160)]
    return x, y, vx, [target(v, c["target"]) for v in vx]


def mse(x, y, p, name):
    return sum((predict(v, p, name)-truth)**2 for v, truth in zip(x, y))/len(x)


def gradient(x, y, p, name):
    """Backpropagation: dL/dy -> dL/dh -> dL/dz -> parameter gradients."""
    g = dict(w1=[0.0]*len(p["w1"]), b1=[0.0]*len(p["w1"]),
             w2=[0.0]*len(p["w1"]), b2=0.0)
    for xi, yi in zip(x, y):
        h = features(xi, p, name)
        y_hat = p["b2"] + sum(w*hj for w, hj in zip(p["w2"], h))
        d = 2*(y_hat-yi)/len(x)  # MSE gradient includes the batch average.
        g["b2"] += d
        for j, hj in enumerate(h):
            dz = d*p["w2"][j]*derivative(p["w1"][j]*xi+p["b1"][j], hj, name)
            g["w2"][j] += d*hj
            g["w1"][j] += dz*xi
            g["b1"][j] += dz
    return g


def step(p, g, lr):
    # Compute every gradient before updating any weight.
    for key in ("w1", "b1", "w2"):
        for j in range(len(p[key])):
            p[key][j] -= lr*g[key][j]
    p["b2"] -= lr*g["b2"]


def linear_fit(x, y):
    mx, my = sum(x)/len(x), sum(y)/len(y)
    slope = sum((v-mx)*(t-my) for v, t in zip(x, y))/sum((v-mx)**2 for v in x)
    return dict(slope=slope, intercept=my-slope*mx)


def train(parameters=None):
    """Generate data, initialize, optimize, and evaluate the completed updates.

    Frames retain parameter snapshots for the website's epoch replay. A numeric
    divergence stops the run and is explicitly reported; a finite loss alone
    does not establish convergence or successful generalization.
    """
    c = settings(parameters)
    x, y, vx, vy = dataset(c)
    p = initialize(c)
    every = max(1, math.ceil(c["epochs"]/80))
    frames, completed, diverged_epoch = [], 0, None

    def record(epoch):
        frames.append(dict(epoch=epoch, train_mse=mse(x, y, p, c["activation"]),
                           validation_mse=mse(vx, vy, p, c["activation"]), params=copy_params(p)))

    record(0)
    for epoch in range(1, c["epochs"]+1):
        candidate = copy_params(p)
        step(candidate, gradient(x, y, p, c["activation"]), c["lr"])
        loss = mse(x, y, candidate, c["activation"])
        if not math.isfinite(loss) or loss > 1e12:
            diverged_epoch = epoch
            break
        p, completed = candidate, epoch
        if epoch % every == 0 or epoch == c["epochs"]:
            record(epoch)
    if frames[-1]["epoch"] != completed:
        record(completed)
    linear = linear_fit(x, y)
    grid = [-1.25*math.pi+2.5*math.pi*i/200 for i in range(201)]
    return dict(ok=diverged_epoch is None, case=c, completed_epochs=completed,
                diverged_epoch=diverged_epoch, parameter_count=3*c["hidden"]+1,
                params=p, frames=frames, train_mse=frames[-1]["train_mse"],
                validation_mse=frames[-1]["validation_mse"],
                linear_validation_mse=sum((linear["slope"]*v+linear["intercept"]-t)**2 for v, t in zip(vx, vy))/len(vx),
                linear=linear, train=dict(x=x, y=y), grid=grid,
                target=[target(v, c["target"]) for v in grid],
                prediction=[predict(v, p, c["activation"]) for v in grid])