# -*- coding: utf-8 -*-
# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .jl
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Julia 1.4.1
#     language: julia
#     name: julia-1.4
# ---

# # A Generic Backprop
#
# The previous notebook, which constructed a neural network with a single hidden layer and a corresponding backprop training function, gave us an idea of how we might generalize these implementations to build and train neural networks of arbitrary width and depth.
#
# In **Part 1**, we'll generalize our model implementation to allow us to easily build fully-connected layers and chain them together into a model. 
#
# In **Part 2**, we'll generalize our backprop implementation to work with these models without having to explicitly hand-code the chain rule multiplications.

# ## Part 1: A generalized model API
#
#    The first thing we'll need is a way to refer to a given function (e.g. softmax or binary cross-entropy) and its derivative in a clear way. We'll do that with a `struct` named `Operation` below. Note that this is just a restructuring of the function definitions from the previous notebook.

# +
# We'll need our clip function again to avoid domain errors when using log.
function clip(x::Real, low=1e-15, hi=1 - 1e-15)
    if x < low
        return low
    elseif x > hi
        return hi
    else
        return x
    end
end

# We'll name it "Operation" because "Function" is already taken by Julia.
struct Operation
    name::String  # The name of this operation
    forward::Function  # The function
    backward::Function # The derivative of this function
end

# Sigmoid activation
function sigmoid(h::Array{Float64})
    return @. 1.0 / (1.0 + exp(-h))
end

function d_sigmoid(h::Array{Float64})
    forward_vals = sigmoid(h)
    return @. forward_vals * (1.0 - forward_vals)
end

Sigmoid = Operation("Sigmoid", sigmoid, d_sigmoid)

# Binary cross entropy
function binary_crossentropy(y::Array, o::Array{Float64})
    # y: the gold-standard labels 0 or 1
    # o: the real-valued output [0,1]
    n = size(o, 1)
    y_hat = clip.(o)
    ces = @. y * log(o) + (1 - y) * log(1 - o)
    return -(1 / n) * sum(ces)
end

function d_binary_crossentropy(y::Array, o::Array{Float64})
    # y: the gold-standard labels 0 or 1
    # o: the real-valued output [0,1]
    y_hat = clip.(o)
    return @. (o - y) / (o * (1 - o))
end

BinaryCrossEntropy = Operation("BinaryCrossEntropy", binary_crossentropy,
                               d_binary_crossentropy)
# -

# Let's try these out to make sure they work as expected. We'll use the same values as in the first notebook to check our answers against. As a reminder, these are
#
# $z = \begin{pmatrix} 0.7 \\ 0.8 \\ 0.4 \\ 0.3 \end{pmatrix}$
#
# $y = \begin{pmatrix} 1 \\ 1 \\ 0 \\ 0 \end{pmatrix} $
#
# $o = \sigma(z) = \begin{pmatrix} 0.67 \\ 0.69 \\ 0.60 \\ 0.57 \end{pmatrix}$
#
#
# $\frac{d{o}}{d{z}} = \frac{d{\sigma}}{d{z}} = \sigma(z)(1-\sigma(z)) =
#     \begin{pmatrix} 0.22 \\ 0.21 \\ 0.24 \\ 0.24 \end{pmatrix}
# $
#
# $ C(y, o) = 0.635 $
#
# <br>
#
# $\frac{\partial{C}}{\partial{o}} = \frac{o - y}{o(1-o)} =
#     \begin{pmatrix} -1.50 \\ -1.45 \\ 2.49 \\ 2.35 \end{pmatrix}
# $
#

# +
z = [0.7, 0.8, 0.4, 0.3]
println("o = σ(z): $(Sigmoid.forward(z))")
println("do = dσ(z): $(Sigmoid.backward(z))")

y = [1, 1, 0, 0]
o = [0.67, 0.69, 0.60, 0.574]
println("C: $(BinaryCrossEntropy.forward(y, o))")
println("∂C: $(BinaryCrossEntropy.backward(y, o))")
# -

