EIT.jl

EIT.jl

using Revise
using Plots
using Images
using Interpolations
using Gridap, GridapGmsh, SparseArrays
using Gridap.FESpaces
using Gridap.Geometry
using Gridap.ReferenceFEs
using LinearAlgebra
using IterativeSolvers
using FFTW
using JLD2

function produce_nonzero_positions(v, atol=1e-8, rtol=1e-5)
    approx_zero(x; atol=atol, rtol=rtol) = isapprox(x, 0; atol=atol, rtol=rtol)
    non_zero_count = count(x -> !approx_zero(x), v)
    non_zero_positions = zeros(Int, non_zero_count)
    non_zero_indices = findall(x -> !approx_zero(x), v)
    g_down = (x) -> x[non_zero_indices]
    g_up = (x) -> begin
        v = zeros(eltype(x), length(v))
        v[non_zero_indices] = x
        return v
    end
    return non_zero_count, non_zero_positions, g_down, g_up
end

struct ModeData
    g::Vector{Float64}
    f::Vector{Float64}
    u_g::Vector{Float64}
    u_f::Vector{Float64}
    U_d::FESpace
    uh_g::FEFunction
    uh_f::FEFunction
    singular_value::Float64
    function ModeData(op,g, singular_value = 0.0, use_LU = false)
        if use_LU
            u_g = op.ND.K_LU \ g
        else
            u_g = op.ND.K \ g
        end
        uh_g = FEFunction(op.ND.V, u_g)
        uh_f = interpolate_everywhere(uh_g, op.DD.V)
        
        U_d = TrialFESpace(op.DD.V, 1.0)
        u_f = uh_f.free_values
        f = op.DD.K * u_f
        #=
        U_d.dirichlet_values .= uh_f.dirichlet_values
        a(u,v) = ∫(∇(u)⋅ ∇(v))op.dΩ
        l(v) = 0
        op_d = AffineFEOperator(a,l,U_d,op.DD.V)
        f = op_d.op.vector
        if use_LU
            u_f = op.DD.K_LU \ f
        else
            u_f = op.DD.K \ f
        end
        =#
        new(g,f,u_g,u_f,U_d, uh_g,uh_f,singular_value)
    end
end

function mean_zero_basis(n::Int, m::Int)
    out = zeros(n)  
    norm = inv(sqrt((m + 1) * m))
    @inbounds for i in 1:m
        out[i] = norm 
    end
    out[m + 1] = -m * norm 
    out
end
function generate_basis_vectors(n::Int, i::Int)
    basis_matrix = zeros(n, i)
    for m in 1:i
        basis_matrix[:, m] = mean_zero_basis(n, m)
    end
    basis_matrix
end
export generate_basis_vectors
function subtract_column_mean!(A::Matrix, n::Int)
    # Calculate the mean of the first n entries of each column
    col_means = mean(A[1:n, :], dims=1)
    # Subtract the mean from each element of the column
    A .-= col_means
    return A
end

mutable struct NeumannData
   V::FESpace
   U::FESpace
   N::Int64
   K::AbstractArray{Float64,2} # force matrix
   K_LU # force matrix after LU decomposition
   m::Int64  # number of nonzeros in stiffness vector
   pos # positions of non zeros in stiffness vector
   down # function that reduces stiffness vector to non zero
   up # function that constructs stiffness vector from non zero
   u        
   v
   γ::FEFunction
   assem
   function NeumannData(mesh, conductivity, Ω, dΩ, Γ, dΓ, reffe)      
      V = TestFESpace(mesh, reffe, conformity=:H1)
      U = TrialFESpace(V)
      γ = interpolate_everywhere(conductivity, V)
      γ.free_values .= max.(γ.free_values, 1e-6)
      a(u,v) = ∫( γ* ∇(v) ⋅ ∇(u))dΩ 
      g_func = (x) -> 1.0
      l(v) = ∫( g_func * v )dΓ
      op_n = AffineFEOperator(a, l, U, V)
      g_m, g_pos, g_down, g_up = produce_nonzero_positions(op_n.op.vector)
      N = length(op_n.op.vector)
      K = op_n.op.matrix
      K_LU = lu(K)
      u = get_trial_fe_basis(U)
      v = get_fe_basis(V)
      assem = SparseMatrixAssembler(U, V)
      new(V,U,N,K,K_LU,g_m,g_pos,g_down,g_up,u,v,γ, assem)
   end
