text:
"""
$(DocStringExtensions.TYPEDSIGNATURES)

A short description of the function

```jldoctest
julia> 1 + 1
2
```
"""
function determine_edge_probabilities(graph, strand)
    kmers = graph.gprops[:kmers]
    outgoing_edge_probabilities = SparseArrays.spzeros(length(kmers), length(kmers))

    for (kmer_index, kmer) in enumerate(kmers)
        if !strand
            kmer = BioSequences.reverse_complement(kmer)
        end

        downstream_neighbor_indices = Int[]
        for neighbor in BioSequences.neighbors(kmer)
            index = get_kmer_index(kmers, BioSequences.canonical(neighbor))
            # kmer must be in our dataset and there must be a connecting edge
            if !isnothing(index) && Graphs.has_edge(graph, ordered_edge(kmer_index, index))
                push!(downstream_neighbor_indices, index)
            end
        end
        sort!(unique!(downstream_neighbor_indices))

        downstream_edge_weights = Int[
            length(get(graph.edge_evidence, ordered_edge(kmer_index, neighbor_index), EdgeEvidence[])) for neighbor_index in downstream_neighbor_indices
        ]

        non_zero_indices = downstream_edge_weights .> 0
        downstream_neighbor_indices = downstream_neighbor_indices[non_zero_indices]
        downstream_edge_weights = downstream_edge_weights[non_zero_indices]

        downstream_edge_likelihoods = downstream_edge_weights ./ sum(downstream_edge_weights)

        for (neighbor_index, likelihood) in zip(downstream_neighbor_indices, downstream_edge_likelihoods)
            outgoing_edge_probabilities[kmer_index, neighbor_index] = likelihood
        end
    end
    return outgoing_edge_probabilities
end

"""
$(DocStringExtensions.TYPEDSIGNATURES)

A short description of the function

```jldoctest
julia> 1 + 1
2
```
"""
function determine_edge_probabilities(graph)
    outgoing_edge_probabilities = determine_edge_probabilities(graph, true)
    incoming_edge_probabilities = determine_edge_probabilities(graph, false)
    return outgoing_edge_probabilities, incoming_edge_probabilities
end

"""
$(DocStringExtensions.TYPEDSIGNATURES)

Get dna (db = "nuccore") or protein (db = "protein") sequences from NCBI
or get fasta directly from FTP site

```jldoctest
julia> 1 + 1
2
```
"""
function oriented_unbranching_walk(kmer_graph, vertex, orientation)
    walk = []
    viable_neighbors = find_unbranched_neighbors(kmer_graph, vertex, orientation)
    while length(viable_neighbors) == 1
#         @show "found a viable neighbor!!"
        viable_neighbor = first(viable_neighbors)
        edge = Graphs.Edge(vertex, viable_neighbor)
        push!(walk, edge)
        vertex = edge.dst
        viable_neighbors = Set{Int}()
        destination_orientations = [o.destination_orientation for o in kmer_graph.eprops[edge][:orientations]]
        for destination_orientation in destination_orientations
            union!(viable_neighbors, find_unbranched_neighbors(kmer_graph, vertex, destination_orientation))
        end
    end
    return walk
end

"""
$(DocStringExtensions.TYPEDSIGNATURES)

Get dna (db = "nuccore") or protein (db = "protein") sequences from NCBI
or get fasta directly from FTP site

```jldoctest
julia> 1 + 1
2
```
"""
function find_unbranched_neighbors(kmer_graph, vertex, orientation)
    downstream_vertices = find_downstream_vertices(kmer_graph, vertex, orientation)
#     backtrack_vertices
#     @show downstream_vertices
    if length(downstream_vertices) == 1
        downstream_vertex = first(downstream_vertices)
#         @show downstream_vertex
        edge = Graphs.Edge(vertex, downstream_vertex)
#         @show edge
        destination_orientations = [o.destination_orientation for o in kmer_graph.eprops[edge][:orientations]]
#         @show destination_orientations
        for destination_orientation in destination_orientations
            backtrack_vertices = find_downstream_vertices(kmer_graph, downstream_vertex, !destination_orientation)
#             @show backtrack_vertices
            # if the only backtrack is the vertex we're on, then we can simplify
            if backtrack_vertices == Set([vertex])
                return downstream_vertices
            end
        end
    end
    return Int[]
end

"""
$(DocStringExtensions.TYPEDSIGNATURES)

Get dna (db = "nuccore") or protein (db = "protein") sequences from NCBI
or get fasta directly from FTP site

```jldoctest
julia> 1 + 1
2
```
"""
function find_downstream_vertices(kmer_graph, vertex, orientation)
    viable_neighbors = Set{Int}()
    for neighbor in Graphs.neighbors(kmer_graph, vertex)
        not_same_vertex = vertex != neighbor
        candidate_edge = Graphs.Edge(vertex, neighbor)
        # palindromes can have multiple viable orientations
        # check each viable orientation individually
        edge_src_orientations = [e.source_orientation for e in kmer_graph.eprops[candidate_edge][:orientations]]
        for edge_src_orientation in edge_src_orientations
#             edge_src_orientation = kmer_graph.eprops[candidate_edge][:orientations].source_orientation
            viable_orientation = edge_src_orientation == orientation
            if not_same_vertex && viable_orientation
                push!(viable_neighbors, neighbor)
            end
        end
    end
    return viable_neighbors
end

# """
# $(DocStringExtensions.TYPEDSIGNATURES)

# A short description of the function

# ```jldoctest
# julia> 1 + 1
# 2
# ```
# """
# function reverse_oriented_path(oriented_path)
#     reversed_path = copy(oriented_path)
#     for (index, state) in enumerate(oriented_path)
#         reversed_path[index] = Mycelia.OrientedKmer(index = state.index, orientation = !state.orientation)
#     end
#     return reverse!(reversed_path)
# end

# """
# $(DocStringExtensions.TYPEDSIGNATURES)

# A short description of the function

# ```jldoctest
# julia> 1 + 1
# 2
# ```
# """
# function maximum_likelihood_walk(graph, connected_component)
#     max_count = maximum(graph.counts[connected_component])
#     max_count_indices = findall(count -> count == max_count, graph.counts[connected_component])
#     initial_node_index = rand(max_count_indices)
#     initial_node = connected_component[initial_node_index]
#     outgoing_edge_probabilities, incoming_edge_probabilities = Mycelia.determine_edge_probabilities(graph)
#     forward_walk = maximum_likelihood_walk(graph, [Mycelia.OrientedKmer(index = initial_node, orientation = true)], outgoing_edge_probabilities, incoming_edge_probabilities)
#     reverse_walk = maximum_likelihood_walk(graph, [Mycelia.OrientedKmer(index = initial_node, orientation = false)], outgoing_edge_probabilities, incoming_edge_probabilities)
#     reversed_reverse_walk = reverse!(
#         [
#             Mycelia.OrientedKmer(index = oriented_kmer.index, orientation = oriented_kmer.orientation)
#             for oriented_kmer in reverse_walk[2:end]
#         ]
#         )
#     full_path = [reversed_reverse_walk..., forward_walk...]
# end

# """
# $(DocStringExtensions.TYPEDSIGNATURES)

# A short description of the function

# ```jldoctest
# julia> 1 + 1
# 2
# ```
# """
# function maximum_likelihood_walk(graph, path::Vector{Mycelia.OrientedKmer}, outgoing_edge_probabilities, incoming_edge_probabilities)
#     done = false
#     while !done
#         maximum_path_likelihood = 0.0
#         maximum_likelihood_path = Vector{Mycelia.OrientedKmer}()
#         for neighbor in Graphs.neighbors(graph.graph, last(path).index)
#             this_path = [last(path).index, neighbor]
#             this_oriented_path, this_path_likelihood =
#                 Mycelia.assess_path(this_path,
#                     graph.kmers,
#                     graph.counts,
#                     last(path).orientation,
#                     outgoing_edge_probabilities,
#                     incoming_edge_probabilities)
#             if this_path_likelihood > maximum_path_likelihood
#                 maximum_path_likelihood = this_path_likelihood
#                 maximum_likelihood_path = this_oriented_path
#             end
#         end
#         if isempty(maximum_likelihood_path) && (maximum_path_likelihood == 0.0)
#             done = true
#         else
#             append!(path, maximum_likelihood_path[2:end])
#         end
#     end
#     return path
# end

# """
# $(DocStringExtensions.TYPEDSIGNATURES)

# A short description of the function

# ```jldoctest
# julia> 1 + 1
# 2
# ```
# """
# function take_a_walk(graph, connected_component)
#     max_count = maximum(graph.counts[connected_component])
#     max_count_indices = findall(count -> count == max_count, graph.counts[connected_component])
#     initial_node_index = rand(max_count_indices)
#     initial_node = connected_component[initial_node_index]
#     outgoing_edge_probabilities, incoming_edge_probabilities = Mycelia.determine_edge_probabilities(graph)

#     # walk forwards from the initial starting node
#     forward_walk = take_a_walk(graph, [Mycelia.OrientedKmer(index = initial_node, orientation = true)], outgoing_edge_probabilities, incoming_edge_probabilities)

#     # walk backwards from the initial starting node
#     reverse_walk = take_a_walk(graph, [Mycelia.OrientedKmer(index = initial_node, orientation = false)], outgoing_edge_probabilities, incoming_edge_probabilities)

#     # we need to reverse everything to re-orient against the forward walk
#     reverse_walk = reverse_oriented_path(reverse_walk)

#     # also need to drop the last node, which is equivalent to the first node of the
#     @assert last(reverse_walk) == first(forward_walk)
#     full_path = [reverse_walk[1:end-1]..., forward_walk...]
# #     @show full_path
# end

# """
# $(DocStringExtensions.TYPEDSIGNATURES)

# A short description of the function

# ```jldoctest
# julia> 1 + 1
# 2
# ```
# """
# function take_a_walk(graph, path::Vector{Mycelia.OrientedKmer}, outgoing_edge_probabilities, incoming_edge_probabilities)
#     done = false
#     while !done
#         maximum_path_likelihood = 0.0
#         maximum_likelihood_path = Vector{Mycelia.OrientedKmer}()
#         for neighbor in Graphs.neighbors(graph.graph, last(path).index)
#             this_path = [last(path).index, neighbor]
#             this_oriented_path, this_path_likelihood =
#                 Mycelia.assess_path(this_path,
#                     graph.kmers,
#                     graph.counts,
#                     last(path).orientation,
#                     outgoing_edge_probabilities,
#                     incoming_edge_probabilities)
#             if this_path_likelihood > maximum_path_likelihood
#                 maximum_path_likelihood = this_path_likelihood
#                 maximum_likelihood_path = this_oriented_path
#             end
#         end
#         if isempty(maximum_likelihood_path) && (maximum_path_likelihood == 0.0)
#             done = true
#         else
#             append!(path, maximum_likelihood_path[2:end])
#         end
#     end
#     return path
# end

# use default cost 1 for breadth first seach
# use default cost 0 for depth first search
function dijkstra_step!(graph, distances, arrival_paths, queue; default_cost = 1)
    current_kmer, cost = DataStructures.dequeue_pair!(queue)
    current_orientation = BioSequences.iscanonical(current_kmer)
    current_canonical_kmer = BioSequences.canonical(current_kmer)
    current_index = graph[current_canonical_kmer, :kmer]

    present_neighbors = Vector{typeof(current_kmer)}()
    for neighbor in BioSequences.neighbors(current_kmer)
        try
            graph[BioSequences.canonical(neighbor), :kmer]
            push!(present_neighbors, neighbor)
        catch
            continue
        end
    end

    true_neighbors =
        filter(neighbor ->
            Graphs.has_edge(
                graph,
                graph[BioSequences.canonical(current_kmer), :kmer],
                graph[BioSequences.canonical(neighbor), :kmer]),
            present_neighbors)

    if !isempty(true_neighbors)
        canonical_true_neighbors = BioSequences.canonical.(true_neighbors)
        true_neighbor_is_canonical = BioSequences.iscanonical.(true_neighbors)
        neighbor_indices = map(canonical_neighbor -> graph[canonical_neighbor, :kmer], canonical_true_neighbors)
        neighbor_weights = map(neighbor_i -> MetaGraphs.get_prop(graph, current_index, neighbor_i, :weight), neighbor_indices)

        # inverse probability!!
        neighbor_costs = 1 .- (neighbor_weights ./ sum(neighbor_weights))
        neighbor_costs .+= default_cost
        # afterwards, 100% likelihood edges cost 1
        # 50% likelihood costs 1.5
        # 0% likelihood is impossible and not included

        for (neighbor, cost) in zip(true_neighbors, neighbor_costs)
            candidate_path = vcat(arrival_paths[current_kmer], current_kmer)
            candidate_path_cost = distances[current_kmer] + cost
            if distances[neighbor] > candidate_path_cost
                arrival_paths[neighbor] = candidate_path
                DataStructures.enqueue!(queue, neighbor => cost)
                distances[neighbor] = candidate_path_cost
            end
        end
    end
    return distances, arrival_paths, queue
end

function evaluate_hits(hits, forward_distances, reverse_distances, forward_arrival_paths, reverse_arrival_paths)

    lowest_cost = Inf
    optimal_path = Int[]

    for hit in hits
        reverse_complement_hit = BioSequences.reverse_complement(hit)
        total_cost = forward_distances[hit] + reverse_distances[reverse_complement_hit]
        if total_cost < lowest_cost
            forward_path = vcat(forward_arrival_paths[hit], hit)
            reverse_path = vcat(reverse_arrival_paths[reverse_complement_hit], reverse_complement_hit)
            reverse_path = reverse!(BioSequences.reverse_complement.(reverse_path))
            full_path = vcat(forward_path, reverse_path[2:end])

            lowest_cost = total_cost
            optimal_path = full_path
        end
    end

    return lowest_cost, optimal_path
end

