#!/usr/bin/env python3
"""Bab 9 — Neural networks from scratch.

Portable standard-library lab: activations, perceptrons, XOR network,
forward pass, gradients, training loop, and SVG decision boundary.
"""
from __future__ import annotations

import math
from pathlib import Path
from typing import Sequence


def sigmoid(z: float) -> float:
    return 1 / (1 + math.exp(-z))


def relu(z: float) -> float:
    return max(0.0, z)


def step(z: float) -> int:
    return 1 if z >= 0 else 0


def dot(a: Sequence[float], b: Sequence[float]) -> float:
    return sum(x * y for x, y in zip(a, b))


def perceptron(x: Sequence[float], w: Sequence[float], b: float) -> int:
    return step(dot(x, w) + b)


def xor_network(x1: int, x2: int) -> int:
    h1 = step(1 * x1 + 1 * x2 - 0.5)       # OR
    h2 = step(-1 * x1 + -1 * x2 + 1.5)     # NAND
    return step(1 * h1 + 1 * h2 - 1.5)     # AND


def sigmoid_xor_probability(x1: int, x2: int) -> float:
    h1 = sigmoid(10 * x1 + 10 * x2 - 5)
    h2 = sigmoid(-10 * x1 - 10 * x2 + 15)
    return sigmoid(10 * h1 + 10 * h2 - 15)


def forward_layer(x: Sequence[float], weights: Sequence[Sequence[float]], bias: Sequence[float], activation=relu) -> list[float]:
    # weights is input_dim x output_dim
    outputs = []
    for j in range(len(bias)):
        z = sum(x[i] * weights[i][j] for i in range(len(x))) + bias[j]
        outputs.append(activation(z))
    return outputs


def mse(preds: Sequence[float], targets: Sequence[float]) -> float:
    return sum((p - y) ** 2 for p, y in zip(preds, targets)) / len(preds)


def train_linear_neuron(xs: Sequence[float], ys: Sequence[float], w: float = 0.0, b: float = 0.0, lr: float = 0.01, epochs: int = 80):
    history = []
    n = len(xs)
    for epoch in range(epochs):
        preds = [w * x + b for x in xs]
        loss = mse(preds, ys)
        grad_w = sum(2 * (p - y) * x for x, y, p in zip(xs, ys, preds)) / n
        grad_b = sum(2 * (p - y) for y, p in zip(ys, preds)) / n
        w -= lr * grad_w
        b -= lr * grad_b
        if epoch % 10 == 0 or epoch == epochs - 1:
            history.append((epoch, loss, w, b, grad_w, grad_b))
    return w, b, history


def write_xor_boundary_svg(path: Path) -> None:
    cells = []
    size = 42
    for i in range(11):
        for j in range(11):
            # map grid to 0..1, threshold for network demonstration
            x1 = 1 if i / 10 >= 0.5 else 0
            x2 = 1 if j / 10 >= 0.5 else 0
            y = xor_network(x1, x2)
            color = "#bbf7d0" if y else "#fecaca"
            x = 90 + i * size
            yy = 90 + (10 - j) * size
            cells.append(f'<rect x="{x}" y="{yy}" width="{size}" height="{size}" fill="{color}" stroke="#ffffff"/>')
    points = [
        (90, 90 + 10 * size, "0"),
        (90, 90, "1"),
        (90 + 10 * size, 90 + 10 * size, "1"),
        (90 + 10 * size, 90, "0"),
    ]
    pt_svg = "".join(f'<circle cx="{x}" cy="{y}" r="14" fill="#0f172a"/><text x="{x-4}" y="{y+5}" font-size="14" fill="white">{label}</text>' for x, y, label in points)
    path.parent.mkdir(parents=True, exist_ok=True)
    path.write_text(f'''<svg xmlns="http://www.w3.org/2000/svg" width="650" height="620" viewBox="0 0 650 620" role="img" aria-label="XOR decision boundary">
<rect width="650" height="620" fill="#f8fafc"/>
<text x="70" y="50" font-family="Arial" font-size="24" font-weight="700">XOR decision boundary dari hidden layer</text>
{''.join(cells)}{pt_svg}
<text x="90" y="580" font-family="Arial" font-size="15">Hijau = prediksi 1, merah = prediksi 0. Pola ini tidak bisa dibuat satu garis lurus.</text>
</svg>''', encoding="utf-8")


def main() -> None:
    print("=== Bab 9 Neural Network Playground ===")
    print("sigmoid(0)=", sigmoid(0))
    print("ReLU(-3), ReLU(4)=", relu(-3), relu(4))

    print("\nPerceptron AND w=[1,1], b=-1.5")
    for x in [(0, 0), (0, 1), (1, 0), (1, 1)]:
        print(x, perceptron(x, [1, 1], -1.5))

    print("\nXOR network step activation")
    xor_outputs = []
    for x in [(0, 0), (0, 1), (1, 0), (1, 1)]:
        y = xor_network(*x)
        xor_outputs.append(y)
        print(x, y, "sigmoid_prob=", round(sigmoid_xor_probability(*x), 3))
    assert xor_outputs == [0, 1, 1, 0]

    print("\nForward layer example")
    h = forward_layer([2, 3], [[1, -1], [2, 1]], [0, 1], relu)
    print("h=", h)

    print("\nGradient training linear neuron for y=3x+1")
    xs = [0, 1, 2, 3, 4]
    ys = [1, 4, 7, 10, 13]
    w, b, history = train_linear_neuron(xs, ys, lr=0.03, epochs=120)
    for epoch, loss, ww, bb, gw, gb in history:
        print(f"epoch={epoch:3d} loss={loss:8.4f} w={ww:6.3f} b={bb:6.3f} grad_w={gw:7.3f} grad_b={gb:7.3f}")
    print("final approx:", round(w, 3), round(b, 3))

    out = (Path(__file__).resolve().parent if "__file__" in globals() else Path.cwd()) / "outputs"
    write_xor_boundary_svg(out / "xor_boundary.svg")
    print("SVG written:", out / "xor_boundary.svg")


if __name__ == "__main__":
    main()