text:
#####
"""
Uses the NearestNeighbors.jl search library to determine which resonances should be grouped together.
"""
function partitionresonancesbyneighbors(x_in::Vector{Vector{T}},
    amplitudes_in::Vector{T},
    threshold_amplitude::T;
    radius = 1e-1,
    Minkowski_parameter = 3.5) where T <: Real

    @assert length(x_in) == length(amplitudes_in)
    @assert !isempty(x_in)

    # take largest amplitude, set as new partition centre.
    amplitudes = copy(amplitudes_in)
    x = copy(x_in)
    inds_buffer = collect(1:length(x))

    out_inds = Vector{Vector{Int}}(undef, 0)

    x_maxs = Vector{Vector{T}}(undef, 0)

    max_amplitude, max_ind = findmax(amplitudes)
    while !isempty(amplitudes) || max_amplitude < threshold_amplitude

        # set up tree.
        X = Matrix{T}(undef, length(x[1]), length(x))
        for r = 1:length(x)
            X[:,r] = x[r]
        end

        balltree = NearestNeighbors.BallTree(X,
            NearestNeighbors.Minkowski(Minkowski_parameter);
            reorder = false)

        inds = NearestNeighbors.inrange(balltree, x[max_ind], radius, true)
        push!(x_maxs, x[max_ind])

        # book keep.
        push!(out_inds, inds_buffer[inds])

        deleteat!(amplitudes, inds)
        deleteat!(x, inds)
        deleteat!(inds_buffer, inds)

        if isempty(amplitudes)
            return out_inds, x_maxs
        end
        max_amplitude, max_ind = findmax(amplitudes)
    end

    return out_inds, x_maxs
end

"""
αs and Ωs must not contain singlet groups.
"""
function partitionresonances(coherence_state_pairs_sys, ms_sys,
    αs::Vector{Vector{T}}, Ωs::Vector{Vector{T}},
    N_spins_sys::Vector{Int};
    α_relative_threshold = 0.05,
    Δc_partition_radius = 1e-1,
    ME::Vector{Vector{Vector{Int}}} = Vector{Vector{Vector{Int}}}(undef, 0),
    simple_coherence_atol::T = -1.2) where T

    N_groups = length(αs)

    part_inds_set = Vector{Vector{Vector{Int}}}(undef, N_groups)
    Δc_m_set = Vector{Vector{Vector{T}}}(undef, N_groups)

    # as = Vector{Vector{Vector{T}}}(undef, N_groups)
    # Fs = Vector{Vector{Vector{T}}}(undef, N_groups)
    as = Vector{Vector{T}}(undef, N_groups)
    Fs = Vector{Vector{T}}(undef, N_groups)

    N_systems = length(coherence_state_pairs_sys)
    @assert N_groups == N_systems == length(N_spins_sys)
    Δc_bar = Vector{Vector{Vector{T}}}(undef, N_systems)

    for i = 1:N_systems

        α_tol = α_relative_threshold*maximum(αs[i])
        inds_amp = findall(xx->(xx>α_tol), αs[i])

        c_states_prune = (coherence_state_pairs_sys[i])[inds_amp]
        αs_i_prune = αs[i][inds_amp]
        Ωs_i_prune = Ωs[i][inds_amp]

        c_m_r = collect( ms_sys[i][r] for (r,s) in c_states_prune )
        c_m_s = collect( ms_sys[i][s] for (r,s) in c_states_prune )
        Δc_m = collect( c_m_r[j] - c_m_s[j] for j = 1:length(c_m_r))

        # println("Δc_m = ", Δc_m)
        # println("ME[i] = ", ME[i])
        # println("N_spins_sys[i] = ", N_spins_sys[i])
        if !isempty(ME)
            if !isempty(ME[i])
                Δc_m = reduceΔc(Δc_m, ME[i], N_spins_sys[i])
            end
        end

        # # check if we should discard the non-simple coherences.
        # if simple_coherence_atol > 0
        #
        #     inds = findsimplecoherences(Δc_m, atol = simple_coherence_atol)
        #
        #     αs_i_prune = αs_i_prune[inds]
        #     Ωs_i_prune = Ωs_i_prune[inds]
        #     Δc_m = Δc_m[inds]
        # end

        part_inds, Δc_centroids = partitionresonancesbyneighbors(Δc_m,
            αs_i_prune, α_tol; radius = Δc_partition_radius)

        # partition_size = length(part_inds)
        # as[i] = Vector{Vector{T}}(undef, partition_size)
        # Fs[i] = Vector{Vector{T}}(undef, partition_size)
        # for m = 1:partition_size
        #
        #     as[i][m] = αs_i_prune[part_inds[m]]
        #     Fs[i][m] = Ωs_i_prune[part_inds[m]]
        # end
        as[i] = αs_i_prune
        Fs[i] = Ωs_i_prune

        part_inds_set[i] = part_inds
        Δc_m_set[i] = Δc_m

        Δc_bar[i] = Δc_centroids
    end

    # for i = 1:N_systems # over elements in a spin group.

    #     N_partition_elements = length(part_inds_compound[i])
    #     Δc_bar[i] = Vector{Vector{T}}(undef, N_partition_elements)

    #     for k = 1:N_partition_elements

    #         inds = part_inds_compound[i][k]

    #         Δc_bar[i][k] = Statistics.mean( Δc_m_compound[i][inds] )

    #         # # weighted mean.
    #         α = as[i][inds]
    #         Ω = Ωs[i][inds]
    #         # tmp = Δc_m_compound[i][inds]
    #         # Δc_bar[i][k] = sum(tmp[l] .* α[l]) / sum(α)
    #     end
    # end

    return as, Fs, part_inds_set, Δc_m_set, Δc_bar
end

function fitproxies(As::Vector{SHType{T}},
    dummy_SSFID::SST,
    λ0::T;
    names::Vector{String} = Vector{String}(undef, 0),
    config_path::String = "",
    Δcs_max_scalar_default = 0.2,
    κ_λ_lb_default = 0.5,
    κ_λ_ub_default = 2.5,
    u_min = Inf,
    u_max = Inf,
    Δr_default = 1.0,
    Δκ_λ_default = 0.05) where {T,SST}

    config_dict = Dict()
    if ispath(config_path)

        # TODO add error-handling if name is not found in the dictionary or filename does not exist.
        config_dict = JSON.parsefile(config_path)
    end

    cores = Vector{FIDModelType{T,SST}}(undef, length(As))

    for n = 1:length(As)
        A = As[n]

        # fit surrogate, save into `core`.
        cores[n] = fitproxy(dummy_SSFID, A, λ0, config_dict;
        compound_name = names[n],
        Δcs_max_scalar_default = Δcs_max_scalar_default,
        κ_λ_lb_default = κ_λ_lb_default,
        κ_λ_ub_default = κ_λ_ub_default,
        u_min = u_min,
        u_max = u_max,
        Δr_default = Δr_default,
        Δκ_λ_default = Δκ_λ_default)
    end

    return cores
end

function fitproxy(dummy_SSFID::SST,
    A::SHType{T},
    λ0::T,
    config_dict;
    compound_name::String = "",
    Δcs_max_scalar_default = 0.2,
    κ_λ_lb_default = 0.5,
    κ_λ_ub_default = 2.5,
    u_min = Inf,
    u_max = Inf,
    Δr_default = 1.0,
    Δκ_λ_default = 0.05)::FIDModelType{T,SST} where {T,SST}

    hz2ppmfunc = uu->(uu - A.ν_0ppm)*A.SW/A.fs
    ppm2hzfunc = pp->(A.ν_0ppm + pp*A.fs/A.SW)

    # allocate `core` data structure.
    N_singlets = length(A.αs_singlets)
    κs_λ_singlets = ones(T, N_singlets)
    κs_β_singlets = zeros(T, N_singlets)
    d_singlets = zeros(T, N_singlets)

    #N_β_vars_sys = A.N_spins_sys # no equivalence used.
    N_β_vars_sys::Vector{Int} = collect( length(A.Δc_bar[i][1]) for i = 1:length(A.Δc_bar) )

    # proxy placeholder.
    #qs = Vector{Vector{Function}}(undef, length(N_spins_sys))

    SSFID_obj = setupSSFIDparams(dummy_SSFID, A.part_inds_compound, N_β_vars_sys)

    # prepare configuration parameters.
    κ_λ_lb = κ_λ_lb_default
    κ_λ_ub = κ_λ_ub_default
    Δκ_λ = Δκ_λ_default
    Δcs_max::Vector{T} = collect( Δcs_max_scalar_default for i = 1:length(A.N_spins_sys))
    Δr = Δr_default

    d_max::Vector{T} = ppm2hzfunc.(Δcs_max) .- ppm2hzfunc(zero(T))

    if !isempty(config_dict) && !isempty(compound_name)
        dict = config_dict[compound_name] # TODO graceful error-handle.

        κ_λ_lb = dict["κ_λ lower bound"]
        κ_λ_ub = dict["κ_λ upper bound"]
        Δκ_λ = dict["κ_λ surrogate sampling step size"]
        if length(dict["Δcs_max"]) == length(A.N_spins_sys)
            Δcs_max = dict["Δcs_max"]
        # else
        #     println("Warning: problem an entry's Δcs_max value in config file. Using default scalar value for Δcs_max")
        end
        Δr = dict["radians surrogate sampling step size"]

        d_max = collect( ppm2hzfunc(Δcs_max[i])-ppm2hzfunc(0.0) for i = 1:length(Δcs_max) )
    end

    # threshold and partition the resonance components.
    if !isfinite(u_min) || !isfinite(u_max)
        Ωs_ppm = hz2ppmfunc.( combinevectors(A.Ωs) ./ (2*π) )
        min_ppm = minimum(Ωs_ppm) - 0.5
        max_ppm = maximum(Ωs_ppm) + 0.5
        u_min = ppm2hzfunc(min_ppm)
        u_max = ppm2hzfunc(max_ppm)
    end

    qs = setupcompoundpartitionitp(d_max,
        SSFID_obj.κs_β,
        A.Δc_bar,
        A.part_inds_compound,
        A.αs, A.Ωs,
        λ0, u_min, u_max;
        κ_λ_lb = κ_λ_lb,
        κ_λ_ub = κ_λ_ub,
        Δr = Δr,
        Δκ_λ = Δκ_λ)

    core = FIDModelType(qs, SSFID_obj, κs_λ_singlets, κs_β_singlets, d_singlets,
        Δcs_max, λ0)

    return core
end

function setupmixtureSH(target_names::Vector{String},
    H_params_path::String,
    dict_compound_to_filename,
    fs::T, SW::T, ν_0ppm::T;
    MEs::Vector{Vector{Vector{Vector{Int}}}} = Vector{Vector{Vector{Vector{Int}}}}(undef, 0),
    config_path = "",
    tol_coherence = 1e-2,
    α_relative_threshold = 0.05,
    Δc_partition_radius = 1e-1,
    simple_coherence_atol::T = -1.2) where {T <: Real, SST}

    ppm2hzfunc = pp->(ν_0ppm + pp*fs/SW)

    N_compounds = length(target_names)
    As = Vector{SHType{T}}(undef, N_compounds)

    for n = 1:N_compounds

        ME = Vector{Vector{Vector{Int}}}(undef, 0)
        if !isempty(MEs)
            ME = MEs[n]
        end

        αs, Ωs, part_inds_compound, Δc_m_compound, Δc_bar, N_spins_sys,
            αs_singlets, Ωs_singlets = setupcompoundSH(target_names[n],
            H_params_path, dict_compound_to_filename, ppm2hzfunc,
            fs, SW, ν_0ppm;
            config_path = config_path,
            ME = ME,
            tol_coherence = tol_coherence,
            α_relative_threshold = α_relative_threshold,
            Δc_partition_radius = Δc_partition_radius,
            simple_coherence_atol = simple_coherence_atol)

        As[n] = SHType(αs, Ωs, Δc_m_compound, part_inds_compound,
            Δc_bar, N_spins_sys, αs_singlets, Ωs_singlets, fs, SW, ν_0ppm)
    end

    return As
end

