Zygote

apt-get update
apt-get install build-essential -y
apt-get install ffmpeg

pkg"up; add Zygote CairoMakie AbstractPlotting PackageCompiler#sd-notomls"

using PackageCompiler

open("/results/precompiles.jl", "w") do io
  PackageCompiler.snoop(abspath("snoop.jl"), io)
end

using Pkg
packages = ["CairoMakie", "AbstractPlotting", "Zygote", "LinearAlgebra"]
ctx = Pkg.Types.Context()
pkgs = PackageSpec.(packages)
packages = PackageCompiler.flat_deps(ctx, packages)
# remove blacklisted packages from full list of packages
imports = PackageCompiler.to_pkgid.(packages)
filter!(x-> !(x.name in ("WinRPM", "HTTPClient")), imports)
usings = join(["const $(x.name) = Base.require($(PackageCompiler.prepr(x)))" for x in imports], "\n")
open("/results/precompiles2.jl", "w") do io
  println(io, usings)
  open(
precompiles.jl
) do io2
  	write(io, io2)
  end
end

ln -fs /lib/x86_64-linux-gnu/libz.so.1 /usr/local/lib/libz.so
ldconfig

syso, sysold = compile_incremental(
precompiles2.jl
)

cp(syso, sysold, force = true)

using AbstractPlotting, CairoMakie, LinearAlgebra
AbstractPlotting.current_backend[] = CairoMakie.CairoBackend("tmp.png")

# We store parameters as tracked one element arrays 
# To easily update them inplace & take the gradient
function chain2points(chain, input)
  points = [input]
  return reduce(chain, init = [input]) do v0, angle
    [v0; next_position(v0[end], angle)]
  end
end

function visualize(input, target)
  points = Node(chain2points(chain, input))
  s = lines(points, color = :gray)
  scatter!(points, markersize = 0.2)
  scatter!([target], color = :red, markersize = 0.3)
  dots = lift(points) do x
    diff = norm(last(x) .- target)
    LinRange(Point(last(x)), Point(target), round(Int, diff / 0.1))
  end
  # CairoMakie doesn't support dotted lines yet -.-
  scatter!(
    dots,
    color = (:indianred1, 0.5), markersize = 0.05,
    limits = FRect(0, -3, 6, 7), scale_plot = false
  )
  display(s)
  return points, s
end
rand_angle() = rand() .* pi

@inline function next_position(position, angle)
  position .+ (sin(angle), cos(angle))
end

# Our "neuronal network" ... or chain of flexible links
@inline function predict(chain, input)
  output = next_position(input,  chain[1]) # Layer 1 
  output = next_position(output, chain[2]) # Layer 2
  output = next_position(output, chain[3]) # Layer 3
  output = next_position(output, chain[4]) # Layer 4
  return output
end

function loss(chain, input, target)
  return sum((predict(chain, input) .- target) .^ 2)
end

chain = [rand_angle() for i in 1:4]

input, target = (0.0, 0.0), (3.0, 3.0)
points, s = visualize(input, target)
using Zygote
function loss_gradient(chain, input, target)
  # first index, to get gradient of first argument
  Zygote.gradient(loss, chain, input, target)[1]
end
record(s, "test.gif", 1:100) do i
  # get gradient of loss function
  angle∇= loss_gradient(chain, input, target)
  # update weights with our optimizer + gradients
  chain .-= 0.01 .* angle∇
  # update visualization
  points[] = chain2points(chain, input)
end