# ## The Layer
#
# A neural network layer is simply a matrix multiplication plus the bias followed by an activation function, i.e. $\sigma(XW + b)$. Thus we'll need variables for $\sigma$, $W$, and $b$. We'll also include placeholders for the weighted inputs $z = WX + b$ and the activations $o = \sigma(z)$, since we'll need these in the backward pass.

# +
import Random
Random.seed!(0)

mutable struct Layer
    name::String            # What we'll call this layer.
    W::Array{Float64}       # The weight matrix
    b::Array{Float64}       # The bias term
    σ::Operation            # The activation function
    δ::Array{Float64}       # Used in backprop
    input::Array{Float64}   # The inputs from the previous layer
    z::Array{Float64}       # Weighted inputs XW + b
    output::Array{Float64}  # The output of this layer, i.e. the actvations

    # Constructor v1: When we have already initialized W and b
    function Layer(name::String, W::Array{Float64}, b::Array{Float64}, σ::Operation)
        if !isa(σ, Operation)
            error("Activation function '$σ' in layer '$name' is not an Operation")
        end
        if size(W, 2) != size(b, 2)
            error("Incompatible W and b shapes in layer '$name'")
        else
            # Since δ, input, and output are computed later, we init as empty arrays.
            new(name, W, b, σ, [], [], [], [])
        end
    end
    
    # Constructor v2: When we want to randomly initialize W and b given the shape of W.
    function Layer(name::String, size::Tuple{Int,Int}, σ::Operation)
        W = Random.rand(Float64, size)
        b = Random.rand(Float64, (1, size[2]))
        # Since δ, input, and output are computed later, we init as empty arrays.
        new(name, W, b, σ, [], [], [], [])
    end
end
# -

# Let's try building a Layer using both constructors.

# +
W = Random.rand(Float64, (2, 2))
b = Random.rand(Float64, (1, 2))
layer1 = Layer("Layer1", W, b, Sigmoid)
println("Layer1 W: $(layer1.W), $(size(layer1.W))")
println("Layer1 b: $(layer1.b), $(size(layer1.b))")

layer2 = Layer("Layer2", (2, 2), Sigmoid)
println("Layer2 W: $(layer2.W), $(size(layer2.W))")
println("Layer2 b: $(layer2.b), $(size(layer2.b))")
# -

# ## Layer forward
#
# We'll now write a function that propagates some inputs $X$ through a layer by doing the matrix multiplication, adding the bias, and applying the activation function.

# The ! after a function name in Julia means that the function has a "side-effect".
# Here, forward! modifies the Layer by storing the values of X, z, and the output.
function forward!(l::Layer, X::Array)
    # It's always a good idea to double check that the shapes are compatible.
    if size(X, 2) != size(l.W, 1)
        error("Input shape $(size(X)) incompatible with layer '$(l.name)' shape $(size(l.W))")
    end
    l.input = X
    l.z = (X * l.W) .+ l.b  # Weighted inputs
    l.output = l.σ.forward(l.z)
    return l.output
end

# Now let's try running some data through a Layer. We'll compute everything by hand here first, and then check the function's output against it.
#
# $ X = \begin{pmatrix} 1 & 3 \\ 2 & 3 \\ 2 & 1 \\ 1 & 1 \end{pmatrix} $
#
# $ W = \begin{pmatrix} 0.1 \\ 0.2 \end{pmatrix} $
#
# $ b = 0.1 $
#
# $ z = XW + b = \begin{pmatrix} .8 \\ .9 \\ .5 \\ .3 \end{pmatrix} $
#
# $ o = \sigma(z) = \begin{pmatrix} .69 \\ .71 \\ .62 \\ .57 \end{pmatrix} $

X = [1 3; 2 3; 2 1; 1 1]
display(X)
W = [0.1, 0.2]
display(W)
b = [0.1]
display(b)
test_layer1 = Layer("test_layer1", W, b, Sigmoid)

# +
o = forward!(test_layer1, X)

# These two should be equal
display(o)
display(test_layer1.output)

display(test_layer1.input)
display(test_layer1.z)
# -

