from ngsolve import *
import numpy as np

alpha_c = 1
n = 20

def rho(u):
    normJ = Norm(u)
    return alpha_c*normJ**(n-1)

def drho_dlambda(u, du):
    normJ = Norm(u)
    return (n-1)*alpha_c*du*normJ**(n-2)

def SimpleNewtonSolve(gfu, a, tol = 1e-6, maxits = 1000):
    fes = gfu.space
    du = fes.TrialFunction()
    # ~ res = gfu.vec.CreateVector()
    er = 1000; it = 0
    while (er > tol) and (it <= maxits):
        print ("Iteration {:3}  ".format(it),end="")
        # ~ a.Apply(gfu.vec, res)
        CF_drho = CoefficientFunction(drho_dlambda(gfu, du))
        res_nl = LinearForm(fes)
        res_nl += InnerProduct(CF_rho * grad(gfu), grad(v)) * dx
        #### Jacobian and assembly
        b = BilinearForm(fes)
        b += InnerProduct(CF_drho * grad(gfu) * du, grad(v)) * dx
        b += InnerProduct(CF_rho * grad(du), grad(v)) * dx
        b.Assemble()
        # ~ res -= f.vec
        res = res_nl.vec - f.vec
        
        gfdu = gfu.vec.CreateVector()
        gfdu.data = b.mat.Inverse(freedofs = fes.FreeDofs()) * res
        
        gfu.vec.data -= gfdu
        #stopping criteria
        er = sqrt(abs(InnerProduct(gfdu, res)))
        it += 1
        print ("iter.: {0:d}; Er.: {1:.6f}".format(it, er))

mesh = Mesh(unit_square.GenerateMesh(maxh=0.2))

fes = H1(mesh, order=0, dirichlet="bottom")

u, v = fes.TnT()

gfu = GridFunction(fes)  # solution

CF_rho = CoefficientFunction(rho(gfu))

a = BilinearForm(fes)
a += InnerProduct(CF_rho * grad(u), grad(v)) * dx
a.Assemble()

f = LinearForm(fes)
f += 1*v*dx
f.Assemble();

gfu.Set(1.0) # initial guess
SimpleNewtonSolve(gfu, a)

vtk = VTKOutput(mesh, coefs = [gfu, grad(gfu)], names = ["u", "grad(u)"], \
    filename = "./solution", subdivision = 1)
vtk.Do()