function bidirectional_dijkstra(graph, a::T, b::T) where T <: BioSequences.AbstractMer
    forward_distances = DataStructures.DefaultDict{T, Float64}(Inf)
    forward_distances[a] = 0

    forward_arrival_paths = Dict{T, Vector{T}}()
    forward_arrival_paths[a] = Int[]

    forward_queue = DataStructures.PriorityQueue{T, Float64}()
    DataStructures.enqueue!(forward_queue, a, 0)

    reverse_b = BioSequences.reverse_complement(b)
    reverse_distances = DataStructures.DefaultDict{T, Float64}(Inf)
    reverse_distances[reverse_b] = 0

    reverse_arrival_paths = Dict{T, Vector{T}}()
    reverse_arrival_paths[reverse_b] = Int[]

    reverse_queue = DataStructures.PriorityQueue{T, Float64}()
    DataStructures.enqueue!(reverse_queue, reverse_b, 0)

    hits = Set{T}()

    while isempty(hits)
        dijkstra_step!(graph, forward_distances, forward_arrival_paths, forward_queue)
        dijkstra_step!(graph, reverse_distances, reverse_arrival_paths, reverse_queue)
        hits = intersect(
                    keys(forward_arrival_paths),
                    BioSequences.reverse_complement.(keys(reverse_arrival_paths)))
    end
    lowest_cost, optimal_path = evaluate_hits(hits, forward_distances, reverse_distances, forward_arrival_paths, reverse_arrival_paths)
    return lowest_cost, optimal_path
end

# # simple point to point
# function dijkstra(graph, a::T, b::T) where {T <: BioSequences.AbstractMer}
#     distances = DataStructures.DefaultDict{T, Float64}(Inf)
#     distances[a] = 0

#     arrival_paths = Dict{T, Vector{T}}()
#     arrival_paths[a] = Int[]

#     queue = DataStructures.PriorityQueue{T, Float64}()
#     DataStructures.enqueue!(queue, a, 0)

#     current_kmer, cost = a, Inf
#     while current_kmer != b
#         distances, arrival_paths, queue = dijkstra_step!(graph, distances, arrival_paths, queue)
#     end
#     shortest_path = vcat(arrival_paths[b], b)
#     return shortest_path, distances[b]
# end

function dijkstra(graph, a::T, targets::Set{T}; search_strategy::Union{Symbol, Missing}=missing) where {T <: BioSequences.AbstractMer}
    if ismissing(search_strategy)
        # local breadth first search to find the single target
        # a-b shortest path
        if length(targets) == 1
            search_strategy = :BFS
        # a - no targets - go find the ends
        # a - series of any target
        else
            search_strategy = :DFS
        end
    end
    if search_strategy == :BFS
        default_cost = 1
    elseif search_strategy == :DFS
        default_cost = 0
    else
        error("please specify either :BFS (breadth-first search) or :DFS (depth-first search)")
    end

    distances = DataStructures.DefaultDict{T, Float64}(Inf)
    distances[a] = 0

    arrival_paths = Dict{T, Vector{T}}()
    arrival_paths[a] = Int[]

    queue = DataStructures.PriorityQueue{T, Float64}()
    DataStructures.enqueue!(queue, a, 0)

    while !isempty(queue) && !(first(DataStructures.peek(queue)) in targets)
        distances, arrival_paths, queue =
            Mycelia.dijkstra_step!(graph, distances, arrival_paths, queue, default_cost = default_cost)
    end
    if !isempty(queue)
        discovered_destination = DataStructures.dequeue!(queue)
        @assert discovered_destination in targets
        shortest_path = vcat(arrival_paths[discovered_destination], discovered_destination)
        return shortest_path, distances[discovered_destination]
    else
        longest_path = Int[]
        distance = Inf
        for (destination, arrival_path) in arrival_paths
            this_distance = distances[destination]
            this_path = vcat(arrival_path, destination)
            if (length(this_path) > length(longest_path)) && (this_distance <= distance)
                longest_path = this_path
                distance = this_distance
            end

        end
        return longest_path, distance
    end
end

function update_remaining_targets(current_walk::AbstractVector{T}, remaining_targets::AbstractSet{T}) where T <: BioSequences.AbstractMer
    # assess whether targets have been hit in the canonical space
    remaining_targets = setdiff(BioSequences.canonical.(remaining_targets), BioSequences.canonical.(current_walk))
    # blow back out into forward and reverse_complement space
    remaining_targets = Set{T}(vcat(remaining_targets, BioSequences.reverse_complement.(remaining_targets)))
    return remaining_targets
end

function assess_downstream_weight(graph, kmer)
    # here we look to see if walking forward or backward from the initial node gets us to heavier weight options
    score = 0
    for neighbor in BioSequences.neighbors(kmer)
        try
            score += MetaGraphs.get_prop(graph, graph[BioSequences.canonical(neighbor), :kmer], :count)
        catch
            continue
        end
    end
    return score
end

# vertices should either be entire graph (by default) or a connected component
# if people want to work on just the connected component, let them induce a subgraph
function find_graph_core(graph; seed=rand(Int))

    Random.seed!(seed)

    T = typeof(MetaGraphs.get_prop(graph, 1, :kmer))

    # targets = [MetaGraphs.get_prop(graph, v, :kmer) for v in Graphs.vertices(graph)]
    # sortperm

    # targets = [
    #         MetaGraphs.get_prop(graph, v_index, :count) for v_index in
    #         sortperm([MetaGraphs.get_prop(graph, v, :count) for v in Graphs.vertices(graph)], rev=true)[1:Int(ceil(Graphs.nv(graph)/10))]]

    targets = [MetaGraphs.get_prop(graph, v, :kmer) for v in Graphs.vertices(graph) if MetaGraphs.get_prop(graph, v, :count) > 1]
            # sortperm([MetaGraphs.get_prop(graph, v, :count) for v in Graphs.vertices(graph)], rev=true)[1:Int(ceil(Graphs.nv(graph)/10))]]

    @show length(targets)
    starting_kmer = first(targets)
    max_degree = 0
    for node in targets
        node_degree = Graphs.degree(graph, graph[node, :kmer])
        if node_degree > max_degree
            max_degree = node_degree
            starting_kmer = node
        end
    end

    current_walk = [starting_kmer]
    prior_walk_length = length(current_walk)
    remaining_targets = update_remaining_targets(current_walk, Set(targets))
    done = isempty(remaining_targets)

    while !done
        # here we look to see if walking forward or backward from the current ends gets us to heavier weight options
        # we want to prioritize walks toward higher coverage nodes
        forward_score = assess_downstream_weight(graph, last(current_walk))
        reverse_score = assess_downstream_weight(graph, BioSequences.reverse_complement(first(current_walk)))
        if reverse_score > forward_score
            current_walk = reverse(BioSequences.reverse_complement.(current_walk))
        end

        forward_source = last(current_walk)
        forward_walk, forward_distance = Mycelia.dijkstra(graph, forward_source, remaining_targets, search_strategy=:DFS)
        current_walk = vcat(current_walk, forward_walk[2:end])
        remaining_targets = update_remaining_targets(current_walk, remaining_targets)
        if isempty(remaining_targets)
            done = true
        else
            reverse_source = BioSequences.reverse_complement(first(current_walk))
            reverse_walk, reverse_distance = Mycelia.dijkstra(graph, reverse_source, remaining_targets, search_strategy=:DFS)
            current_walk = vcat(reverse(BioSequences.reverse_complement.(reverse_walk))[1:end-1], current_walk)
            remaining_targets = update_remaining_targets(current_walk, remaining_targets)
        end
        @show length(current_walk)
        failed_this_expansion = length(current_walk) == prior_walk_length
        prior_walk_length = length(current_walk)
        if isempty(remaining_targets)
            done = true
        elseif failed_this_expansion
            done = true
        end
    end
    return current_walk
end

max_stars_repo_name: cjprybol/Eisenia

max_issues_repo_name: cjprybol/Eisenia

max_forks_repo_name: cjprybol/Eisenia


text:
module RestrictedDiscreteModelsTests

using Gridap.Arrays
using Gridap.ReferenceFEs
using Gridap.Geometry
using Test

domain = (0,2,0,2,0,2)
partition = (10,10,10)
oldmodel = CartesianDiscreteModel(domain,partition)

const R = 0.7

function is_in(coords)
  n = length(coords)
  x = (1/n)*sum(coords)
  d = x[1]^2 + x[2]^2 - R^2
  d < 0
end

oldgrid = get_grid(oldmodel)
oldcell_to_coods = get_cell_coordinates(oldgrid)
cell_to_mask = lazy_map(is_in,oldcell_to_coods)
cell_to_oldcell = findall(collect1d(cell_to_mask))

labels = get_face_labeling(oldmodel)

ne = num_entities(labels)
fluid_entity = ne+1
solid_entity = ne+2
cell_to_entity = get_cell_entity(labels)
cell_to_entity .= fluid_entity
cell_to_entity[cell_to_oldcell] .= solid_entity
add_tag!(labels,"fluid",[fluid_entity])
add_tag!(labels,"solid",[solid_entity])

model = RestrictedDiscreteModel(oldmodel,cell_to_oldcell)
test_discrete_model(model)

trian = get_triangulation(model)
@test isa(trian,RestrictedTriangulation)

trian = Triangulation(model)
@test isa(trian,RestrictedTriangulation)

model = DiscreteModel(oldmodel,cell_to_oldcell)
test_discrete_model(model)
@test isa(model,RestrictedDiscreteModel)

model = DiscreteModel(oldmodel,cell_to_mask)
test_discrete_model(model)
@test isa(model,RestrictedDiscreteModel)

model_fluid = DiscreteModel(oldmodel,tags="fluid")
test_discrete_model(model_fluid)
@test isa(model_fluid,RestrictedDiscreteModel)

model_solid = DiscreteModel(oldmodel,tags="solid")
test_discrete_model(model_solid)
@test isa(model_solid,RestrictedDiscreteModel)

itrian = InterfaceTriangulation(model_fluid,model_solid)

#using Gridap.Visualization
#writevtk(itrian,"itrian",nsubcells=3,cellfields=["normal"=>get_normal_vector(itrian)])

#
#writevtk(model,"model")
#writevtk(trian,"trian")

oldtrian = Triangulation(oldmodel)
trian_s = SkeletonTriangulation(model)
trian_b = BoundaryTriangulation(model)

@test get_background_triangulation(trian_s) === oldtrian
@test get_background_triangulation(trian_b) === oldtrian
@test get_background_triangulation(itrian) === oldtrian

#meas_K_b = cell_measure(trian_b,oldtrian)
#meas_K_sl = cell_measure(trian_s.plus,oldtrian)
#meas_K_sr = cell_measure(trian_s.minus,oldtrian)
#
#oldcell_to_cell = zeros(Int,num_cells(oldmodel))
#oldcell_to_cell[cell_to_oldcell] .= 1:length(cell_to_oldcell)
#
#@test all( oldcell_to_cell[findall(meas_K_b .!= 0)] .!= 0 )
#@test all( oldcell_to_cell[findall(meas_K_sl .!= 0)] .!= 0 )
#@test all( oldcell_to_cell[findall(meas_K_sr .!= 0)] .!= 0 )

#using Gridap.Visualization
#writevtk(trian_b,"trian_b")
#writevtk(trian_s,"trian_s")
#writevtk(trian,"trian")
#writevtk(oldtrian,"oldtrian",celldata=["meas_K_b"=>meas_K_b,"meas_K_sl"=>meas_K_sl,"meas_K_sr"=>meas_K_sr])

end # module

ext: jl


text:
include("SimulationsUtils.jl")
include("MethodUtils.jl")
include("WassersteinVariance.jl")
using Serialization

function compare_wvm_icp(nsim, numS, ne, noisetype="normal", numvar=18, seed=123)
    Random.seed!(seed)
    ######################## SIMULATION ##############################
    # nsim = 100
    # These will store a four tuple for each simulation: nfp, nfn, length(Ŝ)
    errs = []
    iscauses, wvm_pvals, icp_pvals = [], [], []
    simtype = "numS-$(numS)/ne-$(ne)/$(noisetype)"
    for i in 1:nsim
        ############### DATA SIM #################
        # collect simulation data
        X_full, S, nₑs = deserialize("../data/Simulations/$(simtype)/sim-$(i).jls")
        n, p = size(X_full,1), size(X_full,2)
        X, y = X_full[:, 1:end-1], X_full[:, end]

        iscause = zeros(p-1)
        iscause[S] .= 1

        ################# lasso #################
        cf = select_lasso_by_num_var(X, y, numvar)
        select = 1:p-1
        select = filter(x -> cf[x] != 0, select)

        ################# WVM + ICP #################
        αs = [0.1, 0.3, 0.5, 0.7]
        wv = DirectWassersteinVariance(X[:,select], y, nₑs, nₑs / n)
        res = WV_screening(wv, α=αs, solver=LBFGS(m=100), B=50)
        Ŝs_wvm = [select[findall_nonzero(r)] for r in res[2]]
        wvm_errs = [error_count(S, Ŝ_wvm) for Ŝ_wvm in Ŝs_wvm]
        wvm_pval = ones(p-1)
        wvm_pval[select] .= res[4]

        icp_res, icp_pval = icp(X, y, nₑs, αs[1], false, "boosting", true)
        Ŝ_icp = icp_res # findall_nonzero(isnothing(icp_res) ? zeros(p-1) : icp_res)
        err = push!([error_count(S, Ŝ_icp)], wvm_errs...)
        push!(errs, err)

        push!(icp_pvals, icp_pval)
        push!(wvm_pvals, wvm_pval)
        push!(iscauses, iscause)
        println("###### SIMULATION $i IS DONE ! ######")
    end
    pr_curve_res_wvm = hcat(vcat(iscauses...), vcat(wvm_pvals...))
    pr_curve_res_icp = hcat(vcat(iscauses...), vcat(icp_pvals...))
    run(`mkdir -p ../data/results/$(simtype)`)
    # serialize("results/numS-$(numS)/$(noisetype)$(addedinfo)-wvm-and-icp-results.jls", errs)
    serialize("../data/results/$(simtype)/diff-alpha-wvm-and-icp.jls", errs)
    serialize("../data/results/$(simtype)/pr-curve-wvm.jls", pr_curve_res_wvm)
    serialize("../data/results/$(simtype)/pr-curve-icp.jls", pr_curve_res_icp)