end

mutable struct DirichletData
   V::FESpace
   U::FESpace
   N::Int64
   K::AbstractArray{Float64,2} # force matrix
   K_LU # force matrix after LU decomposition
   m::Int64  # number of nonzeros in stiffness vector
   pos # positions of non zeros in stiffness vector
   down # function that reduces stiffness vector to non zero
   up # function that constructs stiffness vector from non zero
   u        
   v
   γ::FEFunction
   assem
        

   function DirichletData(mesh, conductivity, Ω, dΩ, Γ, dΓ, reffe)      
      V = TestFESpace(mesh, reffe, conformity=:H1, dirichlet_tags="boundary")
      f_func = (x) -> 1.0
      U_d = TrialFESpace(V,f_func)
      U = TrialFESpace(V)
      γ = interpolate_everywhere(conductivity, V)
      γ.free_values .= max.(γ.free_values, 1e-6) #ensures values > 0
      a(u,v) = ∫( γ* ∇(v) ⋅ ∇(u))dΩ 
      l(v) = 0
      op_d = AffineFEOperator(a, l, U_d, V)
      f_m, f_pos, f_down, f_up = produce_nonzero_positions(op_d.op.vector)
      N = length(op_d.op.vector)
      K = op_d.op.matrix
      K_LU = lu(K)
      u = get_trial_fe_basis(U)
      v = get_fe_basis(V)
      assem = SparseMatrixAssembler(U, V)
      new(V,U,N,K,K_LU,f_m,f_pos,f_down,f_up,u,v,γ, assem)
   end
end

mutable struct EITOperator
    mesh
    conductivity
    Ω
    dΩ
    Γ
    dΓ
    reffe
    max_modes::Int64 # Maximum number of modes
    ND::NeumannData
    DD::DirichletData
    function EITOperator(mesh,conductivity)
        Ω = Triangulation(mesh)
        dΩ = Measure(Ω, 2)
        Γ = BoundaryTriangulation(mesh, tags= "boundary")
        dΓ = Measure(Γ, 2)
        reffe = ReferenceFE(lagrangian, Float64, 1)
        ND = NeumannData(mesh, conductivity, Ω, dΩ, Γ, dΓ, reffe)
        DD = DirichletData(mesh, conductivity, Ω, dΩ, Γ, dΓ, reffe)
        max_modes = min(ND.m-1, DD.m) 
        new(mesh,conductivity,Ω,dΩ,Γ,dΓ,reffe,max_modes, ND,DD)
    end
end

export EITOperator

function FromImageData(img)
    
    # Convert image to Float64 and ensure positive values
    img_float = Float64.(img)
    img_float[img_float .== 0] .+= 1e-6

    # Define the interpolation
    itp = Interpolations.interpolate(img_float, BSpline(Linear()))
    size_x = size(img, 1)
    size_y = size(img, 2)
    # Define the scaling to map indices to [-1, 1] × [-1, 1]
    x_range = range(-1, 1, length=size_x)  # Map first dimension to [-1, 1]
    y_range = range(-1, 1, length=size_y)  # Map second dimension to [-1, 1]
    img_interp = Interpolations.scale(itp, x_range, y_range)

    domain = (-1,1,-1,1)
    partition = (size_x,size_y)
    mesh = CartesianDiscreteModel(domain,partition)
    return EITOperator(mesh, img_interp)
end
export FromImageData

function interpolate_from_image(img)
        img_float = Float64.(img)
    img_float[img_float .== 0] .+= 1e-6

    # Define the interpolation
    itp = Interpolations.interpolate(img_float, BSpline(Linear()))
    size_x = size(img, 1)
    size_y = size(img, 2)
    # Define the scaling to map indices to [-1, 1] × [-1, 1]
    x_range = range(-1, 1, length=size_x)  # Map first dimension to [-1, 1]
    y_range = range(-1, 1, length=size_y)  # Map second dimension to [-1, 1]
    img_interp = Interpolations.scale(itp, x_range, y_range)
end

