using ForwardDiff
using Tensors
using Ferrite
using MAT
using NLsolve

immutable YeohMaterial{T}
    λ::T
    μ::T
    c2::T
    c3::T
end

function Ψ_Yeoh(C, mp::YeohMaterial)
    μ, λ, c2, c3 = mp.μ, mp.λ, mp.c2, mp.c3
    Ct = convert(Tensor{2, 2}, C)
    J = sqrt(det(Ct))
    Ic = trace(Ct) + 1
    lnJ = log(J)
    return μ/2 * (Ic - 3) + c2*(Ic - 3)^2 + c3 * (Ic - 3)^3 - μ*lnJ + λ/2 * lnJ^2
end

# Parameters used
function get_material()
    μ = 1.267e6
    λ = 1.457e7
    c2 = -μ/7.9
    c3 = μ/41
    return YeohMaterial(μ, λ, c2, c3)
end

# Takes the initial + current configuration, the fe_value object and the material parameters and
# computes the internal forces for the element.
function intf_element{T,Q,dim}(x::Vector{T}, X::Vector{Q}, fe_values::FEValues{dim}, 
                               material_parameters::YeohMaterial)
    # Closures for Forward Diff
    Ψ_Yeohh(C) = Ψ_Yeoh(C, material_parameters)
    ∂Ψ_Yeoh∂C = ForwardDiff.gradient(Ψ_Yeohh)
    S_Yeoh(C) = 2 * ∂Ψ_Yeoh∂C(C)
    
    # Reinterpret x and X as vectors of first order tensors
    n_basefuncs = n_basefunctions(get_functionspace(fe_values))
    @assert length(x) == length(X) == dim * n_basefuncs
    X_vec = reinterpret(Vec{dim, Q}, X, (n_basefuncs,))
    x_vec = reinterpret(Vec{dim, T}, x, (n_basefuncs,))
    
    reinit!(fe_values, X_vec)
    
    fe = [zero(Tensor{1, dim, T}) for i in 1:n_basefuncs]

    for q_point in 1:length(Ferrite.points(get_quadrule(fe_values)))
        F = function_vector_gradient(fe_values, q_point, x_vec)
        C = F' ⋅ F
        S = Tensor{2, 2}(S_Yeoh(vec(C)))
        P = F ⋅ S
        for i in 1:n_basefuncs
            fe[i] += P ⋅ shape_gradient(fe_values, q_point, i) * detJdV(fe_values, q_point)
        end
    end
    
    return reinterpret(T, fe, (dim * n_basefuncs,))
end

immutable CALMesh2D
    coord::Matrix{Float64}
    dof::Matrix{Int}
    edof::Matrix{Int}
    ex::Matrix{Float64}
    ey::Matrix{Float64}
end

# Reads the data from the .mat file and stores it in CALMesh2D as well as returning
# the dofs that are free/prescribed and fixed,
function read_data()
    vars = matread("cass2_mesh_data.mat");
    dof_fixed = convert(Vector{Int}, vec(vars["dof_fixed"]))
    Coord = vars["Coord"]'
    Ex, Ey = vars["Ex"], vars["Ey"]
    dof_prescr = convert(Vector{Int}, vec(vars["dof_prescr"]))
    Edof = convert(Matrix{Int}, vars["Edof"])'[2:end,:]
    Dof = convert(Matrix{Int}, vars["Dof"]')
    dof_free = convert(Vector{Int}, vec(vars["dof_free"]));
    return CALMesh2D(Coord, Dof, Edof, Ex, Ey), dof_prescr, dof_fixed, dof_free
end

# Assembles the global internal forces as well as the reaction forces
# for the prescribed dofs
function internal_forces!(fvec, X, x, mesh, material_parameters, dof_free, dof_fixed, fe_values, f_react)
    f_full = zeros(length(mesh.dof))

    fill!(fvec, 0.0)
    for i in 1:size(mesh.edof, 2)
        dofs = mesh.edof[:, i]
        fe = intf_element(x[dofs], X[dofs], fe_values, material_parameters)
        f_full[dofs] += fe
    end
    f_react[dof_fixed] = f_full[dof_fixed]
    copy!(fvec, f_full[dof_free])
end

function grad!(fvec, X, x, mesh, material_parameters, dof_free, fe_values)
    n_basefuncs = n_basefunctions(get_functionspace(fe_values))

    assembler = start_assemble()
    XX = zeros(2 * n_basefuncs)
    intf(x) = intf_element(x, XX, fe_values, material_parameters)
    grad! = ForwardDiff.jacobian(intf, mutates = true)
    Ke = zeros(2*n_basefuncs, 2*n_basefuncs)
    for i in 1:size(mesh.edof, 2)
        dofs = mesh.edof[:, i]
        copy!(XX, X[dofs])
        grad!(Ke, x[dofs])
        assemble(dofs, assembler, Ke)
    end
    K = end_assemble(assembler)
    return K[dof_free, dof_free]
end

function vtkoutput(pvd, i, mesh, X, x, topology, f_react)
    u = x - X
    nnodes = div(length(mesh.dof), 2)
    nrelem = size(mesh.edof, 2)

    disp = u
    disp = reshape(disp, (2, nnodes))
    disp = [disp; zeros(nnodes)']

    f_react = reshape(f_react, (2, nnodes))
    f_react = [f_react; zeros(nnodes)']


    vtkfile = vtk_grid(topology, mesh.coord, "box_$i")
    vtk_point_data(vtkfile, disp, "displacement")
    vtk_point_data(vtkfile, f_react, "reaction_forces")
    collection_add_timestep(pvd, vtkfile, float(i))
end

function go()
    # Get the data
    mesh, dof_prescr, dof_fixed, dof_free = read_data()
    
    # Get the material
    material_parameters = get_material()

    topology = topologyxtr(mesh.edof, mesh.coord, mesh.dof, 3)
    pvd = paraview_collection("box")
    
    function_space = Lagrange{2, RefTetrahedron, 1}()
    quad_rule = QuadratureRule(Dim{2}, RefTetrahedron(), 1)
    fe_values = FEValues(Float64, quad_rule, function_space)
    
    # Initialize
    X = vec(mesh.coord)
    x = copy(X)
    prev_x = copy(X)
    f_react = zeros(X)
   
    end_displacement = 30
    nsteps = 20
    for i in 1:nsteps
        # Set current config to correct value.
        prev_x[dof_prescr] = X[dof_prescr] + i / nsteps * end_displacement
        
        # Newton guess
        dx0 = zeros(length(dof_free))

        function f!(dx, fvec)
            copy!(x, prev_x)
            x[dof_free] += dx
            internal_forces!(fvec, X, x, mesh, material_parameters, dof_free, dof_fixed, fe_values, f_react)
        end

        function g!(dx, g)
            fvec = zeros(length(dof_free))
            copy!(x, prev_x)
            x[dof_free] += dx
            K = grad!(fvec, X, x, mesh, material_parameters, dof_free, fe_values)
            copy!(g, K)
        end

        println("Timestep $i out of $nsteps")
        df = DifferentiableSparseMultivariateFunction(f!, g!)
        res = nlsolve(df, dx0; ftol = 1e-6, iterations = 20, method=:newton, show_trace=true)
        if !converged(res)
            error("Global equation did not converge")
        end
        dx_conv = res.zero::Vector{Float64}
        # Update converged solution
        prev_x[dof_free] += dx_conv
        vtkoutput(pvd, i, mesh, X, x, topology, f_react)
    end
    vtk_save(pvd)
    return
end

go()