#=
# Weight normalization

The authors of <https://arxiv.org/abs/1602.07868> introduce weight normalization to boost
learning.
Let's try it here.

Preliminaries:
=#

import CairoMakie
import MLDatasets
import RestrictedBoltzmannMachines as RBMs
import WeightNormalizedRBMs as wRBMs
nothing #hide

#=
The following is a convenience function to plot grids of digits.
Given a four dimensional tensor `A` of size `(width, height, ncols, nrows)`
containing `width x height` images in a grid of `nrows x ncols`, this returns
a matrix of size `(width * ncols, height * nrows)`, that can be plotted in a heatmap
to display all images.
=#

function imggrid(A::AbstractArray{<:Any,4})
    reshape(permutedims(A, (1,3,2,4)), size(A,1)*size(A,3), size(A,2)*size(A,4))
end

#=
Load data. Keep only 0, 1 digits for speed.
=#

Float = Float32
train_x, train_y = MLDatasets.MNIST.traindata()
tests_x, tests_y = MLDatasets.MNIST.testdata()
selected_digits = (0, 1)
train_x = Array{Float}(train_x[:, :, train_y .∈ Ref(selected_digits)] .> 0.5)
tests_x = Array{Float}(tests_x[:, :, tests_y .∈ Ref(selected_digits)] .> 0.5)
train_y = train_y[train_y .∈ Ref(selected_digits)]
tests_y = tests_y[tests_y .∈ Ref(selected_digits)]
train_nsamples = length(train_y)
tests_nsamples = length(tests_y)
(train_nsamples, tests_nsamples)

#=
Init RBM and train.
=#

rbm = RBMs.BinaryRBM(zeros(Float,28,28), zeros(Float,400), zeros(Float,28,28,400))
RBMs.initialize!(rbm, train_x)
wrbm = wRBMs.WeightNormRBM(rbm) # weight normalization reparameterization
@time history = RBMs.pcd!(wrbm, train_x; epochs=500, batchsize=256, steps=5)
rbm = RBMs.RBM(wrbm)
nothing #hide

#=
Let's see what the learning curves look like.
=#

fig = CairoMakie.Figure(resolution=(800, 300))
ax = CairoMakie.Axis(fig[1,1], xlabel="train time", ylabel="log(pseudolikelihood)")
CairoMakie.lines!(ax, get(history, :lpl)...)
fig

#=
Let's look at some samples generated by this RBM.
=#

nrows, ncols = 10, 15
fantasy_x = train_x[:, :, rand(1:train_nsamples, nrows * ncols)]
@time fantasy_x .= RBMs.sample_v_from_v(rbm, fantasy_x; steps=10000)
nothing #hide

#

fig = CairoMakie.Figure(resolution=(40ncols, 40nrows))
ax = CairoMakie.Axis(fig[1,1], yreversed=true)
CairoMakie.image!(ax, imggrid(reshape(fantasy_x, 28, 28, ncols, nrows)), colorrange=(1,0))
CairoMakie.hidedecorations!(ax)
CairoMakie.hidespines!(ax)
fig