struct NeumannSolData
    V::FESpace
    U::FESpace
    N::Int64
    m::Int64  # number of nonzeros in stiffness vector
    pos # positions of non zeros in stiffness vector
    down # function that reduces stiffness vector to non zero
    up # function that constructs stiffness vector from non zero
    u        
    v
    T::Union{SparseMatrixCSC{SparseVector{Float64}, Int64},Nothing}
    function NeumannSolData(nd::NeumannData, assem_tensor::Bool = false)
        if assem_tensor 
            T = assemble_tensor(nd.V::FESpace)
            new(nd.V,nd.U,nd.N,nd.m,nd.pos,nd.down,nd.up, nd.u,nd.v, T)
        else
            new(nd.V,nd.U,nd.N,nd.m,nd.pos,nd.down,nd.up, nd.u,nd.v, nothing)
        end
    end
end

struct DirichletSolData
    V::FESpace
    U::FESpace
    N::Int64
    m::Int64  # number of nonzeros in stiffness vector
    pos # positions of non zeros in stiffness vector
    down # function that reduces stiffness vector to non zero
    up # function that constructs stiffness vector from non zero
    u        
    v
    function DirichletSolData(nd::DirichletData)
        
        new(nd.V,nd.U,nd.N,nd.m,nd.pos,nd.down,nd.up, nd.u,nd.v)
    end
end
struct FixedData
    mesh
    Ω
    dΩ
    Γ
    dΓ
    reffe
    max_modes::Int64
    N::NeumannSolData
    D::DirichletSolData
    function FixedData(op::EITOperator, max_modes::Int64, assem_tensor::Bool = false)
        N = NeumannSolData(op.ND, assem_tensor)
        D = DirichletSolData(op.DD)
        new(op.mesh, op.Ω, op.dΩ, op.Γ, op.dΓ, op.reffe, max_modes, N, D) 
    end
end

mutable struct SolutionState
    γ::FEFunction
    K_n::AbstractArray{Float64,2}
    K_d::AbstractArray{Float64,2}
    Kn_LU
    Kd_LU
    ∇J::Union{Nothing,AbstractArray{Float64,2}}
    R2::Union{Nothing,Vector{Float64}} # R2 regularizer
    δTV::Vector{Float64} # TV regularizer
    LR::Union{Nothing,Vector{Float64}} # Later for learned regularization

    function SolutionState(fi::FixedData, γ::FEFunction)
        #=
        a(u,v) = ∫(γ * ∇(u)⋅∇(v))fi.dΩ
        l(v) = 0
        op_n = AffineFEOperator(a,l,fi.ND.U,fi.ND.V)
        K_n = op_n.op.matrix
        op_d = AffineFEOperator(a,l,fi.DD.U,fi.DD.V)
        K_d = op_d.op.matrix
        Kn_LU = lu(K_n)
        Kd_LU = lu(K_d)
        =#

        a(u, v) = ∫( γ * ∇(v) ⋅ ∇(u) )fi.dΩ
        
        assem_n = SparseMatrixAssembler(fi.N.U, fi.N.V)
        matcontribs_n = a(fi.N.u,fi.N.v)
        matdata_n =  collect_cell_matrix(fi.N.U,fi.N.V,matcontribs_n)
        K_n = assemble_matrix(assem_n,matdata_n)
        Kn_LU = lu(K_n)

        assem_d = SparseMatrixAssembler(fi.D.U, fi.D.V)
        matcontribs_d = a(fi.D.u,fi.D.v)
        matdata_d =  collect_cell_matrix(fi.D.U,fi.D.V,matcontribs_d)
        K_d = assemble_matrix(assem_d,matdata_d)
        Kd_LU = lu(K_d)

        δTV = zeros(fi.N.N)
        new(γ, K_n, K_d, Kn_LU, Kd_LU, nothing,nothing, δTV, nothing)
    end 
    function SolutionState(fi::FixedData, conductivity)
        γ= interpolate_everywhere(conductivity,fi.N.V)
        SolutionState(fi,γ)        
    end
    function SolutionState(fi::FixedData, γ_vec::Vector{Float64})
        @assert (length(γ_vec) == length(fi.N.N)) #Sanity check
        γ= FEFunction(fi.N.V,γ_vec)
        SolutionState(fi,γ)        
    end
end


