# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/02_nn.ipynb.

# %% auto 0
__all__ = ['Module', 'Neuron', 'Layer', 'MLP']

# %% ../nbs/02_nn.ipynb 3
from nbdev.showdoc import *
from fastcore.test import *
from fastcore.utils import *
from IPython import display

# %% ../nbs/02_nn.ipynb 5
import random
from .engine import Value

# %% ../nbs/02_nn.ipynb 6
from .utils import draw_dot

# %% ../nbs/02_nn.ipynb 17
class Module:
    """base module"""
    def zero_grad(self):
        for p in self.parameters():
            p.grad = 0

    def parameters(self):
        return []

# %% ../nbs/02_nn.ipynb 20
class Neuron(Module):
    """Artificial Neuron e.g. Perceptron"""
    
    def __init__(self, nin, nonlin=True):
        self.w = [Value(random.uniform(-1,1)) for _ in range(nin)]
        self.b = Value(0)
        self.nonlin = nonlin

    def __call__(self, x):
        act = sum((wi*xi for wi,xi in zip(self.w, x)), self.b)
        return act.relu() if self.nonlin else act

    def parameters(self):
        return self.w + [self.b]

    def __repr__(self):
        return f"{'ReLU' if self.nonlin else 'Linear'}Neuron({len(self.w)})"

# %% ../nbs/02_nn.ipynb 29
class Layer(Module):
    """Layer of Neurons"""
    def __init__(self, nin, nout, nonlin=True, **kwargs):
        self.neurons = [Neuron(nin, nonlin=nonlin, **kwargs) for _ in range(nout)]

    def __call__(self, x):
        out = [n(x) for n in self.neurons]
        return out[0] if len(out) == 1 else out

    def parameters(self):
        return [p for n in self.neurons for p in n.parameters()]

    def __repr__(self):
        return f"Layer of [{', '.join(str(n) for n in self.neurons)}]"

# %% ../nbs/02_nn.ipynb 38
class MLP(Module):
    """Multi-Layer Perceptron (MLP) e.g. Neural Network"""

    def __init__(self, nin, nouts, nonlin=True):
        sz = [nin] + nouts
        if nonlin:
            self.layers = [Layer(sz[i], sz[i+1], nonlin=i!=len(nouts)-1) for i in range(len(nouts))]
        else:
            self.layers = [Layer(sz[i], sz[i+1], nonlin=nonlin) for i in range(len(nouts))]
            
    def __call__(self, x):
        for layer in self.layers:
            x = layer(x)
        return x

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

    def __repr__(self):
        return f"MLP of [{', '.join(str(layer) for layer in self.layers)}]"