end
compare_wvm_icp(100, 6, 100, "normal")
compare_wvm_icp(100, 6, 500, "normal")
compare_wvm_icp(100, 12, 500, "normal")


text:
export atmospherefit, atmospheretable, gravity

#Define Properties Constants
const Tsl = 288.15  #sea level temperature (K)
const Psl = 101325.0 #sea level pressure (Pa)
const musl = 0.0000181206  #sea level viscosity (N-s/m^2)
const Sc = 113.0  #constant in Sutherlands formula (K)
const gamma = 1.4  #ratio of specific heats (air)
const R_M = 287.0520809957016 #Molar Gas Constant: R (gas constant) = 8.31432e3 #N m kmol−1 K−1, M (mean molar mass) = 28.9645 #g/mol, R_M = R/M
const earthradius = 6369.0e3 #radius of the earth (m)
const g = 9.80665 #gravitational constant at mean sea level

"""
atmospherefit(altitude::Real)

Return density (kg/m^3), viscosity (Pa-s), and speed of sound (m/s) functions.

Input altitude must be in meters.
"""
function atmospherefit(altitude::Real)
    if altitude > 47000.0
        @warn("air properties for altitudes above 47000 meters will be innacurate.")
    end
    T, P = temp_presfit(altitude)
    rho = density(T, P)
    mu = viscosity(T)
    a = speedofsound(T)

    return rho, mu, a

end #atmospherefit

"""
temp_presfit(altitude::Real)

Calculate atmospheric temperature and pressure for input altitude (meters) using fits of the standard atmosphere model,
slightly modified from those found in Flight Vehicle Aerodynamics by Mark Drela.

"""
function temp_presfit(altitude::Real)

    #Convert Altitude to km
    altkm = altitude/1000

    T = Tsl - 71.5 + 2*log(1+exp(35.75-3.25*altkm)+exp(-3.0+0.0003*altkm^3))
    P = Psl*exp(-0.118*altkm-(0.0015*altkm^2)/(1-0.018*altkm+0.0011*altkm^2))

    return T, P

end #temp_presfit

"""
atmospheretable(altitude::Real)

Return density (kg/m^3), viscosity (Pa-s), and speed of sound (m/s) functions.

Input altitude must be in meters.
"""
function atmospheretable(altitude::Real)
    if altitude > 86000.0
        @warn("air properties for altitudes above 86000 meters will be innacurate.")
    end

    T, P = temp_prestable(altitude)
    rho = density(T, P)
    mu = viscosity(T)
    a = speedofsound(T)

    return rho, mu, a

end #atmospherefit

"""
temp_prestable(altitude::Real)

Calculate atmospheric temperature and pressure for input altitude (meters)
from 1976 Standard Atmosphere Model.

"""
function temp_prestable(altitude::Real)
    #Convert Geometric Altitude to Geopotential Altitude
    altgeopot = altitude*earthradius/(altitude+earthradius)

    #table values
    altitudetable = [0.0; 11.0; 20.0; 32.0; 47.0; 51.0; 71.0; 84.5]*1e3
    temperaturetable = [288.150; 216.650; 216.650; 228.650; 270.650; 270.650; 214.650; 187.650]
    temperaturegradient = [-0.0065; 0.0; 0.001; 0.0028; 0.0; -0.0028; -0.002; 0.0]
    pressuretable = [101325.0; 22632.0; 5474.8; 868.01; 110.90; 66.938; 3.9564; 0.39814]

    #find relavent index in tables
    idx = findall(altitudetable.<=altgeopot)[end]

    #find temperature
    T = temperaturetable[idx] + (altgeopot-altitudetable[idx])*temperaturegradient[idx]

    #find pressure
    if temperaturegradient[idx] == 0.0
        P = pressuretable[idx]*exp(-g*(altgeopot-altitudetable[idx])/R_M/temperaturetable[idx])
    else
        P = pressuretable[idx]*(temperaturetable[idx]/T)^(g/R_M/temperaturegradient[idx])
    end

    return T, P

end

"""
density(Temperature::Real, Pressure::Real)

Return Air Density (kg/m^3) from Ideal Gas Law.

Input Temperature must be in Kelvin, and Input Pressure must be in Pascals.
"""
function density(Temperature::Real, Pressure::Real)
    #Ideal Gas Law
    rho = Pressure/(R_M*Temperature)
    return rho
end #density

"""
speedofsound(Temperature::Real)

Return Speed of Sound (m/s) from Ideal Gas Law.

Input Temperature must be in Kelvin.
"""
function speedofsound(Temperature::Real)
    #speed of sound (ideal gas law)
    a = sqrt(gamma*R_M*Temperature)
    return a
end #speedofsound

"""
viscosity(Temperature::Real)

Return Dynamic Viscosity (Pa-s) from Sutherlands Equation.

Input Temperature must be in Kelvin.
"""
function viscosity(Temperature::Real)
    mu = musl*(Temperature/Tsl)^(3.0/2)*(Tsl+Sc)/(Temperature+Sc)
    return mu
end #viscosity

"""
gravity(altitude::Real)

Calculate adjusted gravitational acceleration at altitude (meters) and latitude (degrees).

1984 Ellipsoidal Gravity Formula.

"""
function gravity(altitude::Real, latitude::Real=45.0)

    #latitude adjustment of mean sea level gravitational acceleration.
    gravadjlat = 9.7803253359*(1 + 0.00193185265241*sind(latitude)^2)/sqrt(1 - 0.00669437999013*sind(latitude)^2)
    #altitude adjustment of gravitational constant.
    gravadjalt = gravadjlat*(earthradius/(altitude+earthradius))^2

    return gravadjalt
end #gravity


text:
using Sads
using Test

@testset "Sads.jl" begin
    # Write your own tests here.
end


text:
#
# Stuff to read FRS and HBAI files into frames and merge them
# a lot of this is just to overcome the weird way I have these stored
# and would need adaption by anyone with a sensible layout who had all the
# files in the one format
# SINGLE (LATEST) year only
#
using DataFrames
using StatFiles
using TableTraits
using CSV
using Utils
using Definitions
using CSVFiles

function loadfrs(which::AbstractString, year::Integer)::DataFrame
    filename = "$(FRS_DIR)/$(year)/tab/$(which).tab"
    loadtoframe(filename)
end

## FIXME Move these and stb versions to libs

function infer_hours_of_care(hourtot::Integer)::Real
    hrs = Dict(
        0 => 0.0,
        1 => 2.0,
        2 => 7.0,
        3 => 14.0,
        4 => 27.5,
        5 => 42.5,
        6 => 75.0,
        7 => 100.0,
        8 => 10.0,
        9 => 27.5,
        10 => 50.0
    )
    h = 0.0
    if hourtot in keys(hrs)
        h = hrs[hourtot]
    end
    h
end

function remapRegion( r :: Integer )
        if r == 112000001 # North East
                r = 1
        elseif r == 112000002 # North West
                r = 2
        # no region 3
        elseif r == 112000003 # Yorks and the Humber
                r = 4
        elseif r == 112000004 # East Midlands
                r = 5
        elseif r == 112000005 # West Midlands
                r = 6
        elseif r == 112000006
                r = 7  # East of England
        elseif r == 112000007 # | London
                r = 8
        elseif r == 112000008 # South East
                r = 9
        elseif r == 112000009 # South West
                r = 10
        elseif r == 399999999 #  wales
                r = 11
        elseif r == 299999999
                r = 12        # Scotland
        elseif r == 499999999 # Northern Ireland
                r = 13
        end
        return r
end

function initialise_person( n::Integer )::DataFrame
    DataFrame(
        frs_year    = Vector{Union{Int64,Missing}}(missing,n),
        household_number  = Vector{Union{Int64,Missing}}(missing,n),
        benefit_unit = Vector{Union{Int8,Missing}}(missing,n),
        person  = Vector{Union{Int8,Missing}}(missing,n),
        government_region  = Vector{Union{Int8,Missing}}(missing,n),
        tenure_type  = Vector{Union{Int8,Missing}}(missing,n), # f enums OK
        household_income = Vector{Union{Real,Missing}}(missing,n),
        benefit_unit_income = Vector{Union{Real,Missing}}(missing,n),
        num_children = Vector{Union{Int8,Missing}}(missing,n),
        num_adults = Vector{Union{Int8,Missing}}(missing,n),
        age = Vector{Union{Int8,Missing}}(missing,n), # age80 - 2005-2012 ; iagegr2 full
        employment_status = Vector{Union{Int8,Missing}}(missing,n),
        sex = Vector{Union{Int8,Missing}}(missing,n),     # full
        marital_status = Vector{Union{Int8,Missing}}(missing,n), # full
        ethnic_group  = Vector{Union{Int8,Missing}}(missing,n), # ethgr3
        years_in_full_time_work =  Vector{Union{Int64,Missing}}(missing,n), # ftwk
        age_completed_full_time_education = Vector{Union{Int64,Missing}}(missing,n), # tea
        socio_economic_grouping = Vector{Union{Real,Missing}}(missing,n), # FIXME Really a Decimal nssc
        highest_qualification = Vector{Union{Int8,Missing}}(missing,n), # 2008-11
        occupational_classification = Vector{Union{Int64,Missing}}(missing,n),
        registered_blind_or_deaf = Vector{Union{Int8,Missing}}(missing,n), # full years
        any_disability = Vector{Union{Int8,Missing}}(missing,n), # disd01..09
        health_status = Vector{Union{Int8,Missing}}(missing,n), # heathad
        hours_of_care_received = Vector{Union{Real,Missing}}(missing,n), # full
        hours_of_care_given = Vector{Union{Real,Missing}}(missing,n), # f
        is_informal_carer = Vector{Union{Int8,Missing}}(missing,n), # carefl
        employment_earnings = Vector{Union{Real,Missing}}(missing,n), # 07-08 missing
        self_employment_income = Vector{Union{Real,Missing}}(missing,n),
        in_poverty = Vector{Union{Int8,Missing}}(missing,n), # low60ahc hbai adult
        happiness = Vector{Union{Int8,Missing}}(missing,n), # happywb adult
        in_debt_now = Vector{Union{Int8,Missing}}(missing,n),
        in_debt_in_last_year = Vector{Union{Int8,Missing}}(missing,n),
        total_hours_worked  = Vector{Union{Real,Missing}}(missing,n),
    )
end

function create_adults(
        year         :: Integer,
        hhld         :: DataFrame,
        benunit      :: DataFrame,
        frs_adults   :: DataFrame,
        hbai_adults  :: DataFrame )

        num_adults = size(frs_adults)[1]
        adult_model = initialise_person(num_adults)
        adno = 0
        hbai_year = year - 1993
        println("hbai_year $hbai_year")
        for pn in 1:num_adults
            if pn % 1000 == 0
                println("adults: on year $year, pno $pn")
            end

            single_person = frs_adults[pn, :]
            sernum = single_person.sernum
            matching_hbai = hbai_adults[((hbai_adults.year.==hbai_year).&(hbai_adults.sernum.==sernum).&(hbai_adults.person.==single_person.person).&(hbai_adults.benunit.==single_person.benunit)), :]
            matching_hhld = hhld[ single_person.sernum .==  hhld.sernum,:]
            matching_benunit = benunit[ ((single_person.sernum .==  benunit.sernum).&(single_person.benunit.==benunit.benunit)),:]
            # print( "matching_benunit=$(matching_benunit)\n" )
            @assert size( matching_benunit)[1] == 1
            @assert size( matching_hhld )[1] == 1
            nhbai = size(matching_hbai)[1]
            @assert nhbai in [0, 1]

            if nhbai == 1 # only non-missing in HBAI
                adno += 1
                    ## also for children
                output_adult = adult_model[adno, :]
                # 'single' here just extracts 1 row and allows us to treat
                # it like a tuple/struct
                single_hbai = matching_hbai[1,:]
                single_hhld = matching_hhld[1,:]
                single_benunit = matching_benunit[1,:]
                output_adult.tenure_type = safe_assign( single_hhld.tentyp2 )
                output_adult.government_region = remapRegion( single_hhld.gvtregn )
                output_adult.frs_year = year
                output_adult.benefit_unit = single_benunit.benunit
                output_adult.household_income = single_hbai.esninchh
                output_adult.benefit_unit_income = single_hbai.esnincbu
                output_adult.person = single_person.person
                output_adult.household_number = single_person.sernum
                output_adult.household_number = single_person.sernum
                output_adult.num_children = single_hbai.depchldh
                output_adult.num_adults =  single_hbai.adulth
                is_hbai_spouse = ( single_hbai.personsp == single_hbai.person )
                is_hbai_head = ( single_hbai.personhd == single_hbai.person )

                output_adult.age = single_person.age80
                output_adult.sex = safe_assign(single_person.sex)
                output_adult.ethnic_group = safe_assign(single_person.ethgr3)
                # plan 'B' wages and SE from HBAI; first work out hd/spouse so we can extract right ones
                ## adult only

                output_adult.marital_status = safe_assign(single_person.marital)
                output_adult.highest_qualification = safe_assign(single_person.dvhiqual)

                output_adult.socio_economic_grouping = safe_assign(Integer(trunc(single_person.nssec)))
                output_adult.age_completed_full_time_education = safe_assign(single_person.tea)
                output_adult.years_in_full_time_work = safe_inc(0, single_person.ftwk)
                output_adult.employment_status = safe_assign(single_person.empstati)
                output_adult.occupational_classification = safe_assign(single_person.soc2010)
                hbai_wages = coalesce( is_hbai_head ? single_hbai.esgjobhd : single_hbai.esgjobsp, 0.0 )
                hbai_se = coalesce( is_hbai_head ? single_hbai.esgrsehd : single_hbai.esgrsesp, 0.0 )
                output_adult.employment_earnings = hbai_wages
                output_adult.self_employment_income = hbai_se
                ## also for child
                output_adult.registered_blind_or_deaf =
                    ((single_person.spcreg1 == 1 ) ||
                     (single_person.spcreg2 == 1 ) ||
                     (single_person.spcreg3 == 1)) ? 1 : 0

                output_adult.any_disability = (
                    (single_person.disd01 == 1) || # cdisd kids ..
                    (single_person.disd02 == 1) ||
                    (single_person.disd03 == 1) ||
                    (single_person.disd04 == 1) ||
                    (single_person.disd05 == 1) ||
                    (single_person.disd06 == 1) ||
                    (single_person.disd07 == 1) ||
                    (single_person.disd08 == 1) ||
                    (single_person.disd09 == 1)) ? 1 : 0

                # dindividual_savings_accountbility_other_difficulty = Vector{Union{Real,Missing}}(missing, n),
                output_adult.health_status = safe_assign(single_person.heathad)
                output_adult.hours_of_care_received = safe_inc(0.0, single_person.hourcare)
                output_adult.hours_of_care_given = infer_hours_of_care(single_person.hourtot) # also kid

                output_adult.is_informal_carer = (single_person.carefl == 1 ? 1 : 0) # also kid
                output_adult.in_poverty =  single_hbai.low60ahc == 1

                output_adult.in_debt_now = (
                    (single_benunit.debt01 == 1 ) ||
                    (single_benunit.debt02 == 1 ) ||
                    (single_benunit.debt03 == 1 ) ||
                    (single_benunit.debt04 == 1 ) ||
                    (single_benunit.debt05 == 1 ) ||
                    (single_benunit.debt06 == 1 ) ||
                    (single_benunit.debt07 == 1 ) ||
                    (single_benunit.debt08 == 1 ) ||
                    (single_benunit.debt09 == 1 ) ||
                    (single_benunit.debt10 == 1 ) ||
                    (single_benunit.debt11 == 1 ) ||
                    (single_benunit.debt12 == 1 ) ||
                    (single_benunit.debt13 == 1 )) ? 1 : 0
                output_adult.in_debt_in_last_year = (
                    (single_benunit.debtar01 == 1 ) ||
                    (single_benunit.debtar02 == 1 ) ||
                    (single_benunit.debtar03 == 1 ) ||
                    (single_benunit.debtar04 == 1 ) ||
                    (single_benunit.debtar05 == 1 ) ||
                    (single_benunit.debtar06 == 1 ) ||
                    (single_benunit.debtar07 == 1 ) ||
                    (single_benunit.debtar08 == 1 ) ||
                    (single_benunit.debtar09 == 1 ) ||
                    (single_benunit.debtar10 == 1 ) ||
                    (single_benunit.debtar11 == 1 ) ||
                    (single_benunit.debtar12 == 1 ) ||
                    (single_benunit.debtar13 == 1 )) ? 1 : 0
                output_adult.happiness =  safe_assign(single_person.happywb)
                output_adult.total_hours_worked = safe_assign(single_person.tothours )
            end # if in HBAI
        end # adult loop
        println("final adno $adno")
        adult_model[1:adno, :]

