#
"""
 FWD Pass: NonLin layer -> Lin Layer
 BWD Pass: GD on NonLin layer, Lin layer

"""

using Revise
using diffMesh
import Zygote
using NNlib, Flux, LinearAlgebra, Plots
#--------------------------------------#
np  = 2001 # data points
ng  = 31   # FEM nodes

function udata(x)
    u   = 2 .+ zero(x)
    u .+= @. sin(1*pi*x) * (-1 <= x <= 0)
    u .+= @. sin(4*pi*x) * ( 0 <  x <= 1)
    return u
end
#--------------------------------------#
function addEndPts(xe::AbstractArray
                  ,x0::Number,x1::Number)

    n = length(xe) + 2
    A,R = Zygote.ignore() do
        A = zeros(n,2)
        A[1,1] = 1
        A[n,2] = 1
        R = restrict(n)
        return A,R
    end

    xg = A * [x0,x1] + R'*xe

    return xg
end
#--------------------------------------#
function model(xe,ug,x)

    xg = addEndPts(xe,-1.0,1.0)

    return reluHat(xg,ug,x)
end
#--------------------------------------#
function llsq(xg)

    NN = reluHatFunc(xg)
    A  = NN(xp')'
    linOp(u) = reluHat(NN,u,xp)

#   ug = udata(xg) # wrong   least sq. for mse loss minimization
    ug = A \ up    # correct least sq. for mse loss minimization

    return ug
end
#--------------------------------------#
function loss()
    fp   = model(xe,ug,xp)
    loss = Flux.Losses.mse(fp,up)
    return loss
end
#--------------------------------------#
function makeplot()
    fp = model(xe,ug,xp)

    xg = addEndPts(xe,-1,1)
    NN = reluHatFunc(xg)
#   ug;= llsq(xg) # defined in global scope

    plt = plot()
    plt = plot!(title="Gradient Descent",xlabel="x",ylabel="y")
    plt = plot!(ylims=(-1,4))

    plt = plot!(xp,up,width=3,label="data")
    plt = plot!(xp,fp,width=3,label="prediction")
#   plt = scatter!(xg,ug,width=4,color="red",label=false)
    
    # hat functions
    NNp = NN(xp')
    for i=1:ng
        scale = ug[i]
#       scale = 1.0
        global plt = plot!(xp,scale .* NNp[i,:]
                          ,width=2,color="black",label=false)
    end
    
    #savefig(plt,"diffMesh.png")
    display(plt)
    return plt
end
#--------------------------------------#
function callback(;doplot=true)

    l2 = loss()
    fp = model(xe,ug,xp)
    er = norm(up .- fp,Inf)
    println("Error ∞ norm: $er, loss: $l2")

    if(doplot) makeplot() end
    return false
end
#--------------------------------------#

# data
xp = Array(range(-1,stop=1,length=np))
up = udata(xp)

# initialize
xg = sort(2*rand(ng).-1) 
ug = rand(length(xg))
#xg = Array(range(-1,stop=1,length=ng))
#ug = udata(xg)

xe = xg[2:end-1] # w/o end points
fp = model(xe,ug,up)

callback(doplot=true)

ps = Flux.params(xe,ug)
data = Iterators.repeated((),100)
opt = ADAM(1e-2)

gs = gradient(()->loss(),ps)

#anim = Animation()
for d in data
    global gs .= gradient(()->loss(),ps)
    Flux.Optimise.update!(opt,ps,gs)
    callback(doplot=true)
    #frame(anim)
end
#gif(anim,"examples/GD.gif",fps=30)
#--------------------------------------#
nothing
