"""An educational scalar autograd engine. Run with Python 3; no dependencies.

Inspired by Andrej Karpathy's micrograd lecture. This implementation stores
local derivatives on edges and clears the reachable graph on each backward().
"""

import math
import random


class Value:
    def __init__(self, data, edges=()):
        self.data = float(data)
        self.grad = 0.0
        self.edges = tuple(edges)  # (parent, local derivative) pairs

    @staticmethod
    def wrap(value):
        return value if isinstance(value, Value) else Value(value)

    def __add__(self, other):
        other = Value.wrap(other)
        return Value(self.data + other.data, ((self, 1.0), (other, 1.0)))

    def __mul__(self, other):
        other = Value.wrap(other)
        return Value(self.data * other.data,
                     ((self, other.data), (other, self.data)))

    def __pow__(self, exponent):
        if not isinstance(exponent, int) or exponent < 1:
            raise ValueError("This engine supports positive integer powers only")
        return Value(self.data ** exponent,
                     ((self, exponent * self.data ** (exponent - 1)),))

    def tanh(self):
        output = math.tanh(self.data)
        return Value(output, ((self, 1.0 - output * output),))

    def __neg__(self):
        return self * -1.0

    def __sub__(self, other):
        return self + -Value.wrap(other)

    def __rsub__(self, other):
        return Value.wrap(other) + -self

    __radd__ = __add__
    __rmul__ = __mul__

    def backward(self):
        order, seen = [], set()

        def visit(node):
            if node in seen:
                return
            seen.add(node)
            for parent, _ in node.edges:
                visit(parent)
            order.append(node)

        visit(self)
        for node in order:
            node.grad = 0.0
        self.grad = 1.0
        for node in reversed(order):
            for parent, local_derivative in node.edges:
                parent.grad += node.grad * local_derivative


class Neuron:
    def __init__(self, n_inputs, rng, nonlinear=True):
        self.weights = [Value(rng.uniform(-1, 1)) for _ in range(n_inputs)]
        self.bias = Value(0)
        self.nonlinear = nonlinear

    def __call__(self, inputs):
        total = sum((w * x for w, x in zip(self.weights, inputs)), self.bias)
        return total.tanh() if self.nonlinear else total

    def parameters(self):
        return self.weights + [self.bias]


class MLP:
    def __init__(self, sizes, seed=42):
        rng = random.Random(seed)
        self.layers = [
            [Neuron(n_in, rng, nonlinear=i < len(sizes) - 2)
             for _ in range(n_out)]
            for i, (n_in, n_out) in enumerate(zip(sizes, sizes[1:]))
        ]

    def __call__(self, inputs):
        for layer in self.layers:
            inputs = [neuron(inputs) for neuron in layer]
        return inputs

    def parameters(self):
        return [p for layer in self.layers for neuron in layer
                for p in neuron.parameters()]


def single_example():
    w, b = Value(1), Value(0)
    prediction = w * 2 + b
    loss = (prediction - 5) ** 2
    loss.backward()
    print(f"loss={loss.data:.1f}, dL/dw={w.grad:.1f}, dL/db={b.grad:.1f}")

    # Update only w to match the article's first gradient-descent example.
    w.data -= 0.1 * w.grad
    new_loss = (w * 2 + b - 5) ** 2  # a fresh forward pass
    print(f"w={w.data:.1f}, new loss={new_loss.data:.2f}")


def train_line(steps=1000, learning_rate=0.05):
    data = [(0, 1), (1, 3), (2, 5), (3, 7)]
    w, b = Value(-1), Value(0)
    for _ in range(steps):
        loss = sum((w * x + b - y) ** 2 for x, y in data) * (1 / len(data))
        loss.backward()  # also clears this graph's old gradients
        for parameter in (w, b):
            parameter.data -= learning_rate * parameter.grad
    final_loss = sum((w.data * x + b.data - y) ** 2 for x, y in data) / len(data)
    return w.data, b.data, final_loss


def train_xor(steps=1000, learning_rate=0.1):
    # Same inputs twice equal -1; different inputs equal +1.
    examples = [([0, 0], -1), ([0, 1], 1), ([1, 0], 1), ([1, 1], -1)]
    model = MLP([2, 4, 1], seed=42)
    parameters = model.parameters()
    for _ in range(steps):
        loss = sum((model(x)[0] - y) ** 2 for x, y in examples) * 0.25
        loss.backward()
        for parameter in parameters:
            parameter.data -= learning_rate * parameter.grad
    predictions = [model(x)[0].data for x, _ in examples]
    final_loss = sum((p - y) ** 2 for p, (_, y) in zip(predictions, examples)) / 4
    return predictions, final_loss


if __name__ == "__main__":
    single_example()
    w, b, loss = train_line()
    print(f"line: w={w:.4f}, b={b:.4f}, loss={loss:.8f}")
    predictions, loss = train_xor()
    print("XOR predictions:", [round(p, 4) for p in predictions])
    print(f"XOR loss={loss:.8f}")