function setupcompoundSH(name, base_path, dict_compound_to_filename,
    ppm2hzfunc,
    fs::T, SW::T, ν_0ppm::T;
    ME::Vector{Vector{Vector{Int}}} = Vector{Vector{Vector{Int}}}(undef, 0),
    config_path::String = "",
    tol_coherence = 1e-2,
    α_relative_threshold = 0.05,
    Δc_partition_radius = 1e-1,
    simple_coherence_atol::T = -1.2) where T <: Real

    # TODO add error-handling if name is not found in the dictionary, or filename does not exist.
    load_path = joinpath(base_path, dict_compound_to_filename[name]["file name"])
    H_IDs, H_css, J_IDs, J_vals = loadcouplinginfojson(load_path)

    if ispath(config_path)

        # TODO add error-handling if name is not found in the dictionary or filename does not exist.
        db_dict = JSON.parsefile(config_path)
        dict = db_dict[name]

        tol_coherence = dict["coherence tolerance"]
        α_relative_threshold = dict["relative amplitude threshold"]
        Δc_partition_radius = dict["maximum Δc deviation"]
        simple_coherence_atol = dict["simple coherence absolute tolerance"]
    end

    # SH.
    # p_cs_sys, css_sys, cs_singlets, J_vals_sys, J_IDs_sys, intermediates_sys,
    # cs_LUT, cs_singlets_compact, N_spins_singlet, css, p_compound,
    # cs_compound = fetchSHparameters(H_IDs, H_css, J_IDs, J_vals)
    J_inds_sys, J_inds_sys_local, J_IDs_sys, J_vals_sys, H_inds_sys,
        cs_sys, H_inds_singlets, cs_singlets, H_inds, J_inds,
        g = setupcsJ(H_IDs, H_css, J_IDs, J_vals)

    N_spins_singlet = length.(H_inds_singlets)

    N_spins_sys = collect( length(cs_sys[m]) for m = 1:length(cs_sys) )
    intermediates_sys = prepcouplingalgorithm(N_spins_sys)

    # css_sys becomes cs_sys
    # cs_singlets_compact becomes cs_singlets.

    αs_inp, Ωs_inp, coherence_mat_sys, eig_vals_sys, Q_sys,
    coherence_state_pairs_sys, H_sys, states_sys, ms_sys,
    M_sys = setupcompoundspectrum!(cs_sys,
        J_vals_sys, J_inds_sys_local, intermediates_sys,
        ppm2hzfunc, cs_singlets, N_spins_singlet, fs, SW;
        tol_coherence = tol_coherence)

    k = findfirst(xx->length(xx)==1, αs_inp)
    αs_spin_sys = copy(αs_inp)
    Ωs_spin_sys = copy(Ωs_inp)
    αs_singlets = Vector{T}(undef, 0)
    Ωs_singlets = Vector{T}(undef, 0)
    if typeof(k) == Int
        # the case there are singlet groups in this compound.
        αs_spin_sys = αs_spin_sys[1:k-1]
        Ωs_spin_sys = Ωs_spin_sys[1:k-1]

        αs_singlets = collect( αs_inp[l][1] for l = k:length(αs_inp) )
        Ωs_singlets = collect( Ωs_inp[l][1] for l = k:length(Ωs_inp) )
    end

    # spectra eval func for spin systems.
    N_spins_sys = collect( length(cs_sys[i]) for i = 1:length(cs_sys))

    αs, Ωs, part_inds_compound,
    Δc_m_compound, Δc_bar = partitionresonances(coherence_state_pairs_sys,
    ms_sys, αs_spin_sys, Ωs_spin_sys, N_spins_sys;
    ME = ME,
    α_relative_threshold = α_relative_threshold,
    Δc_partition_radius = Δc_partition_radius,
    simple_coherence_atol = simple_coherence_atol)

    # f = uu->evalcLcompoundviapartitions(uu, d,
    # αs, Ωs, κs_λ, κs_β, λ0, Δc_m_compound, part_inds_compound)

    return αs, Ωs, part_inds_compound, Δc_m_compound, Δc_bar, N_spins_sys,
    αs_singlets, Ωs_singlets
end

function setupSSFIDparams(dummy_SSFID::SpinSysParamsType1{T}, part_inds_compound, N_β_vars_sys)::SpinSysParamsType1{T} where T
    L = length(part_inds_compound)

    κs_λ = ones(T, L)
    κs_β = collect( zeros(T, N_β_vars_sys[i]) for i = 1:length(N_β_vars_sys))
    #d = rand(length(αs))
    d = zeros(T, L)

    return constructorSSFID(dummy_SSFID, κs_λ, κs_β, d)
end

function setupSSFIDparams(dummy_SSFID::SpinSysParamsType2{T}, part_inds_compound::Vector{Vector{Vector{Int}}}, N_β_vars_sys)::SpinSysParamsType2{T} where T

    N_sys = length(part_inds_compound)
    κs_λ = Vector{Vector{T}}(undef, N_sys)
    d = Vector{Vector{T}}(undef, N_sys)

    for i = 1:length(d)
        N_partition_elements = length(part_inds_compound[i])

        κs_λ[i] = ones(T, N_partition_elements)
        d[i] = zeros(T, N_partition_elements)
    end

    κs_β = collect( zeros(T, N_β_vars_sys[i]) for i = 1:length(N_β_vars_sys))

    return constructorSSFID(dummy_SSFID, κs_λ, κs_β, d)
end

function getNβ(A::CompoundFIDType{T,SST}) where {T,SST}

    counter_sys = 0
    for i = 1:length(A.κs_β)
        counter_sys += length(A.κs_β[i])
    end

    N_β = counter_sys + length(A.β_singlets)

    return N_β
end


text:
df_show(args...) = Base.show(args...)

df_show_type(io::IO, ::Type{T}) where {T} = dfprintln(io, crayon"red", "$T:", crayon"reset")
df_show_type(io::IO, x) = df_show_type(io, typeof(x))

function df_show(io::IO, block::InputData)
    s = """Block name: $(block.name)
    Block option: $(block.option)
    Block data:
    """
    dfprintln(io, s)
    dfprintln(io, string(block.data) * "\n\n")
    return
end

function df_show(io::IO, band::Band)
    df_show_type(io, band)
    string = """
    k_points of length $(length(band.k_points_cart)):
    cart:    $(band.k_points_cart[1]) -> $(band.k_points_cart[end])
    """
    !isempty(band.k_points_cryst) &&
        (string *= "cryst:   $(band.k_points_cryst[1]) -> $(band.k_points_cryst[end])\n")
    string *= """eigvals: $(band.eigvals[1]) -> $(band.eigvals[end])
    extra:   $(band.extra)
    """
    dfprintln(io, string)
    return
end

function df_show(io::IO, job::Job)
    reset = crayon"reset"
    fieldns = [:name, :version]
    fs = string.(filter(x -> !isempty(x), getfield.((job,), fieldns)))
    fns = string.(fieldns)
    insert!(fns, 2, "dir")
    insert!(fs, 2, Jobs.main_job_dir(job))
    push!(fns, "versions")
    versions = Client.versions(job)
    push!(fs, join(string.(versions), ", "))
    timestamp = Client.submission_time(job)
    if timestamp != 0
        push!(fns, "last submission")
        push!(fs, string(round(Dates.unix2datetime(timestamp), Second)))
    end
    push!(fns, "state")
    state = Client.state(job)
    push!(fs, string(state))
    lfns = maximum(length.(fns)) + maximum(length.(fs)) + 4
    line = "+"
    for i in 1:div(lfns, 2) + 1
        line *= "-"
    end
    dfprint(io, crayon"cyan", line[1:end-2])
    dfprint(io, crayon"cyan", "DFJOB")
    for i in 1:div(lfns, 2)
        line *= "-"
    end
    line *= "+"
    totlen = length(line)
    dfprintln(io, crayon"cyan", line[div(lfns, 2)+6:end], reset)
    lfns = maximum(length.(string.(fns)))
    for (fn, f) in zip(fns, fs)
        isname = fn == "name"
        l = length(fn)
        fn *= ":"
        for i in 1:lfns-l
            fn *= " "
        end
        dfprint(io, crayon"cyan", "|", reset)
        isname ? dfprint(io, " $fn ", crayon"magenta", "$f", reset) :
        dfprint(io, " $fn $f", reset)
        for i in 1:totlen-(length(fn)+4+length(f))
            dfprint(io, " ")
        end
        dfprintln(io, crayon"cyan", "|", reset)
    end
    is = job.calculations
    last = !isempty(versions) ? Client.last_running_calculation(job) : -1
    if !isempty(is)
        dfprintln(io, crayon"cyan", line, reset)
        dfprintln(io, reset, "(", crayon"green", "scheduled", reset, ", ", crayon"red",
                  "not scheduled", reset, ")")
        ln = maximum(length.(string.(map(x -> x.name, is))))
        for i in is
            n = i.name
            cr = i.run ? crayon"green" : crayon"red"
            dfprint(io, cr,
                    i == last ? (state == Jobs.Pending || state == Jobs.Running ? "\t$n <- running\n" : "\t$n <- ran last\n") :
                    "\t$n\n")
        end
    end
    dfprint(io, reset)
    return
end

function write_flags(io, fls, prestr = "")
    fl = maximum(length.(string.(keys(fls))))
    for (f, v) in fls
        fs = string(f)
        l = length(fs)
        for i in 1:fl-l
            fs *= " "
        end
        dfprintln(io, crayon"cyan", prestr * "\t$fs", crayon"yellow", " => ",
                  crayon"magenta", "$v")
    end
end

function df_show(io::IO, c::Calculation)
    df_show_type(io, c)
    s = """name  = $(c.name)
    exec = $(c.exec.exec)
    run   = $(c.run)
    data  = $([e.name for e in c.data])
    flags:"""
    dfprintln(io, s)

    if !isempty(c.flags)
        if eltype(c) == QE
            namelist_flags = Dict{Symbol,Dict{Symbol,Any}}()
            info = Calculations.qe_calculation_info(c)
            if info !== nothing
                flag_keys = keys(c.flags)
                for i in info.control
                    for f in i.flags
                        if f.name ∈ flag_keys
                            if haskey(namelist_flags, i.name)
                                namelist_flags[i.name][f.name] = c[f.name]
                            else
                                namelist_flags[i.name] = Dict{Symbol,Any}()
                                namelist_flags[i.name][f.name] = c[f.name]
                            end
                        end
                    end
                end
                for (inf, flgs) in namelist_flags
                    dfprintln(io, crayon"green", "\t&$inf", crayon"reset")
                    write_flags(io, flgs, "\t\t")
                end
            else
                write_flags(io, c.flags)
            end
        else
            write_flags(io, c.flags)
        end
    end
    dfprint(io, crayon"reset")
    return
end

function df_show(io::IO, flag_info::Calculations.QEFlagInfo{T}) where {T}
    df_show_type(io, flag_info)
    dfprintln(io, crayon"yellow", "name : $(flag_info.name)")
    dfprintln(io, crayon"cyan", "description:", crayon"reset")
    dfprint(io, "\t" * replace(flag_info.description, "\n" => "\n\t"))
    return
end

function df_show(io::IO, flag_info::Calculations.ElkFlagInfo{T}) where {T}
    df_show_type(io, flag_info)
    dfprintln(io, crayon"yellow", "name : $(flag_info.name)")
    dfprintln(io, crayon"green", "default:", crayon"reset")
    if flag_info.default != nothing
        dfprintln(io, "\t" * "$(flag_info.default)")
    else
        dfprintln(io, "\t" * "nothing")
    end
    dfprintln(io, crayon"cyan", "description:", crayon"reset")
    dfprintln(io, "\t" * replace(flag_info.description, "\n" => "\n\t"))
    return
end

function df_show(io::IO, info::Calculations.ElkControlBlockInfo)
    dfprintln(io, crayon"yellow", "name = $(info.name)")
    dfprintln(io, crayon"cyan", "flags:", crayon"reset")
    for flag in info.flags
        dfprintln(io, "\t$(flag.name)")
    end
    dfprintln(io, crayon"cyan", "description:", crayon"reset")
    dfprintln(io, "\t" * replace(info.description, "\n" => "\n\t"))
    dfprint(io, crayon"reset")
    return
end

function df_show(io::IO, el::Element)
    for f in fieldnames(typeof(el))
        dfprintln(io, crayon"red", "$f: ", crayon"reset", "$(getfield(el, f))")
    end
    return dfprint(io, crayon"reset")
end

function df_show(io::IO, str::Structure)
    dfprintln(io, crayon"cyan", "Structure", crayon"reset")
    dfprintln(io, crayon"red", "    cell parameters:")
    dfprint(io, crayon"reset",
            "\t a = $((str.cell[:,1]...,))\n\t b = $((str.cell[:,2]...,))\n\t c = $((str.cell[:,3]...,))\n")
    dfprintln(io, crayon"red", "    nat:", crayon"reset", " $(length(str.atoms))")
    dfprintln(io, crayon"red", "    ntyp:", crayon"reset", " $(length(unique(str.atoms)))")
    show(io, str.atoms)
    return dfprintln(io, crayon"reset")
end

function df_show(io::IO, at::Atom)
    dfprintln(io)
    dfprintln(io, crayon"cyan", "Atom")
    dfprintln(io, crayon"red", "    name: ", crayon"reset", "$(at.name)")
    for f in fieldnames(typeof(at))[3:end-1]
        f == :pseudo && continue
        fld = getfield(at, f)
        if f in (:position_cart, :position_cryst)
            dfprintln(io, crayon"red", "    $f: ", crayon"reset", "$((fld...,))")
        else
            dfprintln(io, crayon"red", "    $f: ", crayon"reset", "$fld")
        end
    end
    dfprintln(io, crayon"red", "    dftu:", crayon"reset")
    df_show(io, at.dftu)
    dfprintln(io, crayon"reset")
    return
end

function df_show(io::IO, proj::Projection)
    println(io, crayon"cyan", "Orbital: ", crayon"reset", "$(proj.orbital.name)")
    println(io, crayon"red", "start index: ", crayon"reset", "$(proj.start)")
    return println(io, crayon"red", "last index: ", crayon"reset", "$(proj.last)")
end

ext: jl


text:
#Running mean
"""
    running_mean!(m, x)

Compute (in-place) the mean, `m`, of a set of vectors, `x`.
* `x` must be a matrix of Floats, with each column of `x` corresponding to a vector.
* `m` must be a vector of Floats of appropriate length (equal to the number of rows in `x`).
"""
function running_mean!(m::AbstractVector{T}, x::AbstractMatrix{T}) where T<:AbstractFloat
    k = zero(T)
    fill!(m,zero(T))
    @inbounds for j ∈ axes(x,2)
        k = 1/j
        @turbo for i ∈ axes(x,1)
            temp = (x[i,j] - m[i])
            m[i] += temp*k
        end
    end
end