# ## Our first Network
#
# Looks good! Let's create another `struct` called `Network`, which will allow us to connect an arbitrary number of layers together. In the constructor, we'll iterate through the layers in order, checking that the shape of the previous layer's output is compatible with the shape of the current layer's input for the matrix multiplication.

struct Network
    name::String
    layers::Array{Layer}
    loss::Operation
    η::Real  # The learning rate
    
    # Constructor: Check that all layer shapes are compatible
    function Network(name::String, layers::Array{Layer}, loss::Operation, η::Real)
        l1 = layers[1]
        for l2 in layers[2:end]
            l1_outdim = size(l1.W, 2)
            l2_indim = size(l2.W, 1)
            if l1_outdim != l2_indim
                error("""Output dimension $(size(l1.W, 2)) of layer
                      '$(l1.name)' incompatible with input dimension
                      $(size(l2.W, 1)) of layer '$(l2.name)'.""")
            end
        end
        new(name, layers, loss, η)
    end
end

# ### The forward pass
#
# We have a `forward!` function for a Layer, but to implement the forward pass through the entire network we'll have to string these together. This is pretty simple: the input to `forward!` for `Layer` $l$ is simply the output of `forward!` from `Layer` $l-1$. 

# The forward! function applied to a Network applies the forword! function to each layer.
# Julia allows duplicate function names as long as the argument signatures are unique.
function forward!(network::Network, X::Array)
    layer_input = X
    for l in network.layers
        layer_input = forward!(l, layer_input)
    end
    return layer_input
end

# Let's try out our new Network by passing the XOR data through it.

# +
import Random
Random.seed!(0)

X = [0 0; 1 1; 0 1; 1 0]
y = [0, 0, 1, 1]

layer1 = Layer("layer1", (2, 3), Sigmoid)
layer2 = Layer("layer2", (3, 1), Sigmoid)

nn = Network("FFNN", [layer1, layer2], BinaryCrossEntropy, 0.1)
# -

o = forward!(nn, X)
loss = nn.loss.forward(y, o)
println("Output: $o")
println("Loss: $loss")

# ## Part 2: A generalized backward pass
#
# In the previous notebook, we casually introduced the $\delta$ term when computing the gradient of the final layer.
#
# $\delta^2 = \frac{\partial{C}}{\partial{o^2}} \frac{d{o^2}}{d{z^2}}$
#
# We also casually defined an analogous term for the first layer.
#
# $\delta^1 = \delta^2 {W^2}^T \odot \frac{d{o^1}}{d{z^1}}$
#
# A keen eye may notice a pattern. Specifically, the $\delta$ value of a given layer can be defined recursively in terms of the $\delta$ value of the subsequent layer.
#
# $\delta^l = \delta^{l+1} {W^{(l+1)}}^T \odot \frac{d{o^l}}{d{z^l}}$
#
# The base case of this recursion (i.e. the $\delta$ of the last layer) is the partial derivative of the loss wrt the output activations.
#
# $\delta^L = \frac{\partial{C}}{\partial{o^L}} \frac{d{o^L}}{d{z^L}}$
#
# Practically, we'll compute $\delta$ using a dynamic programming approach, rather than pure recursion. We know that our base case is always the last layer of the network, so all we have to do is iterate backwards through the layers. We implement this computation below in `backward!` using the helper function `compute_δ`, which takes care of the base case.

# +
function compute_δ(network::Network, y::Array, i::Int)
    l = network.layers[i]
    dσ_dz = l.σ.backward(l.z)
    # The base case of the recursion: the last layer in the network.
    if i == length(network.layers)
        o = network.layers[i].output
        δ = network.loss.backward(y, o) .* dσ_dz
        return δ
    end
    l_next = network.layers[i+1]
    # l_next.δ should have already been computed because we're iterating
    # backwards through the network's layers in backward!() below.
    δ = (l_next.δ * transpose(l_next.W)) .* dσ_dz
    return δ
end

function backward!(net::Network, y::Array)
    num_layers = length(net.layers)
    for i in num_layers:-1:1  # Iterate backwards through the network
        l = net.layers[i]
        l.δ = compute_δ(net, y, i)
        ∂C_∂W = transpose(l.δ) * l.input
        l.W = l.W .- (net.η * transpose(∂C_∂W))
    end