function update_gamma!(fi::FixedData,sol::SolutionState, γ::FEFunction, clip::Bool = false; λ::Float64=1e-8)
    sol.γ = γ
    if clip
        sol.γ.free_values .= min.(max.(sol.γ.free_values, 1e-6),1)
    else    
        sol.γ.free_values .= max.(sol.γ.free_values, 1e-6)
    end
    a(u, v) = ∫( γ * ∇(v) ⋅ ∇(u) )fi.dΩ
        
    assem_n = SparseMatrixAssembler(fi.N.U, fi.N.V)
    matcontribs_n = a(fi.N.u,fi.N.v)
    matdata_n =  collect_cell_matrix(fi.N.U,fi.N.V,matcontribs_n)
    sol.K_n = assemble_matrix(assem_n,matdata_n)
    sol.K_n += λ * I
    assem_d = SparseMatrixAssembler(fi.D.U, fi.D.V)
    matcontribs_d = a(fi.D.u,fi.D.v)
    matdata_d =  collect_cell_matrix(fi.D.U,fi.D.V,matcontribs_d)
    sol.K_d = assemble_matrix(assem_d,matdata_d)
    sol.K_d += λ * I
    return nothing
end
function update_gamma!(fi::FixedData,sol::SolutionState, γ::Vector{Float64}, clip::Bool = false; λ::Float64 = 1e-8)
    update_gamma!(fi,sol, FEFunction(fi.N.V,γ), clip, λ = λ)
end

mutable struct OptMode 
    singular_value::Float64
    # There vectors are fixed and should not be changed. They represent the force vectors for the FEM problems:
    f::Vector{Float64} # force vector for the dirichlet problem
    g::Vector{Float64} # force vector for the neumann problem
    f_dirichlet::Vector{Float64} # truncated force vector for the dirichlet problem according to the FESpace. Do not use with matrix.
    U_d::Union{FESpace,Nothing} # FESpace for the dirichlet problem
    # These are the solutions to the FEM problems as vectors:
    u_f::Vector{Float64} # solution vector for the dirichlet problem
    u_g::Vector{Float64} # solution vector for the neumann problem
    # These are the solutions to the FEM problems as FEFunctions:
    uh_f::Union{FEFunction,Nothing} # solution FEFunction for the dirichlet problem
    uh_g::Union{FEFunction,Nothing} # solution FEFunction for the neumann problem
    
    ud_g::Union{FEFunction,Nothing} # solution FEFunction for the neumann problem interpolated on the dirichlet FESpace
    un_f::Union{FEFunction,Nothing} # solution FEFunction for the dirichlet problem interpolated on the neumann FESpace
    ∇J_n::Union{Vector{Float64},Nothing} # gradient of the objective function
    ∇J_d::Union{Vector{Float64},Nothing}
    w_d::Union{FEFunction,Nothing} # difference between the dirichlet and neumann solutions on the dirichlet FESpace
    w_n::Union{FEFunction,Nothing} # difference between the dirichlet and neumann solutions on the neumann FESpace
    error_d::Vector{Float64} # error between the dirichlet and neumann solutions on the dirichlet FESpace
    error_n::Vector{Float64} # error between the dirichlet and neumann solutions on the neumann FESpace
    λ_d::Union{FEFunction,Nothing} # Adjoint solution (if used)
    λ_n::Union{FEFunction,Nothing}

    λd::Vector{Float64} # Adjoint solution vector
    λn::Vector{Float64} # Adjoint solution vector
    n_vec::Union{Vector{Float64},Nothing}
    d_vec::Union{Vector{Float64},Nothing}

    Δu


    function OptMode(fi::FixedData,f,g,f_dirichlet,singular_value)
        u_f = zeros(fi.D.N)
        u_g = zeros(fi.N.N)
        λd = copy(u_f)
        λn = copy(u_g)
        U_d = TrialFESpace(fi.D.V,1.0)
        U_d.dirichlet_values .= f_dirichlet
        error_d = zeros(fi.D.m)
        error_n = zeros(fi.N.m)
        new(singular_value,f,g,f_dirichlet,U_d, u_f,u_g,nothing,nothing,nothing,nothing, nothing, nothing, nothing, nothing, error_d,error_n, nothing,nothing, λd,λn , nothing, nothing,0.0)
    end
    function OptMode(fi::FixedData,mode::ModeData)
        u_f = zeros(fi.D.N)
        u_g = zeros(fi.N.N)
        λd = copy(u_f)
        λn = copy(u_g)
        U_d = TrialFESpace(fi.D.V,1.0)
        U_d.dirichlet_values .= mode.uh_f.dirichlet_values
        error_d = zeros(fi.D.m)
        error_n = zeros(fi.N.m)
        new(mode.singular_value,mode.f,mode.g,mode.uh_f.dirichlet_values,U_d, u_f,u_g,nothing,nothing,nothing, nothing,nothing, nothing, nothing, nothing, error_d,error_n,nothing, nothing,λd,λn,nothing, nothing ,0.0)
    end