end

hbai_adults = loadtoframe("$(HBAI_DIR)/tab/i1718_all.tab")
output_adults = initialise_person(0)

for year in 2017:2017

    print("on year $year ")

    ## hbf = HBAIS[year]
    ## hbai_adults = loadtoframe("$(HBAI_DIR)/tab/$hbf")
    househol =loadfrs("househol", year)
    benunit = loadfrs("benunit", year)
    adult = loadfrs("adult", year)

    output_adults_yr = create_adults(
        year,
        househol,
        benunit,
        adult,
        hbai_adults
    )
    append!(output_adults, output_adults_yr)
end

CSVFiles.save( File( format"CSV", "/mnt/data/teaching/frs/2019J/data/dd309_frs_adults_2017.tab" ),
    output_adults, delim = "\t",  nastring="")


text:
const GLOBAL_CI_CACHE = GPUCompiler.CodeCache()

function __init__()
    ccall((:brutus_init, "libbrutus"), Cvoid, (Any,), @__MODULE__)
end

max_stars_repo_name: stephenneuendorffer/brutus

max_issues_repo_name: stephenneuendorffer/brutus

max_forks_repo_name: stephenneuendorffer/brutus


text:
using TensorToolbox
using LinearMaps
using KrylovKit

include("util.jl")

function cov_Tail(T, d, ε, τ)
    #N = (d * log(d/ε/τ))^6 / ε^2
    #12 * exp(-T) + 3ε/(d*log(N/τ))
    if T <= 10 * log(1/ε)
        return 1
    end
    return 3ε/(T * log(T))^2
end

function Q(G, P)
    # Computes the mean of x'Px - tr(P) for x ~ N(0, I)

    #χ = Chisq(1)
    #return mean((sum(v*rand(χ) for v in c) - tr(P))^2 for _ in 1:10000

    #P = (P + P') ./ 2
    #Z = cholesky(G.Σ)
    #M = Z.U*P*Z.L
    #U = eigvecs(M)'
    #c = diag(U*M*U')
    #return 2*norm(c)^2 + sum(c*c') - 2*tr(P)*sum(c) + tr(P)^2

    return 2norm(P)^2
end