"""
    running_mean(x)

Compute the mean of a set of vectors `x`.
* `x` must be a matrix of Floats, with each column of `x` corresponding to a vector.
"""
function running_mean(x::AbstractMatrix{T}) where T<:AbstractFloat
    m = zeros(T,size(x,1))
    running_mean!(m, x)
    return m
end

#Update running mean with a nth data point where n is the running index
"""
    update_running_mean!(m, x, n)

Update the mean, `m`, of some set of `n-1` vectors to be the mean of the set of `n` vectors where `x` is the ``n^th`` vector.
* `m` and `x` must be vectors of equal length.
* Typically used when the mean of a part of a dataset has been previously computed and needs to be updated with a new data point (vector).
"""
function update_running_mean!(mean::AbstractVector{T}, data::AbstractVector{T}, n::Int) where T<:AbstractFloat
    k = 1/n
    @inbounds @turbo for i ∈ eachindex(mean)
        temp = (data[i] - mean[i])*k
        mean[i] += temp
    end
end

#Running mean and variance
"""
    running_meanvar!(m, v, x)

Compute (in-place) the mean, `m`, and variance along each dimension, `v`, of a set of vectors, `x`.
* `x` must be a matrix of Floats, with each column of `x` corresponding to a vector.
* `m` must be a vector of Floats of appropriate length (equal to the number of rows in `x`).
* `v` must be a vector of Floats of appropriate length (equal to the number of rows in `x`).
"""
function running_meanvar!(m::AbstractVector{T}, v::AbstractVector{T}, x::AbstractMatrix{T}) where T<:AbstractFloat
    k = zero(T)
    fill!(m,zero(T))
    fill!(v,zero(T))
    @inbounds for j ∈ axes(x,2)
        k = 1/j
        @inbounds for i ∈ axes(x,1)
            temp = (x[i,j] - m[i])
            m[i] += temp*k
            v[i] += temp*(x[i,j] - m[i])
        end
    end
    N = 1/(size(x,2)-1)
    @turbo for i ∈ eachindex(v)
        v[i] *= N
    end
end

"""
    running_meanvar(x)

Compute the mean and variance along each dimension of a set of vectors, `x`.
* `x` must be a matrix of Floats, with each column of `x` corresponding to a vector.
"""
function running_meanvar(x::AbstractMatrix{T}) where T<:AbstractFloat
    m = zeros(T,size(x,1))
    v = zeros(T,size(x,1))
    running_meanvar!(m,v,x)
    return m, v
end

#Update running mean and variance with a nth data point where n is the running index
"""
    update_running_meanvar!(m, v, x, n)

Update the mean, `m`, and variance along each dimension, `v`, of some set of `n-1` vectors to be the mean and variance of the set of `n` vectors where `x` is the ``n^th`` vector (for `n>=3`).
* `m`, `v` and `x` must be vectors of equal length.
* Typically used when the mean and variance of a part of a dataset has been previously computed and needs to be updated with a new data point (vector).

!!! note
    `n` must be 3 or greater.
    [`running_meanvar()`](@ref) or [`running_meanvar!()`](@ref) can be used to obtain an initial estimate of `m` and `v` from at least two vectors.
"""
function update_running_meanvar!(m::AbstractVector{T}, v::AbstractVector{T}, x::AbstractVector{T}, n::Int) where T<:AbstractFloat
    n >=3 || error("n must be greater than 2. Use running_meanvar() to obtain an initial estimate of mean and variance from at least two vectors before using update_running_meanvar() to update the estimates.")
    k = 1/n
    k1 = 1/(n-1)
    @inbounds for i ∈ eachindex(m,v,x)
        temp = (x[i] - m[i])
        m[i] += temp*k
        v[i] *= (n-2)
        v[i] += temp*(x[i] - m[i])
        v[i] *= k1
    end
end

#running mean and covariance
"""
    running_meancov!(m, C, x)

Compute (in-place) the mean, `m`, and covariance matrix, `C`, of a set of vectors, `x`.
* `x` must be a matrix of Floats, with each column of `x` corresponding to a vector.
* `m` must be a vector of Floats of appropriate length (equal to the number of rows in `x`).
* `C` must be a square matrix of Floats of appropriate size (side equal to the number of rows in `x`).
"""
function running_meancov!(m::AbstractVector{T}, C::AbstractMatrix{T}, x::AbstractMatrix{T}) where T<:AbstractFloat
    k = zero(T)
    temp = Vector{T}(undef,length(m))
    fill!(m,zero(T))
    fill!(C,zero(T))
    @inbounds for n ∈ axes(x,2)
        k = 1/n
        @turbo for i ∈ axes(x,1)
            temp[i] = (x[i,n] - m[i])
            m[i] += temp[i]*k
        end
        @turbo for i ∈ axes(x,1), j ∈ axes(x,1)
            C[j,i] += (x[j,n] - m[j])*temp[i]
        end
    end
    N = 1/(size(x,2)-1)
    @turbo for i ∈ eachindex(C)
        C[i] *= N
    end
    C += C'
    C /= 2
end

"""
    running_meancov(x)

Compute the mean and covariance matrix a set of vectors, `x`.
* `x` must be a matrix of Floats, with each column of `x` corresponding to a vector.
"""
function running_meancov(x::AbstractMatrix{T}) where T<:AbstractFloat
    m = Vector{T}(undef,size(x,1))
    C = Matrix{T}(undef,size(x,1),size(x,1))
    running_meancov!(m,C,x)
    return m, C
end

#Update running mean and covariance with a nth data point where n is the running index
"""
    update_running_meancov!(m, C, x, n)

Update the mean, `m`, and covariance, `C`, of some set of `n-1` vectors to be the mean and variance of the set of `n` vectors where `x` is the ``n^th`` vector (for `n>=3`).
* `m` and `x` must be vectors of equal length.
* `C` must be a square matrix of side equal to the length of `x`.
* Typically used when the mean and covariance of a part of a dataset has been previously computed and needs to be updated with a new data point (vector).

!!! note
    `n` must be 3 or greater.
    [`running_meancov()`](@ref) or [`running_meanvcov!()`](@ref) can be used to obtain an initial estimate of `m` and `v` from at least two vectors.
"""
function update_running_meancov!(m::AbstractVector{T}, C::AbstractMatrix{T}, x::AbstractVector{T}, n::Int) where T<:AbstractFloat
    n >=3 || error("n must be greater than 2. Use running_meanvar() to obtain an initial estimate of mean and variance from at least two vectors before using update_running_meanvar() to update the estimates.")
    k = 1/n
    k1 = 1/(n-1)
    temp = Vector{T}(undef,length(m))
    @turbo for i ∈ eachindex(x)
        temp[i] = (x[i] - m[i])
        m[i] += temp[i]*k
    end
    @turbo for i ∈ eachindex(x), j ∈ eachindex(x)
        C[j,i] *= (n-2)
        C[j,i] += (x[j] - m[j])*temp[i]
        C[j,i] *= k1
    end
    C += C'
    C /= 2
end


text:
#=======================================================================

We now use the PermLTCSC structure to solve systems.
The other code is just for testing and recording other things we tried.

=========================================================================#
#=
function forwardsolve!(lt::PermLTCSC{Ti,StaticArrays.SArray{Tuple{D,D},T,2,L}},
    b::AbstractArray{SVector{D, T}}) where {Ti, D, T, L}
    n = length(lt.colptr)-1

    for j in 1:n
        col = lt.rowval[lt.colptr[j]]

        b[col] = lt.nzval[lt.colptr[j]] \ b[col]

        for ii in (lt.colptr[j]+1):(lt.colptr[j+1]-1)
            i = lt.rowval[ii]
            b[i] -= lt.nzval[ii] * b[col]
        end
    end
end

function forwardsolve!(lt::PermLTCSC{Ti,Tv}, b::Vector{Tv}) where {Ti,Tv}
    n = length(lt.colptr)-1

    @inbounds for j in 1:n
        col = lt.rowval[lt.colptr[j]]

        b[col] /= lt.nzval[lt.colptr[j]]

        for ii in (lt.colptr[j]+1):(lt.colptr[j+1]-1)
            i = lt.rowval[ii]
            b[i] -= b[col] * lt.nzval[ii]
        end
    end
end
=#

import Laplacians.pcg

function forwardsolve!(lt::PermLTCSC, b)
    n = length(lt.colptr)-1

    @inbounds for j in 1:n
        col = lt.rowval[lt.colptr[j]]

        b[col] = lt.nzval[lt.colptr[j]] \ b[col]

        for ii in (lt.colptr[j]+1):(lt.colptr[j+1]-1)
            i = lt.rowval[ii]
            b[i] -= lt.nzval[ii] * b[col]
        end
    end
end

# this works only when L has been permuted to lower tri.
function forwardsolve!(L::LTCSC, b::Vector)

    n = length(b)

    @inbounds for j in 1:n

        b[j] = L.nzval[L.colptr[j]] \ b[j]

        for ii in (L.colptr[j]+1):(L.colptr[j+1]-1)
            i = L.rowval[ii]
            b[i] -= L.nzval[ii] * b[j]
        end
    end

end

function backsolve!(lt::PermLTCSC, b)
    n = length(lt.colptr)-1

    @inbounds for j in n:-1:1
        col = lt.rowval[lt.colptr[j]]

        for ii in (lt.colptr[j]+1):(lt.colptr[j+1]-1)
            b[col] -= lt.nzval[ii]' * b[lt.rowval[ii]]
        end
        b[col] = lt.nzval[lt.colptr[j]] \ b[col]
    end
end

# this works only when L has been permuted to lower tri.
function backsolve!(L::LTCSC, b::Vector)
    n = length(b)

    @inbounds for j in n:-1:1
        for ii in (L.colptr[j]+1):(L.colptr[j+1]-1)
            b[j] -= L.nzval[ii]' * b[L.rowval[ii]]
        end
        b[j] = L.nzval[L.colptr[j]] \ b[j]
    end
end

# For vectors y of block vectors, and block matrix lt.
function solver!(lt::AbsLTCSC{Ti,StaticArrays.SArray{Tuple{D,D},T,2,L}},
    y::AbstractArray{SVector{D, T}}) where {Ti, D, T, L}

    n = length(lt.colptr)-1

    @assert length(y) == n

    forwardsolve!(lt, y)
    @inbounds for i in 1:length(y)
        y[lt.rowval[lt.colptr[i]]] = lt.nzval[lt.colptr[i]] * y[lt.rowval[lt.colptr[i]]]
    end
    backsolve!(lt, y)
end

# for real entry  lt, and real vectors
function solver!(lt::AbsLTCSC{Ti,Tv},
    y::Vector{Tv}) where {Ti, Tv}

    n = length(lt.colptr)-1

    @assert length(y) == n

    forwardsolve!(lt, y)
    @inbounds for i in 1:length(y)
        y[lt.rowval[lt.colptr[i]]] = lt.nzval[lt.colptr[i]] * y[lt.rowval[lt.colptr[i]]]
    end
    backsolve!(lt, y)
end

# For vectors y of block vectors, and block matrix lt.
function solver(lt::AbsLTCSC{Ti,StaticArrays.SArray{Tuple{D,D},T,2,L}},
    b::AbstractArray{SVector{D, T}}) where {Ti, D, T, L}

    y = copy(b)
    solver!(lt, y)
    return y
end

# For vectors y of real entries, and block matrix lt.
function solver(lt::AbsLTCSC{Ti,StaticArrays.SArray{Tuple{D,D},T,2,L}},
    b::Vector{T}) where {Ti, D, T, L}
    y = copy(reinterpret(SVector{D,T},b))
    solver!(lt, y)
    return collect(reinterpret(T,y))
end

# for real entry  lt, and real vectors
function solver(lt::AbsLTCSC{Ti,Tv},
    b::Vector{Tv}) where {Ti, Tv}
    y = copy(b)
    solver!(lt, y)
    return y
end

function pcg(mat::BlockCSC{D,Tv,L}, b::Vector{Tv}, pre::Function;
    tol::Real=1e-6, maxits=Inf, maxtime=Inf, verbose=false, pcgIts=Int[],
    stag_test::Integer=0)::Vector{Tv} where {D,Tv,L}

    FlatVec(
        pcg(mat, BlockVec(D,b), pre,
        tol=tol, maxits=maxits, maxtime=maxtime,
        pcgIts=pcgIts,
        verbose=verbose,
        stag_test = stag_test)
     )

end

function pcg_inplace(mat::BlockCSC{D,Tv,L}, b::Vector{Tv}, pre::Function;
    tol::Real=1e-6, maxits=Inf, maxtime=Inf, verbose=false, pcgIts=Int[],
    stag_test::Integer=0)::Vector{Tv} where {D,Tv,L}

    FlatVec(
        pcg_inplace(mat, BlockVec(D,b), pre,
        tol=tol, maxits=maxits, maxtime=maxtime,
        pcgIts=pcgIts,
        verbose=verbose,
        stag_test = stag_test)
     )

end

#import Laplacians.pcg
import Laplacians.axpy2!