end

function solve_with_guess_cg!(A, b, x0, maxiter=500)
    cg!(x0,A, b; maxiter=maxiter)
end
function solve_with_guess_gmres!(A, b, x0, maxiter=500)
    gmres!(x0,A, b;  maxiter=maxiter)
end


function adj_gradient_from_mode!(fi::FixedData,sol::SolutionState, mode::OptMode, update_dirichlet::Bool=false, max_iter::Int64=500)
    # Solve dirichlet state equation

    #mode.u_f = sol.K_d \ mode.f
    if update_dirichlet
        solve_with_guess_gmres!(sol.K_d, mode.f, mode.u_f, max_iter)
        mode.uh_f = FEFunction(mode.U_d,mode.u_f)
        mode.un_f = interpolate_everywhere(mode.uh_f, fi.N.V)
    end
    #mode.u_g = sol.K_n \ mode.g
    solve_with_guess_cg!(sol.K_n, mode.g, mode.u_g, max_iter)

    mode.uh_g = FEFunction(fi.N.V,mode.u_g)
    mode.ud_g = interpolate_everywhere(mode.uh_g, fi.D.V)


    # normalize u_g:

    g_mean = mean(mode.ud_g.dirichlet_values)
    mode.ud_g.dirichlet_values .-= g_mean
    mode.ud_g.free_values .-= g_mean
    mode.uh_g.free_values .-= g_mean

    mode.w_d = interpolate_everywhere(mode.uh_g - mode.uh_f, fi.D.V)
    mode.w_n = interpolate_everywhere(mode.uh_g - mode.un_f, fi.N.V)

    a(u,v) = ∫(∇(u) ⋅ ∇(v))fi.dΩ
    l_n(v) = ∫(mode.w_n * v)fi.dΓ
    #l_d(v) = 0
    
    #U = TrialFESpace(fi.D.V, 1.0)
    #U.dirichlet_values .= mode.w_d.dirichlet_values
    #Hopefully the compiler optimizes this:
    mode.n_vec = AffineFEOperator(a,l_n, fi.N.U,fi.N.V).op.vector
    # This is false:
    #mode.d_vec = AffineFEOperator(a,l_d, U,fi.D.V).op.vector
    
    #λd = sol.K_d \ (2*mode.d_vec)
    #solve_with_guess_gmres!(λd, sol.K_d, 2*mode.d_vec)
    #λn = sol.K_n \ (2*mode.n_vec)

    solve_with_guess_cg!(sol.K_n, 2*mode.n_vec, mode.λn, max_iter)
    
    #mode.λ_d = FEFunction(U, λd)
    mode.λ_n = FEFunction(fi.N.U, mode.λn)

    mode.∇J_n =  interpolate_everywhere(∇(mode.uh_g)⋅∇(mode.λ_n), fi.N.V).free_values
    #mode.∇J_n = evaluate_bilinear_map(fi.N.T,mode.u_g,mode.λn)
    #mode.∇J_d =  interpolate_everywhere(∇(mode.uh_f)⋅∇(mode.λ_d), fi.N.V).free_values


    mode.error_d = mode.w_d.dirichlet_values    
    #mode.error_n = fi.N.down(sol.K_n * mode.w_n.free_values)
    mode.Δu = norm(mode.error_n)/length(mode.error_n)
    return mode.Δu 
end

function conductivity_from_image_data(img)
    
    # Convert image to Float64 and ensure positive values
    img_float = Float64.(img)
    img_float[img_float .== 0] .+= 1e-6

    # Define the interpolation
    itp = Interpolations.interpolate(img_float, BSpline(Linear()))
    size_x = size(img, 1)
    size_y = size(img, 2)
    # Define the scaling to map indices to [-1, 1] × [-1, 1]
    x_range = range(-1, 1, length=size_x)  # Map first dimension to [-1, 1]
    y_range = range(-1, 1, length=size_y)  # Map second dimension to [-1, 1]
    img_interp = Interpolations.scale(itp, x_range, y_range)

end

