using Revise
using AdFem
using LinearAlgebra
using PyPlot
using DelimitedFiles
using MAT 
using Optim

function f(x, y)
    10/(1+x^2) + x * y + 10*y^2
end


function calc_residual_and_jacobian(θ, u)
    E = θ
    μ = E / (2*(1 + nu))
    λ = E*nu/((1 + nu)*(1 - 2*nu))

    ψ, J = neo_hookean(u, μ, λ, mmesh)
    B = compute_fem_source_term(zeros(get_ngauss(mmesh)), -0.5*ones(get_ngauss(mmesh)), mmesh)
    ψ = ψ - B
    J, ψ = impose_Dirichlet_boundary_conditions(J, ψ, bdnode, zeros(length(bdnode)))
    ψ, J
end


mmesh = Mesh(joinpath(PDATA, "twoholes_large.stl"))
mmesh = Mesh(mmesh.nodes * 10, mmesh.elems)

# E = eval_f_on_gauss_pts(f, mmesh)
# θ = Variable(10ones(get_ngauss(mmesh)))

θ = Variable(ones(mmesh.nelem))
θ_ = repeat(θ, 1, get_ngauss(mmesh)÷mmesh.nelem)
θ_ = reshape(θ_, (-1,)) * 10


nu = constant(0.3 * ones(get_ngauss(mmesh)))
bdnode = bcnode((x,y)->(x<1e-5) || (x>0.49-1e-5), mmesh)
bdnode = [bdnode; bdnode .+ mmesh.ndof]
nr = newton_raphson_with_grad(calc_residual_and_jacobian, zeros(2mmesh.ndof), θ_)
u = matread("fenics/fwd2.mat")["x"]
loss = sum((nr-u)^2) * 1e10
g = gradients(loss, θ)
sess = Session(); init(sess)

_loss = Float64[]
cb = (x, iter, loss)->begin 
    global _loss
    push!(_loss, loss)
    printstyled("[#iter $iter] loss=$loss\n", color=:green)
    nu = run(sess, θ_, θ=>x)
    close("all")
    visualize_scalar_on_gauss_points(nu, mmesh)
    if mod(iter, 10)==1
        matwrite("fenics/bwd4-$iter.mat", Dict("iter"=>iter,"loss"=>_loss, "E"=>nu))
        savefig("fenics/bwd4_pixel$iter.png")
    end
end

# run(sess, loss)
loss_ = BFGS!(sess, loss, g, θ, callback = cb)