function pcg(mat::BlockCSC{D,Tv,L},
    b::AbstractVector{SArray{Tuple{D},Tv,1,D}},
    pre::Function;
    tol::Real=1e-6,
    maxits=Inf,
    maxtime=Inf,
    verbose=false,
    pcgIts=Int[],
    stag_test::Integer=0) where {D,Tv,L}

    local al::Tv

    n = size(mat,2)

    nb = norm(b)

    # If input vector is zero, quit
    if nb == 0
    return zeros(size(b))
    end

    x = zeros(SVector{D,Tv} ,n)
    bestx = zeros(SVector{D,Tv}, n)
    bestnr = one(Tv)

    r = copy(b)
    z = pre(r)
    p = copy(z)

    rho = dot(r, z)
    best_rho = rho
    stag_count = 0

    t1 = time()

    itcnt = 0
    while itcnt < maxits
        itcnt = itcnt+1

        q = mat*p

        pq = dot(p,q)

        if (pq < eps(Tv) || isinf(pq))
            if verbose
                println("PCG Stopped due to small or large pq")
            end
            break
        end

        al = rho/pq

        # the following line could cause slowdown
        if al*norm(p) < eps(Tv)*norm(x)
            if verbose
                println("PCG: Stopped due to stagnation.")
            end
            break
        end

        axpy2!(al,p,x)
        # x = x + al * p
        #=
        @inbounds @simd for i in 1:n
            x[i] += al*p[i]
        end
        =#
        #axpy

        axpy2!(-al,q,r)
        #r .= r .- al.*q
        #=
        @inbounds @simd for i in 1:n
            r[i] -= al*q[i]
        end
        =#

        nr = norm(r)/nb
        if nr < bestnr
            bestnr = nr
            @inbounds @simd for i in 1:n
                bestx[i] = x[i]
            end
        end

        if nr < tol #Converged?
            break
        end

        # here is the top of the code in numerical templates

        z = pre(r)

        oldrho = rho
        rho = dot(z, r) # this is gamma in hypre.

        if rho < best_rho*(1-1/stag_test)
            best_rho = rho
            stag_count = 0
        else
            if stag_test > 0
                if best_rho > (1-1/stag_test)*rho
                    stag_count += 1
                    if stag_count > stag_test
                        println("PCG Stopped by stagnation test ", stag_test)
                        break
                    end
                end
            end
        end

        if (rho < eps(Tv) || isinf(rho))
            if verbose
                println("PCG Stopped due to small or large rho")
            end
            break
        end

        # the following would have higher accuracy
        #       rho = sum(r.^2)

        beta = rho/oldrho
        if (beta < eps(Tv) || isinf(beta))
            if verbose
                println("PCG Stopped due to small or large beta")
            end
            break
        end

        bzbeta!(beta,p,z)
        #=
        # p = z + beta*p
        @inbounds @simd for i in 1:n
            p[i] = z[i] + beta*p[i]
        end
        =#

        if (time() - t1) > maxtime
            if verbose
                println("PCG New stopped at maxtime.")
            end
            break
        end

    end

    if verbose
        println("PCG stopped after: ", round((time() - t1),digits=3), " seconds and ", itcnt, " iterations with relative error ", (norm(r)/norm(b)), ".")
    end

    if length(pcgIts) > 0
        pcgIts[1] = itcnt
    end

    return bestx
end

# uses a preconditioner that acts in place
function pcg_inplace(mat::BlockCSC{D,Tv,L},
    b::AbstractVector{SArray{Tuple{D},Tv,1,D}},
    pre!::Function;
    tol::Real=1e-6,
    maxits=Inf,
    maxtime=Inf,
    verbose=false,
    pcgIts=Int[],
    stag_test::Integer=0) where {D,Tv,L}

    local al::Tv

    n = size(mat,2)

    nb = norm(b)

    # If input vector is zero, quit
    if nb == 0
    return zeros(size(b))
    end

    x = zeros(SVector{D,Tv} ,n)
    bestx = zeros(SVector{D,Tv}, n)
    bestnr = one(Tv)

    r = copy(b)
    z = copy(r)
    pre!(z)
    p = copy(z)
    q = similar(z)

    rho = dot(r, z)
    best_rho = rho
    stag_count = 0

    t1 = time()

    itcnt = 0
    while itcnt < maxits
        itcnt = itcnt+1

        mul!(q, mat, p)
        #q = mat*p

        pq = dot(p,q)

        if (pq < eps(Tv) || isinf(pq))
            if verbose
                println("PCG Stopped due to small or large pq")
            end
            break
        end

        al = rho/pq

        # the following line could cause slowdown
        if al*norm(p) < eps(Tv)*norm(x)
            if verbose
                println("PCG: Stopped due to stagnation.")
            end
            break
        end

        axpy2!(al,p,x)

        axpy2!(-al,q,r)

        nr = norm(r)/nb
        if nr < bestnr
            bestnr = nr
            @inbounds @simd for i in 1:n
                bestx[i] = x[i]
            end
        end

        if nr < tol #Converged?
            break
        end

        # here is the top of the code in numerical templates

        copy!(z,r) # copy r in to z
        pre!(z)

        oldrho = rho
        rho = dot(z, r) # this is gamma in hypre.

        if rho < best_rho*(1-1/stag_test)
            best_rho = rho
            stag_count = 0
        else
            if stag_test > 0
                if best_rho > (1-1/stag_test)*rho
                    stag_count += 1
                    if stag_count > stag_test
                        println("PCG Stopped by stagnation test ", stag_test)
                        break
                    end
                end
            end
        end

        if (rho < eps(Tv) || isinf(rho))
            if verbose
                println("PCG Stopped due to small or large rho")
            end
            break
        end

        # the following would have higher accuracy
        #       rho = sum(r.^2)

        beta = rho/oldrho
        if (beta < eps(Tv) || isinf(beta))
            if verbose
                println("PCG Stopped due to small or large beta")
            end
            break
        end

        bzbeta!(beta,p,z)
        #=
        # p = z + beta*p
        @inbounds @simd for i in 1:n
            p[i] = z[i] + beta*p[i]
        end
        =#

        if (time() - t1) > maxtime
            if verbose
                println("PCG New stopped at maxtime.")
            end
            break
        end

    end

    if verbose
        println("PCG stopped after: ", round((time() - t1),digits=3), " seconds and ", itcnt, " iterations with relative error ", (norm(r)/norm(b)), ".")
    end

    if length(pcgIts) > 0
        pcgIts[1] = itcnt
    end

    return bestx
end

function axpy2!(al,p::AbstractArray{SVector{D,T}},
        x::AbstractArray{SVector{D,T}}) where {D, T}
    n = length(x)
    @inbounds @simd for i in 1:n
        x[i] = x[i] + al*p[i]
    end
  end

  function bzbeta!(beta,p::AbstractArray{SVector{D,T}},
    z::AbstractArray{SVector{D,T}}) where {D, T}
    n = length(p)
    @inbounds @simd for i in 1:n
        p[i] = z[i] + beta*p[i]
    end
  end

#=============================================================

The routines that do the solve: forwards, backwards, and diagonal inversion.
Note that they have special versions for D=1, and probably should for D=2 and 3.

These all use the ops structure.  It should be replaced.

=============================================================#

mapD(D,i) = (D*(i-1)+1):(D*i)

# should hard-code D=2 and D=3
function forward!(ops, D, y)
    if D == 1
        forward1!(ops, y)
    else
        forwardD!(ops, D, y)
    end
end

# For D = 1
function forward1!(ops, y)

    for i in 1:length(ops)
        op = ops[i]

        y[op.j] += op.wji * y[op.i]
        y[op.i] = op.wii*y[op.i]
    end
end

#=
# This is the idealized function, but it is not fast enough.
=#
#=
function forwardD!(ops, D, y)

    for i in 1:length(ops)
        op = ops[i]
        di = mapD(D,op.i)
        dj = mapD(D,op.j)

        y[dj] = y[dj] + op.wji * y[di]
        y[di] = op.wii*y[di]
    end
end
=#

function forwardD!(ops, D, y)

    tmp = similar(y[1:D])

    @inbounds for i in 1:length(ops)
        op = ops[i]

        di = D*(op.i-1)
        dj = D*(op.j-1)

        # The following lines simulate:
        # y[dj] = y[dj] + op.wji * y[di]

        for j in 1:D
            for k in 1:D
                y[dj+j] += op.wji[j,k] * y[di+k]
            end
        end

        # The following lines simulate:
        # y[di] = op.wii*y[di]

        tmp .= 0.0

        for j in 1:D
            for k in 1:D
                tmp[j] += op.wii[j,k] * y[di+k]
            end
        end

        for j in 1:D
            y[di+j] = tmp[j]
        end
    end
end

#=
Invert diag d[i] on ith block of y, for each i.
Ideal routine, which is too slow, is:

    for i in 1:(length(d))

        if is_non_zero(d[i])
            di = mapD(D,i)
            y[di] = inv(d[i])*y[di]
        end

    end

=#
function diag!(d, D, y)
    if D == 1
        diag1!(d, y)
    else
        diagD!(d, D, y)
    end
end

function diag1!(d, y)
    for i in 1:length(d)
        y[i] *= d[i]
    end
end

function diagD!(d, D, y)

    tmp = similar(y[1:D])

    for i in 1:length(d)
        Di = D*(i-1)
        y[(Di+1):(Di+D)] = d[i]*y[(Di+1):(Di+D)]
    end
end

#=
You would think we should use the following, but it is slower
=#
function diagD2!(d, D, y)

    tmp = similar(y[1:D])

    @inbounds for i in 1:length(d)
        Di = D*(i-1)

        tmp .= 0.0

        for j in 1:D
            for k in 1:D
                tmp[j] += d[i][j,k] * y[Di+k]
            end
        end

        for j in 1:D
            y[Di+j] = tmp[j]
        end

        #y[(Di+1):(Di+D)] = d[i]*y[(Di+1):(Di+D)]
    end
end

# should hard-code D=2 and D=3
function backward!(ops, D, y)
    if D == 1
        backward1!(ops, y)
    else
        backwardD!(ops, D, y)
    end
end

# For D = 1
function backward1!(ops, y)

    for i in length(ops):-1:1
        op = ops[i]

        y[op.i] = op.wii*y[op.i] + op.wji*y[op.j]
    end
end

#=
# This is the idealized function, but it is not fast enough.
=#
#=
function backwardD!(ops, D, y)

    for i in length(ops):-1:1
        op = ops[i]
        di = D*(op.i-1)
        dj = D*(op.j-1)

        y[(di+1):(di+D)] = op.wii*y[(di+1):(di+D)] + op.wji'*y[(dj+1):(dj+D)]
    end

end
=#

function backwardD!(ops, D, y)

    tmp = similar(y[1:D])

    for i in length(ops):-1:1
        op = ops[i]

        di = D*(op.i-1)
        dj = D*(op.j-1)

        tmp .= 0

        for j in 1:D
            for k in 1:D
                tmp[j] += op.wii[j,k] * y[di+k] + op.wji[k,j]'*y[dj+k]
            end
        end

        for j in 1:D
            y[di+j] = tmp[j]
        end
    end
end

# function solver{Tv}(ops::Array{IJop,1}, d::Array, b::Array{Tv,1})
function solver(ops, d, b)
    y = copy(b)
    D = size(ops[1].wii,1)

    forward!(ops, D, y)

    diag!(d, D, y)

    backward!(ops, D, y)

    return y
end

#=======================

  Experimental Solvers

==========================#

"""
    cnt = count_nonzero_diags(a::SparseMatrixCSC)

Count number of nonzero diagonal elements in a.
"""
function count_nonzero_diags(a::SparseMatrixCSC)

    @assert a.m == a.n
    cnt = 0
    for j in 1:a.n
        for i in a.colptr[j]:(a.colptr[j+1]-1)
            cnt += (a.rowval[i] == j) && (!iszero(a.nzval[i]))
        end
    end
    return cnt
end

"""
    D, R = diags_and_rest(a::SparseMatrixCSC{Tv,Ti})

Create a diagonal matrix D and a sparse matrix R with zero diagonal
such that D+R = a.  Requires a to be square.
Actually returns the vector of the diagonal D.
So, if you want a matrix, use Diagonal(D)
"""
function diags_and_rest(a::SparseMatrixCSC{Tv,Ti}) where {Tv,Ti}

    @assert a.m == a.n
    cnt = count_nonzero_diags(a)
    dv = zeros(Tv, a.n)

    colptr = similar(a.colptr)
    rowval = Vector{Ti}(undef, length(a.rowval) - cnt)
    nzval = Vector{Tv}(undef, length(a.rowval) - cnt)

    ptr = one(Ti)

    for j in 1:a.n
        colptr[j] = ptr
        for i in a.colptr[j]:(a.colptr[j+1]-1)
            if a.rowval[i] == j
                dv[j] = a.nzval[i]
            else
                rowval[ptr] = a.rowval[i]
                nzval[ptr] = a.nzval[i]
                ptr += 1
            end
        end
    end

    colptr[end] = ptr

    R = SparseMatrixCSC(a.m, a.n, colptr, rowval, nzval)

    return dv, R
end

#================================================

  Triangular Solves for matrices triangular up to an order.
  Will need to compare to removing diagonals
  and permutation followed by triangular Solves

  =================================================#

function backsolve!(L::SparseMatrixCSC, b::Vector)

    @assert length(b) == L.n
    @inbounds for j in L.n:-1:1
        for ii in (L.colptr[j]+1):(L.colptr[j+1]-1)
            b[j] -= L.nzval[ii] * b[L.rowval[ii]]
        end
        b[j] /= L.nzval[L.colptr[j]]
    end
end

function backsolve(L::SparseMatrixCSC, b::Vector)
    x = copy(b)
    backsolve!(L, x)
    return x
end

function backsolve(L::SparseMatrixCSC, b::Vector, p::Vector, ip::Vector)
    x = b[p]
    backsolve!(L, x)
    return x[ip]
end

function forwardsolve(L::SparseMatrixCSC, b::Vector, p::Vector, ip::Vector)
    x = b[p]
    forwardsolve!(L, x)
    return x[ip]
