#
using DiffMesh, DiffEqFlux, LinearAlgebra, Plots
using GalacticOptim, ParameterSchedulers
#using GLMakie
#----------------------------------------------------------------------------#
using Descartes
ci = translate([0.0,0.0])Circle(1.0)
sq = translate([0.0,0.0])Square([3.0,3.0];center=true)
ss = diff(sq,ci)
sdf = x -> FRep(ss,x)

#sdf = (x,y) -> @. (x^2 + y^2) - 1

#nprt = 2
#shape = x -> @. sin(pi/2*x)
#pou = femPOU2D(nprt,shape,sdf)

nprt = 6
pou  = nnPOU2D(nprt,sdf;depth=4,act=Flux.tanh)

nx = ny = 4
bsp = TensorProdCosines(nprt,nx,ny)
net = MeshNet2D(pou,bsp)

p,re = Flux.destructure(net)
#--------------------------------------#
function utrue(x,y)

    u  = similar(x)
    r2 = @. x*x + y*y

    u .= @. 1 - r2
    u .= @. sin(3*pi*x) + cos(3*pi*y)*exp(x*y)
    u .= @. x*y^3 + sin(r2)  + exp(x^2*y)
    return u
end

# train
ne = 1000
Re = rand(1,ne)
θe = 2*pi*rand(1,ne)

xe = @. Re * cos(θe)
ye = @. Re * sin(θe)

xe = rand(1,1000) .* 2.5
ye = rand(1,1000) .* 2.5

ue = utrue(xe,ye)

# test
np = 50
xp = Array(range(-2f5,2f5,length=np))
yp = Array(range(-2f5,2f5,length=np))

xp = xp * ones(1,np)
yp = ones(np) * yp'
up = utrue(xp,yp)
Ωp = sdf(xp,yp) .< 1e-8

xpp = reshape(xp,np*np)' # flattened
ypp = reshape(yp,np*np)'
upp = reshape(up,np*np)'
Ωpp = reshape(Ωp,np*np)'

#----------------------------------------------------------------------------#
it = 0 # iter
#fig = Figure()
#ax1 = Axis(fig[1,1])
#ax2 = Axis(fig[1,2])
#ax3 = Axis(fig[1,3])
#ax1.title = "Pred"
#ax2.title = "True"
#ax3.title = "Err"
#fig

#--------------------------------------#
function loss(p,xe)
    phi = re(p)

    fe = phi(xe,ye)
    ee = ue .- fe
    loss = Flux.Losses.mse(fe,ue)
    return loss
end
#--------------------------------------#
function makeplot(phi)

    fpp = phi(xpp,ypp)
    fp  = reshape(fpp,np,np)
    ep  = @. abs(up-fp)

#   heatmap!(ax1,Ωp.*fp)
#   heatmap!(ax2,Ωp.*up)
#   heatmap!(ax3,Ωp.*ep)
#   fig

    p1 = heatmap(Ωp.*fp)
    p2 = heatmap(Ωp.*up)
    p3 = heatmap(Ωp.*ep)

    return
end
#--------------------------------------#
function cb1(p,l)
    phi = re(p)

    leastSq!(phi,xe,ye,ue)

    fpp = phi(xpp,ypp)
    epp = (upp .- fpp) .* Ωpp
    er  = norm(epp,Inf)

    global it += 1
    println("it: $it, loss: $l, pointwise er: $er")
    (it % 10) == 0 && makeplot(phi)
    (l < 1e-6) && return true
    return false
end
#--------------------------------------#
function cb2(p,l)
    phi = re(p)

    fpp = phi(xpp,ypp)
    epp = (upp .- fpp) .* Ωpp
    er  = norm(epp,Inf)

    global it += 1
    println("it: $it, loss: $l, pointwise er: $er")
    (it % 10) == 0 && makeplot(phi)
    (l < 1e-6) && return true
    return false
end
#----------------------------------------------------------------------------#
# training

func = OptimizationFunction(loss,GalacticOptim.AutoZygote())
prob = OptimizationProblem(func,p,xe)

sch  = ParameterSchedulers.SinExp(λ0=1e-0,λ1=1e-2,γ=5e-1,period=50)
opt1 = ParameterSchedulers.Scheduler(sch,ADAM())
opt1 = ADAM(1e-2)
opt2 = BFGS()

res  = GalacticOptim.solve(prob,opt1,cb=cb1,maxiters=200)
prob = remake(prob,u0=res.minimizer)
res  = GalacticOptim.solve(prob,opt2,cb=cb1,maxiters=200)
prob = remake(prob,u0=res.minimizer)
res  = GalacticOptim.solve(prob,opt2,cb=cb2,maxiters=100)

mzr = res.minimizer
phi = re(mzr)
#----------------------------------------------------------------------------#
return