#This is one of the absolute bottel necks of the whole computation. For whatever reason i have not found a better way of extracting the values of the FEFunction on a grid, that runs reasonably fast.
function evaluate_fe_on_grid(u::FEFunction, n::Int)
    # Generate grid points over [-1,1] × [-1,1]
    xs = range(-1.0, 1.0; length=n)
    ys = range(-1.0, 1.0; length=n)

    # Prepare the array to store values
    values = Array{Float64}(undef, n, n)

    # Evaluate u at each point
    for i in 1:n, j in 1:n
        x = Point(xs[j], ys[i])  # Note: (x,y) but Julia arrays are row (y)-major
        values[j, i] = u(x)  # u must be callable at a Point
    end

    return values
end
export evaluate_fe_on_grid

function to_grayscale_image(values::AbstractMatrix) #Looks unnecessay
    min_val = minimum(values) #Think about this again: it should be about the same vaue as values contains. Else it would shift grayscale values to be darker/lighter than original.
    max_val = maximum(values)
    normalized = (values .- min_val) ./ (max_val - min_val + eps()) # avoid div-by-zero
    img = Gray.(normalized)
    return img
end
export to_grayscale_image

function get_TV_regularization!(fi::FixedData,sol::SolutionState,ϵ = 1e-6,max_iter::Int64=500)
    dΩ = fi.dΩ
    V = fi.N.V
    γ = sol.γ
    norm_sq_grad_γ = ∇(γ) ⋅ ∇(γ)
    norm_sq_grad_γ_ie = interpolate_everywhere(norm_sq_grad_γ, V)
    norm_sq_γ_ϵ_array = sqrt.(norm_sq_grad_γ_ie.free_values )
    TV_error = sum(norm_sq_γ_ϵ_array)
    norm_sq_γ_ϵ_array= 1 ./ (norm_sq_γ_ϵ_array.+ ϵ)
    norm_sq_γ_ϵ = FEFunction(V, norm_sq_γ_ϵ_array)
    l(v) = ∫(norm_sq_γ_ϵ*(∇(γ)⋅∇(v)))dΩ
    a(u, v) = ∫(u⋅v)dΩ
    op  = AffineFEOperator(a, l, V, V)
    A = op.op.matrix
    b = op.op.vector
    cg!(sol.δTV, A, b;  maxiter = max_iter)
    #δTV = A \ b
    return TV_error
end

function get_L1_regularization(fi::FixedData,sol::SolutionState)
    return norm(sol.γ.free_values,1), -sol.γ.free_values./abs.(sol.γ.free_values)
end
function get_R2_regularization(fi::FixedData,sol::SolutionState)
    return norm(sol.γ.free_values,2), -sol.γ.free_values
end

# This delivers worse optimization than Gauss Newton, despite knowing real conductivity:
function cheat_step(fi::FixedData,sol::SolutionState, γ_true::Vector{Float64}, J::AbstractArray{Float64,2})
    Δγ = γ_true - sol.γ.free_values
    grad_γ = J \ Δγ
    return J * grad_γ, grad_γ
end
# I'm not sure abou that type check: Rejects transpose and adjoint