end

function solver(sdd_tri::SparseMatrixCSC, d_for_ldl::Vector, p::Vector, ip::Vector, b)
    y = forwardsolve(sdd_tri, b, p, ip)

    diag!(d, D, y)

    backward!(ops, D, y)

    return y
end

function forwardsolve!(L::SparseMatrixCSC, b::Vector)

    @assert length(b) == L.n
    @inbounds for j in 1:L.n

        b[j] /= L.nzval[L.colptr[j]]

        for ii in (L.colptr[j]+1):(L.colptr[j+1]-1)
            i = L.rowval[ii]
            b[i] -= b[j] * L.nzval[ii]
        end
    end

end

function disordered_sparse(L, order)
    colptr = similar(L.colptr)
    rowval = similar(L.rowval)
    nzval = similar(L.nzval)
    n = L.n
    ptr = 1
    for i in 1:n
        col = order[i]
        colptr[i] = ptr
        lead = ptr
        rowval[lead] = col
        ptr += 1
        for j in L.colptr[col]:(L.colptr[col+1]-1)
            row = L.rowval[j]
            if row != col
                rowval[ptr] = L.rowval[j]
                nzval[ptr] = L.nzval[j]
                ptr += 1
            else
                nzval[lead] = L.nzval[j]
            end
        end

    end
    colptr[n+1] = ptr
    return colptr, rowval, nzval
end

#=
Solves with the order built in to leading entries of rowval:
the semantics are not exactly what they are for sparse matrices.
rowval[colptr[i]] is the ith column in the elimination order.
nzval[colptr[i]] is the corresponding diagonal entry.
rowval[colptr[i]+1] up to the next are the other entries in that column.
=#
function forwardsolve!(colptr::Vector{Ti}, rowval::Vector{Ti},
    nzval::Vector{Tv}, b::Vector{Tv}) where {Ti,Tv}
    n = length(colptr)-1

    @assert length(b) == n
    @inbounds for j in 1:n
        col = rowval[colptr[j]]

        b[col] /= nzval[colptr[j]]

        for ii in (colptr[j]+1):(colptr[j+1]-1)
            i = rowval[ii]
            b[i] -= b[col] * nzval[ii]
        end
    end
end

function backsolve!(colptr::Vector{Ti}, rowval::Vector{Ti},
    nzval::Vector{Tv}, b::Vector{Tv}) where {Ti,Tv}
    n = length(colptr)-1

    @assert length(b) == n
    @inbounds for j in n:-1:1
        col = rowval[colptr[j]]

        for ii in (colptr[j]+1):(colptr[j+1]-1)
            b[col] -= nzval[ii] * b[rowval[ii]]
        end
        b[col] /= nzval[colptr[j]]
    end
end

"""
    x = backsolve(L, order, b)

Does a backsolve (as in upper triangular version) in a matrix that
is lower-triangular under the order b.

Not necessarily an efficient routine, but one we will try to start.
"""
function backsolve!(L::SparseMatrixCSC, order::Vector, b::Vector)

    @assert length(b) == L.n == length(order)
    @inbounds for jj in L.n:-1:1
        j = order[jj]
        for ii in L.colptr[j]:(L.colptr[j+1]-1)
            i = L.rowval[ii]
            if j != i
                b[j] -= L.nzval[ii] * b[i]
            end
        end
        b[j] /= L[j,j]
    end
end

function backsolve(L::SparseMatrixCSC, order::Vector, b::Vector)
    x = copy(b)
    backsolve!(L, order, x)
    return x
end

"""
    x = forwardsolve(L, order, b)

Does a forwardsolve (as in upper triangular version) in a matrix that
is lower-triangular under the order b.

Not necessarily an efficient routine, but one we will try to start.
"""
function forwardsolve!(L::SparseMatrixCSC, order, b)

    @assert length(b) == L.n == length(order)
    @inbounds for jj in 1:L.n
        j = order[jj]
        b[j] /= L[j,j]

        for ii in L.colptr[j]:(L.colptr[j+1]-1)
            i = L.rowval[ii]
            if i != j
               b[i] -= b[j] * L.nzval[ii]
            end
        end
    end
end

function forwardsolve(L::SparseMatrixCSC, order, b)
    x = copy(b)
    forwardsolve!(L, order, x)
    return x
end

"""
For separted diagonal and lower triangular part, R
"""
function forwardsolve!(D::Vector, L::SparseMatrixCSC, order, b)

    @assert length(b) == L.n == length(order) == length(D)
    @inbounds for jj in 1:L.n
        j = order[jj]
        b[j] /= D[j]

        for ii in L.colptr[j]:(L.colptr[j+1]-1)
            i = L.rowval[ii]
            b[i] -= b[j] * L.nzval[ii]
        end
    end
end

function forwardsolve(D::Vector, L::SparseMatrixCSC, order, b)
    x = copy(b)
    forwardsolve!(D, L, order, x)
    return x
end

function backsolve!(D::Vector, L::SparseMatrixCSC, order, b)

    @assert length(b) == L.n == length(order) == length(D)
    @inbounds for jj in L.n:-1:1
        j = order[jj]
        for ii in L.colptr[j]:(L.colptr[j+1]-1)
            i = L.rowval[ii]
            b[j] -= L.nzval[ii] * b[i]
        end
        b[j] /= D[j]
    end
end

function backsolve(D::Vector, L::SparseMatrixCSC, order, b)
    x = copy(b)
    backsolve!(D, L, order, x)
    return x
end


text:
using HorizonSideRobots

"""
    moves!(r::Robot, side::HorizonSide)

-- Перемещает робота до упора в заданном направлении
"""
function moves!(r::Robot, side::HorizonSide)
    while isborder(r, side) == false
        move!(r, side)
    end
end

"""
    moves!(r::Robot, side::HorizonSide, num_steps::Int)

-- Перемещает робота в направлении side на заданное число шагов
"""
function number_move_put!(r::Robot, side::HorizonSide, num_steps::Int)
    for _ in 1:num_steps
        move!(r, side)
        putmarker!(r)
    end
end

"""
    putmarkers!(r::Robot, side::HorizonSide)

-- Ставит маркеры всюду в заданном направлении, пока робот не встретит перегородку,
   исключая начальную позицию
"""
function putmarkers!(r::Robot, side::HorizonSide)
    while isborder(r, side) == false
        move!(r, side)
        putmarker!(r)
    end
end
function move_by_markers(r::Robot, side::HorizonSide)
    while ismarker(r) == true
        move!(r, side)
    end
end

"""
    inverse(side::HorizonSide)

-- Возвращает направление, противоположное заданному
"""
inverse(side::HorizonSide) = HorizonSide(mod(Int(side) + 2, 4))

"""

    cornerwalk(r::Robot, coordinates(название любого массива, в этом случае такой))

-- Идет в угол
"""
function cornerwalk(r::Robot, coordinates::Vector{Any})
    while !(isborder(r, West) && isborder(r, Sud))
        if isborder(r, West) == false
            move!(r, West)
            push!(coordinates, West)
        end
        if isborder(r, Sud) == false
            move!(r, Sud)
            push!(coordinates, Sud)
        end
    end
end

"""

    go_back(r::Robot, coordinates)

---возвращается обратно из угла, дополнение к прошлой функции, ставить после всех махинаций
"""
function go_back!(r::Robot, coordinates)
    while length(coordinates) > 0
        side = pop!(coordinates)
        move!(r, inverse(side))
    end
end

"""

    mark_perimetr!(r::Robot, side::HorizonSide)

---маркирует фигуру (квадрад, прямоугольник, линия) по кругу с заданным начальным направлением
"""

function mark_perimetr!(r::Robot, side::HorizonSide)
    if side == Nord
        mov = (Nord, Ost, Sud, West)
        bord = (Ost, Sud, West, Nord)
    else
        mov = (Sud, Ost, Nord, West)
        bord = (Ost, Nord, West, Sud)
    end
    for i in 1:4
        while isborder(r, bord[i]) == true
            putmarker!(r)
            move!(r, mov[i])
        end
        putmarker!(r)
        if i < 4
            move!(r, mov[i+1])
        else
            move!(r, Sud)
        end
    end
end

"""

find_mrk_spiral(r::Robot)

--ходит спиралью и ищет маркер
"""

function find_mrk_spiral(r)

    counter = 1
    while ismarker(r)==false

        for i in 1:counter
            move!(r, Ost)
            if ismarker(r)==true
                break
            end
        end
        if ismarker(r)==true
            break
        end

        for i in 1:counter
            move!(r, Nord)
            if ismarker(r)==true
                break
            end
        end
        if ismarker(r)==true
            break
        end

        counter += 1

        for i in 1:counter
            move!(r, West)
            if ismarker(r)==true
                break
            end
        end
        if ismarker(r)==true
            break
        end

        for i in 1:counter
            move!(r, Sud)
            if ismarker(r)==true
                break
            end
        end
        if ismarker(r)==true
            break
        end

        counter += 1

    end
end

"""
    go_to_corner!(x::Robot, mvs::AbstractArray)
-- Перемещает Робота в любой угол, обходя препядствия, и записывает путь в массив
"""
function go_to_corner!(x::Robot, mvs::AbstractArray)
    while (!isborder(x, side1) || !isborder(x, side2))
        if !isborder(x, side1)
            moves!(x, side1, mvs)
        else
            moves!(x, side2, mvs)

        end
    end
end

function moves1!(r::Robot,side::HorizonSide,num_steps::Int)
    q=0
    while q<num_steps
        if isborder(r,Nord)==true
            while isborder(r,side)==true
                move!(r,Ost)
            end
            move!(r,Nord)
            q+=1
            if isborder(r,West)==true
                while isborder(r,West)==true
                    move!(r,Nord)
                    q+=1
                end
            end
            move!(r,West)
        end
        if q<num_steps
        move!(r,side)
        q+=1
        end
    end
end

"""
obhod(r::Robot,side::HorizonSide)
-- Обходит препядствие на своем пути
"""

function obhod(r::Robot,side::HorizonSide)
    k=0
    j=0
    while isborder(r,side)==true && isborder(r,next(side))==false
        move!(r,next(side))
        k+=1
    end
    if isborder(r,side)==true && isborder(r,next(side))==true
        for _ in 1:k
            move!(r,inverse(next(side)))
        end
        k=0
    end
    while isborder(r,side)==true && isborder(r,inverse(next(side)))==false
        move!(r,inverse(next(side)))
        j+=1
    end
    if isborder(r,side)==true && isborder(r,inverse(next(side)))==true
        for _ in 1:j
            move!(r,next(side))
        end
    end
    if k==0
        k=j
        side1 = next(side)
    else
        side1 = inverse(next(side))
    end
    if isborder(r,side)==false
        print(side)
        move!(r,side)
        while (isborder(r,next(inverse(side)))==true &&(isborder(r,next(side))==false || isborder(r,inverse(next(side)))==false))
            move!(r,side)
        end
        for _ in 1:k
            move!(r,side1)
        end
    end
end

"""
go_and_obhod(r::Robot,side::HorizonSide,num_steps::Int)
-- Будет делать треугольник и обходить препядствия на своем пути
"""

function go_and_obhod(r::Robot,side::HorizonSide,num_steps::Int)
    q=0
    k=0
    while q<num_steps
        k=0
       putmarker!(r)
       if isborder(r,side)==false
            move!(r,side)
            q+=1
       else isborder(r,side)==true
            while isborder(r,side)==true
                move!(r,Nord)
                k+=1
            end
            move!(r,side)
            q+=1
            if isborder(r,Sud)==true
                move!(r,side)
                q+=1
            end
            for _ in 1:k
                move!(r,Sud)
            end
       end
    end
end

"""
go_to_border_and_return_path!(r::Robot, side::HorizonSide; go_around_barriers::Bool = false, markers = false)::Array{Tuple{HorizonSide,Int64},1}
Перемещает робота до границы и возвращает путь для обраного следования в виде массива пар типа (направление, количество шагов)
Именованные параметры:
go_around_barriers - робот будет по пути обходить перегородки, если такие имеются;
markers - робот будет ставить маркеры по пути (не ставит их при обходе перегородок).
"""
function go_to_border_and_return_path!(r::Robot, side::HorizonSide; go_around_barriers::Bool = false, markers = false)::Array{Tuple{HorizonSide,Int64},1}
    my_ans = [ (Nord, 0) ]
    if go_around_barriers
        steps = 0
        if markers
            putmarker!(r)
        end
        if !isborder(r, side)
            move!(r, side)
            steps = 1
            push!(my_ans, (inverse_side(side), 1) )
        else
            path = go_around_barrier_and_return_path!(r, side)
            steps = get_path_length_in_direction(path, side)
            for i in path
                push!(my_ans, i)
            end
        end
        if markers
            putmarker!(r)
        end
        while steps > 0
            if !isborder(r, side)
                move!(r, side)
                steps = 1
                push!(my_ans, (inverse_side(side), 1) )
                if markers
                    putmarker!(r)
                end
            else
                path = go_around_barrier_and_return_path!(r, side)
                steps = get_path_length_in_direction(path, side)
                for i in path
                    push!(my_ans, i)
                end
                if markers
                    putmarker!(r)
                end
            end
        end

    else
        steps=0
        steps_now = go!(r,side; markers)
        while steps_now > 0
            steps += steps_now
            steps_now = go!(r,side; markers)
        end
        push!(my_ans, (inverse_side(side), steps) )
    end
    return my_ans
end