end
# -

# ### Training a model
#
# These new `forward!` and `backward!` functions significantly simplify the training loop compared to the previous notebooks, as we no longer have to hand-code the computations. All we have to do is call the functions in turn for a number of training steps. We can also wrap the training loop into a convenient `fit!` function.

# +
function fit!(net::Network, X::Array, y::Array, epochs::Int)
    for i=1:epochs
        o = forward!(net, X)
        backward!(net, y)
        if i == 1 || i % 100 == 0
            loss = net.loss.forward(y, o)
            println("Loss ($i): $(loss)")
        end
    end
end

function decision(o)
    return Int.(o .>= 0.5)
end

function accuracy(y, y_hat)
    return sum(y .== y_hat) / size(y, 1)
end

# +
# Create the training dataset
X = [0 0; 1 1; 0 1; 1 0]
y = [0, 0, 1, 1]

# Instantiate the network
layer1 = Layer("layer1", (2, 3), Sigmoid)
layer2 = Layer("layer2", (3, 1), Sigmoid)
nn = Network("FFNN", [layer1, layer2], BinaryCrossEntropy, 0.2)

# Fit the network to the training data
fit!(nn, X, y, 1000)

# Compute the accuracy of the network on the training data.
logits = forward!(nn, X)
predictions = decision(logits)
acc = accuracy(y, predictions)
println("Logits: $(logits)")
println("Predictions: $(predictions)")
println("Gold Labels: $(y)")
println("Accuracy: $(acc)")
# -

# ## Extensibility
#
# Our generic model API allows us to easily implement extensions. For example, we implement two additional activations functions, ReLU and Tanh, below.
#
# #### ReLU: Rectified Linear Unit
# $$ \text{ReLU}(x) = max\{0, x\} $$
#
# $$ \frac{d}{dx} \text{ReLU}(x) = max\{0, 1\} $$
#
#
# #### Tanh
# $$ \text{tanh}(x) = \frac{sinh(x)}{cosh(x)} = \frac{e^x - e^{-x}}{e^x + e^{-x}} $$
# which is just a scaled and shifted form of the sigmoid function over $[-1, 1]$.
#
# $$ \frac{d}{dx} \text{tanh}(x) = 1 - \text{tanh}(x)^2 $$

# +
function relu(x)
    return replace(a -> a<0 ? 0 : a, x)
end

function d_relu(x)
    return replace(a -> a<0 ? 0 : 1, x)
end

ReLU = Operation("ReLU", relu, d_relu)


function tanh(x)
    ex = exp.(x)
    enx = exp.(-x)
    return @. (ex - enx) / (ex + enx)
end

function d_tanh(x)
    return @. 1 - (tanh(x)^2)
end

Tanh = Operation("Tanh", tanh, d_tanh)
# -

# With our API, we can just plug these new activation functions into a model, call `fit!`, and we're done!

# +
X = [0 0; 1 1; 0 1; 1 0]
y = [0, 0, 1, 1]

layer1 = Layer("layer1", (2, 3), Tanh)
layer2 = Layer("layer2", (3, 3), ReLU)
layer3 = Layer("layer3", (3, 1), Sigmoid)
loss_fn = BinaryCrossEntropy
η = 0.1

nn = Network("FFNN2", [layer1, layer2, layer3], loss_fn, η)
# -

fit!(nn, X, y, 1000)
predictions = decision(forward!(nn, X))
println("Accuracy: $(accuracy(y, predictions))")

# Finally, it would be nice to summarize a network without seeing all the individual weights. Let's write a function to do that, which overloads the built-in `display`.

# +
import Base.display

function display(nn::Network)
    println("Network: $(nn.name)")
    for layer in nn.layers
        println("$(layer.name) : $(size(layer.W)) : $(layer.σ.name)")
    end
    println("Loss: $(nn.loss.name)")
end

display(nn)
# -



StarcoderdataJupyterScriptsDedupFiltered