function cov_estimation_filter(S′, ε, τ=0.1; limit=nothing, method=:krylov)
    d, n = size(S′)
    C = 10
    C′ = 0
    Σ′ = S′*S′' ./ n
    G′ = MvNormal(Σ′)
    invsqrtΣ′ = Symmetric(Σ′)^(-1/2)
    Y = invsqrtΣ′ * S′
    xinvΣ′x = [y'y for y in eachcol(Y)]
    mask = xinvΣ′x .>= C*d*log(n/τ)
    if any(mask)
        println("early filter")
        if limit == nothing
            return .! mask
        else
            return .! mask .| k_lowest_ind(xinvΣ′x, max(0, n - limit))
        end
    end
    if method == :arpack
        Z = mapcols(y->kron(y, y), Y)
        Id♭ = ♭(Matrix(I, d, d))
        TS′ = Symmetric(-Id♭ * Id♭' + Z * Z' ./ n)
        (λ,), v = eigs(TS′; nev=1)
    else
        Z = LinearMap(v->krtv(Y, Y, v), v->tkrtv(Y', Y', v), d^2, n)
        Id♭ = LinearMap(reshape(♭(Matrix(I, d, d)), :, 1))
        TS′ = -Id♭ * Id♭' + Z * Z' / n
        (λ,), (v,) = eigsolve(TS′, ones(d^2), issymmetric=true)
    end
    if λ <= (1 + C*ε*log(1/ε)^2)*Q(collect(invsqrtΣ′) * G′, ♯(v)) / 2
        return G′
    end
    V = Symmetric(♯(v) + ♯(v)')/2
    ps = [1/√2 * (y'V*y - tr(V)) for y in eachcol(Y)]
    μ = median(ps)
    diffs = abs.(ps .- μ)
    for (i, diff) in enumerate(sort(diffs))
        shift = 3
        if diff < shift
            continue
        end
        T = diff - shift
        if T <= C′
            continue
        end
        if i/n >= cov_Tail(T, d, ε, τ)
            if limit == nothing
                return diffs .<= T
            else
                return (diffs .<= T) .| k_lowest_ind(diffs, max(0, n - limit))
            end
        end
    end
end

function cov_estimation_iterate(S′, ε, τ=0.1, k=nothing; iters=nothing, limit=nothing)
    _, n = size(S′)
    idxs = 1:n
    i = 0
    if limit != nothing
        orig_limit = limit
        p = Progress(limit, 1)
    end
    while true
        if iters != nothing && i >= iters
            break
        end
        if k == nothing
            S′k = S′
        else
            S′k, _ = pca(S′, k)
        end
        select = cov_estimation_filter(S′k, ε, τ, limit=limit)
        if select isa MvNormal
            println("Terminating early $(i) success...")
            break
        end
        if select == nothing
            println("Terminating early $(i) fail...")
            break
        end
        if limit != nothing
            limit -= length(select) - sum(select)
            @assert limit >= 0
            update!(p, orig_limit - limit)
        end
        S′ = S′[:, select]
        idxs = idxs[select]
        i += 1
        if limit == 0
            break
        end
    end
    select = falses(n)
    for i in idxs
        select[i] = 1
    end
    return select
end

function rcov(S′, ε, τ=0.1, k=nothing; iters=nothing, limit=nothing)
    select = cov_estimation_iterate(S′, ε, τ, k; iters=iters, limit=limit)
    selected = S′[:, select]
    return selected*selected'
end

function rpca(S′, ε, τ=0.1, k=100; iters=nothing, limit=nothing)
    d, n = size(S′)
    perm = randperm(n)
    S′paired = S′[:, perm][:, 1:div(n, 2)] - S′[:, perm][:, div(n, 2)+1:end-n%2]
    if limit != nothing
        # TODO: Is this correction right?
        limit = round(Int, limit - limit^2/2/n)
    end
    selected = cov_estimation_iterate(S′paired, ε, τ, iters=iters, limit=limit)
    _, U = pca(S′paired[:, selected], k)
    S′selected = U*S′paired[:, selected]
    return U * S′, U, cov(S′paired[:, selected]'U', corrected=false) ./ 2
end

mean_Tail(T, d, ε, δ, τ, ν=1) = 8exp(-T^2/(2ν)) + 8ε/(T^2*log(d*log(d/(ε*τ))))

function mean_estimation_filter(S′, ε, τ=0.1, ν=1; limit=nothing)
    d, n = size(S′)
    μ, Σ = mean(S′, dims=2), cov(S′', corrected=false)
    (λ,), v = eigs(Σ; nev=1)
    if λ - 1 <= ε * log(1/ε)
        return
    end
    δ = 3sqrt(ε * (λ - 1))
    λmags = abs.((S′ .- μ)'v)[:]
    for (i, mag) in enumerate(sort(λmags))
        if mag < δ
            continue
        end
        T = mag - δ
        if (n - i)/n > mean_Tail(T, d, ε, δ, τ, ν)
            if limit == nothing
                return λmags .<= mag
            else
                return (λmags .<= mag) .| k_lowest_ind(λmags, max(0, n - limit))
            end
        end
    end
end

function mean_estimation_iterate(A, ε, τ=0.1, ν=1; iters=nothing, limit=nothing)
    d, n = size(A)
    idxs = 1:n
    i = 0
    if limit != nothing
        orig_limit = limit
        p = Progress(limit, 1)
    end
    while true
        if iters != nothing && i >= iters
            break
        end
        select = mean_estimation_filter(A, ε, τ, ν, limit=limit)
        if select == nothing
            println("Terminating early $(i)...")
            break
        end
        if limit != nothing
            limit -= length(select) - sum(select)
            @assert limit >= 0
            update!(p, orig_limit - limit)
        end
        A = A[:, select]
        idxs = idxs[select]
        i += 1
    end
    select = falses(n)
    for i in idxs
        select[i] = 1
    end
    return select
end

function rmean(A, ε, τ=0.1, ν=1; iters=nothing, limit=nothing)
    select = mean_estimation_iterate(A, ε, τ, ν; iters=iters, limit=limit)
    return mean(A[:, select], dims=2)
end

max_stars_repo_name: TDteach/spectre-defense

max_issues_repo_name: TDteach/spectre-defense

max_forks_repo_name: TDteach/spectre-defense


text:
# Use baremodule to shave off a few KB from the serialized `.ji` file
baremodule HDF5_jll
using Base
using Base: UUID
import JLLWrappers

JLLWrappers.@generate_main_file_header("HDF5")
JLLWrappers.@generate_main_file("HDF5", UUID("0234f1f7-429e-5d53-9886-15a909be8d59"))
end  # module HDF5_jll


text:
using ..Continuations: AbstractContinuationProblem, AbstractContinuationSolver,
    AbstractContinuationCache,
    as, SweepSetup, ContinuationSweep, ContinuationSolution,
    ContinuationCache, ContinuationOptions, ContinuationSolver,
    residual_jacobian!

using ..BifurcationsBase: timekind, Continuous, Discrete
import ..BifurcationsBase: TimeKind

using ..BifurcationsBase
using ..BifurcationsBase: SpecialPoint, SpecialPointInterval,
    BifurcationSweep, BifurcationSolution, BifurcationSolver,
    BifurcationProblem, BifurcationCache,
    eigvals_prototype, allocate_sweep!, check_sweep_length, record!
import ..BifurcationsBase: analyze!, re_analyze!

using ..Codim1LimitCycle: set_reference!

module PointTypes
@enum(
    PointType,
    none,
)
end  # module
using .PointTypes: PointType

const Codim2LCSweep{tkind <: TimeKind, ptType <: PointType} =
    BifurcationSweep{tkind, ptType}
const Codim2LCSolution{S <: ContinuationSolution, W <: Codim2LCSweep} =
    BifurcationSolution{S, W}
# The constraints `<: TimeKind` and `<: ContinuationSolution` are
# important.
# See: [[../codim1/solver.jl::Codim1Sweep]]

BifurcationsBase.point_type_type(::Codim2LCProblem) = PointType

BifurcationsBase.regular_point(::Type{PointType}) = PointTypes.none

function BifurcationsBase.eigvals_prototype(prob::Codim2LCProblem,
                                            cache::ContinuationCache)
    n = length(prob.super.de_prob.u0)
    return similar(cache.u[1:1] .+ 1im, n)
end

mutable struct Codim2LCCache{P, C <: ContinuationCache{P},
                             JType, eType,
                             } <: BifurcationCache{P}
    super::C
    J::JType
    eigvals::eType
    point_type::PointType
end

function Codim2LCCache(super::C,
                       J::JType,
                       eigvals::eType,
                       point_type = PointTypes.none,
                       ) where {P, C <: ContinuationCache{P},
                                JType, eType}
    return Codim2LCCache{P, C, JType, eType}(super, J, eigvals, point_type)
end
# TODO: Remove this constructor after removing the type parameter `P`.

Codim2LCCache(prob::Codim2LCProblem, super::ContinuationCache) =
    Codim2LCCache(
        super,
        let v = eigvals_prototype(prob, super)
            v * v'
        end,
        # ds_jacobian(super),
        copy(eigvals_prototype(prob, super)),
    )

BifurcationsBase.BifurcationCache(prob::Codim2LCProblem,
                                  super::ContinuationCache) =
    Codim2LCCache(prob, super)

const Codim2LCSolver{
        R <: ContinuationSolver,
        P <: Codim2LCProblem,
        C <: Codim2LCCache,
        S <: Codim2LCSolution,
        } =
    BifurcationSolver{R, P, C, S}

# TODO: trait
BifurcationsBase.contkind(::Codim2LCProblem) = FoldLimitCycleCont()

function re_analyze!(solver::Codim2LCSolver, u::AbstractVector)
    residual_jacobian!(as(solver.cache, ContinuationCache), u)
    analyze!(solver.cache, solver.opts)

    # Suppress special point recording:
    # It's a bit ugly hack... (communicate by sharing!)
    solver.cache.point_type = PointTypes.none  # TODO: FIX!

    record!(solver.sol, solver.cache)
end

function analyze!(cache::Codim2LCCache, opts)
    #=
    cache.J = J = ds_jacobian(cache)
    eigvals = ds_eigvals(timekind(cache), J)
    cache.point_type = guess_point_type(timekind(cache), cache, eigvals, opts)
    cache.eigvals = eigvals
    =#

    # This works since the "layout" of the head of cache.super.u is
    # identical to `LimitCycleCache` case.  Maybe better to not rely
    # on the layout implicitly:
    set_reference!(cache)
end


text:
# Author: Mathias Louboutin, mlouboutin3@gatech.edu
# Date: December 2020

export isLegal, lbfgsUpdate, lbfgsHvFunc2, ssbin, solveSubProblem!, subHv, result, soft_thresholding

mutable struct result{T}
    x::AbstractArray{T}
    g::AbstractArray{T}
    ϕ::T
    ϕ_trace::Vector{T}
    x_trace::Vector{AbstractArray{T}}
    n_project::Integer
    n_ϕeval::Integer
    n_geval::Integer
end

function update!(r::result; x=nothing, ϕ=nothing, g=nothing, iter=1, store_trace=false)
    ~isnothing(x) && copyto!(r.x, x)
    ~isnothing(ϕ) && (r.ϕ = ϕ)
    ~isnothing(g) && copyto!(r.g, g)
    (~isnothing(x) && length(r.x_trace) == iter-1 && store_trace) && (push!(r.x_trace, deepcopy(x)))
    (~isnothing(ϕ) && length(r.ϕ_trace) == iter-1) && (push!(r.ϕ_trace, ϕ))
end

function result(init_x::AbstractArray{T}; ϕ0=0, ϕeval=0, δϕeval=0) where T
    return result(deepcopy(init_x), T(0)*init_x, T(ϕ0), Vector{T}(), Vector{AbstractArray{T}}(), 0, ϕeval, δϕeval)
end

function isLegal(v::AbstractArray{T}) where T
    nv = norm(v)
    return !isnan(nv) && !isinf(nv)
end

isLegal(v::T) where T = (!isnan(v) && !isinf(v))

function terminate(options, optCond, t, d, ϕ, ϕ_old)
    ~isLegal(ϕ) && return true
    ~isLegal(d) && return true
    # Check optimality
    if optCond < options.optTol
        options.verbose >= 1 &&  @printf("First-Order Optimality Conditions %2.2e Below optTol %2.2e\n",
                                         optCond, options.optTol)
        return true
    end

    if norm(t*d, Inf) < options.progTol
        options.verbose >= 1 && @printf("Step size: %2.2e below progTol: %2.2e\n",
                                        norm(t*d, Inf), options.progTol)
        return true
    end

    if norm(d, Inf) < options.progTol
        options.verbose >= 1 && @printf("Update direction smaller than progTol (%2.2e < %2.2e)\n",
                                         norm(d, Inf), options.progTol)
        return true
    end
    t == 0 && return true
    return false
end

"""
    PQN lbfgs functions
"""
function lbfgsUpdate(y::AbstractArray{T}, s::AbstractArray{T}, corrections::Integer,
                     old_dirs::AbstractArray{T, 2}, old_stps::AbstractArray{T, 2},
                     Hdiag) where T
    ys = dot(y,s)
    if ys > 1e-10 || size(old_dirs,2)==0
        numCorrections = size(old_dirs, 2)
        if numCorrections < corrections
            # Full Update
            old_dirs = [old_dirs s]
            old_stps = [old_stps y]
        else
            # Limited-Memory Update
            old_dirs = [old_dirs[:, 2:corrections] s]
            old_stps = [old_stps[:, 2:corrections] y]
        end

        # Update scale of initial Hessian approximation
        Hdiag = ys/dot(y, y)
    end
    return old_dirs, old_stps, Hdiag
end

function lbfgsHvFunc2(v::AbstractArray{T}, Hdiag, N::AbstractArray{T, 2}, M::AbstractArray{T, 2}) where T
    if cond(M)>(1/(eps(T)))
        pr = ssbin(M, 500)
        L = Diagonal(vec(pr))
        Hv = v/Hdiag - N*L*((L*M*L)\(L*(N'*v)))
    else
        Hv = v/Hdiag - N*(M\(N'*v))
    end

    return Hv
end

function ssbin(A::AbstractArray{T, N}, nmv) where {T, N}
    # Stochastic matrix-free binormalization for symmetric real A.
    # x = ssbin(A,nmv,n)
    #   A is a symmetric real matrix or function handle. If it is a
    #     function handle, then v = A(x) returns A*x.
    #   nmv is the number of matrix-vector products to perform.
    #   [n] is the size of the matrix. It is necessary to specify n
    #     only if A is a function handle.
    #   diag(x) A diag(x) is approximately binormalized.

    # Jan 2010. Algorithm and code by Andrew M. Bradley (ambrad@stanford.edu).
    # Aug 2010. Modified to record and use dp. New omega schedule after running
    #   some tests.
    # Jul 2011. New strategy to deal with reducible matrices: Use the old
    #   iteration in the early iterations; then switch to snbin-like behavior,
    #   which deals properly with oscillation.
    n = size(A, 1)
    d = ones(T, n)
    dp = deepcopy(d)
    tmp = zeros(T, n)
    for k = 1:nmv
      # Approximate matrix-vector product
      u = randn(T, n)
      s = u ./ sqrt.(dp)
      y = A*s
      # omega^k
      alpha = (k - 1)/nmv
      omega = (1 - alpha)/2 + alpha/nmv
      # Iteration
      d .= (1-omega)*d/sum(d) + omega*y.^2/sum(y.^2)
      if (k < min(32,floor(nmv/2)))
        # First focus on making d a decent approximation
        dp .= d
      else
        # This block makes ssbin behave like snbin except for omega
        tmp .= dp
        dp .= d
        d .= tmp
      end
    end

    return 1 ./(d.*dp).^(T(.25))
end

function solveSubProblem!(p::AbstractArray{T}, x::AbstractArray{T}, g::AbstractArray{T}, H,
                         funProj::Function, options, x_init::AbstractArray{T}) where T
    # Uses SPG to solve for projected quasi-Newton direction
    funObj(γ) = subHv(γ, x[1:end], g[1:end], H)
    sol = spg(funObj, x_init, funProj, options)
    p .= sol.x
end

function subHv(p::AbstractArray{T}, x::AbstractArray{T}, g::AbstractArray{T}, HvFunc::Function) where T
    d = p - x
    Hd = HvFunc(d)
    f = dot(g, d) + dot(d, Hd) / 2
    g = g + Hd
    return f, g
end

# THresholding
soft_thresholding(x::AbstractArray{Complex{T}}, λ::T) where {T} = exp.(angle.(x)im) .* max.(abs.(x) .- convert(T, λ), T(0))
soft_thresholding(x::AbstractArray{Complex{T}}, λ::Array{T}) where {T} = exp.(angle.(x)im) .* max.(abs.(x) .- convert(Array{T}, λ), T(0))
soft_thresholding(x::AbstractArray{T}, λ::T) where {T} = sign.(x) .* max.(abs.(x) .- convert(T, λ), T(0))
soft_thresholding(x::AbstractArray{T}, λ::Array{T}) where {T} = sign.(x) .* max.(abs.(x) .- convert(Array{T}, λ), T(0))


text:
using COSMOAccelerators, Test, Random, LinearAlgebra, SparseArrays

@testset "AndersonAccelerator" begin

    rng = Random.MersenneTwister(2)

    @testset "Constructors" begin
        dim = 10
        # default: Float64, Type 2, QR decomposition, NoRegularizer, RestartedMemory
        aa = AndersonAccelerator(dim, mem = 10)
        @test typeof(aa) == AndersonAccelerator{Float64, Type2{QRDecomp}, RestartedMemory, NoRegularizer}

        aa = AndersonAccelerator{Float32}(dim)
        @test typeof(aa) == AndersonAccelerator{Float32, Type2{QRDecomp}, RestartedMemory, NoRegularizer}

        aa = AndersonAccelerator{Type1}(dim)
        @test typeof(aa) == AndersonAccelerator{Float64, Type1, RestartedMemory, NoRegularizer}

        aa = AndersonAccelerator{TikonovRegularizer}(dim)
        @test typeof(aa) == AndersonAccelerator{Float64,  Type2{NormalEquations}, RestartedMemory, TikonovRegularizer}

        aa = AndersonAccelerator{FrobeniusNormRegularizer}(dim)
        @test typeof(aa) == AndersonAccelerator{Float64,  Type2{NormalEquations}, RestartedMemory, FrobeniusNormRegularizer}

        aa = AndersonAccelerator{RollingMemory}(dim)
        @test typeof(aa) == AndersonAccelerator{Float64,  Type2{NormalEquations}, RollingMemory, NoRegularizer}

        @test_throws ArgumentError AndersonAccelerator(dim, mem = 0)
        @test_throws ArgumentError AndersonAccelerator(dim, min_mem = 0)
        @test_throws ArgumentError AndersonAccelerator(dim, λ = -1.)
        @test_throws ArgumentError AndersonAccelerator{Float64, Type2{QRDecomp}, RollingMemory, NoRegularizer}(dim)
        @test_throws ArgumentError AndersonAccelerator{Float64, Type2{QRDecomp}, RestartedMemory, TikonovRegularizer}(dim)
        @test_throws ArgumentError AndersonAccelerator{Float64, Type2{QRDecomp}, RestartedMemory, FrobeniusNormRegularizer}(dim)
    end

    @testset "Anderson Acceleration (RollingMemory)" begin
        # add a history of vectors
        n = 2
        m = 4
        hist_length = 10
        G = randn(rng, n + m, hist_length)
        X = randn(rng, n + m, hist_length)
        F = X - G
        mem = 5
        aa = AndersonAccelerator{Float64, Type1, RollingMemory, NoRegularizer}(m + n, mem = mem);

        # add more vectors than memory, it should start appending at the beginning of the matrix caches again
        for k = 1:mem+2
            update!(aa, G[:, k], X[:, k], k)
        end
        X_ref = X[:,7] - X[:, 6]
        F_ref = F[:,7] - F[:, 6]
        G_ref = G[:,7] - G[:, 6]
        @test aa.X[:, 1] == X_ref
        @test aa.G[:, 1] == G_ref
        @test aa.F[:, 1] == F_ref
        return nothing

    end

    @testset "Anderson Acceleration (RestartedMemory)" begin

        # add a history of vectors
        n = 2
        m = 4
        hist_length = 10
        G = randn(rng, n + m, hist_length)
        X = randn(rng, n + m, hist_length)
        F = X - G
        mem = 5
        aa = AndersonAccelerator{Float64, Type2{NormalEquations}, RestartedMemory, NoRegularizer}(m + n, mem = mem);

        # lets add a bunch of inputs and double check that everything is stored correctly
        @test aa.init_phase

        # add first vector pair
        update!(aa, G[:, 1], X[:, 1], 1)
        @test !aa.init_phase
        @test aa.iter == 0
        @test aa.x_last == X[:, 1]
        @test aa.g_last == G[:, 1]
        @test aa.f_last == X[:, 1] - G[:, 1]
        @test iszero(aa.F) && iszero(aa.G) && iszero(aa.F) && iszero(aa.M)

        # add second vector
        update!(aa, G[:, 2], X[:, 2], 2)
        @test aa.iter == 1
        @test aa.F[:, 1] ==  F[:, 2] - F[:, 1]
        @test aa.X[:, 1] ==  X[:, 2] - X[:, 1]
        @test aa.G[:, 1] ==  G[:, 2] - G[:, 1]

        # add third and fourth vector pairs
        k = 3
        update!(aa, G[:, k], X[:, k], 3)
        k += 1
        update!(aa, G[:, k], X[:, k], 4)
        @test aa.iter == 3
        # let's check the matrices
        X_ref = [X[:,2] - X[:, 1] X[:, 3] - X[:, 2] X[:, 4] - X[:, 3] ]
        F_ref = [F[:,2] - F[:, 1] F[:, 3] - F[:, 2] F[:, 4] - F[:, 3] ]
        G_ref = [G[:,2] - G[:, 1] G[:, 3] - G[:, 2] G[:, 4] - G[:, 3] ]
        @test aa.X[:,1:3] ==  X_ref
        @test aa.F[:,1:3] ==  F_ref
        @test aa.G[:,1:3] ==  G_ref

        # let's test the accelerated vector
        v = copy(G[:, k])
        v_prev = copy(X[:, k])
        f = F[:, k]
        eta = F_ref' * f
        eta = F_ref' * F_ref \ eta
        v_acc_ref = v - (X_ref - F_ref) * eta
        accelerate!(v, v_prev, aa, k)
        @test isapprox(v, v_acc_ref, atol = 1e-8)

        # let's add more vectors to exceed memory and check that memory gets flushed
        k += 1
        update!(aa, G[:, k], X[:, k], k)
        k += 1
        update!(aa, G[:, k], X[:, k], k)
        # memory should now be full
        @test aa.iter % aa.mem == 0

        # add one more and double check it is added in first position
        k += 1
        update!(aa, G[:, k], X[:, k], k)

        @test aa.X[:, 1] == X[:,7] - X[:, 6]
        @test aa.F[:, 1] == F[:,7] - F[:, 6]
        @test aa.G[:, 1] == G[:,7] - G[:, 6]
        @test aa.iter == 1
        @test aa.init_phase == false
        return nothing

    end

    @testset "Gram-Schmidt Orthogonalisation method" begin
        F = LinearAlgebra.qr(rand(rng, 8, 4))
        V = Matrix(F.Q)

        # now test our method
        V2 = [V zeros(8)]

        COSMOAccelerators.qr!(V2, zeros(5, 5), v2, 5)
        v2 = V2[:, 5]
        # is new vector orthogonal to all other vectors in V
        for k = 1:4
            @test dot(V[:, k], v2) <= 1e-12
        end
        @test isapprox(norm(v2), 1., atol = 1e-8)
        return nothing

    end

end

max_issues_repo_name: oxfordcontrol/COSMOAccelerators

max_forks_repo_name: oxfordcontrol/COSMOAccelerators


text:
"Integrated projection of leaf area for a single leaf inclination of θₗ, assumes azimuthally uniform distribution"
function A(θ::FT, θₗ::FT) where FT<:Real
    a = cos(θ) * cos(θₗ)
    # Eq. 14.24 in Bonan et al.
    if θₗ ≤ FT(π/2) - θ
        return a
    else
        b = sin(θ) * sin(θₗ)
        c = sqrt(sin(θₗ)^2  - cos(θ)^2)
        return FT(2/π)*(c + a * asin(a/b))
    end
end

"""
    $(FUNCTIONNAME)(μ::Array{FT}, LD::AbstractLeafDistribution; nLeg=20)

Returns the integrated projection of leaf area in the direction of μ, assumes azimuthally uniform distribution and a LD distribution for leaf polar angle θ.
This function is often referred to as the function O(B) (Goudriaan 1977) or G(Ζ) (Ross 1975,1981), see Bonan modeling book, eqs. 14.21-14.26.

# Arguments
- `μ` an array of cos(θ) (directions [0,1])
- `LD` an [`AbstractLeafDistribution`](@ref) type struct, includes a leaf distribution function
- `nLeg` an optional parameter for the number of legendre polynomials to integrate over the leaf distribution (default=20)

# Examples
```julia-repl
julia> μ,w = CanopyOptics.gauleg(10,0.0,1.0);       # Create 10 quadrature points in μ
julia> LD  = CanopyOptics.spherical_leaves()        # Create a default spherical leaf distribution
julia> G   = CanopyOptics.G(μ, LD)                  # Compute G(μ)
10-element Vector{Float64}:
 0.5002522783000879
 0.5002715115149204
 0.5003537989277846
 0.5004432798701134
 0.5005134448870893
 0.5003026448466977
 0.4999186257540982
 0.4994511190721635
 0.49907252201082375
 0.49936166823681594
```
"""
function G(μ::Array{FT}, LD::AbstractLeafDistribution; nLeg=20) where FT
    θₗ,w = gauleg(nLeg,FT(0),FT(π/2))
    Fᵢ = pdf.(LD.LD,2θₗ/π) * LD.scaling
    res = similar(μ);
    θ = acos.(μ)
    for i in eachindex(μ)
        res[i] =  sum(w .* Fᵢ .* A.(θ[i],θₗ))
    end
    return res
end

"""
    $(FUNCTIONNAME)(μ::Array{FT,1},μꜛ::Array{FT,1}, r,t, LD::AbstractLeafDistribution; nLeg = 20)

Computes the azimuthally-averaged area scattering transfer function following Shultis and Myneni (https://doi.org/10.1016/0022-4073(88)90079-9), Eq 43:

``Γ(μ' -> μ) = \\int_0^1 dμ_L g_L(μ_L)[t_L Ψ⁺(μ, μ', μ_L) + r_L Ψ⁻(μ, μ', μ_L)]``

assuming an azimuthally uniform leaf angle distribution.
# Arguments
- `μ::Array{FT,1}` : Quadrature points incoming direction (cos(θ))
- `μꜛ::Array{FT,1}`: Quadrature points outgoing direction (cos(θ))
- `r` : Leaf lambertian reflectance
- `t` : Leaf lambertian transmittance
- `LD` a [`AbstractLeafDistribution`](@ref) struct that describes the leaf angular distribution function.
- `nLeg = 20`: number of quadrature points used for integration over all leaf angles (default is 20).
"""
function compute_lambertian_Γ(μ::Array{FT,1},μꜛ::Array{FT,1}, r,t, LD::AbstractLeafDistribution; nLeg = 20) where FT
    Γ = zeros(length(μ), length(μ))
    θₗ,w = gauleg(nLeg,FT(0),FT(π/2))
    for i in eachindex(θₗ)
        Ψ⁺, Ψ⁻ = compute_Ψ(μ,μꜛ, cos(θₗ[i]));
        Γ += pdf.(LD.LD,2θₗ[i]/π) * LD.scaling * w[i] * (t * Ψ⁺ + r * Ψ⁻)
    end
    return Γ
end

"""
    $(FUNCTIONNAME)(mod::BiLambertianCanopyScattering, μ::Array{FT,1}, LD::AbstractLeafDistribution, m::Int)

Computes the single scattering Z matrices (𝐙⁺⁺ for same incoming and outgoing sign of μ, 𝐙⁻⁺ for a change in direction). Internally computes the azimuthally-averaged area scattering transfer function following Shultis and Myneni (https://doi.org/10.1016/0022-4073(88)90079-9), Eq 43::

``Γ(μ' -> μ) = \\int_0^1 dμ_L g_L(μ_L)[t_L Ψ⁺(μ, μ', μ_L) + r_L Ψ⁻(μ, μ', μ_L)]``

assuming an azimuthally uniform leaf angle distribution. Normalized Γ as 𝐙 = 4Γ/(ϖ⋅G(μ)).
Returns 𝐙⁺⁺, 𝐙⁻⁺

# Arguments
- `mod` : A bilambertian canopy scattering model [`BiLambertianCanopyScattering`](@ref), uses R,T,nQuad from that model.
- `μ::Array{FT,1}`: Quadrature points ∈ [0,1]
- `LD` a [`AbstractLeafDistribution`](@ref) struct that describes the leaf angular distribution function.
- `m`: Fourier moment (for azimuthally uniform leave distributions such as here, only m=0 returns non-zero matrices)
"""
function compute_Z_matrices(mod::BiLambertianCanopyScattering, μ::Array{FT,1}, LD::AbstractLeafDistribution, m::Int) where FT
    @unpack R,T,nQuad = mod
    # Transmission (same direction)
    𝐙⁺⁺ = zeros(length(μ), length(μ))
    # Reflection (change direction)
    𝐙⁻⁺ = zeros(length(μ), length(μ))

    # skip everything beyond m=0
    if m>0
        return 𝐙⁺⁺, 𝐙⁻⁺
    end
    # Ross kernel
    G = CanopyOptics.G(μ, LD)
    # Single Scattering Albedo (should make this a vector too)
    ϖ = R+T

    θₗ,w = gauleg(nQuad,FT(0),FT(π/2));
    for i in eachindex(θₗ)
        Ψ⁺, Ψ⁻ = compute_Ψ(μ,μ, cos(θₗ[i]));
        𝐙⁺⁺ += pdf.(LD.LD,2θₗ[i]/π) * LD.scaling * w[i] * (T * Ψ⁺ + R * Ψ⁻)
        Ψ⁺, Ψ⁻ = compute_Ψ(μ,-μ, cos(θₗ[i]));
        𝐙⁻⁺ += pdf.(LD.LD,2θₗ[i]/π) * LD.scaling * w[i] * (T * Ψ⁺ + R * Ψ⁻)
    end
    return 4𝐙⁺⁺ ./(G*ϖ), 4𝐙⁻⁺ ./(G*ϖ)
end

# Page 20, top of Knyazikhin and Marshak
# Example
# ϕ = range(0.0, 2π,  length=200)
# θ = range(0.0, π/2, length=150)
# dirs = [dirVector(a,b) for a in θ, b in ϕ];
# R = CanopyOptics.compute_specular_reflection.([dirs[10,1]],dirs, [1.5], [0.3], [LD])
function compute_reflection(mod::SpecularCanopyScattering, Ωⁱⁿ::dirVector{FT}, Ωᵒᵘᵗ::dirVector{FT}, LD) where FT
    @unpack nᵣ,κ = mod
    Ωstar = getSpecularΩ(Ωⁱⁿ, Ωᵒᵘᵗ)
    #θstar = min(abs(Ωstar.θ), (π-abs(Ωstar.θ))) # min(abs(Ωstar.θ), abs(π+Ωstar.θ))
    θstar = Ωstar.θ;
    #if Ωⁱⁿ.θ ≈ Ωᵒᵘᵗ.θ && Ωⁱⁿ.ϕ ≈ Ωᵒᵘᵗ.ϕ
    #    θstar = Ωⁱⁿ.θ
    #end
    # Still needs to be implemented!
    # incident angle on leaf surface (half of in and out angle):
    sa = Ωⁱⁿ ⋅ Ωᵒᵘᵗ
    sa > 1 ? sa = FT(1) : nothing
    αstar = acos(abs(sa))/2
    #@show Ωstar.ϕ, Ωstar.θ
    #a = (Ωⁱⁿ ⋅ Ωstar) * (Ωᵒᵘᵗ ⋅ Ωstar)
    return FT(1/8) * pdf(LD.LD,2θstar/π) * LD.scaling * K(κ, αstar) * Fᵣ(nᵣ,αstar)

end

function compute_reflection(mod::SpecularCanopyScattering,Ωⁱⁿ::dirVector_μ{FT}, Ωᵒᵘᵗ::dirVector_μ{FT}, LD) where FT
    @unpack nᵣ,κ = mod
    Ωstar, αstar = getSpecularΩ(Ωⁱⁿ, Ωᵒᵘᵗ)
    # Can change this later as well do have the pdf in μ, not theta!
    θstar = acos(abs(Ωstar.μ));
    # Eq. 2.39 in "Discrete Ordinates Method for Photon Transport in Leaf Canopies", page 59
    return FT(1/8) * pdf(LD.LD,2θstar/π) * LD.scaling * K(κ, αstar) * Fᵣ(nᵣ,αstar)
end

function compute_Z_matrices(mod::SpecularCanopyScattering, μ::Array{FT,1}, LD::AbstractLeafDistribution, m::Int) where FT
    @unpack nᵣ, κ,nQuad = mod
    # Transmission (same direction)
    𝐙⁺⁺ = zeros(length(μ), length(μ))
    # Reflection (change direction)
    𝐙⁻⁺ = zeros(length(μ), length(μ))

    # Quadrature points in the azimuth:
    ϕ, w_azi = gauleg(nQuad,FT(0),FT(2π));
    # Fourier weights (cosine decomposition)
    f_weights = cos.(m*ϕ)
    for i in eachindex(μ)
        # Incoming beam at ϕ = 0
        Ωⁱⁿ = dirVector_μ(μ[i], FT(0));
        # Create outgoing vectors in θ and ϕ
        dirOutꜛ = [dirVector_μ(a,b) for a in μ, b in ϕ];
        dirOutꜜ = [dirVector_μ(a,b) for a in -μ, b in ϕ];
        # Compute over μ and μ_azi:
        Zup   = compute_reflection.([mod],[Ωⁱⁿ],dirOutꜛ, [LD]);
        Zdown = compute_reflection.([mod],[Ωⁱⁿ],dirOutꜜ, [LD]);
        # integrate over the azimuth:
        𝐙⁻⁺[i,:] = Zup   * (w_azi .* f_weights)
        𝐙⁺⁺[i,:] = Zdown * (w_azi .* f_weights)
    end
    return 𝐙⁺⁺, 𝐙⁻⁺
end

"The reduction factor proposed by Nilson and Kuusk, κ ≈ 0.1-0.3, returns exp(-κ * tan(abs(α))"
function K(κ::FT, α::FT) where FT
    exp(-κ * tan(abs(α)));
end


text:
func(x) = [0.5*sqrt(abs(x[1] + x[2])), 1.5*x[1] + 0.5*x[2]]
Inputs = [0.3,900.0]
fp_simple   = fixed_point(func, Inputs; Algorithm = :Simple)
fp_simple.Convergence_ < 1e-10
fp_anderson = fixed_point(func, Inputs; Algorithm = :Anderson)
fp_anderson.Convergence_ < 1e-10
fp_aitken   = fixed_point(func, Inputs; Algorithm = :Aitken)
fp_aitken.Convergence_ < 1e-10
fp_newton   = fixed_point(func, Inputs; Algorithm = :Newton)
fp_newton.Convergence_ < 1e-10
fp_VEA      = fixed_point(func, Inputs; Algorithm = :VEA)
fp_VEA.Convergence_ < 1e-10
fp_SEA      = fixed_point(func, Inputs; Algorithm = :SEA)
fp_SEA.Convergence_ < 1e-10
fp_MPE      = fixed_point(func, Inputs; Algorithm = :MPE)
fp_MPE.Convergence_ < 1e-10
fp_RRE      = fixed_point(func, Inputs; Algorithm = :RRE)
fp_RRE.Convergence_ < 1e-10

# Now trying a function with a typeswitch from Int to Float
Inputs = [1,4]
fp_simple   = fixed_point(func, Inputs; Algorithm = :Simple)
fp_simple.Convergence_ < 1e-10
fp_anderson = fixed_point(func, Inputs; Algorithm = :Anderson)
fp_anderson.Convergence_ < 1e-10

# And finally printing a couple of them to test printing.
fp_anderson2 = fixed_point(func, Inputs; Algorithm = :Anderson, PrintReports = true)
fp_anderson2.Convergence_ < 1e-10
fp_simple2   = fixed_point(func, Inputs; Algorithm = :Simple, PrintReports = true)
fp_simple2.Convergence_ < 1e-10

# Testing function that returns a missing
f(x) = [missing, missing]
fp_simple3   = fixed_point(f, Inputs; Algorithm = :Simple, PrintReports = true)
fp_simple3.TerminationCondition_ == :InvalidInputOrOutputOfIteration
f(x) = [1, missing]
fp_simple4   = fixed_point(f, Inputs; Algorithm = :Simple, PrintReports = true)
fp_simple4.TerminationCondition_ == :InvalidInputOrOutputOfIteration

# Testing the outputting of a tuple
f(x) = (a = [1.0, 1.0], b = :goodProgressInFunction)
@test_throws ErrorException("The Fixedpoint function can only return a vector or a tuple  of which the first entry is the vector for which a fixedpoint is sought and the second is a namedtuple (the contents of which are output for the user but are not used in fixed point acceleration).") fixed_point(f, Inputs; Algorithm = :Simple, PrintReports = true)
# And doing side effects properly.
f(x) = ([1.0, 1.0], (b = :goodProgressInFunction,))
fp_simple5 = fixed_point(f, Inputs; Algorithm = :Simple, PrintReports = true)
fp_simple5.Convergence_ < 1e-10


text:
const Container{T} = Union{NTuple{N, T}, Vector{T}} where N

function getmask(ls::Container{<:Container})
    lens = map(length, ls)
    m = zeros(Float32, maximum(lens), length(lens))

    for (i, l) ∈ enumerate(ls)
        selectdim(selectdim(m, 2, i), 1, 1:length(l)) .= 1
    end
    reshape(m, (1, size(m)...))
end

getmask(m1::A, m2::A) where A <: Abstract3DTensor = permutedims(m1, [2,1,3]) .* m2


text:
using NIPALS_PCA
using CSV
using DataFrames
using Pipe

isnumcol(type) = type <: Union{Missing,Number}

irisdir = @pipe pathof(NIPALS_PCA) |> splitpath |> _[1:end - 2] |> joinpath(_..., "test", "data", "iris")

#x_df = CSV.File(joinpath(irisdir, "iris.csv"), header=false) |> DataFrame
#select!(x_df, eltype.(eachcol(x_df)) |> coltypes -> isnumcol.(coltypes) |> findall)
#xdataset = parseDataFrame(x_df) |> dataset -> filterDataset(dataset, filters=[dataset -> dataset.mvs .< 0.25]) |> normalize

#pca = calcPCA(xdataset, 3)

#calcVariances(xdataset, pca)


text:
using HorizonSideRobots
include("BackPath.jl")
include("Functions.jl")

function paint_border(r::Robot, side1::HorizonSide, side2::HorizonSide)
    while isborder(r, side1)
        putmarker!(r)
        move!(r, side2)
    end
    move!(r, side1)
end

function paint_perimeter(r::Robot)
    back_path = BackPath(r, (Nord, Ost))
    higth = steps_to_border!(r, Sud)

    num_side = 0
    cur = 0

    while cur != higth + 1
        while !isborder(r, HorizonSide(num_side))
            move!(r, HorizonSide(num_side))
            cur += 1
        end
        if cur != higth
            cur = higth + 1
        else
            move!(r, West)
            num_side = (num_side + 2) % 4
        end
    end

    cur_side = num_side
    for i in 0:3
        paint_border(r, HorizonSide(cur_side % 4), HorizonSide((cur_side + 3) % 4))
        cur_side = (cur_side + 1) % 4
    end

    to_nord_ost_corner = BackPath(r, (Nord, Ost))
    back!(r, back_path)
end

max_stars_repo_name: Mananhina/mirea-progs

max_issues_repo_name: Mananhina/mirea-progs

max_forks_repo_name: Mananhina/mirea-progs


text:
stress_scale = 1.0e5
strain_scale = 1

include("nnutil.jl")

force_scale = 5.0
force_scales = [5.0]

testtype = "NeuralNetwork2D"
nntype = "piecewise"

# ! define H0
# Trained with nx, ny = 10, 5
# H0 = [1.04167e6  2.08333e5  0.0
#       2.08333e5  1.04167e6  0.0
#       0.0        0.0        4.16667e5]/stress_scale

H0 = [1.0406424793819175e6 209077.08366547766         0.0
      209077.08366547766   1.0411467691352057e6       0.0
      0.0                  0.0                   419057.32049008965]/stress_scale

n_data = [100, 101, 102, 103, 104, 200, 201, 202, 203, 204]

porder = 2
prop = Dict("name"=> testtype, "rho"=> 4.5, "nn"=>nn)

T = 200.0
NT = 200

# DNS computaional domain
fiber_size = 2
# nx_f, ny_f = 40*fiber_size, 20*fiber_size
nx_f, ny_f = 10*fiber_size, 5*fiber_size
Lx, Ly = 1.0, 0.5
# nx_f, ny_f = 12, 4

# homogenized computaional domain
# number of elements in each directions
nx, ny = 10, 5
# number of subelements in one element in each directions
sx_f, sy_f = div(nx_f,nx), div(ny_f,ny)

ndofs = 2
fine_to_coarse = zeros(Int64, ndofs*(nx*porder+1)*(ny*porder+1))
for idof = 1:ndofs
    for iy = 1:ny*porder+1
        for ix = 1:nx*porder+1
            fine_to_coarse[ix + (iy - 1)*(nx*porder+1) + (idof-1)*(nx*porder+1)*(ny*porder+1)] =
            1 + (ix - 1) * sx_f + (iy - 1) * (nx_f*porder + 1) * sy_f + (nx_f*porder + 1)*(ny_f*porder + 1)*(idof - 1)
        end
    end
end

function compute_loss(tid, force_scale)
    nodes, EBC, g, gt, FBC, fext, ft, npoints, node_to_point = BoundaryCondition(tid, nx, ny, porder, Lx, Ly; force_scale=force_scale )
    domain = Domain(nodes, elements, ndofs, EBC, g, FBC, fext)
    state = zeros(domain.neqs)
    ∂u = zeros(domain.neqs)
    globdat = GlobalData(state,zeros(domain.neqs), zeros(domain.neqs),∂u, domain.neqs, gt, ft)
    assembleMassMatrix!(globdat, domain)
    # full_state_history, full_fext_history = read_data("$(@__DIR__)/Data/order$porder/$(tid)_$(force_scale)_$(fiber_size).dat")
    full_state_history, full_fext_history = read_data("$(@__DIR__)/Data/order$porder/$(tid)_$(force_scale)_$(fiber_size).dat")

    #update state history and fext_history on the homogenized domain
    state_history = [x[fine_to_coarse] for x in full_state_history]

    fext_history = []
    setConstantNodalForces!(domain, FBC, fext)
    for i = 1:NT
        globdat.time = Δt*i
        updateTimeDependentEssentialBoundaryCondition!(domain, globdat)
        fext = getExternalForce(domain, globaldat)
        push!(fext_history, fext)
    end
    DynamicMatLawLoss(domain, globdat, state_history, fext_history, nn, Δt)
end

Δt = T/NT

nodes, _, _, _, _, _, _,_,_ = BoundaryCondition(n_data[1], nx, ny, porder, Lx, Ly; force_scale=force_scale)
elements = []
for j = 1:ny
    for i = 1:nx
        n = (nx*porder+1)*(j-1)*porder + (i-1)porder+1
        #element (i,j)
        if porder == 1
            #   4 ---- 3
            #
            #   1 ---- 2

            elnodes = [n, n + 1, n + 1 + (nx + 1), n + (nx + 1)]
        elseif porder == 2
            #   4 --7-- 3
            #   8   9   6
            #   1 --5-- 2
            elnodes = [n, n + 2, n + 2 + 2*(2*nx+1),  n + 2*(2*nx+1), n+1, n + 2 + (2*nx+1), n + 1 + 2*(2*nx+1), n + (2*nx+1), n+1+(2*nx+1)]
        else
            error("polynomial order error, porder= ", porder)
        end
        coords = nodes[elnodes,:]
        push!(elements,SmallStrainContinuum(coords,elnodes, prop, 3))
    end
end

losses = Array{PyObject}(undef, length(n_data)*length(force_scales))
k = 1
for i in n_data
    global k
    for force_scale in force_scales
        losses[k] = compute_loss(i, force_scale)/stress_scale
        k += 1
    end
end

@show stress_scale^2
loss = sum(losses)

sess = Session(); init(sess)

start_id = 4
ADCME.load(sess, "$(@__DIR__)/Data/$nntype/NNPreLSfit_$(idx)_$(H_function)_$(start_id).mat")
#vars = get_collection()
for i = 1:100
    println("************************** Outer Iteration = $i ************************** ")
    BFGS!(sess, loss, 1000)
    #BFGS!(sess, loss, gradients(loss,vars), vars, iterations=1000)
    ADCME.save(sess, "Data/$nntype/NN_Train_$(idx)_from_$(start_id)_$(H_function)_ite$(i).mat")
end


text:
abstract type AbstractPotential end
struct ConstantPotential{T <: Number} <: AbstractPotential
    _::T
end
struct CoeffPotential{P <: AbstractPotential, T <: Number} <: AbstractPotential
    potential::P
    coefficient::T
end
struct PotentialFunction{P <: AbstractPotential, F <: Function} <: AbstractPotential
    potential::P
    func::F
end
struct PotentialExponent{P <: AbstractPotential, T <: Number} <: AbstractPotential
    potential::P
    exponent::T
end
struct PotentialFold{POTS <: NTuple, O <: Function} <: AbstractPotential
    _::POTS
    op::O
end
size(::AbstractPotential) = 1

import Base: *, +, -, /, //, ^, .+, .*, ./
/(p::AbstractPotential, n::Number) = CoeffPotential(p, 1/n)
//(p::AbstractPotential, n::Number) = CoeffPotential(p, 1//n)
*(n::Number, p::AbstractPotential) = CoeffPotential(p, n)
*(n::Number, p::CoeffPotential) = CoeffPotential(p.potential, p.coefficients * n)
*(n::Number, p::PotentialFold) = PotentialFold(tuple((n * u for u in p._)...), +)
*(p::AbstractPotential, n::Number) = n * p
for op in (:+, :.+, :*, :.*, :/, :./, ://, :.//)
    @eval begin
        $op(a::AbstractPotential, b::AbstractPotential) = PotentialFold((a, b), $op)
        $op(a::PotentialFold, b::AbstractPotential) = PotentialFold((a._..., b), $op)
        $op(a::AbstractPotential, b::PotentialFold) = PotentialFold((a, b._...), $op)
    end
end
for op in (:+, :-)
    @eval begin
        $op(a::AbstractPotential, b::Number) = $op(a, ConstantPotential(b))
        $op(b::Number, a::AbstractPotential) = $op(ConstantPotential(b), a)
        $op(a::ConstantPotential, b::ConstantPotential) = ConstantPotential($op(a._, b._))
    end
end
-(a::AbstractPotential, b::AbstractPotential) = PotentialFold((a, -b), +)
-(a::PotentialFold, b::AbstractPotential) = PotentialFold((a._..., -b), +)
-(a::AbstractPotential, b::PotentialFold) = PotentialFold((a, b._...), +)
.-(a::AbstractPotential, b::AbstractPotential) = PotentialFold((a, -b), .+)
.-(a::PotentialFold, b::AbstractPotential) = PotentialFold((a._..., -b), .+)
.-(a::AbstractPotential, b::PotentialFold) = PotentialFold((a, b._...), .+)
-(a::AbstractPotential) = CoeffPotential(a, -1)
^(a::AbstractPotential, n::Integer) = PotentialExponent(a, n)
^(a::AbstractPotential, n::Number) = PotentialExponent(a, n)

for func in [:sin, :cos, :exp, :log10, :log]
    @eval Base.$func(a::AbstractPotential) = PotentialFunction(a, $func)
end

# potential(pot::AbstractPotential, ρ::AxisArrays) = potential!(pot, ρ, similar_potential(ρ))


text:
using QuantEcon
using LaTeXStrings
using Plots
using Plots.PlotMeasures
pyplot()

wp = JvWorker(grid_size=25)
v_init = wp.x_grid .* 0.5

f(x) = bellman_operator(wp, x)
V = compute_fixed_point(f, v_init, max_iter=300)

s_policy, ϕ_policy = bellman_operator(wp, V, ret_policies=true)

# === plot solution === #
p = plot(wp.x_grid, [ϕ_policy s_policy V],
    title = ["ϕ policy" "s policy" "value function"],
    color = [:orange :blue :green], width = 3,
    xaxis = ("x", (0.0, maximum(wp.x_grid))),
    yaxis = ((-0.1, 1.1)),
    size = (800, 800),
    legend = false,
    layout = (3, 1), bottom_margin = Measures.Length(:mm, 20), show = false)


text:
##############   EXAMPLE-0    ##############
# No errors in X nor in Y

using LinearFitXYerrors
using Plots; gr()

# Example-0a
# INPUT DATA:
# http://pba.ucdavis.edu/files/45007.pdf

X = [16.4, 17.2, 17.6, 18.0, 18.2, 18.5]
Y = [2.67, 2.75, 2.99, 3.14, 3.88, 4.23]

# COMPUTE and PLOT:
stxy = linearfitxy(X,Y, isplot=true);

# Example-0b
# INPUT DATA:
# Altman, D. and Gardner, M. [1988] Statistics in Medicine: Calculating confidence intervals for regression and correlation. British Medical Journal, 296(6631), pp.1238–1242.

# Table I
X = [91., 104, 107, 107, 106, 100, 92, 92, 105, 108]
Y = [9.8, 7.4, 7.9, 8.3, 8.3, 9.0, 9.7, 8.8, 7.6, 6.9]

# COMPUTE and PLOT:
stxy = linearfitxy(X,Y, isplot=true);    # σX =  σY = r = 0

savefig("Example0b_LinearFitXYerrors.png")
########################################


text:
using LazIO

function countsgrid(bbox, cellsizes)
    min_x, min_y, min_z, max_x, max_y, max_z = bbox
    rows = Int(cld(max_x - min_x, cellsizes[1]))
    cols = Int(cld(max_y - min_y, cellsizes[2]))
    # heights = Int(cld(max_z - min_z, cellsizes[3]))
    zeros(Int64, rows, cols) #, heights)
end


text:
using Test
using PolynomialBasis
using ImplicitDomainQuadrature
# using Revise
using CutCellDG
include("../useful_routines.jl")

function test_incoherent_interface_simple_tension()
    L = 1.0
    W = 1.0
    lambda1, mu1 = 1.0, 2.0
    lambda2, mu2 = 3.0, 4.0
    penaltyfactor = 1e2
    nelmts = 1
    dx = L / nelmts
    dy = W / nelmts
    penalty = penaltyfactor / dx * (lambda1 + mu1 + lambda2 + mu2) * 0.5
    eta = 1

    applydisplacement = 0.01

    stiffness = CutCellDG.HookeStiffness(lambda1, mu1, lambda2, mu2)
    polyorder = 1
    numqp = 2

    elasticitybasis = LagrangeTensorProductBasis(2, polyorder)
    levelsetbasis = HermiteTensorProductBasis(2)
    dim, nf = size(interpolation_points(levelsetbasis))
    refpoints = interpolation_points(elasticitybasis)

    cgmesh = CutCellDG.CGMesh([0.0, 0.0], [L, W], [nelmts, nelmts], nf)
    mesh = CutCellDG.DGMesh([0.0, 0.0], [L, W], [nelmts, nelmts], refpoints)

    x0 = [0.5, 0.0]
    normal = [1.0, 0.0]
    levelset = CutCellDG.LevelSet(
        x -> plane_distance_function(x, normal, x0),
        cgmesh,
        levelsetbasis,
    )

    cutmesh = CutCellDG.CutMesh(mesh, levelset)
    cellquads = CutCellDG.CellQuadratures(cutmesh, levelset, numqp)
    interfacequads = CutCellDG.InterfaceQuadratures(cutmesh, levelset, numqp)
    facequads = CutCellDG.FaceQuadratures(cutmesh, levelset, numqp)

    mergedmesh = cutmesh

    onbottomboundary(x) = x[2] ≈ 0.0 ? true : false
    onrightboundary(x) = x[1] ≈ L ? true : false
    ontopboundary(x) = x[2] ≈ W ? true : false
    onleftboundary(x) = x[1] ≈ 0.0 ? true : false

    sysmatrix = CutCellDG.SystemMatrix()
    sysrhs = CutCellDG.SystemRHS()

    CutCellDG.assemble_displacement_bilinear_forms!(
        sysmatrix,
        elasticitybasis,
        cellquads,
        stiffness,
        mergedmesh,
    )
    CutCellDG.assemble_incoherent_interface_condition!(
        sysmatrix,
        elasticitybasis,
        interfacequads,
        stiffness,
        mergedmesh,
        penalty,
        eta,
    )
    CutCellDG.assemble_penalty_displacement_component_bc!(
        sysmatrix,
        sysrhs,
        x -> 0.0,
        elasticitybasis,
        facequads,
        stiffness,
        mergedmesh,
        x -> onleftboundary(x),
        [1.0, 0.0],
        penalty,
    )
    CutCellDG.assemble_penalty_displacement_component_bc!(
        sysmatrix,
        sysrhs,
        x -> 0.0,
        elasticitybasis,
        facequads,
        stiffness,
        mergedmesh,
        x -> onbottomboundary(x),
        [0.0, 1.0],
        penalty,
    )
    CutCellDG.assemble_penalty_displacement_component_bc!(
        sysmatrix,
        sysrhs,
        x -> applydisplacement,
        elasticitybasis,
        facequads,
        stiffness,
        mergedmesh,
        x -> onrightboundary(x),
        [1.0, 0.0],
        penalty,
    )

    matrix = CutCellDG.sparse_displacement_operator(sysmatrix, mergedmesh)
    rhs = CutCellDG.displacement_rhs_vector(sysrhs, mergedmesh)

    solution = matrix \ rhs
    nodaldisplacement = reshape(solution, 2, :)

    E1 = mu1 * (lambda1 + mu1) / (lambda1 + 2mu1)
    E2 = mu2 * (lambda2 + mu2) / (lambda2 + 2mu2)

    uI = E1 / (E1 + E2) * applydisplacement

    u10 = uI + (applydisplacement - uI) / 0.5 * (-0.5)
    Δx = applydisplacement

    u12 = -lambda1 / (lambda1 + 2mu1) * 2 * (Δx - uI)
    u22 = -lambda2 / (lambda2 + 2mu2) * 2 * uI

    testdisplacement = [
        u10 u10 Δx Δx 0.0 0.0 2uI 2uI
        0.0 u12 0.0 u12 0.0 u22 0.0 u22
    ]

    @test allapprox(nodaldisplacement, testdisplacement, 1e2eps())
end

test_incoherent_interface_simple_tension()


text:
export DDPGPolicy

using Random
using Flux

mutable struct DDPGPolicy{
    BA<:NeuralNetworkApproximator,
    BC<:NeuralNetworkApproximator,
    TA<:NeuralNetworkApproximator,
    TC<:NeuralNetworkApproximator,
    P,
    R<:AbstractRNG,
} <: AbstractPolicy

    behavior_actor::BA
    behavior_critic::BC
    target_actor::TA
    target_critic::TC
    γ::Float32
    ρ::Float32
    batch_size::Int
    start_steps::Int
    start_policy::P
    update_after::Int
    update_every::Int
    act_limit::Float64
    act_noise::Float64
    step::Int
    rng::R
end

"""
    DDPGPolicy(;kwargs...)

# Keyword arguments

- `behavior_actor`,
- `behavior_critic`,
- `target_actor`,
- `target_critic`,
- `start_policy`,
- `γ = 0.99f0`,
- `ρ = 0.995f0`,
- `batch_size = 32`,
- `start_steps = 10000`,
- `update_after = 1000`,
- `update_every = 50`,
- `act_limit = 1.0`,
- `act_noise = 0.1`,
- `step = 0`,
- `rng = Random.GLOBAL_RNG`,
"""
function DDPGPolicy(;
    behavior_actor,
    behavior_critic,
    target_actor,
    target_critic,
    start_policy,
    γ = 0.99f0,
    ρ = 0.995f0,
    batch_size = 32,
    start_steps = 10000,
    update_after = 1000,
    update_every = 50,
    act_limit = 1.0,
    act_noise = 0.1,
    step = 0,
    rng = Random.GLOBAL_RNG,
)
    copyto!(behavior_actor, target_actor)  # force sync
    copyto!(behavior_critic, target_critic)  # force sync
    DDPGPolicy(
        behavior_actor,
        behavior_critic,
        target_actor,
        target_critic,
        γ,
        ρ,
        batch_size,
        start_steps,
        start_policy,
        update_after,
        update_every,
        act_limit,
        act_noise,
        step,
        rng,
    )
end

# TODO: handle Training/Testing mode
function (p::DDPGPolicy)(env)
    p.step += 1

    if p.step <= p.start_steps
        p.start_policy(env)
    else
        D = device(p.behavior_actor)
        s = get_state(env)
        s = Flux.unsqueeze(s, ndims(s) + 1)
        action = p.behavior_actor(send_to_device(D, s)) |> vec |> send_to_host
        clamp(action[] + randn(p.rng) * p.act_noise, -p.act_limit, p.act_limit)
    end
end

function RLBase.update!(p::DDPGPolicy, traj::CircularCompactSARTSATrajectory)
    length(traj[:terminal]) > p.update_after || return
    p.step % p.update_every == 0 || return

    inds = rand(p.rng, 1:(length(traj[:terminal])-1), p.batch_size)
    s = select_last_dim(traj[:state], inds)
    a = select_last_dim(traj[:action], inds)
    r = select_last_dim(traj[:reward], inds)
    t = select_last_dim(traj[:terminal], inds)
    s′ = select_last_dim(traj[:next_state], inds)

    A = p.behavior_actor
    C = p.behavior_critic
    Aₜ = p.target_actor
    Cₜ = p.target_critic

    γ = p.γ
    ρ = p.ρ

    # !!! we have several assumptions here, need revisit when we have more complex environments
    # state is vector
    # action is scalar
    a′ = Aₜ(s′)
    qₜ = Cₜ(vcat(s′, a′)) |> vec
    y = r .+ γ .* (1 .- t) .* qₜ
    a = Flux.unsqueeze(a, 1)

    gs1 = gradient(Flux.params(C)) do
        q = C(vcat(s, a)) |> vec
        mean((y .- q) .^ 2)
    end

    update!(C, gs1)

    gs2 = gradient(Flux.params(A)) do
        -mean(C(vcat(s, A(s))))
    end

    update!(A, gs2)

    # polyak averaging
    for (dest, src) in zip(Flux.params([Aₜ, Cₜ]), Flux.params([A, C]))
        dest .= ρ .* dest .+ (1 - ρ) .* src
    end
end


text:
# Implements passing in Jacobians (and possibly other functions) via
# function overloading:
#
# - f(...) - objective function
# - f(Val{:jac}, ...) - Jacobian of objective function
# - the details of what `...` needs to be depends on the
#   AbstractODEProblem subtype

# Method_exists does not work:
#
# julia> f(::Val{:jac}, a, b, c) = 5
# f (generic function with 1 method)
#
# julia> method_exists(f, Tuple{Val{:jac}, Vararg})
# false
#
# Thus hand-code it:
check_first_arg(f,T::Type) = check_first_arg(typeof(f),T)
function check_first_arg{F}(::Type{F}, T::Type)
    typ = Tuple{Any, T, Vararg}
    typ2 = Tuple{Any, Type{T}, Vararg} # This one is required for overloaded types
    method_table = Base.MethodList(F.name.mt) # F.name.mt gets the method-table
    for m in method_table
        (m.sig<:typ || m.sig<:typ2) && return true
    end
    return false
end
# Standard
@traitdef HasJac{F}
@traitdef HastGrad{F}
__has_jac(f) = check_first_arg(f, Val{:jac})
__has_tgrad(f) = check_first_arg(f, Val{:tgrad})
@generated SimpleTraits.trait{F}(::Type{HasJac{F}}) = __has_jac(F) ? :(HasJac{F}) : :(Not{HasJac{F}})
@generated SimpleTraits.trait{F}(::Type{HastGrad{F}}) = __has_tgrad(F) ? :(HastGrad{F}) : :(Not{HastGrad{F}})
has_jac{T}(f::T)   = istrait(HasJac{T})
has_tgrad{T}(f::T) = istrait(HastGrad{T})

# Performance
@traitdef HasExpJac{F}
@traitdef HasInvJac{F}
@traitdef HasInvW{F}
@traitdef HasInvW_t{F}
__has_expjac(f) = check_first_arg(f, Val{:expjac})
__has_invjac(f) = check_first_arg(f, Val{:invjac})
__has_invW(f) = check_first_arg(f, Val{:invW})
__has_invW_t(f) = check_first_arg(f, Val{:invW_t})
@generated SimpleTraits.trait{F}(::Type{HasExpJac{F}}) = __has_expjac(F) ? :(HasExpJac{F}) : :(Not{HasExpJac{F}})
@generated SimpleTraits.trait{F}(::Type{HasInvJac{F}}) = __has_invjac(F) ? :(HasInvJac{F}) : :(Not{HasInvJac{F}})
@generated SimpleTraits.trait{F}(::Type{HasInvW{F}}) = __has_invW(F) ? :(HasInvW{F}) : :(Not{HasInvW{F}})
@generated SimpleTraits.trait{F}(::Type{HasInvW_t{F}}) = __has_invW_t(F) ? :(HasInvW_t{F}) : :(Not{HasInvW_t{F}})
has_expjac{T}(f::T) = istrait(HasExpJac{T})
has_invjac{T}(f::T) = istrait(HasInvJac{T})
has_invW{T}(f::T)   = istrait(HasInvW{T})
has_invW_t{T}(f::T) = istrait(HasInvW_t{T})

# Hessians
@traitdef HasHes{F}
@traitdef HasInvHes{F}
__has_hes(f) = check_first_arg(f, Val{:hes})
__has_invhes(f) = check_first_arg(f, Val{:invhes})
@generated SimpleTraits.trait{F}(::Type{HasHes{F}}) = __has_hes(F) ? :(HasHes{F}) : :(Not{HasHes{F}})
@generated SimpleTraits.trait{F}(::Type{HasInvHes{F}}) = __has_invhes(F) ? :(HasInvHes{F}) : :(Not{HasInvHes{F}})
has_hes{T}(f::T)      = istrait(HasHes{T})
has_invhes{T}(f::T)   = istrait(HasInvHes{T})

# Parameter-Based
@traitdef HasParamDeriv{F}
@traitdef HasParamJac{F}
__has_paramderiv(f) = check_first_arg(f, Val{:deriv})
__has_paramjac(f) = check_first_arg(f, Val{:paramjac})
@generated SimpleTraits.trait{F}(::Type{HasParamDeriv{F}}) = __has_paramderiv(F) ? :(HasParamDeriv{F}) : :(Not{HasParamDeriv{F}})
@generated SimpleTraits.trait{F}(::Type{HasParamJac{F}}) = __has_paramjac(F) ? :(HasParamJac{F}) : :(Not{HasParamJac{F}})
has_paramderiv{T}(f::T) = istrait(HasParamDeriv{T})
has_paramjac{T}(f::T)   = istrait(HasParamJac{T})

## Parameter Names Check
@traitdef HasSyms{F}
__has_syms(f) = :syms ∈ fieldnames(f)
@generated SimpleTraits.trait{F}(::Type{HasSyms{F}}) = __has_syms(F) ? :(HasSyms{F}) : :(Not{HasSyms{F}})
has_syms{T}(f::T) = istrait(HasSyms{T})

## Analytical Solution Check
@traitdef HasAnalytic{F}
__has_analytic(f) = check_first_arg(f, Val{:analytic})
@generated SimpleTraits.trait{F}(::Type{HasAnalytic{F}}) = __has_analytic(F) ? :(HasAnalytic{F}) : :(Not{HasAnalytic{F}})
has_analytic{T}(f::T) = istrait(HasAnalytic{T})

# now a trait methods can dispatch on this:
# @traitfn fn(g::::HasJac, ...) = ...
# @traitfn fn(g::::(!HasJac), ...) = ...


text:
"""

```
inbulk(d, R::Result)
```

Function that returns if a distance is in the bulk region or not, according
to the options

"""
function inbulk(d, R::Result)
  if R.options.usecutoff
    if d >= R.dbulk && d < R.cutoff
      return true
    end
  else
    if d >= R.dbulk
      return true
    end
  end
  return false
end