function go!(r::Robot, side::HorizonSide; steps::Int = 1, go_around_barriers::Bool = false, markers = false)::Int
    my_ans = 0
    if markers
        putmarker!(r)
    end
    if (go_around_barriers)
        path = around_move_return_path!(r, side; steps, markers)
        my_ans = get_path_length_in_direction(path, side)
    else
        for i ∈ 1:steps

            if (markers)
                putmarker!(r)
            end

            if !isborder(r, side)
                move!(r, side)
                my_ans += 1
            else
                for i ∈ 1:my_ans
                    move!(r, inverse_side(side))
                end
                my_ans = 0
                break
            end
        end
        if (markers)
            putmarker!(r)
        end
    end

    return my_ans
end


text:
#=
These unit tests for the 3 types of imaging condition are a little bit heuristic. We use two
half-space models with the upper half-plane velocity greater in the model we compute the
gradients in, so that we expect all positive gradient.

Tests:

1. Ensure gradients from different imaging conditions are different
    a. test: norm(g_STD .- g_FWI) > ϵ
    b. test: norm(g_FWI .- g_RTM) > ϵ

2. For FWI, we expect longer wavelengths, so a larger sum after maximum value normalization
    a. normalize to maximum value unity
    b. test: sum(sign(g_FWI)) > sum(sign(g_STD))

3. For RTM, we test the correlation of an image with that image after high-pass filter.
   With the RTM imaging condition, there should be less low wavenumber energy, hence this
   correlation should be larger for RTM IC.
   a. model in water velocity and subtract to remove direct wave
   b. test: cc(g_RTM,highpass * g_RTM) > cc(g_STD,highpass * g_STD)

Note these tests are skipped currently as serialized wavefields for the imaging condition
unit tests are too large for github CI.
=#
using WaveFD, JetPackWaveFD, JetPack, Jets, Random, LinearAlgebra, Test

fpeak = 20
dtmod,dtrec = 0.001,0.005
ntrec,nsponge,reportinterval = 251,40,0
z,y,x,δz,δy,δx = 600.0,700.0,800.0,10.0,10.0,10.0
nz,ny,nx = round(Int,z/δz)+1,round(Int,y/δy)+1,round(Int,x/δx)+1

