#
"""
 Functional based mesh adaptation

 FWD Pass: NonLin layer -> least sq. on linear layer
 BWD Pass: GD all the way through (takes adjoint on least sq.)

 Task: replace least sq. with PDE solve
       test problem: steady advection diffusion eqn at large
                     Peclet # (thin thin boundary 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,xp,up) # grid points, eval points, values to fit

    xg = addEndPts(xe,-1,1)
    NN = reluHatFunc(xg)
    ug = llsq(NN,xp,up)

    return reluHat(NN,ug,xp)
end
#--------------------------------------#
function llsq(NN,xp,up)

    A  = NN(xp)
    linOp(u) = reluHat(NN,u,xp)

#   u = udata(x)
    u = A \ up
    return u
end
#--------------------------------------#
function loss()
    fp   = model(xe,xp,up)
    loss = Flux.Losses.mse(fp,up)
    return loss
end
#--------------------------------------#
function makeplot()
    fp = model(xe,xp,up)

    xg = addEndPts(xe,-1,1)
    NN = reluHatFunc(xg)
    ug = llsq(NN,xp,up)

    plt = plot()
    plt = plot!(title="FEM AutoMesh",xlabel="x",ylabel="y")

    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)

    global xe .= sort(xe)
    l2 = loss()
    fp = model(xe,xp,up)
    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 = Array(range(-1,stop=1,length=ng))
xg = sort(2*rand(ng).-1) 

xe = xg[2:end-1]
fp = model(xe,xp,up)

callback(doplot=true)

ps = Flux.params(xe)
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/AutoMesh.gif",fps=30)

#--------------------------------------#
nothing