function gauss_newton_lm_step_svd(J::Matrix{Float64}, r::Vector{Float64}, λ::Float64=1e-3)
    U, Σ, V = svd(J, full=false)
    n = length(Σ)
    Σ_damped = zeros(n)
    for i in 1:n
        Σ_damped[i] = Σ[i] / (Σ[i]^2 + λ) # Levenberg-Marquardt regularization
    end
    Jtr = J' * r
    δ = V * (Σ_damped .* (U' * r))
    return δ
end

function gaussian_smoother(F::Matrix{Float64},n::Int64,λ::Float64,s::Float64 = 1.0)
    C  = dct(dct(F,1),2)
    k  = repeat(0:n-1, 1, n)
    ℓ  = repeat((0:n-1)', n, 1)
    C .= C .* exp.(-λ * real((k.^2 + ℓ.^2)^s))
    Fsm = idct(idct(C,1),2)
    return Fsm
end

function full_step!(fi::FixedData,sol::SolutionState,opt_modes::Dict{Int64,OptMode}, steps::Int64, num_nodes::Int64,γ_true::Vector{Float64}; use_TV::Bool = false, β::Float64 = 1e-6, λ_CG::Float64=1e-8, λ_LM::Float64= 1e-3, clip::Bool = false, ϵ::Float64 = 1e-6, s::Float64 = 2.0, max_iter::Int64 = 1000)
    @assert length(opt_modes) >= num_nodes
    @assert s >= 0
    if use_TV
        J = zeros(fi.N.N,num_nodes+1) 
        err = zeros(num_nodes+1)
    else
        J = zeros(fi.N.N,num_nodes) 
        err = zeros(num_nodes)
    end
    for step in 1:steps
    @time begin
        if use_TV
            tv_task = Threads.@spawn begin
                return TV_error = get_TV_regularization!(fi, sol,ϵ)
            end, max_iter
        end
        if step == 1
            Threads.@threads for i in 1:num_nodes
                adj_gradient_from_mode!(fi,sol, opt_modes[i],true, max_iter)
            end
        else
            Threads.@threads for i in 1:num_nodes
                adj_gradient_from_mode!(fi,sol, opt_modes[i],false, max_iter)
            end
        end
        for i in 1:num_nodes
            J[:,i] = opt_modes[i].∇J_n
            err[i] = norm(opt_modes[i].error_d)
        end
        error = norm(γ_true - sol.γ.free_values)
        if use_TV
            TV_error = fetch(tv_task)
            err[num_nodes+1] = β * TV_error
            J[:,num_nodes+1] = β * sol.δTV
        end

        δγ = gauss_newton_lm_step_svd( copy(transpose(J)), err.^s,λ_LM)
        
        update_gamma!(fi,sol, sol.γ.free_values + δγ, clip,λ = λ_CG)
    end
        if use_TV
            println("Step $step done", " with error: $error  and TV_error: $TV_error")
        else
            println("Step $step done", " with error: $error")
        end
    end
end




#Remove Inefficiencies:
# This is for testing on a small mesh! DO NOT USE ELSEWHERE!!!!
# Try pkill -f julia if you're that stupid to try.
function assemble_tensor_test(V::FESpace)
    n = num_free_dofs(V)
    #dict = Dict{Tuple{Int, Int}, Vector{Float64}}()
    T = spzeros(SparseVector{Float64}, n, n)
    α = FEFunction(V, zeros(n))
    β = FEFunction(V, zeros(n))
    zero_vec = zeros(n)
    test_vec = zeros(n)
    for i in 1:n
        α.free_values[i] = 1.0
        for j in i:n
            β.free_values[j] = 1.0
            test_vec = interpolate_everywhere(∇(α) ⋅ ∇(β), V).free_values
            if !(zero_vec ≈ test_vec)
               #dict[(i, j)] = test_vec
               T[i, j] = sparsevec(test_vec)
               T[j, i] = T[i,j]
            end
            β.free_values[j] = 0.0
        end
        α.free_values[i] = 0.0
    println("step $i of $n done")
    end
    return T
end
function assemble_tensor(V::FESpace) # Still need a faster version
    n = num_free_dofs(V)
    #dict = Dict{Tuple{Int, Int}, Vector{Float64}}()
    T = spzeros(SparseVector{Float64}, n, n)
    α = FEFunction(V, zeros(n))
    β = FEFunction(V, zeros(n))
    cell_dof_ids = get_cell_dof_ids(V)
    m = length(cell_dof_ids[1])
    zero_vec = zeros(n)
    test_vec = zeros(n)
    for cell in cell_dof_ids
        for s in 1:m
            for t in s:m
                i = cell[s]
                j = cell[t]
                α.free_values[i] = 1.0
                β.free_values[j] = 1.0
                test_vec = interpolate_everywhere(∇(α) ⋅ ∇(β), V).free_values
                if !(test_vec ≈ zero_vec)
                    T[i,j] = sparsevec(test_vec)
                    T[j,i] = T[i,j]
                end
                α.free_values[i] = 0.0
                β.free_values[j] = 0.0
            end
        end
    end
    return T
end

function evaluate_bilinear_map(T::SparseMatrixCSC{SparseVector{Float64}, Int}, u::Vector{Float64}, v::Vector{Float64})
    n = size(T, 1)
    @assert length(u) == n && length(v) == n "Vectors u and v must have length $n"
    
    result = zeros(Float64, n)  # Initialize result vector
    # Iterate over non-zero entries of T
    for j in 1:n
        for i in nzrange(T, j)
            row = rowvals(T)[i]
            sparse_vec = nonzeros(T)[i]  # Get the sparse vector T[i,j]
            # Add contribution u[row] * v[j] * T[i,j] to result
            for (k, val) in zip(findnz(sparse_vec)[1], findnz(sparse_vec)[2])
                result[k] += u[row] * v[j] * val
            end
        end
    end
    return result
end