function makeF(; dim::Int=2, physics::String="ISO", imgcondition::String="standard")
    if dim == 2
        sz = δz * 1;
        sx = δx * div(nx,2);
        rz = δz .* ones(Float32, nx);
        rx = δx .* Float32[0:(nx-1);];
    else
        sz = δz * 1;
        sy = δy * div(ny,2);
        sx = δx * div(nx,2);
        rz = δz .* ones(Float32, nx*ny);
        ry = (δx .* ones(Float32, nx) * Float32[0:ny-1;]')[:];
        rx = (δx .* Float32[0:nx-1;] * ones(Float32, ny)')[:];
    end

    if dim == 2 && physics == "ISO"
        b = 0.001f0 .* ones(Float32,nz,nx)

        JopNlProp2DAcoIsoDenQ_DEO2_FDTD(; b = b,
            comptype = UInt32, compscale = 1e-1,
            sz = sz, sx = sx, rz = rz, rx = rx, dz = δz, dx = δx,
            dtmod = dtmod, dtrec = dtrec, ntrec = ntrec, nsponge = nsponge,
            wavelet = WaveletCausalRicker(f=fpeak),
            imgcondition = imgcondition, reportinterval = reportinterval)

    elseif dim == 2 && physics == "VTI"
        b = 0.001f0 .* ones(Float32,nz,nx)
        ϵ = 0.1f0 .* ones(Float32,nz,nx)
        η = 0.2f0 .* ones(Float32,nz,nx)

        JopNlProp2DAcoVTIDenQ_DEO2_FDTD(; b = b, ϵ = ϵ, η = η,
            comptype = UInt32, compscale = 1e-1,
            sz = sz, sx = sx, rz = rz, rx = rx, dz = δz, dx = δx,
            dtmod = dtmod, dtrec = dtrec, ntrec = ntrec, nsponge = nsponge,
            wavelet = WaveletCausalRicker(f=fpeak),
            imgcondition = imgcondition, reportinterval = reportinterval)

    elseif dim == 2 && physics == "TTI"
        b = 0.001f0 .* ones(Float32,nz,nx)
        ϵ = 0.1f0 .* ones(Float32,nz,nx)
        η = 0.2f0 .* ones(Float32,nz,nx)
        θ = Float32(π/8) .* ones(Float32,nz,nx)

        JopNlProp2DAcoTTIDenQ_DEO2_FDTD(; b = b, ϵ = ϵ, η = η, θ = θ,
            comptype = UInt32, compscale = 1e-1,
            sz = sz, sx = sx, rz = rz, rx = rx, dz = δz, dx = δx,
            dtmod = dtmod, dtrec = dtrec, ntrec = ntrec, nsponge = nsponge,
            wavelet = WaveletCausalRicker(f=fpeak),
            imgcondition = imgcondition, reportinterval = reportinterval)

    elseif dim == 3 && physics == "ISO"
        b = 0.001f0 .* ones(Float32,nz,ny,nx)

        JopNlProp3DAcoIsoDenQ_DEO2_FDTD(; b = b,
            comptype = UInt32, compscale = 1e-1,
            sz = sz, sy = sy, sx = sx, rz = rz, ry = ry, rx = rx, dz = δz, dy = δy, dx = δx,
            dtmod = dtmod, dtrec = dtrec, ntrec = ntrec, nsponge = nsponge,
            wavelet = WaveletCausalRicker(f=fpeak),
            imgcondition = imgcondition, reportinterval = reportinterval)

    elseif dim == 3 && physics == "VTI"
        b = 0.001f0 .* ones(Float32,nz,ny,nx)
        ϵ = 0.1f0 .* ones(Float32,nz,ny,nx)
        η = 0.2f0 .* ones(Float32,nz,ny,nx)

        JopNlProp3DAcoVTIDenQ_DEO2_FDTD(; b = b, ϵ = ϵ, η = η,
            comptype = UInt32, compscale = 1e-1,
            sz = sz, sy = sy, sx = sx, rz = rz, ry = ry, rx = rx, dz = δz, dy = δy, dx = δx,
            dtmod = dtmod, dtrec = dtrec, ntrec = ntrec, nsponge = nsponge,
            wavelet = WaveletCausalRicker(f=fpeak),
            imgcondition = imgcondition, reportinterval = reportinterval)

    elseif dim == 3 && physics == "TTI"
        b = 0.001f0 .* ones(Float32,nz,ny,nx)
        ϵ = 0.1f0 .* ones(Float32,nz,ny,nx)
        η = 0.2f0 .* ones(Float32,nz,ny,nx)
        θ = Float32(π/8) .* ones(Float32,nz,ny,nx)
        ϕ = Float32(π/6) .* ones(Float32,nz,ny,nx)

        JopNlProp3DAcoTTIDenQ_DEO2_FDTD(; b = b, ϵ = ϵ, η = η, θ = θ, ϕ = ϕ,
            comptype = UInt32, compscale = 1e-1,
            sz = sz, sy = sy, sx = sx, rz = rz, ry = ry, rx = rx, dz = δz, dy = δy, dx = δx,
            dtmod = dtmod, dtrec = dtrec, ntrec = ntrec, nsponge = nsponge,
            wavelet = WaveletCausalRicker(f=fpeak),
            imgcondition = imgcondition, reportinterval = reportinterval)
    end
end

@testset "ImgCondition" begin
    @testset "Imaging Condition 2D tests, $(physics)" for physics in ("ISO", "VTI", "TTI")
        write(stdout, "\n")
        write(stdout, "Imaging conditions 2D, physics=$(physics)\n")

        opStd = makeF(;dim=2, physics=physics, imgcondition="standard");
        opFwi = makeF(;dim=2, physics=physics, imgcondition="FWI");
        opRtm = makeF(;dim=2, physics=physics, imgcondition="RTM");

        vw = 1500 .* ones(domain(opStd)); # water velocity model
        va = 1500 .* ones(domain(opStd));
        vb = 1505 .* ones(domain(opStd));
        va[div(nz,2):end,:,:] .= 2500;
        vb[div(nz,2):end,:,:] .= 2500;

        d1a = opStd * va;
        d1w = opStd * vw;

        # standard FWI
        d1b = opStd * vb;
        r1  = d1b .- d1a;
        J1 = jacobian!(opStd,vb);
        g1 = J1' * r1;

        # wavefield separation FWI
        d2b = opFwi * vb;
        r2  = d2b .- d1a;
        J2 = jacobian!(opFwi,vb);
        g2 = J2' * r2;

        # standard RTM
        d3b = opRtm * vb;
        r3 = d1a .- d1w
        J3 = jacobian!(opStd,vb);
        g3 = J3' * r3;

        # wavefield separation RTM
        d4b = opRtm * vb;
        r4 = d1a .- d1w
        J4 = jacobian!(opRtm,vb);
        g4 = J4' * r4;

        close(opStd)
        close(opFwi)
        close(opRtm)

        # remove pixels near source to eliminate very large amplitudes
        nzero = 5
        g1[1:nzero,:,:] .= 0
        g2[1:nzero,:,:] .= 0
        g3[1:nzero,:,:] .= 0
        g4[1:nzero,:,:] .= 0

        # FWI gradients differ
        @test norm(g1 .- g2) > eps(Float32)

        # RTM gradients differ
        @test norm(g3 .- g4) > eps(Float32)

        # FWI gradient longer wavelength (more positive) than standard gradient
        @show sum(sign.(g1)), sum(sign.(g2))
        @test sum(sign.(g2)) > sum(sign.(g1))

        # cc(g_RTM,highpass,g_RTM) > cc(g_STD,highpass,g_STD)
        op = JopLaplacian(JetSpace(Float32,nz,nx)) ∘ JopReshape(domain(opStd), JetSpace(Float32,nz,nx))
        h3 = op * g3;
        h4 = op * g4;
        cc3 = dot(abs.(g3[:]),abs.(h3[:])) / sqrt(dot(abs.(g3[:]),abs.(g3[:])) * dot(abs.(h3[:]),abs.(h3[:])))
        cc4 = dot(abs.(g4[:]),abs.(h4[:])) / sqrt(dot(abs.(g4[:]),abs.(g4[:])) * dot(abs.(h4[:]),abs.(h4[:])))
        @show cc3, cc4
        @test cc4 > cc3
    end

    @testset "Imaging Condition 3D tests, $(physics)" for physics in ("ISO", "VTI", "TTI")
        write(stdout, "\n")
        write(stdout, "Imaging conditions 3D, physics=$(physics)\n")

        opStd = makeF(;dim=3, physics=physics, imgcondition="standard");
        opFwi = makeF(;dim=3, physics=physics, imgcondition="FWI");
        opRtm = makeF(;dim=3, physics=physics, imgcondition="RTM");

        vw = 1500 .* ones(domain(opStd)); # water velocity model
        va = 1500 .* ones(domain(opStd));
        vb = 1505 .* ones(domain(opStd));
        va[div(nz,2):end,:,:,:] .= 2500;
        vb[div(nz,2):end,:,:,:] .= 2500;

        d1a = opStd * va;
        d1w = opStd * vw;

        # standard FWI
        d1b = opStd * vb;
        r1  = d1b .- d1a;
        J1 = jacobian!(opStd,vb);
        g1 = J1' * r1;

        # wavefield separation FWI
        d2b = opFwi * vb;
        r2  = d2b .- d1a;
        J2 = jacobian!(opFwi,vb);
        g2 = J2' * r2;

        # standard RTM
        d3b = opRtm * vb;
        r3 = d1a .- d1w
        J3 = jacobian!(opStd,vb);
        g3 = J3' * r3;

        # wavefield separation RTM
        d4b = opRtm * vb;
        r4 = d1a .- d1w
        J4 = jacobian!(opRtm,vb);
        g4 = J4' * r4;

        close(opStd)
        close(opFwi)
        close(opRtm)

        # remove pixels near source to eliminate very large amplitudes there
        nzero = 5
        g1[1:nzero,:,:,:] .= 0
        g2[1:nzero,:,:,:] .= 0
        g3[1:nzero,:,:,:] .= 0
        g4[1:nzero,:,:,:] .= 0

       # FWI gradients differ
       @test norm(g1 .- g2) > eps(Float32)

       # RTM gradients differ
       @test norm(g3 .- g4) > eps(Float32)

       # FWI gradient longer wavelength (more positive) than standard gradient
       @show sum(sign.(g1)), sum(sign.(g2))
       @test sum(sign.(g2)) > sum(sign.(g1))

       # cc(g_RTM,highpass,g_RTM) > cc(g_STD,highpass,g_STD)
       op = JopLaplacian(JetSpace(Float32,nz,ny,nx)) ∘ JopReshape(domain(opStd), JetSpace(Float32,nz,ny,nx))
       h3 = op * g3;
       h4 = op * g4;
       cc3 = dot(abs.(g3[:]),abs.(h3[:])) / sqrt(dot(abs.(g3[:]),abs.(g3[:])) * dot(abs.(h3[:]),abs.(h3[:])))
       cc4 = dot(abs.(g4[:]),abs.(h4[:])) / sqrt(dot(abs.(g4[:]),abs.(g4[:])) * dot(abs.(h4[:]),abs.(h4[:])))
       @show cc3, cc4
       @test cc4 > cc3
    end
end

max_stars_repo_name: ChevronETC/JetPackWaveFD

max_issues_repo_name: ChevronETC/JetPackWaveFD

max_forks_repo_name: ChevronETC/JetPackWaveFD


text:
using JuliaDB
using MXNet
using PyPlot
using ProgressMeter
using JLD
using LaTeXStrings
include("self_confidence.jl")

mutable struct SQ_model
    input_sch::ML.Schema
    output_sch::ML.Schema
    X3_1net::MXNet.mx.FeedForward #neural network
    X3_2net::MXNet.mx.FeedForward #neural network
    range::Array{Float64}
end

function data_source(batchsize::Int64,train_in,train_out,valid_in,valid_out)
  train = mx.ArrayDataProvider(
    :data => train_in,
    :label => train_out,
    batch_size = batchsize,
    shuffle = true,
    )
  valid = mx.ArrayDataProvider(
    :data => valid_in,
    :label => valid_out,
    batch_size = batchsize,
    shuffle = true,
    )
  println("DONE")
  return train, valid
end

function make_nn_SQ_model(train_fname::String,valid_fname::String,log_fname::String; input_dict::Dict=Dict(),
                          output_dict::Dict=Dict(),nn_epoc::Int64=150,nn_batch_size::Int64=-1,
                          train_subsample::Int64=1,valid_subsample::Int64=1)

    training_in, training_out, training_table, input_sch, output_sch = return_data(train_fname,inputs=input_dict,outputs=output_dict,subsample=train_subsample)
    valid_in, valid_out, valid_table, _notused, _notused = return_data(valid_fname,inputs=input_dict,outputs=output_dict,subsample=valid_subsample)

    output_list = collect(keys(output_dict))
    input_list = collect(keys(input_dict))
    info("*****outputs: $output_list\n")
    info("*****inputs: $input_list")

    # create a DNN
    data = mx.Variable(:data)
    label = mx.Variable(:label)
    net = @mx.chain     mx.Variable(:data) =>
                        mx.FullyConnected(num_hidden=10) =>
                        mx.Activation(act_type=:relu) =>
                        mx.FullyConnected(num_hidden=10) =>
                        mx.Activation(act_type=:relu) =>
                        mx.FullyConnected(num_hidden= 1) =>
                        mx.LinearRegressionOutput(mx.Variable(:label))

    # final network definition, don't change, except if using gpu
    X3_1net = mx.FeedForward(net, context=mx.cpu())
    X3_2net = mx.FeedForward(net, context=mx.cpu())

    # set up the optimizer: select one, explore parameters, if desired
    #  optimizer = mx.SGD(η=0.01, μ=0.9, λ=0.00001)
    optimizer = mx.ADAM()

    # get data providers
    # Array(training_out[1,:]') -- this is some crazy conversion so the output is the right type....
    if nn_batch_size < 0
        # no value was provided
        nn_batch_size = round(Int,0.25*length(training_in))
    end

    X3_1trainprovider, X3_1evalprovider = data_source(#= batchsize =# nn_batch_size,training_in,Array(training_out[1,:]'),valid_in,Array(valid_out[1,:]'))
    X3_2trainprovider, X3_2evalprovider = data_source(#= batchsize =# nn_batch_size,training_in,Array(training_out[2,:]'),valid_in,Array(valid_out[2,:]'))

    # train, reporting loss for training and evaluation sets
    info("Training X3_1")
    mx.train(X3_1net, optimizer, X3_1trainprovider,
           initializer = mx.NormalInitializer(0.0, 0.1),
           eval_metric = mx.MSE(),
           eval_data = X3_1evalprovider,
           n_epoch = nn_epoc,
           callbacks = [mx.speedometer(),mx.do_checkpoint(string(log_fname,"_net1"),frequency=nn_epoc)])
    info("Training X3_2")
    mx.train(X3_2net, optimizer, X3_2trainprovider,
           initializer = mx.NormalInitializer(0.0, 0.1),
           eval_metric = mx.MSE(),
           eval_data = X3_2evalprovider,
           n_epoch = nn_epoc,
           callbacks = [mx.speedometer(),mx.do_checkpoint(string(log_fname,"_net2"),frequency=nn_epoc)])

    #  println(size(training_out))
    training_range = [minimum(training_out,2) maximum(training_out,2)]

    # put into a SQ_model for return
    SQ = SQ_model(input_sch,output_sch,X3_1net,X3_2net,training_range)
    return SQ
end

function SQ_predict(SQ::SQ_model,inputs::Array;use_eng_units::Bool=false)
    println("inputs:")
    display(inputs)
    X3_1plotprovider = mx.ArrayDataProvider(:data => inputs, shuffle=false)
    X3_2plotprovider = mx.ArrayDataProvider(:data => inputs, shuffle=false)
    X3_1fit = mx.predict(SQ.X3_1net, X3_1plotprovider)
    X3_2fit = mx.predict(SQ.X3_2net, X3_2plotprovider)
    println("Predicted Output:")
    display(X3_1fit)
    display(X3_2fit)
    if use_eng_units
        return restore_eng_units(inputs,SQ.input_sch), restore_eng_units([X3_1fit; X3_2fit],SQ.output_sch)
    else
        return inputs, [X3_1fit; X3_2fit]
    end
end

# convert data back to original units
function restore_eng_units(ary::Array,sch::ML.Schema)
    info("Restoring Engineering Units")
    println(size(ary))
    arrays = Dict()
    i = 1
    for entry in sch
        if entry.second != nothing
            println("key: $(entry.first), value: $(entry.second)")
            entry_mean = mean(entry.second)
            entry_std = std(entry.second)
            arrays[entry.first] = ary[i,:].*entry_std + entry_mean
            i += 1
        end
    end
    return arrays
    #  return [x.second for x in arrays]
end

function return_data(fname::String;inputs::Dict=Dict(),outputs::Dict=Dict(),
                     schema::Array=[],subsample::Int64=1)

    function my_splitschema(xs::ML.Schema,ks::Vector{Symbol})
        return filter((k,v) -> k ∉ ks, xs), filter((k,v) -> k ∈ ks, xs)
    end

    println("Processing $fname")
    table = loadtable(fname,escapechar='"')
    input_sch = nothing
    output_sch = nothing
    schema_dict = Dict()

    if isempty(schema)
        #  p = Progress(length(colnames(table)),dt=0.1,desc="Processing",output=STDOUT)
        for itm_name in colnames(table)
            if itm_name ∈ keys(inputs)
                #  ProgressMeter.next!(p,showvalues=[itm_name])
                key_type = nothing
                if contains(inputs[itm_name],"Continuous")
                    key_type = ML.Continuous
                elseif contains(inputs[itm_name],"Categorical")
                    key_type = ML.Categorical
                end
                schema_dict[itm_name] = key_type
            elseif itm_name ∈ keys(outputs)
                key_type = nothing
                if contains(outputs[itm_name],"Continuous")
                    key_type = ML.Continuous
                elseif contains(outputs[itm_name],"Categorical")
                    key_type = ML.Categorical
                end
                #  ProgressMeter.next!(p,showvalues=[itm_name])
                schema_dict[itm_name] = key_type
            else
                #  ProgressMeter.next!(p,showvalues=[itm_name])
                # ignore variables that aren't inputs or outputs
                schema_dict[itm_name] = nothing
            end
        end

        output_list = collect(keys(outputs))
        input_list = collect(keys(inputs))

        sch = ML.schema(table, hints=schema_dict)
        input_sch, output_sch = my_splitschema(sch,output_list)
    else
        input_sch = schema[1]
        output_sch = schema[2]
    end

    # allow subsampling data, subsample=1 means every data point is used
    # subsample=2 means every other point is used, etc...
    in_data = ML.featuremat(input_sch, table)[:,1:subsample:end]
    out_data = ML.featuremat(output_sch, table)[:,1:subsample:end]

    if any(isnan.(in_data)) || any(isnan.(out_data))
        println(input_sch)
        display(in_data)
        display(output_sch)
        display(out_data)
        # this can happen from mis-spelled keys, OR data like a zero in the standard deviation
        error("NaN in data!!!")
    end

    return in_data, out_data, table, input_sch, output_sch
end

function add_sq_scatter3d_annotation(ax1::PyCall.PyObject,ax2::PyCall.PyObject,in::Array,out_eng_ary::Dict,i1::Symbol,i2::Symbol,poi::Array,SQmodel::SQ_model;yval::Symbol=:X3_1,yval_std::Symbol=:X3_2,marker::AbstractString="*",s::Int64=100,color=:black,fontsize::Int64=12)

    in_eng, pred_outputs = SQ_predict(SQmodel,in,use_eng_units=true)
    display(in_eng)

    x_idx = nearest_to_x(in,poi,val_and_idx=false)
    println("#####################")
    display(in_eng)
    println(size(in_eng[i1]))
    println(x_idx)
    println(poi)
    println("[$(in_eng[i1][x_idx]) $(in_eng[i2][x_idx])]")
    println("#####################")
    #  error()

    limits = restore_eng_units(SQmodel.range,SQmodel.output_sch)
    #  println(limits)
    rwd_rng = [limits[yval][1],limits[yval][2]]

    #  println(in_eng)
    #  println(pred_outputs)
    sq = X3(Normal(out_eng_ary[yval][x_idx],out_eng_ary[yval_std][x_idx]),Normal(pred_outputs[yval][x_idx],pred_outputs[yval_std][x_idx]),global_rwd_range=rwd_rng)
    sq_txt = @sprintf("%0.3f",sq)

    println("!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!")
    println("vals: $([out_eng_ary[yval][x_idx], out_eng_ary[yval_std][x_idx]])")
    println("preds: $([pred_outputs[yval][x_idx], pred_outputs[yval_std][x_idx]])")
    println("pred_outputs: $pred_outputs")
    println("!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!")
    ax2[:errorbar](0.1,pred_outputs[yval][x_idx],pred_outputs[yval_std][x_idx],color=:blue,label="Trusted",fmt="o",capsize=3,alpha=0.6)
    ax2[:errorbar](-0.1,out_eng_ary[yval][x_idx],out_eng_ary[yval_std][x_idx],color=:red,fmt="o",capsize=3,alpha=0.6,label="Candidate")
    ax2[:set_xlim]([-0.2,0.2])
    ax2[:set_xticklabels]("")
    #  ax2[:set_ylim](rwd_rng)
    ax2[:axhline](limits[:X3_1][2])
    ax2[:axhline](limits[:X3_1][1])
    ax2[:text](0.8*0.2,limits[:X3_1][2],L"r_H",fontsize=fontsize,va="top",ha="right")
    ax2[:text](0.8*0.2,limits[:X3_1][1],L"r_L",fontsize=fontsize,va="bottom",ha="right")
    ax2[:set_title](string("SQ at ",marker,": $sq_txt"),fontsize=fontsize)
    ax2[:set_ylabel]("Reward")
    ax2[:set_adjustable]("box-forced")
    #  ax2[:set_aspect](1/(5*(limits[:X3_1][2]-limits[:X3_2][1])))
end
function norm2(ary::Array)
    n = NaN*ones(1,size(ary,2))
    for i = 1:size(ary,2)
        n[i] = norm(ary[:,i])
    end
    return n
end
function nearest_to_x(ary::Array,x::Float64;val_and_idx::Bool=false)
    resid = abs.(ary-x)
    @assert size(resid,2) == size(ary,2)
    idx_min = indmin(resid)
    if val_and_idx
        return idx_min,ary[idx_min]
    else
        return idx_min
    end
end
function nearest_to_x(ary::Array,x::Array;val_and_idx::Bool=false)
    resid = norm2(ary.-x)
    idx_min = indmin(resid)
    println(idx_min)
    if val_and_idx
        return idx_min,ary[:,idx_min]
    else
        return idx_min
    end
end

function add_sq_annotation(ax::PyCall.PyObject,x_ary::Dict,y_ary::Dict,preds::Dict,idx::Int64,xvar::Symbol,yvar_mean::Symbol,yvar_std::Symbol,SQmodel::SQ_model;fontsize::Int64=12,txt_loc::Array=[])
    #  println("########## Debug Annotation ##########")
    #  idx, idx_val = nearest_to_x(tst_in_eng[xvar],target_val,val_and_idx=true)
    limits = restore_eng_units(SQmodel.range,SQmodel.output_sch)

    note_x = x_ary[xvar][idx]

    ex_c_m = y_ary[yvar_mean][idx]
    ex_c_s = y_ary[yvar_std][idx]
    ex_c = Normal(ex_c_m,ex_c_s)

    ex_t_m = preds[yvar_mean][idx]
    ex_t_s = preds[yvar_std][idx]
    ex_t = Normal(ex_t_m,ex_t_s)

    ex_X3 = X3(ex_c,ex_t,global_rwd_range=[limits[yvar_mean][1],limits[yvar_mean][2]])

    if isempty(txt_loc)
        # auto calculate the text location
        x_lims = (minimum(x_ary[xvar]),maximum(x_ary[xvar]))
        y_lims = (limits[yvar_mean][1],limits[yvar_mean][2])

        rng_xlims = x_lims[2]-x_lims[1]
        rng_ylims = y_lims[2]-y_lims[1]
        note_pct_x = (note_x-x_lims[1])/rng_xlims*1.1
        note_pct_y = (ex_c_m-y_lims[1])/rng_ylims*1.1
        txt_loc = [note_pct_x*rng_xlims+x_lims[1],note_pct_y*rng_ylims+y_lims[1]]
        #  println("##########")
        #  println("limits: $limits")
        #  println("nx: $note_x, xlims: $x_lims")
        #  println("ny: $ex_c_m, ylims: $y_lims")
        #  #  println("pct_coords: $")
        #  println("arrow coords: $([note_x,ex_c_m])")
        #  println("##########")
    end

    ax[:annotate](@sprintf("SQ:%0.2f",ex_X3),xy=[note_x,ex_c_m],xytext=txt_loc,arrowprops=Dict(:arrowstyle=>"->"),fontsize=fontsize)
    #  ax[:annotate](".",xy=[note_x,ex_c_m],xytext=[0.8,0.],arrowprops=Dict(:arrowstyle=>"->"),fontsize=fontsize)
    #  ax[:errorbar](note_x, ex_t_m,yerr=ex_t_s/2,fmt="o",capsize=3,color=:black,alpha=0.6,label="")
    #  println("########## END Debug Annotation ##########")
end

function return_cat_stats(y::Dict,yval::Symbol,x::Dict,xval::Symbol,yval_std::Symbol;return_std::Bool=false)
    unq_x = unique(x[xval])
    num_xs = length(unq_x)
    ymat = zeros(num_xs)
    t_ucl = zeros(num_xs)
    t_lcl = zeros(num_xs)
    std_val = NaN
    for u_x in unq_x
        u_x_ind = convert(Int,round(u_x))
        vals = y[yval][x[xval].==u_x]
        if length(unique(vals)) == 1
            #using model to predict
            ymat[u_x_ind] = mean(vals)
            std_val = y[yval_std][u_x_ind]/2
            t_ucl[u_x_ind] = ymat[u_x_ind]+std_val
            t_lcl[u_x_ind] = ymat[u_x_ind]-std_val
        else
            #  println(vals)
            ymat[u_x_ind] = mean(vals)
            #  t_ucl[u_x_ind] = ymat[u_x_ind]+y[yval_std][u_x_ind]/2
            #  t_lcl[u_x_ind] = ymat[u_x_ind]-y[yval_std][u_x_ind]/2
            std_val = std(vals)
            t_ucl[u_x_ind] = ymat[u_x_ind]+std_val
            t_lcl[u_x_ind] = ymat[u_x_ind]-std_val
        end
    end
    if return_std
        return ymat,unq_x,std_val
    else
        return ymat,unq_x,t_ucl,t_lcl
    end
end

function scatter_with_conf_bnds(ax::PyCall.PyObject,x::Dict,y::Dict,xval::Symbol,yval::Symbol,yval_std::Symbol,color::Symbol;subsample::Array{Int64}=collect(1:length(x[xval])),label::String="",bar::Bool=false)
    unq_x = unique(x[xval])
    num_xs = length(unq_x)
    if bar
        x_srt2 = sortperm(x[xval])[subsample]
        ax[:errorbar](x[xval][subsample], y[yval][subsample],yerr=y[yval_std][subsample],fmt="o",capsize=3,color=color,alpha=0.6,label=label)
    else
        x_srt = sortperm(x[xval])[subsample]

        ax[:plot](x[xval][x_srt],y[yval][x_srt],label=label,color=color)
        t_ucl = y[yval]+y[yval_std]
        t_lcl = y[yval]-y[yval_std]
        #  println(t_ucl)
        #  println(x_srt)

        ax[:scatter](x[xval][x_srt],t_ucl[x_srt],label="",s=0)
        ax[:scatter](x[xval][x_srt],t_lcl[x_srt],label="",s=0)

        ax[:fill_between](x[xval][x_srt],y[yval][x_srt],t_ucl[x_srt],alpha=0.2,color=color)
        ax[:fill_between](x[xval][x_srt],y[yval][x_srt],t_lcl[x_srt],alpha=0.2,color=color)
    end

end
function make_label_from_keys(d::Dict)
    z = ""
    for x in collect(keys(d))
        z = string(z,x)
    end
    return z
end
function make_label_from_keys(a::Array)
    d = Dict()
    for z in a
        d[z] = ""
    end
    return make_label_from_keys(d)
end

function make_sq_model(net_type::String,inpts::Array{Symbol};num_epoc::Int64=250,subsample::Int64=1,
                       trusted_fname::String="",net_folder::String="nn_logs")
    train_fname = trusted_fname
    test_fname = trusted_fname #see if this works... I don't this is important for training

    #  inputs = Dict(:tprob=>"ML.Continuous",:e_mcts=>"ML.Continuous")
    inputs = Dict()
    for i in inpts
        inputs[i] = "ML.Continuous"
    end
    outputs = Dict(:X3_1=>"ML.Continuous",:X3_2=>"ML.Continuous")

    log_fname = joinpath(net_folder,"$(net_type)_$(make_label_from_keys(inputs))")

    training_subsample = subsample
    # make model
    SQmodel = make_nn_SQ_model(train_fname,test_fname,log_fname,input_dict=inputs,output_dict=outputs,nn_epoc=num_epoc,train_subsample=training_subsample)

    info("Writing variable to file:")
    println(log_fname)
    jldopen(log_fname*"_SQmodel.jld","w") do file
        JLD.addrequire(file,MXNet)
        JLD.addrequire(file,JuliaDB)
        write(file,"train_fname",train_fname)
        write(file,"test_fname",test_fname)
        write(file,"input_sch",SQmodel.input_sch)
        write(file,"output_sch",SQmodel.output_sch)
        write(file,"inputs",inputs)
        write(file,"outputs",outputs)
        write(file,"range",SQmodel.range)
    end
end

max_stars_repo_name: COHRINT/FaMSeC

max_issues_repo_name: COHRINT/FaMSeC

max_forks_repo_name: COHRINT/FaMSeC


text:
using Random: Xoshiro

const MosiRNG = Xoshiro

new_rng(seed) = MosiRNG(seed)

restore_rng(state::NTuple{4, UInt64}) = MosiRNG(state...)

rng_state(rng::MosiRNG) = (rng.s0, rng.s1, rng.s2, rng.s3)

make_seed() = UInt64(Random.make_seed()[1])
make_seed(n::Integer) = UInt64(Random.make_seed(n)[1])

macro _with_default_rng(func_def)
    @capture(func_def, function f_(__) _ end) ||
        @capture(func_def, f_(__) = _) ||
            error("Could not parse function definition")
    insert!(func_def.args[1].args, 2, :(rng::AbstractRNG))
    esc(quote
        $func_def
        $f(args...; kwargs...) = $f(GLOBAL_RNG, args...; kwargs...)
    end)
end


text:
using Test
using StaticArrays
using IntervalArithmetic
# using Revise
using HDGElasticity

function allapprox(v1,v2)
    return all(v1 .≈ v2)
end

@test HDGElasticity.linear_map_slope(1.,2.,[1.,2.],[3.,4.]) ≈ [2.,2.]

map = HDGElasticity.LineMap(0.,1.,[0.,0.],[1.,1.])
@test map.xiL ≈ 0.0
@test map.xiR ≈ 1.0
@test allapprox(map.xL,[0.,0.])
@test allapprox(map.xR,[1.,1.])
@test allapprox(map.slope,[1.,1.])
@test HDGElasticity.dimension(map) == 2
@test allapprox(map(0.),[0.,0.])
@test allapprox(map(0.5),[0.5,0.5])

map = HDGElasticity.LineMap([0.5,0.5],[1.,0.5])
@test map.xiL ≈ -1.0
@test map.xiR ≈ 1.0
@test allapprox(map.xL,[0.5,0.5])
@test allapprox(map.xR,[1.0,0.5])
@test allapprox(map.slope,[0.5,0.0]/2)

map = HDGElasticity.LineMap([0.,0.],[1.,1.]*sqrt(2))
@test allapprox(map(-1),[0.,0.])
@test allapprox(map(1),[1,1]*sqrt(2))
@test allapprox(HDGElasticity.jacobian(map),[1,1]/sqrt(2))
@test allapprox(HDGElasticity.determinant_jacobian(map),1.)
@test allapprox(HDGElasticity.inverse_jacobian(map),[1,1]*sqrt(2))

xiL = [0.,0.]
xiR = [1.,1.]
xL = [0.,0.]
xR = [2.,1.]
map = HDGElasticity.CellMap(xiL,xiR,xL,xR)
@test allapprox(map.xiL,xiL)
@test allapprox(map.xiR,xiR)
@test allapprox(map.xL,xL)
@test allapprox(map.xR,xR)
@test HDGElasticity.dimension(map) == 2
@test allapprox(map([0.5,0.5]),[1.,0.5])
@test allapprox(map([0.25,0.75]),[0.5,0.75])
@test allapprox(HDGElasticity.jacobian(map),[2.,1.])
@test allapprox(HDGElasticity.inverse_jacobian(map),[0.5,1.])
@test HDGElasticity.determinant_jacobian(map) ≈ 2.0
@test allapprox(HDGElasticity.face_determinant_jacobian(map),[2.,1.,2.,1.])

HDGElasticity.update_range!(map,[1.,1.],[2.,2.])
@test allapprox(map([0.5,0.5]),[1.5,1.5])
@test allapprox(map([0.75,0.25]),[1.75,1.25])
@test allapprox(HDGElasticity.jacobian(map),[1.,1.])
@test HDGElasticity.determinant_jacobian(map) ≈ 1.0

map = HDGElasticity.CellMap(2)
@test allapprox(map.xiL,[-1.,-1.])
@test allapprox(map.xiR,[+1.,+1.])
@test allapprox(map.xL,[-1.,-1.])
@test allapprox(map.xR,[+1.,+1.])
@test allapprox(map.slope,[1.,1.])
HDGElasticity.update_range!(map,[0.,0.],[1.,1.])
@test allapprox(map.xL,[0.,0.])
@test allapprox(map.xR,[1.,1.])
@test allapprox(map.slope,[0.5,0.5])

box = IntervalBox(1..2,3..5)
HDGElasticity.update_range!(map,box)
@test allapprox(map.xL,[1.,3.])
@test allapprox(map.xR,[2.,5.])
@test allapprox(map.slope,[0.5,1.])

map = HDGElasticity.CellMap([1.,2.],[2.,4.])
@test allapprox(map.xiL,[-1.,-1.])
@test allapprox(map.xiR,[1.,1.])
@test allapprox(map.xL,[1.,2.])
@test allapprox(map.xR,[2.,4.])
@test allapprox(HDGElasticity.jacobian(map),[0.5,1.])
@test HDGElasticity.determinant_jacobian(map) ≈ 0.5

xi = IntervalBox(-1..1,2)
x = IntervalBox(0..1,2)
map = HDGElasticity.CellMap(xi,x)
@test allapprox(map.xiL,[-1.,-1.])
@test allapprox(map.xiR,[1.,1.])
@test allapprox(map.xL,[0.,0.])
@test allapprox(map.xR,[1.,1.])

map = HDGElasticity.CellMap(IntervalBox(0..2,1..2))
@test allapprox(map.xiL,[-1.,-1.])
@test allapprox(map.xiR,[1.,1.])
@test allapprox(map.xL,[0.,1.])
@test allapprox(map.xR,[2.,2.])

fmaps = HDGElasticity.reference_cell_facemaps(2)
@test allapprox(fmaps[1](-1.0),[-1.,-1.])
@test allapprox(fmaps[1](+1.0),[+1.,-1.])

@test allapprox(fmaps[2](-1.0),[+1.,-1.])
@test allapprox(fmaps[2](+1.0),[+1.,+1.])

@test allapprox(fmaps[3](-1.0),[-1.,+1.])
@test allapprox(fmaps[3](+1.0),[+1.,+1.])

@test allapprox(fmaps[4](-1.0),[-1.,-1.])
@test allapprox(fmaps[4](+1.0),[-1.,+1.])


text:
import Base: size, getindex, *, +, -, ^, \, adjoint, ones
import LinearAlgebra: inv, eigen

# Assuming we work with real matrices...
struct Circulant{T} <: AbstractArray{T, 2}
    co::Vector{T}
    fft::Vector{Complex{T}}
end

Circulant(co) = Circulant(co, fft(co))

size(c::Circulant) = length(c.co), length(c.co)
getindex(c::Circulant, i::Int, j::Int) = c.co[mod(i - j, length(c.co))+1]

function *(a::Circulant, b::Circulant)
    return Circulant(real(ifft!(a.fft  .* b.fft)), a.fft  .* b.fft)
end

function *(a::Circulant, x::Vector)
     return ifft!(a.fft .* fft(x))
end

function *(a::Circulant, x::Vector{T}) where T <: Real
     return real(ifft!(a.fft .* fft(x)))
end

function +(a::Circulant, b::Circulant)
    return Circulant(a.co + b.co, a.fft + b.fft)
end

function -(a::Circulant, b::Circulant)
    return Circulant(a.co - b.co, a.fft - b.fft)
end

function ^(a::Circulant, n::Int)
    return Circulant(real(ifft!(a.fft.^n)), a.fft.^n)
end

function \(a::Circulant, x::Vector{T}) where T <: Real
     return real(ifft!(fft(x) ./ a.fft))
end

function inv(a::Circulant)
     return Circulant( real(ifft!(1 ./ a.fft )), 1 ./ a.fft)
end

function adjoint(a::Circulant)
    c = a.co
    c_new = vcat(c[1], c[end:-1:2])
    return Circulant(c_new)
end

function eigen(a::Circulant)
    λ = a.fft
    n = length(a.co)
    ω = [exp(2*im*π*k*j/n) for j in 0:n-1, k in 1:n]
    return λ, ω
end

function circ_ones(N::Int)
    return Circulant(ones(N))
end

function circ_id(N::Int)
    return Circulant([1, zeros(N-1)...])
end