text:
#!/usr/bin/env bash
    #=
    exec julia --project="$(realpath $(dirname $(dirname $0)))" --color=yes --startup-file=no -e "include(popfirst!(ARGS))" \
    "${BASH_SOURCE[0]}" "$@"
    =#

#=
Gauss-Jordan elimination over finite fields, modulo n
=#

function swaprows!(A::Matrix{T}, i::Int, j::Int) where T
    for k in axes(A, 2)
        A[i, k], A[j, k] = A[j, k], A[i, k]
    end

    return A
end

function swapcols!(A::Matrix{T}, i::Int, j::Int) where T
    for k in axes(A, 1)
        A[k, i], A[k, j] = A[k, j], A[k, i]
    end

    return A
end

"""
```julia
rref!(A::Matrix{Int}, n::Int; colswap::Bool=false, verbose::Bool=false, vverbose::Bool=false) -> Matrix{Int}
```

Performs Gauss-Jordan Elimination on a matrix `A`.  *This directly changes the matrix A.  Use `rref` for a non-mutating version of this function.*

Parameters:
  - `A::Matrix{Int}`: A matrix of Ints you wish to perform Gauss-Jordan elimiation on.
  - `n::Int`: The modulus of the finite field you are working under.
  - `colswap::Bool` (kwarg): Whether or not you allow for column swapping.
  - `verbose::Bool` (kwarg): Print the row operations.
  - `vverbose::Bool` (kwarg): Print the intermediate matrices of the algorithm.

Returns:
  - `Matrix{Int}`: a matrix in row echelon form.

---

### Examples

```julia
julia> rref!([1 1 0 2 3 1; 2 0 1 3 4 1; 1 2 2 1 4 3], 5, colswap=true) # gauss-jordan elimitation modulo 5 with column swapping
3×6 Array{Int64,2}:
 1  0  0  3  2  2
 0  1  0  2  1  1
 0  0  1  0  0  4
```
"""
function rref!(A::Matrix{Int},
    n::Int;
    colswap::Bool = false,
    verbose::Bool = false,
    vverbose::Bool = false)

    i = j = 1

    while i ≤ size(A, 1) && j ≤ size(A, 2)
        # Rule 1: Swap zero rows if out of order and ensure leading ones cascade down diagonally.
        s = findfirst(!iszero, A[i:end, :])
        isnothing(s) && break
        swaprows!(A, i, i + s[1] - 1)
        if verbose
            println("r$(i) ⟷  r$(i + s[1] - 1)")
            if vverbose
                displaymatrix(A); println()
            end
        end

        # Rule 2: Normalize each row
        s = findfirst(!iszero, A[i,:])
        isnothing(s) && break
        α = invmod(A[i, s], n)
        A[i,:] .= mod.(A[i,:] .* α, n)
        if verbose
            if ! iszero(α)
                isone(α) || println("r$(i) ⟵  r$(i) × $(α)")
                if vverbose && ! isone(α)
                    displaymatrix(A); println()
                end
            end
        end

        # Rule 3: Subtract it from the others
        s = findfirst(!iszero, A[i,:])
        isnothing(s) && break
        for k in axes(A, 1)
            if i ≠ k
                β = A[k, s]
                A[k, :] .= mod.(A[k, :] .- β .* A[i, :], n)
                if verbose
                    if ! iszero(β)
                        if isone(β)
                            println("r$(k) ⟵  r$(k) - r$(i)")
                        else
                            println("r$(k) ⟵  r$(k) - $(β)r$(i)")
                        end
                        if vverbose
                            displaymatrix(A); println()
                        end
                    end
                end
            end
        end

        # Rule 4: Swap columns if needed (and allowed)
        if colswap
            if iszero(A[i, j])
                s = findfirst(!iszero, A[i, :])
                swapcols!(A, j, s)
                if verbose
                    println("c$(j) ⟷  c$(s)")
                    if vverbose
                        displaymatrix(A); println()
                    end
                end
            end
        end

        # increment row and column counter respectively
        i += 1
        j += 1
    end

    if verbose
        println()
    end

    return A
end

"""
```julia
rref(A::Matrix{Int}, n::Int; colswap::Bool=false, verbose::Bool=false, vverbose::Bool=false) -> Matrix{Int}
```

Performs Gauss-Jordan Elimination on a matrix `A`.

Parameters:
  - `A::Matrix{Int}`: A matrix of Ints you wish to perform Gauss-Jordan elimiation on.
  - `n::Int`: The modulus of the finite field you are working under.
  - `colswap::Bool` (kwarg): Whether or not you allow for column swapping.
  - `verbose::Bool` (kwarg): Print the row operations.
  - `vverbose::Bool` (kwarg): Print the intermediate matrices of the algorithm.

Returns:
  - `Matrix{Int}`: a matrix in row echelon form.

---

### Examples

```julia
julia> rref([1 1 0 2 3 1; 2 0 1 3 4 1; 1 2 2 1 4 3], 5, colswap=true) # gauss-jordan elimitation modulo 5 with column swapping
3×6 Array{Int64,2}:
 1  0  0  3  2  2
 0  1  0  2  1  1
 0  0  1  0  0  4
```
"""
rref(A::Matrix{Int},
    n::Int;
    colswap::Bool = false,
    verbose::Bool = false,
    vverbose::Bool = false
    ) = rref!(copy(A), n::Int; colswap = colswap, verbose = verbose, vverbose = vverbose)


text:
mutable struct SequentialRoomsDirected{T, R} <: AbstractGridWorld
    world::GridWorldBase{Tuple{Agent, Wall, Goal}}
    agent_pos::CartesianIndex{2}
    agent_dir::AbstractDirection
    reward::T
    rng::R
    terminal_reward::T
    num_rooms::Int
    room_length_range::UnitRange{Int}
    rooms::Array{Room, 1}
    done::Bool
end

@generate_getters(SequentialRoomsDirected)
@generate_setters(SequentialRoomsDirected)

mutable struct SequentialRoomsUndirected{T, R} <: AbstractGridWorld
    world::GridWorldBase{Tuple{Agent, Wall, Goal}}
    agent_pos::CartesianIndex{2}
    reward::T
    rng::R
    terminal_reward::T
    num_rooms::Int
    room_length_range::UnitRange{Int}
    rooms::Array{Room, 1}
    done::Bool
end

@generate_getters(SequentialRoomsUndirected)
@generate_setters(SequentialRoomsUndirected)

#####
# Directed
#####

function SequentialRoomsDirected(; T = Float32, num_rooms = 3, room_length_range = 4:6, rng = Random.GLOBAL_RNG)
    objects = (AGENT, WALL, GOAL)
    big_n = 2 * num_rooms * room_length_range.stop
    world = GridWorldBase(objects, big_n, big_n)
    agent_pos = CartesianIndex(2, 2)
    agent_dir = RIGHT
    reward = zero(T)
    terminal_reward = one(T)
    done = false

    env = SequentialRoomsDirected(world, agent_pos, agent_dir, reward, rng, terminal_reward, num_rooms, room_length_range, Room[], done)

    RLBase.reset!(env)

    return env
end

RLBase.state_space(env::SequentialRoomsDirected, ::RLBase.Observation, ::RLBase.DefaultPlayer) = nothing
const SEQUENTIAL_ROOMS_DIRECTED_LAYERS = SA.SVector(2, 3)
RLBase.state(env::SequentialRoomsDirected, ::RLBase.Observation, ::RLBase.DefaultPlayer) = get_grid(get_world(env), get_agent_pos(env), get_agent_dir(env), get_half_size(env), SEQUENTIAL_ROOMS_DIRECTED_LAYERS)

RLBase.state_space(env::SequentialRoomsDirected, ::RLBase.InternalState, ::RLBase.DefaultPlayer) = nothing
RLBase.state(env::SequentialRoomsDirected, ::RLBase.InternalState, ::RLBase.DefaultPlayer) = (copy(get_grid(env)), get_agent_dir(env))

RLBase.action_space(env::SequentialRoomsDirected, ::RLBase.DefaultPlayer) = DIRECTED_NAVIGATION_ACTIONS
RLBase.reward(env::SequentialRoomsDirected, ::RLBase.DefaultPlayer) = get_reward(env)
RLBase.is_terminated(env::SequentialRoomsDirected) = get_done(env)

function RLBase.reset!(env::SequentialRoomsDirected{T}) where {T}
    big_n = 2 * env.num_rooms * env.room_length_range.stop
    world = GridWorldBase(get_objects(env), big_n, big_n)
    set_world!(env, world)

    rng = get_rng(env)

    env.rooms = Room[]

    room = generate_first_room(env)
    place_room!(env, room)
    push!(env.rooms, room)

    tries = 1

    while tries < env.num_rooms
        candidate_rooms = generate_candidate_rooms(env)

        if length(candidate_rooms) > 0
            room = rand(rng, candidate_rooms)
            place_room!(env, room)
            push!(env.rooms, room)

            door_pos = rand(rng, intersect(env.rooms[end - 1].region, room.region)[2:end-1])
            world[WALL, door_pos] = false
        end

        tries += 1
    end

    # create a compact world
    centered_world = get_centered_world(world)
    shift_rooms!(env)
    set_world!(env, centered_world)
    world = get_world(env)

    # add the GOAL randomly in the last room
    goal_pos = rand(rng, get_interior(env.rooms[end]))
    while goal_pos == get_agent_pos(env)
        goal_pos = rand(rng, get_interior(env.rooms[end]))
    end
    world[GOAL, goal_pos] = true

    # add the agent randomly in the first room
    agent_start_pos = rand(rng, get_interior(env.rooms[1]))
    set_agent_pos!(env, agent_start_pos)
    world[AGENT, agent_start_pos] = true

    set_agent_dir!(env, rand(rng, DIRECTIONS))

    set_reward!(env, zero(T))
    set_done!(env, false)

    return nothing
end

function (env::SequentialRoomsDirected{T})(action::AbstractTurnAction) where {T}
    new_dir = turn(action, get_agent_dir(env))
    set_agent_dir!(env, new_dir)
    world = get_world(env)

    if world[GOAL, get_agent_pos(env)]
        set_done!(env, true)
        set_reward!(env, get_terminal_reward(env))
    else
        set_done!(env, false)
        set_reward!(env, zero(T))
    end

    return nothing
end

function (env::SequentialRoomsDirected{T})(action::AbstractMoveAction) where {T}
    world = get_world(env)

    agent_pos = get_agent_pos(env)
    dest = move(action, get_agent_dir(env), agent_pos)

    if !world[WALL, dest]
        world[AGENT, agent_pos] = false
        set_agent_pos!(env, dest)
        world[AGENT, dest] = true
    end

    if world[GOAL, get_agent_pos(env)]
        set_done!(env, true)
        set_reward!(env, get_terminal_reward(env))
    else
        set_done!(env, false)
        set_reward!(env, zero(T))
    end

    return nothing
end

#####
# Undirected
#####

function SequentialRoomsUndirected(; T = Float32, num_rooms = 3, room_length_range = 4:6, rng = Random.GLOBAL_RNG)
    objects = (AGENT, WALL, GOAL)
    big_n = 2 * num_rooms * room_length_range.stop
    world = GridWorldBase(objects, big_n, big_n)
    agent_pos = CartesianIndex(2, 2)
    reward = zero(T)
    terminal_reward = one(T)
    done = false

    env = SequentialRoomsUndirected(world, agent_pos, reward, rng, terminal_reward, num_rooms, room_length_range, Room[], done)

    RLBase.reset!(env)

    return env
end

RLBase.state_space(env::SequentialRoomsUndirected, ::RLBase.Observation, ::RLBase.DefaultPlayer) = nothing
const SEQUENTIAL_ROOMS_UNDIRECTED_LAYERS = SA.SVector(2, 3)
RLBase.state(env::SequentialRoomsUndirected, ::RLBase.Observation, ::RLBase.DefaultPlayer) = get_grid(get_world(env), get_agent_pos(env), get_half_size(env), SEQUENTIAL_ROOMS_UNDIRECTED_LAYERS)

RLBase.state_space(env::SequentialRoomsUndirected, ::RLBase.InternalState, ::RLBase.DefaultPlayer) = nothing
RLBase.state(env::SequentialRoomsUndirected, ::RLBase.InternalState, ::RLBase.DefaultPlayer) = copy(get_grid(env))

RLBase.action_space(env::SequentialRoomsUndirected, player::RLBase.DefaultPlayer) = UNDIRECTED_NAVIGATION_ACTIONS
RLBase.reward(env::SequentialRoomsUndirected, ::RLBase.DefaultPlayer) = get_reward(env)
RLBase.is_terminated(env::SequentialRoomsUndirected) = get_done(env)

function RLBase.reset!(env::SequentialRoomsUndirected{T}) where {T}
    big_n = 2 * env.num_rooms * env.room_length_range.stop
    world = GridWorldBase(get_objects(env), big_n, big_n)
    set_world!(env, world)

    rng = get_rng(env)

    env.rooms = Room[]

    room = generate_first_room(env)
    place_room!(env, room)
    push!(env.rooms, room)

    tries = 1

    while tries < env.num_rooms
        candidate_rooms = generate_candidate_rooms(env)

        if length(candidate_rooms) > 0
            room = rand(rng, candidate_rooms)
            place_room!(env, room)
            push!(env.rooms, room)

            door_pos = rand(rng, intersect(env.rooms[end - 1].region, room.region)[2:end-1])
            world[WALL, door_pos] = false
        end

        tries += 1
    end

    # create a compact world
    centered_world = get_centered_world(world)
    shift_rooms!(env)
    set_world!(env, centered_world)
    world = get_world(env)

    # add the GOAL randomly in the last room
    goal_pos = rand(rng, get_interior(env.rooms[end]))
    while goal_pos == get_agent_pos(env)
        goal_pos = rand(rng, get_interior(env.rooms[end]))
    end
    world[GOAL, goal_pos] = true

    # add the agent randomly in the first room
    agent_start_pos = rand(rng, get_interior(env.rooms[1]))
    set_agent_pos!(env, agent_start_pos)
    world[AGENT, agent_start_pos] = true

    set_reward!(env, zero(T))
    set_done!(env, false)

    return nothing
end

function (env::SequentialRoomsUndirected{T})(action::AbstractMoveAction) where {T}
    world = get_world(env)

    agent_pos = get_agent_pos(env)
    dest = move(action, agent_pos)

    if !world[WALL, dest]
        world[AGENT, agent_pos] = false
        set_agent_pos!(env, dest)
        world[AGENT, dest] = true
    end

    if world[GOAL, get_agent_pos(env)]
        set_done!(env, true)
        set_reward!(env, get_terminal_reward(env))
    else
        set_done!(env, false)
        set_reward!(env, zero(T))
    end

    return nothing
end

#####
# Common
#####

#####
# Room generation and placement
#####

function generate_first_room(env::Union{SequentialRoomsDirected, SequentialRoomsUndirected})
    big_n = get_width(env)
    rng = get_rng(env)
    origin = CartesianIndex(big_n ÷ 2 + 1, big_n ÷ 2 + 1)
    height = rand(rng, env.room_length_range)
    width = rand(rng, env.room_length_range)
    room = Room(origin, height, width)
    return room
end

is_valid_room(env::Union{SequentialRoomsDirected, SequentialRoomsUndirected}, room::Room) = !any(x -> is_intersecting(room, x), env.rooms)

function generate_candidate_rooms(env::Union{SequentialRoomsDirected, SequentialRoomsUndirected})
    rooms = Room[]
    for height in env.room_length_range, width in env.room_length_range, dir in DIRECTIONS
        push!(rooms, generate_candidate_rooms(env, height, width, dir)...)
    end
    return rooms
end

function generate_candidate_rooms(env::Union{SequentialRoomsDirected, SequentialRoomsUndirected}, height::Int, width::Int, dir::AbstractDirection)
    origins = generate_candidate_origins(env.rooms[end], height, width, dir)
    rooms = map(x -> Room(x, height, width), origins)
    valid_rooms = filter(x -> is_valid_room(env, x), rooms)
end

function generate_candidate_origins(room::Room, height::Int, width::Int, dir::Up)
    i = room.region.indices[1].start - height + 1
    jj = room.region.indices[2].start - width + 3 : room.region.indices[2].stop - 2
    return CartesianIndices((i:i, jj))
end

function generate_candidate_origins(room::Room, height::Int, width::Int, dir::Down)
    i = room.region.indices[1].stop
    jj = room.region.indices[2].start - width + 3 : room.region.indices[2].stop - 2
    return CartesianIndices((i:i, jj))
end

function generate_candidate_origins(room::Room, height::Int, width::Int, dir::Left)
    ii = room.region.indices[1].start - height + 3 : room.region.indices[1].stop - 2
    j = room.region.indices[2].start - width + 1
    return CartesianIndices((ii, j:j))
end

function generate_candidate_origins(room::Room, height::Int, width::Int, dir::Right)
    ii = room.region.indices[1].start - height + 3 : room.region.indices[1].stop - 2
    j = room.region.indices[2].stop
    return CartesianIndices((ii, j:j))
end

#####
# World centering
#####

function get_bounding_region(world::GridWorldBase)
    all_wall_pos = findall(world[WALL, :, :])
    top_extreme = minimum(pos -> pos.I[1], all_wall_pos)
    bottom_extreme = maximum(pos -> pos.I[1], all_wall_pos)
    left_extreme = minimum(pos -> pos.I[2], all_wall_pos)
    right_extreme = maximum(pos -> pos.I[2], all_wall_pos)
    CartesianIndices((top_extreme:bottom_extreme, left_extreme:right_extreme))
end

function get_padding(world::GridWorldBase)
    bounding_region = get_bounding_region(world)
    small_n = get_width(world) ÷ 2
    top_padding = (small_n - size(bounding_region, 1)) ÷ 2
    left_padding = (small_n - size(bounding_region, 2)) ÷ 2
    CartesianIndex(top_padding, left_padding)
end

function get_new_global_origin(world::GridWorldBase)
    padding = get_padding(world)
    bounding_region = get_bounding_region(world)
    return get_origin(bounding_region) - padding
end

function shift_rooms!(env::Union{SequentialRoomsDirected, SequentialRoomsUndirected})
    world = get_world(env)
    new_global_origin = get_new_global_origin(world)
    for (i, room) in enumerate(env.rooms)
        env.rooms[i] = shifted_room(room, new_global_origin)
    end
end

function shifted_room(room::Room, new_global_origin::CartesianIndex{2})
    height = size(room.region, 1)
    width = size(room.region, 2)
    new_room_origin = get_origin(room) - new_global_origin + CartesianIndex(1, 1)
    Room(new_room_origin, height, width)
end

function get_centered_world(world::GridWorldBase)
    small_n = get_width(world) ÷ 2
    new_global_origin = get_new_global_origin(world)
    new_global_region = CartesianIndices((new_global_origin.I[1] : new_global_origin.I[1] + small_n - 1, new_global_origin.I[2] : new_global_origin.I[2] + small_n - 1))
    centered_world = GridWorldBase(get_objects(world), small_n, small_n)
    centered_world[:, :, :] .= world[:, new_global_region]
    return centered_world
end

ext: jl


text:
# ------------ Block alias and APIs ------------ #

const Block = MLIR.API.MlirBlock

create_block(n::Int, ts::Vector{Type}) = MLIR.API.mlirBlockCreate(n, ts)
create_block(ts::Vector{Type}) = create_block(length(ts), ts)
destroy!(b::Block) = MLIR.API.mlirBlockDestroy(b)
is_null(b::Block) = MLIR.API.mlirBlockIsNull(b)
Base.:(==)(b1::Block, b2::Block) = MLIR.API.mlirBlockEqual(b1, b2)
get_arg(b::Block, pos::Int) = MLIR.API.mlirBlockGetArgument(b, pos)
get_num_args(b::Block) = MLIR.API.mlirBlockGetNumArguments(b)
get_next_in_region(b::Block) = MLIR.API.mlirBlockGetNextInRegion(b)
get_first_operation(b::Block) = MLIR.API.mlirBlockGetFirstOperation(b)
get_terminator(b::Block) = MLIR.API.mlirBlockGetTerminator(b)
insert!(b::Block, pos::Int, op::Operation) = MLIR.API.mlirBlockInsertOwnedOperation(b, pos, unwrap(op))
insertbefore!(b::Block, ref::Operation, op::Operation) = MLIR.API.mlirBlockInsertOwnedOperationAfter(b, unwrap(ref), unwrap(op))
insertafter!(b::Block, ref::Operation, op::Operation) = MLIR.API.mlirBlockInsertOwnedOperation(b, unwrap(ref), unwrap(op))

function Base.getindex(blk::Block, pos::Int)
    arg = get_arg(blk, pos)
    unwrap(arg) == C_NULL && error("Pointer to SSA argument at block position $pos is NULL.")
    arg
end

# Constructor.
Block(ts::Type...) = create_block(collect(ts))

@doc(
"""
const Block = MLIR.API.MlirBlock
""", Block)


text:
doc"""
    output(result, nv, ne)

Print Graph500 Benchmark statistics(http://graph500.org/?page_id=12#sec-9_3)
"""
function output(result::Graph500Results,nv::Integer,ne::Integer)
    io = IOContext(stdout, :color=>true)
    printstyled(io, "\n", "Benchmark Info", "\n", color=:green)
    println("Scale : ", result.scale)
    println("Number of vertices : ", nv)
    println("Number of Edges : ", ne)
    println("Number of Bfs : ", result.nbfs, "\n")

    # kernel 1 Construction time
    printstyled(io, "\n", "Graph construction time", "\n", color=:green)
    println("Graph Construction Time : ", result.kernel_1_time, "\n")

    # bfs time statistics
    bfs_time_stats = summarystats(result.kernel_2_time)
    printstyled(io, "\n", "BFS time statistics", "\n", color=:green)
    println("BFS min time : ", bfs_time_stats.min)
    println("BFS first quartile time : ", bfs_time_stats.q25)
    println("BFS median time : ", bfs_time_stats.median)
    println("BFS third quartile time : ", bfs_time_stats.q75)
    println("BFS max time : ", bfs_time_stats.max)
    println("BFS mean time : ", bfs_time_stats.mean)
    println("BFS standard deviation time : ", std(result.kernel_2_time), "\n")

    # bfs statistics based on number of edges
    bfs_edgetraversed_stats = summarystats(result.kernel_2_nedge)
    printstyled(io, "\n", "Edges traversed per BFS statistics", "\n", color=:green)
    println("BFS min nedge : ", bfs_edgetraversed_stats.min)
    println("BFS first quartile nedge : ", bfs_edgetraversed_stats.q25)
    println("BFS median nedge : ", bfs_edgetraversed_stats.median)
    println("BFS third quartile nedge : ", bfs_edgetraversed_stats.q75)
    println("BFS max nedge : ", bfs_edgetraversed_stats.max)
    println("BFS mean nedge : ", bfs_edgetraversed_stats.mean)
    println("BFS standard deviation nedge : ", std(result.kernel_2_nedge), "\n")

    # statistics on traversed edges per second(teps)
    bfs_teps = result.kernel_2_nedge ./ result.kernel_2_time
    bfs_n = length(bfs_teps)
    teps_stats = summarystats(bfs_teps)
    printstyled(io, "\n", "Edges traversed per second statistics", "\n", color=:green)
    println("BFS min TEPS : ", teps_stats.min)
    println("BFS first quartile TEPS : ", teps_stats.q25)
    println("BFS median TEPS : ", teps_stats.median)
    println("BFS third quartile TEPS : ", teps_stats.q75)
    println("BFS max TEPS : ", teps_stats.max)

    # harmomnic mean of teps
    k2tmp = 1.0 ./ bfs_teps
    k2tmp = k2tmp .- (1/harmmean(bfs_teps))
    harmonic_std = (sqrt(sum(k2tmp.^2)) / (bfs_n-1)) * harmmean(bfs_teps)^2
    println("BFS harmonic mean TEPS : ", harmmean(bfs_teps))
    println("BFS harmonic standard deviation TEPS : ", harmonic_std, "\n")
end


text:
# # Example 1: Inverted Coax Detector

using Plots; pyplot(fmt = :png);
using SolidStateDetectors
using Unitful

T = Float32
simulation = Simulation{T}(SSD_examples[:InvertedCoax])

plot(simulation.detector, size = (700, 700))

# One can also have a look at how the initial conditions look like on the grid (its starts with a very coarse grid):

apply_initial_state!(simulation, ElectricPotential) # optional
plot(
    plot(simulation.electric_potential), # initial electric potential (boundary conditions)
    plot(simulation.point_types), # map of different point types: fixed point / inside or outside detector volume / depleted/undepleted
    plot(simulation.ρ), # charge density distribution
    plot(simulation.ϵ), # dielectric distribution
    layout = (1, 4), size = (1400, 700)
)

# Next, calculate the electric potential:

calculate_electric_potential!( simulation,
                               max_refinements = 4)

plot(
    plot(simulation.electric_potential, φ = 20), # initial electric potential (boundary conditions)
    plot(simulation.point_types), # map of different point types: fixed point / inside or outside detector volume / depleted/undepleted
    plot(simulation.ρ), # charge density distribution
    plot(simulation.ϵ), # dielectric distribution
    layout = (1, 4), size = (1400, 700)
)

# SolidStateDetectors.jl supports active (i.e. depleted) volume calculation:

get_active_volume(simulation.point_types) # approximation (sum of the volume of cells marked as depleted)

# ## Partially depleted detectors

# SolidStateDetectors.jl can also calculate the electric potential of a partially depleted detector:

detector_undep = deepcopy(simulation.detector)
detector_undep.contacts[end].potential = 500; # V  <-- Bias Voltage of Mantle

simulation_undep = Simulation(detector_undep);

calculate_electric_potential!( simulation_undep,
                               depletion_handling = true,
                               max_refinements = 4,
                               verbose = false)

plot(
    plot(simulation_undep.electric_potential),
    plot(simulation_undep.point_types),
    layout = (1, 2), size = (800, 700)
)

# Compare both volumes:

println("Depleted:   ", get_active_volume(simulation.point_types))
println("Undepleted: ", get_active_volume(simulation_undep.point_types));

# ## Electric field calculation

# Calculate the electric field of the fully depleted detector, given the already calculated electric potential:

calculate_electric_field!(simulation, n_points_in_φ = 72)

plot_electric_field(simulation, size = (350, 500))

# ## Drift field calculation

# Given the electric field and a charge drift model, calculate drift fields for electrons and holes. Precalculating the drift fields saves time during charge drift simulation:

# And apply the charge drift model to the electric field:

apply_charge_drift_model!(simulation)

# Now, let's create an "random" (multiside) event:

starting_positions = [ CylindricalPoint{T}( 0.020, deg2rad(10), 0.015 ),
                       CylindricalPoint{T}( 0.015, deg2rad(20), 0.045 ),
                       CylindricalPoint{T}( 0.025, deg2rad(30), 0.025 ) ]
energy_depos = T[1460, 609, 1000] * u"keV" # are needed later in the signal generation

event = Event(starting_positions, energy_depos);

time_step = 5u"ns"
drift_charges!(event, simulation, Δt = time_step)

plot(simulation.detector, size = (700, 700))
plot!(event.drift_paths)

# ## Weighting potential calculation

# We need weighting potentials to simulate the detector charge signal induced by drifting charges. We'll calculate the weighting potential for the point contact and the outer shell of the detector:

for contact in simulation.detector.contacts
    calculate_weighting_potential!(simulation, contact.id, max_refinements = 4, n_points_in_φ = 2, verbose = false)
end

plot(
    plot(simulation.weighting_potentials[1]),
    plot(simulation.weighting_potentials[2]),
    size = (900, 700)
)

# ## Detector waveform generation

# ### Single-event simulation

# Given an interaction at an arbitrary point in the detector, we can now simulate the charge drift and the resulting detector charge signals (e.g. at the point contact):

simulate!(event, simulation) # drift_charges + signal generation of all channels

p_pc_signal = plot( event.waveforms[1], lw = 1.5, xlims = (0, 2000), xlabel = "Time / ns",
                    legend = false, tickfontsize = 12, ylabel = "Energy / eV", guidefontsize = 14)


text:
import Distributions
using Statistics

include(joinpath(@__DIR__, "simulation_multi_target_localization.jl"))

function getMetricHistory(s_history)
    # Evaluation Metric: Worst Entropy
    MetricArray = Float64[];
    for t = 1:length(s_history)
        bVec = Distributions.MvNormal.(s_history[t].b)
        push!(MetricArray,maximum(Distributions.entropy.(bVec)));
    end
    return MetricArray
end

function evaluate()
    data_dir = normpath(joinpath(@__DIR__, "../data",
                                 "multi_target_localization"));
    # For ergodic control, the ErgodicControl.jl package is not compatible with Julia v1.3
    # Therefore, we use previously collected data with Julia v0.6 for the WAFR paper.
    ergodic_data = FileIO.load(joinpath(data_dir, "statistical_data_localization_openloop_ergodic.jld2"));

    safe_rngnums = ergodic_data["safe_rngnums"]

    num_successful_cases = 0;

    cost_greedy = Array{Float64}[];
    tcalc_greedy = Array{Float64}[];
    cost_sacbp = Array{Float64}[];
    tcalc_sacbp = Array{Float64}[];
    cost_mcts = Array{Float64}[];
    tcalc_mcts = Array{Float64}[];

    time_idx_array = [];

    for rngnum in safe_rngnums
        println("------- RNG Seed: $(rngnum) & #Cases Tested: $(num_successful_cases) -------")
        println("1. Simulation with Greedy Controller")
        #try
        y_history,tcalc_true_history,s_history,u_history,U_pool =
        simulate_main("greedy",rng_num=rngnum, animate=false, plotfig=false, verbose=false);
        push!(cost_greedy,getMetricHistory(s_history));
        push!(tcalc_greedy,tcalc_true_history);
        if isempty(time_idx_array)
            push!(time_idx_array, unique([s.t for s in s_history]))
        end
        #catch
        #    rngnum += 1
        #    continue
        #end
        println("2. Simulation with SACBP Controller")
        #try
        y_history,tcalc_true_history,s_history,u_history,U_pool =
        simulate_main("sacbp",rng_num=rngnum, animate=false, plotfig=false, verbose=false);
        push!(cost_sacbp,getMetricHistory(s_history));
        push!(tcalc_sacbp,tcalc_true_history);
        if isempty(time_idx_array)
            push!(time_idx_array, unique([s.t for s in s_history]))
        end
        #catch
        #    pop!(cost_greedy)
        #    pop!(tcalc_greedy)
        #    rngnum += 1
        #    continue
        #end
        println("3. Simulation with MCTS Controller")
        #try
        y_history,tcalc_true_history,s_history,u_history,U_pool =
        simulate_main("mcts",rng_num=rngnum, animate=false, plotfig=false, verbose=false);
        push!(cost_mcts,getMetricHistory(s_history));
        push!(tcalc_mcts,tcalc_true_history);
        if isempty(time_idx_array)
            push!(time_idx_array, unique([s.t for s in s_history]))
        end
        #catch
        #    pop!(cost_greedy)
        #    pop!(tcalc_greedy)
        #    pop!(cost_sacbp)
        #    pop!(tcalc_sacbp)
        #    rngnum += 1
        #    continue
        #end
        #append!(safe_rngnums,rngnum)
        #rngnum += 1;
        num_successful_cases += 1;
    end

    tcalc_greedy = vcat(tcalc_greedy...);
    tcalc_sacbp = vcat(tcalc_sacbp...);
    tcalc_mcts = vcat(tcalc_mcts...);

    cost_ergodic, tcalc_ergodic = ergodic_data["cost_ergodic"], ergodic_data["tcalc_ergodic"]

    FileIO.save(joinpath(data_dir, "statistical_data_localization_openloop.jld2"),
                "cost_greedy",cost_greedy,"tcalc_greedy",tcalc_greedy,
                "cost_sacbp",cost_sacbp,"tcalc_sacbp",tcalc_sacbp,
                "cost_mcts",cost_mcts,"tcalc_mcts",tcalc_mcts,
                "cost_ergodic",cost_ergodic,"tcalc_ergodic",tcalc_ergodic)

    # Plot Stats Data
    cost_std_mcts = [std([cost[ii] for cost in cost_mcts]) for ii = 1:length(cost_mcts[1])]
    cost_std_sacbp = [std([cost[ii] for cost in cost_sacbp]) for ii = 1:length(cost_sacbp[1])]
    cost_std_greedy = [std([cost[ii] for cost in cost_greedy]) for ii = 1:length(cost_greedy[1])];
    cost_std_ergodic = [std([cost[ii] for cost in cost_ergodic]) for ii = 1:length(cost_ergodic[1])];

    MetricHistoryGreedy = mean(cost_greedy)
    plot(time_idx_array,MetricHistoryGreedy[1:end-1],xlabel="Time [s]",ribbon=cost_std_greedy,fillalpha=0.3,ylabel="Worst Entropy Value",label="Greedy",color=:dimgrey,linewidth=2.,linestyle=:dot)
    MetricHistoryMCTS = mean(cost_mcts)
    plot!(time_idx_array,MetricHistoryMCTS[1:end-1],label="MCTS-DPW",ribbon=cost_std_mcts,color=:lightcoral,fillalpha=0.3,linewidth=2.,linestyle=:dash)
    MetricHistoryErgodic = mean(cost_ergodic)
    plot!(time_idx_array,MetricHistoryErgodic[1:end-1],label="Ergodic",ribbon=cost_std_ergodic,color=:darkorange,fillalpha=0.3,linewidth=2.,linestyle=:dashdot)
    MetricHistorySACBP = mean(cost_sacbp)
    plot!(time_idx_array,MetricHistorySACBP[1:end-1],label="SACBP",ribbon=cost_std_sacbp,fillalpha=0.3,color=:darkblue,linewidth=2.,size=(400,400))
    savefig(joinpath(data_dir, "tracking_costs.pdf"))

    # Plot Computation Time Comparison
    tcalcTrueArray = [tcalc_greedy,tcalc_mcts,tcalc_ergodic,tcalc_sacbp]
    meanTcalcTrueArray = mean.(tcalcTrueArray);
    stdTcalcTrueArray = std.(tcalcTrueArray);
    bar([1,2,3,4],meanTcalcTrueArray,yerror=stdTcalcTrueArray,linecolor=:black,marker=stroke(4.0,:black, :dash),c=[:dimgrey,:lightcoral,:darkorange,:darkblue],label="",size=(400,400));
    xticks!([1,2,3,4],["Greedy","MCTS-DPW","Ergodic","SACBP"]);
    plot!([0.5,4.5],[3/4*0.2,3/4*0.2],width=2,label="Targeted",color=:darkred,ylim=(-0.03,0.4))
    ylabel!("Control Computation Time [s]")
    savefig(joinpath(data_dir, "tracking_time.pdf"))
end

if abspath(PROGRAM_FILE) == @__FILE__
    evaluate();
end


text:
# JET doesn't conretize this by default, but just analyze its type
const GLOBAL_CODE_STORE = Dict()

macro with_code_record(a)
    GLOBAL_CODE_STORE[__source__] = a # record the code location in the global store
    esc(a)
end

# here JET will try to actually expand `@with_code_record`,
# but since `GLOBAL_CODE_STORE` didn't get concretized (i.e. instantiated), JET analysis fails at this point
@with_code_record foo(a) = identity(a)

foo(10) # top-level callsite, abstracted away


text:
module Smg2s

using LinearAlgebra
using SparseArrays
using Random

export Nilpotent, Nilp
export Spectrum!, checkSpectrum
export initMat!
export nonherm, nonsym
export NilpxM, MxNilp, MNilpM, NilpMat

include("nilpotent.jl")
include("spectrum.jl")
include("init.jl")
include("nonsym.jl")
include("nonherm.jl")

end #endmodule


text:
# SPDX-License-Identifier: MIT

module SpeechFeatures

using PaddedViews
using FFTW

include("utils.jl")
include("features.jl")

export filterbank, stft, mfcc, add_deltas

end

max_stars_repo_name: lucasondel/SpeechFeatures

max_issues_repo_name: BUTSpeechFIT/SpeechFeatures

max_forks_repo_name: BUTSpeechFIT/SpeechFeatures


text:
# Tests for observation-based learning
function t_observation_networklearner()

#############################################
# Test the observation-based NetworkLearner #
#############################################
Ntrain = 100						# Number of training observations
Ntest = 10						# Number of testing observations
inferences = [:ic, :rl, :gs, :unknown]			# Collective inferences
rlearners = [:rn, :wrn, :bayesrn, :cdrn, :unknown]	# Relational learners
nAdj = 2						# Number of adjacencies to generate

nlmodel=[]

#####################
# Column-major case #
#####################

# Initializations
X = rand(1,Ntrain); 					# Training data
ft=x->vec(x)
y = vec(sin.(X));
Xo = zeros(1,Ntest)

# Train and test methods for local model
fl_train = (x)->mean(x[1]);
fl_exec=(m,x)->x.-m;

# Train and test methods for relational model
fr_train=(x)->sum(x[1], dims=2);
fr_exec=(m,x)->sum(x.-m, dims=1)
amv = sparse.(Symmetric.([sprand(Float64, Ntrain,Ntrain, 0.5) for i in 1:nAdj]));
adv = adjacency.(amv);

for infopt in inferences
	for rlopt in rlearners
		Test.@test try
			# Train NetworkLearner
			nlmodel=fit(NetworkLearnerObs, X, y,
				   adv, fl_train, fl_exec,fr_train,fr_exec;
				   learner=rlopt,
				   inference=infopt,
				   use_local_data=rand(Bool),
				   f_targets=ft,
				   normalize=false, maxiter = 5
			)

			# Test NetworkLearner
			Xtest = rand(1,Ntest)

			# Add adjacency
			amv_t = sparse.(Symmetric.([sprand(Float64, Ntest,Ntest, 0.7) for i in 1:nAdj]));
			adv_t = adjacency.(amv_t);
			add_adjacency!(nlmodel, adv_t)

			#Run NetworkLearner
			predict!(Xo, nlmodel, Xtest); # in-place
			predict(nlmodel, Xtest);      # creates output matrix
			true
		catch
			false
		end
	end
end

##################
# Row-major case #
##################

# Initializations
X = rand(Ntrain,1); 					# Training data
ft=x->vec(x)
y = vec(sin.(X));
Xo = zeros(Ntest,1)

# Train and test methods for local model
fl_train = (x)->mean(x[1]);
fl_exec=(m,x)->x.-m;

# Train and test methods for relational model
fr_train=(x)->sum(x[1], dims=1);
fr_exec=(m,x)->sum(x.-m, dims=2)
amv = sparse.(Symmetric.([sprand(Float64, Ntrain,Ntrain, 0.5) for i in 1:nAdj]));
adv = adjacency.(amv);

for infopt in inferences
	for rlopt in rlearners
		Test.@test try
			# Train NetworkLearner
			nlmodel=fit(NetworkLearnerObs, X, y,
				   adv, fl_train, fl_exec,fr_train,fr_exec;
				   learner=rlopt,
				   inference=infopt,
				   use_local_data=rand(Bool),
				   f_targets=ft,
				   normalize=false, maxiter = 5,
				   obsdim=1
			)

			# Test NetworkLearner
			Xtest = rand(Ntest,1)

			# Add adjacency
			amv_t = sparse.(Symmetric.([sprand(Float64, Ntest,Ntest, 0.7) for i in 1:nAdj]));
			adv_t = adjacency.(amv_t);
			add_adjacency!(nlmodel, adv_t)

			#Run NetworkLearner
			predict!(Xo, nlmodel, Xtest); # in-place
			predict(nlmodel, Xtest);      # creates output matrix
			true
		catch
			false
		end
	end
end

buf = IOBuffer()
Test.@test try
	show(buf,nlmodel)
	true
catch
	false
end

end


text:
"""
    singletonof(::Type{T}, x) :: T
    singletonof(::T, x) :: T

Create a singleton container of type `T`.

# Examples
```jldoctest
julia> using BangBang

julia> @assert singletonof(Vector, 1) == [1]

julia> @assert singletonof(Dict, :a => 1) == Dict(:a => 1)

julia> @assert singletonof(Set, 1) == Set([1])

julia> using StructArrays: StructVector

julia> @assert singletonof(StructVector, (a=1, b=2)) == StructVector(a=[1], b=[2])

julia> using TypedTables: Table

julia> @assert singletonof(Table, (a=1, b=2)) == Table(a=[1], b=[2])

julia> using StaticArrays: SArray, SVector

julia> @assert singletonof(SArray, 1) === SVector(1)

julia> @assert singletonof(SVector, 1) === SVector(1)
```
"""
singletonof(::Type{T}, x) where T = T(SingletonVector(x))

function singletonof(::T, x) where T
    C = constructorof(T)
    return singletonof(C isa Type ? C : T, x)
end


text:
include("attributes.jl")

function PMM1_leader(x, y)
    k = min(length(x), length(y)) ÷ 2

    x1 = x[1:k]; x2 = x[k+1:end];
    y1 = y[1:k]

    q1 = sum( (y1 - x1.^3/100).^2 )
    q2 = 10 + y[k+1] + 1e6sum( y[k+2:end].^2 )

    r1 = -sum(x1.^2)
    r2 = -sum(x2.^2)

    p1 = q1 + r1
    p2 = q2 + r2
    F = abs(p1) + abs(p2)

    F, zeros(0), zeros(0)
end

function PMM1_follower(x, y)
    k = min(length(x), length(y)) ÷ 2

    x1 = x[1:k]; x2 = x[k+1:end];
    y1 = y[1:k]

    q1 = sum( (y1 - x1.^3/100).^2 )
    q2 = 10 + y[k+1] + 1e6sum( y[k+2:end].^2 )

    r1 = -sum(x1.^2)
    r2 = -sum(x2.^2)

    p1 = q1 + r1
    p2 = q2 + r2

    f = p1 + p2

    f, zeros(0), zeros(0)
end

function PMM2_leader(x, y)
    k = min(length(x), length(y)) ÷ 2

    ###########################
    q1 = (y[1]- x[1] *sin(x[1]) ).^2 + 1e6sum((y[2:k]- x[2:k] .* sin.(x[2:k]) ).^2)

    m = k+1:length(y)
    v = (y[m] - sqrt.(m))
    q2 = sum(v.^2) + sum(0.5v)^2 + sum(0.5v)^4

    x1 = x[1:k]; x2 = x[k+1:end];
    r1 = -sum(x1.^2)
    r2 = -x2[1]^2 - 1e6sum(x2[2:end].^2)
    ###########################

    p1 = q1 + r1
    p2 = q2 + r2

    F = abs(p1) + abs(p2)

    F, zeros(0), zeros(0)
end

function PMM2_follower(x, y)
    k = min(length(x), length(y)) ÷ 2

    ###########################
    q1 = (y[1]- x[1] *sin(x[1]) ).^2 + 1e6sum((y[2:k]- x[2:k] .* sin.(x[2:k]) ).^2)

    m = k+1:length(y)
    v = (y[m] - sqrt.(m))
    q2 = sum(v.^2) + sum(0.5v)^2 + sum(0.5v)^4

    x1 = x[1:k]; x2 = x[k+1:end];
    r1 = -sum(x1.^2)
    r2 = -x2[1]^2 - 1e6sum(x2[2:end].^2)
    ###########################

    p1 = q1 + r1
    p2 = q2 + r2

    f = p1 + p2

    f, zeros(0), zeros(0)
end

function PMM3_leader(x, y)
    k = min(length(x), length(y)) ÷ 2

    ###########################
    x1 = x[1:k]; x2 = x[k+1:end];
    y1 = y[1:k]; y2 = y[k+1:end];

    q1 = sum( (y1 - x1.^3/100).^2 )
    q2 = sum( 100(y2[1:end-1].^2 - y2[2:end]).^2 + (y2[1:end-1] .- 1).^2 )

    r1 = -sum(x1.^2)
    r2 = -sum(x2.^2)
    ###########################

    p1 = q1 + r1
    p2 = q2 + r2

    F = abs(p1) + abs(p2)

    F, zeros(0), zeros(0)
end

function PMM3_follower(x, y)
    k = min(length(x), length(y)) ÷ 2

    ###########################
    x1 = x[1:k]; x2 = x[k+1:end];
    y1 = y[1:k]; y2 = y[k+1:end];

    q1 = sum( (y1 - x1.^3/100).^2 )
    q2 = sum( 100(y2[1:end-1].^2 - y2[2:end]).^2 + (y2[1:end-1] .- 1).^2 )

    r1 = -sum(x1.^2)
    r2 = -sum(x2.^2)
    ###########################

    p1 = q1 + r1
    p2 = q2 + r2

    f = p1 + p2

    f, zeros(0), zeros(0)
end

function PMM4_leader(x, y)
    k = min(length(x), length(y)) ÷ 2

    ###########################
    x1 = x[1:k]; x2 = x[k+1:end];
    y1 = y[1:k]; y2 = y[k+1:end];

    q1 = sum( (y1 - 10 ./ (1 .+ 2.5*x1.^2)).^2 )
    q2 = 10length(y2) + sum(y2.^2 - 10cos.(2π*y2))

    r1 = -sum(x1.^2)
    r2 = -sum(x2.^2)
    ###########################

    p1 = q1 + r1
    p2 = q2 + r2

    F = abs(p1) + abs(p2)

    F, zeros(0), zeros(0)
end

function PMM4_follower(x, y)
    k = min(length(x), length(y)) ÷ 2

    ###########################
    x1 = x[1:k]; x2 = x[k+1:end];
    y1 = y[1:k]; y2 = y[k+1:end];

    q1 = sum( (y1 - 10 ./ (1 .+ 2.5*x1.^2)).^2 )
    q2 = 10length(y2) + sum(y2.^2 - 10cos.(2π*y2))

    r1 = -sum(x1.^2)
    r2 = -sum(x2.^2)
    ###########################

    p1 = q1 + r1
    p2 = q2 + r2

    f = p1 + p2

    f, zeros(0), zeros(0)
end

function PMM5_leader(x, y)
    k = min(length(x), length(y)) ÷ 2

    ###########################
    x1 = x[1:k]; x2 = x[k+1:end];
    y1 = y[1:k]; y2 = y[k+1:end];

    q1 = 10k + sum( (y1-x1).^2 - 10cos.(2π*abs.(y1-x1)))
    q2 = sum( abs.(y2).^(2:length(y2)+1) )

    r1 = -sum(x1.^2)
    r2 = -sum(x2.^2)
    ###########################

    p1 = q1 + r1
    p2 = q2 + r2

    F = abs(p1) + abs(p2)

    F, zeros(0), zeros(0)
end

function PMM5_follower(x, y)
    k = min(length(x), length(y)) ÷ 2

    ###########################
    x1 = x[1:k]; x2 = x[k+1:end];
    y1 = y[1:k]; y2 = y[k+1:end];

    q1 = 10k + sum( (y1-x1).^2 - 10cos.(2π*abs.(y1-x1)))
    q2 = sum( abs.(y2).^(2:length(y2)+1) )

    r1 = -sum(x1.^2)
    r2 = -sum(x2.^2)
    ###########################

    p1 = q1 + r1
    p2 = q2 + r2

    f = p1 + p2

    f, zeros(0), zeros(0)
end

function PMM6_leader(x, y)
    k = min(length(x), length(y)) ÷ 2

    ###########################
    x1 = x[1:k]; x2 = x[k+1:end];
    y1 = y[1:k]; y2 = y[k+1:end];

    q1 = 1 + (1/4)*sum((y1 - x1).^2) - prod(cos.(10(y1 - x1) / sqrt.(1:k)))
    q2 = 10length(y2) + sum(y2.^2 - 10cos.(2π*y2))

    r1 = -sum(x1.^2)
    r2 = -10length(x2) - sum(x2.^2 - 10cos.(2π*x2))
    ###########################

    p1 = q1 + r1
    p2 = q2 + r2

    F = abs(p1) + abs(p2)

    F, zeros(0), zeros(0)
end

function PMM6_follower(x, y)
    k = min(length(x), length(y)) ÷ 2

    ###########################
    x1 = x[1:k]; x2 = x[k+1:end];
    y1 = y[1:k]; y2 = y[k+1:end];

    q1 = 1 + (1/4)*sum((y1 - x1).^2) - prod(cos.(10(y1 - x1) / sqrt.(1:k)))
    q2 = 10length(y2) + sum(y2.^2 - 10cos.(2π*y2))

    r1 = -sum(x1.^2)
    r2 = -10length(x2) - sum(x2.^2 - 10cos.(2π*x2))
    ###########################

    p1 = q1 + r1
    p2 = q2 + r2

    f = p1 + p2

    f, zeros(0), zeros(0)
end

function PMM1_optimum(uldim, lldim)
    k = min(uldim, lldim) ÷ 2
    x = zeros(uldim)
    y = zeros(lldim)
    y[k+1] = -10
    return x, y
end

function PMM2_optimum(uldim, lldim)
    k = min(uldim, lldim) ÷ 2
    x = zeros(uldim)
    y = zeros(lldim)
    y[k+1:length(y)] = sqrt.(k+1:length(y))
    return x, y
end

function PMM3_optimum(uldim, lldim)
    k = min(uldim, lldim) ÷ 2
    x = zeros(uldim)
    y = zeros(lldim)
    y[k+1:length(y)] .= 1
    return x, y
end

function PMM4_optimum(uldim, lldim)
    k = min(uldim, lldim) ÷ 2
    x = zeros(uldim)
    y = zeros(lldim)
    y[1:k] .= 10
    return x, y
end

function PMM5_optimum(uldim, lldim)
    k = min(uldim, lldim) ÷ 2
    x = zeros(uldim)
    y = zeros(lldim)
    return x, y
end

function PMM6_optimum(uldim, lldim)
    k = min(uldim, lldim) ÷ 2
    x = zeros(uldim)
    y = zeros(lldim)
    return x, y
end

#=
const PMM6_optimum = PMM_optimum
PMM1_optimum = PMM2_optimum = PMM3_optimum = PMM4_optimum = PMM5_optimum = PMM6_optimum
=#


text:
@recipe function f(t::PointCharge{T}) where {T}

    seriescolor --> :blue

    @series begin
        seriescolor --> seriescolor
        seriestype --> :scatter
        markersize --> 6
        label --> "PointCharge"
        t.locations
    end
end

@recipe function f(nbc::NBodyChargeCloud{T, N, SH}; connect = true, markersize = 10) where {T, N, SH}

    seriescolor --> :blue
    points = nbc.locations

    #Center point
    @series begin
        markersize := markersize
        label --> "NBodyChargeCloud"
        seriescolor --> seriescolor
        PointCharge{T}(points[1])
    end

    vertex_no = 1

    #Shells
    for shell in Base.OneTo(N)
        vertices = get_vertices(nbc.shell_structure)
        @series begin
            markersize := markersize * exp(-(shell))
            label --> ""
            connect --> connect
            seriescolor --> seriescolor
            nbc.shell_structure(points[vertex_no+1:vertex_no+vertices])
        end
        vertex_no += vertices
    end
end

@recipe function f(nbc::NBodyChargeCloud{T, N, NTuple{N, Int}}; connect = true, markersize = 10) where {T, N}

    seriescolor --> :blue
    points = nbc.locations
    seriestype := :scatter

    @series begin
        markersize := markersize
        label --> "NBodyChargeCloud"
        seriescolor --> seriescolor
        PointCharge{T}(points[1])
    end

    vertex_no = 1

    #Shells
    for shell in Base.OneTo(N)
        vertices = nbc.shell_structure[shell]
        @series begin
            markersize := markersize * exp(-(shell))
            label --> ""
            connect --> connect
            seriescolor --> seriescolor
            points[vertex_no+1:vertex_no+vertices]
        end
        vertex_no += vertices
    end
end

@recipe function f(t::Tetrahedron{T}; connect = true) where {T}

    linestyle --> :dot
    seriescolor --> :blue

    @series begin
        seriescolor --> seriescolor
        seriestype --> :scatter
        markersize --> 6
        label --> "Tetrahedron"
        t.locations
    end

    if connect
        @series begin
            seriescolor --> seriescolor
            seriestype --> :line3d
            label --> ""
            linewidth --> 1
            t.locations[vcat(1,2,3,4,2)]
        end

        @series begin
            seriescolor --> seriescolor
            seriestype --> :line3d
            label --> ""
            linewidth --> 1
            t.locations[vcat(3,1,4)]
        end
    end
end

@recipe function f(o::Octahedron{T}; connect = true) where {T}

    linestyle --> :dot
    seriescolor --> :blue

    @series begin
        seriescolor --> seriescolor
        seriestype --> :scatter
        markersize --> 6
        label --> "Octahedron"
        o.locations
    end

    if connect
        @series begin
            seriescolor --> seriescolor
            seriestype --> :line3d
            label --> ""
            linewidth --> 1
            o.locations[vcat(1,2,3,4,5,2,6,5,1,3,6,4,1)]
        end
    end
end

@recipe function f(h::Hexahedron{T}; connect = true) where {T}

    linestyle --> :dot
    seriescolor --> :blue

    @series begin
        seriescolor --> seriescolor
        seriestype --> :scatter
        markersize --> 6
        label --> "Hexahedron"
        h.locations
    end

    if connect
        @series begin
            seriescolor --> seriescolor
            seriestype --> :line3d
            label --> ""
            linewidth --> 1
            h.locations[vcat(1,2,5,3,6,4,7,8,5)]
        end

        @series begin
            seriescolor --> seriescolor
            seriestype --> :line3d
            label --> ""
            linewidth --> 1
            h.locations[vcat(3,1,4)]
        end

        @series begin
            seriescolor --> seriescolor
            seriestype --> :line3d
            label --> ""
            linewidth --> 1
            h.locations[vcat(2,7)]
        end

        @series begin
            seriescolor --> seriescolor
            seriestype --> :line3d
            label --> ""
            linewidth --> 1
            h.locations[vcat(6,8)]
        end
    end
end

@recipe function f(i::Icosahedron{T}; connect = true) where {T}

    linestyle --> :dot
    seriescolor --> :blue

    @series begin
        seriescolor --> seriescolor
        seriestype --> :scatter
        markersize --> 6
        label --> "Icosahedron"
        i.locations
    end

    if connect
        @series begin
            seriescolor --> seriescolor
            seriestype --> :line3d
            label --> ""
            linewidth --> 1
            i.locations[vcat(1,2,7,8,9,10,11,7,3,8,4,9,5,10,6,11,2,3,4,5,6,2)]
        end

        @series begin
            seriescolor --> seriescolor
            seriestype --> :line3d
            label --> ""
            linewidth --> 1
            i.locations[vcat(3,1,6)]
        end

        @series begin
            seriescolor --> seriescolor
            seriestype --> :line3d
            label --> ""
            linewidth --> 1
            i.locations[vcat(5,1,4)]
        end

        @series begin
            seriescolor --> seriescolor
            seriestype --> :line3d
            label --> ""
            linewidth --> 1
            i.locations[vcat(7,12,8)]
        end

        @series begin
            seriescolor --> seriescolor
            seriestype --> :line3d
            label --> ""
            linewidth --> 1
            i.locations[vcat(9,12,10)]
        end

        @series begin
            seriescolor --> seriescolor
            seriestype --> :line3d
            label --> ""
            linewidth --> 1
            i.locations[vcat(11,12)]
        end
    end
end

@recipe function f(d::Dodecahedron{T}; connect = true) where {T}

    linestyle --> :dot
    seriescolor --> :blue

    @series begin
        seriescolor --> seriescolor
        seriestype --> :scatter
        markersize --> 6
        label --> "Dodecahedron"
        d.locations
    end

    if connect
        @series begin
            seriescolor --> seriescolor
            seriestype --> :line3d
            label --> ""
            linewidth --> 1
            d.locations[vcat(1,2,3,4,5,1,6,11,7,12,8,13,9,14,10,15,6)]
        end

        @series begin
            seriescolor --> seriescolor
            seriestype --> :line3d
            label --> ""
            linewidth --> 1
            d.locations[vcat(15,20,16,17,18,19,20)]
        end

        for j in [2,3,4,5,11,12,13,14]
            @series begin
                seriescolor --> seriescolor
                seriestype --> :line3d
                label --> ""
                linewidth --> 1
                d.locations[vcat(j,j+5)]
            end
        end
    end
end


text:
GQL_FRAGMENT_VARIABLES = """
  fragment ppe_fields on Ppe {
    solveKey
    suggested
    max
    mean
    lastUpdatedTimestamp
  }
  fragment solverdata_fields on SolverData {
    solveKey
    BayesNetOutVertIDs
    BayesNetVertID
    dimIDs
    dimbw
    dims
    dimval
    dontmargin
    eliminated
    infoPerCoord
    initialized
    ismargin
    separator
    solveInProgress
    solvedCount
    variableType
    vecbw
    vecval
    _version
  }
  fragment variable_skeleton_fields on Variable {
    label
    tags
  }
  fragment variable_summary_fields on Variable {
    timestamp
    ppes {
      ...ppe_fields
    }
    variableType
    _version
  }
  fragment variable_full_fields on Variable{
    smallData
    solvable
    solverData
    {
      ...solverdata_fields
    }
  }
  """

GQL_GETVARIABLE = """
  query sdk_get_variable(
      \$userId: ID!,
      \$robotId: ID!,
      \$sessionId: ID!,
      \$label: ID!) {
    users(where:{id:\$userId}){
      robots(where:{id: \$robotId}) {
        sessions(where:{id: \$sessionId}) {
          variables(where:{label: \$label}) {
            ...variable_skeleton_fields
            ...variable_summary_fields
            ...variable_full_fields
          }
        }
      }
    }
  }"""

GQL_GETVARIABLES = """
  query sdk_get_variables(
      \$userId: ID!,
      \$robotId: ID!,
      \$sessionId: ID!,
      \$fields_summary: Boolean! = false,
      \$fields_full: Boolean! = false){
    users(where:{id: \$userId}) {
      name
      robots(where:{id: \$robotId}) {
        name
        sessions(where:{id: \$sessionId}){
          name
          variables {
            ...variable_skeleton_fields
            ...variable_summary_fields @include(if: \$fields_summary)
            ...variable_full_fields @include(if: \$fields_full)
          }
        }
      }
    }
  }"""

GQL_GETVARIABLESFILTERED = """
  query sdk_get_variables_filtered(
      \$userId: ID!,
      \$robotIds: [ID!]!,
      \$sessionIds: [ID!]!,
      \$variable_label_regexp: String = ".*",
      \$variable_tags: [String] = ["VARIABLE"],
      \$solvable: Int! = 0,
      \$fields_summary: Boolean! = false,
      \$fields_full: Boolean! = false){
    users(where:{id: \$userId}) {
      name
      robots(where:{id_in: \$robotIds}) {
        name
        sessions(where:{id_in: \$sessionIds}){
          name
          variables(where:{
              label_MATCHES: \$variable_label_regexp,
              tags: \$variable_tags,
              solvable_GTE: \$solvable}) {
            ...variable_skeleton_fields
            ...variable_summary_fields @include(if: \$fields_summary)
            ...variable_full_fields @include(if: \$fields_full)
          }
        }
      }
    }
  }"""


text:
# LSRTM of Marmousi using stochastic gradient descent
# Author: Philipp Witte, pwitte@eoas.ubc.ca
# Date: January 2017

using PyCall,PyPlot, HDF5, opesciSLIM.TimeModeling, opesciSLIM.LeastSquaresMigration, SeisIO, JLD

# Fetch data and starting model from ftp server
~isfile("marmousi_2D.segy") && run(`wget ftp://ftp.slim.gatech.edu/data/SoftwareRelease/Imaging.jl/2DLSRTM/marmousi_2D.segy`)
~isfile("marmousi_migration_velocity.h5") && run(`wget ftp://ftp.slim.gatech.edu/data/SoftwareRelease/Imaging.jl/2DLSRTM/marmousi_migration_velocity.h5`)

# Load migration velocity model
n,d,o,m0 = read(h5open("marmousi_migration_velocity.h5","r"), "n", "d", "o", "m0")

# Load data
block = segy_read("marmousi_2D.segy")
dD = joData(block)

# Set up model structure
model0 = Model((n[1],n[2]), (d[1],d[2]), (o[1],o[2]), m0)

# Set up wavelet
src_geometry = Geometry(block; key="source", segy_depth_key="SourceDepth")
wavelet = source(src_geometry.t[1],src_geometry.dt[1],0.03)	# 30 Hz wavelet
q = joData(src_geometry,wavelet)

# Set up info structure
ntComp = get_computational_nt(q.geometry,dD.geometry,model0)	# no. of computational time steps
info = Info(prod(model0.n),dD.nsrc,ntComp)

###################################################################################################

# Setup operators
F = joModeling(info,model0,q.geometry,dD.geometry)
J = joJacobian(F,q)

# Right-hand preconditioners (model topmute)
Mr = opTopmute(model0.n,42,10)

# Stochastic gradient
x = zeros(Float32,info.n)
batchsize = 10
niter = 32
fval = zeros(Float32,niter)

# Main loop
for j=1:niter
	println("Iteration: ", j)

	# Select batch and set up left-hand preconditioner
	idx = randperm(dD.nsrc)[1:batchsize]
	Jsub = subsample(J,idx)
	dsub = subsample(dD,idx)
	Ml = opMarineTopmute(30,dsub.geometry)	# data topmute

	# Compute residual and gradient
	r = Ml*Jsub*Mr*x - Ml*dsub
	g = Mr'*Jsub'*Ml'*r

	# Step size and update variable
	fval[j] = .5*norm(r)^2
	t = norm(r)^2/norm(g)^2
	x -= t*g

	save("snapshot.jld","x",reshape(x,model0.n))

end

# Save final velocity model, function value and history
h5open("lsrtm_marmousi_2D_result.h5", "w") do file
	write(file, "x", reshape(x,model0.n), "fval", fval)
end

max_stars_repo_name: SINBADconsortium/opesciSLIM

max_issues_repo_name: SINBADconsortium/opesciSLIM

max_forks_repo_name: SINBADconsortium/opesciSLIM


text:
# Julia code
# Copyright (c) 2022 the SLICOTMath.jl developers
# Portions extracted from SLICOT-Reference distribution:
# Copyright (c) 2002-2020 NICONET e.V.

function run_mb03fz(datfile, io=stdout)
    NIN = 5
    NOUT = 6
    NMAX = 50
    LDB  = NMAX
    LDC =   NMAX
    LDD = NMAX
    LDFG = NMAX
    LDQ = 2*NMAX
    LDU = NMAX
    LDWORK = 18*NMAX*NMAX + NMAX + 3 + max( 2*NMAX, 24 )
    LDZ  = NMAX
    LIWORK = 2*NMAX + 9
    LZWORK = 8*NMAX + 28
    Z = Array{ComplexF64,2}(undef, LDZ,NMAX)
    B = Array{ComplexF64,2}(undef, LDB,NMAX)
    FG = Array{ComplexF64,2}(undef, LDFG,NMAX)
    D = Array{ComplexF64,2}(undef, LDD,NMAX)
    C = Array{ComplexF64,2}(undef, LDC,NMAX)
    Q = Array{ComplexF64,2}(undef, LDQ,2*NMAX)
    U = Array{ComplexF64,2}(undef, LDU,2*NMAX)
    ALPHAR = Array{Float64,1}(undef, NMAX)
    ALPHAI = Array{Float64,1}(undef, NMAX)
    BETA = Array{Float64,1}(undef, NMAX)
    BWORK = Array{BlasBool,1}(undef, NMAX)
    ZWORK = Array{ComplexF64,1}(undef, LZWORK)
    DWORK = Array{Float64,1}(undef, LDWORK)
    IWORK = Array{BlasInt,1}(undef, LIWORK)
    f = open(datfile,"r")
    readline(f)
    vs = split(readline(f))
    COMPQ =  vs[1][1]
    COMPU =  vs[2][1]
    ORTH =  vs[3][1]
    N = parse(BlasInt, vs[4])
if ( N<0 || N>NMAX || mod( N, 2 )!=0 )
else
    M = N÷2

    vs = String[]
    _isz,_jsz = (N,N)
    while length(vs) < _isz*_jsz
        append!(vs, replace.(split(readline(f)),'D'=>'E'))
    end
    for i in 1:_isz
       _i0 = (i-1)*_jsz
       Z[i,1:_jsz] .= parsex.(ComplexF64, vs[_i0+1:_i0+_jsz])
    end

    vs = String[]
    _isz,_jsz = (M,M)
    while length(vs) < _isz*_jsz
        append!(vs, replace.(split(readline(f)),'D'=>'E'))
    end
    for i in 1:_isz
       _i0 = (i-1)*_jsz
       B[i,1:_jsz] .= parsex.(ComplexF64, vs[_i0+1:_i0+_jsz])
    end

    vs = String[]
    _isz,_jsz = (M,M+1)
    while length(vs) < _isz*_jsz
        append!(vs, replace.(split(readline(f)),'D'=>'E'))
    end
    for i in 1:_isz
       _i0 = (i-1)*_jsz
       FG[i,1:_jsz] .= parsex.(ComplexF64, vs[_i0+1:_i0+_jsz])
    end
# interp call 1

    NEIG, INFO = SLICOT.mb03fz!(COMPQ, COMPU, ORTH, N, Z, B, FG, D, C, Q, U, ALPHAR, ALPHAI, BETA, LIWORK, BWORK)
    println(io, "NEIG = $NEIG")
    @test INFO == 0
    INFO == 0 || return

if ( INFO!=0 )
else
# interp output 1
    println(io,"Z:")
    _nc = N
    _nr = N
    show(io,"text/plain",Z[1:_nr,1:_nc])
    println(io,)

if ( LSAME( COMPQ, 'C' ) || LSAME( COMPU, 'C' ) )
# interp output 2
    println(io,"D:")
    _nc = N
    _nr = N
    show(io,"text/plain",D[1:_nr,1:_nc])
    println(io,)

# interp output 3
    println(io,"C:")
    _nc = N
    _nr = N
    show(io,"text/plain",C[1:_nr,1:_nc])
    println(io,)

# interp output 4
    println(io,"B:")
    _nc = N
    _nr = N
    show(io,"text/plain",B[1:_nr,1:_nc])
    println(io,)

# interp output 5
    println(io,"FG:")
    _nc = N
    _nr = N
    show(io,"text/plain",FG[1:_nr,1:_nc])
    println(io,)

end # if
# interp output 6
    println(io,"ALPHAR:")
    _nr = N
    show(io,"text/plain",ALPHAR[1:_nr])
    println(io,)

# interp output 7
    println(io,"ALPHAI:")
    _nr = N
    show(io,"text/plain",ALPHAI[1:_nr])
    println(io,)

# interp output 8
    println(io,"BETA:")
    _nr = N
    show(io,"text/plain",BETA[1:_nr])
    println(io,)

if ( LSAME( COMPQ, 'C' ) && NEIG>0 )
# interp output 9
    println(io,"Q:")
    _nc = NEIG
    _nr = N
    show(io,"text/plain",Q[1:_nr,1:_nc])
    println(io,)

end # if
if ( LSAME( COMPU, 'C' ) && NEIG>0 )
# interp output 10
    println(io,"U:")
    _nc = NEIG
    _nr = N
    show(io,"text/plain",U[1:_nr,1:_nc])
    println(io,)

end # if
end # if
end # if
    close(f)
end # run_X()


text:
using OrdinaryDiffEq, DiffEqDevTools, Test, LinearAlgebra
using DiffEqProblemLibrary.ODEProblemLibrary: importodeproblems; importodeproblems()
import DiffEqProblemLibrary.ODEProblemLibrary: prob_ode_linear, prob_ode_2Dlinear, van

testTol = 0.2

for prob in [prob_ode_linear, prob_ode_2Dlinear]
  sim21 = test_convergence(1 .//2 .^(6:-1:3),prob,RadauIIA5())
  @test sim21.𝒪est[:final] ≈ 5 atol=testTol
end

# test adaptivity
for iip in (true, false)
  if iip
    vanstiff = ODEProblem{iip}(van, [0;sqrt(3)], (0.0,1.0), 1e6)
  else
    vanstiff = ODEProblem{false}((u,p,t)->van(u,p,t), [0;sqrt(3)], (0.0,1.0), 1e6)
  end
  sol = solve(vanstiff, RadauIIA5())
  if iip
    @test sol.destats.naccept + sol.destats.nreject > sol.destats.njacs # J reuse
    @test sol.destats.njacs < sol.destats.nw # W reuse
  end
  @test length(sol) < 150
  @test length(solve(remake(vanstiff, p=1e7), RadauIIA5())) < 150
  @test length(solve(remake(vanstiff, p=1e7), reltol=[1e-4, 1e-6], RadauIIA5())) < 170
  @test length(solve(remake(vanstiff, p=1e7), RadauIIA5(), reltol=1e-9, abstol=1e-9)) < 870
  @test length(solve(remake(vanstiff, p=1e9), RadauIIA5())) < 170
  @test length(solve(remake(vanstiff, p=1e10), RadauIIA5())) < 190
end

##Tests fot RadauIIA3

f = (u,p,t)->cos(t)
prob_ode_sin = ODEProblem(ODEFunction(f; analytic=(u0,p,t)->sin(t)), 0.,(0.0,1.0))

f = (u,p,t)-> 1.01*u
prob_ode_exp = ODEProblem(ODEFunction(f;analytic=(u0,p,t)->u0*exp(1.01*t)), 1.01,(0.,1.0))

probs_oop = [prob_ode_sin, prob_ode_exp]

alg = RadauIIA3()

dts = 1 ./2 .^(8:-1:1)

for prob in [probs_oop; prob_ode_2Dlinear]
  sim = test_convergence(dts, prob, alg)
  @test sim.𝒪est[:final] ≈ 3 atol = 0.25
end


text:
using RPOMDPs
using RPOMDPModels
# using POMDPToolbox

# ambiguity set examples
using JuMP, Clp
mins = [0.05, 0.6, 0.05]
maxs = [0.3, 0.9, 0.4]
m = Model(solver = ClpSolver())
@variable(m, mins[i] <= p[i = 1:3] <= maxs[i])
@constraint(m, sum(p) == 1)
@objective(m, Min, 5 * p[1] + 10 * p[2] - 3 * p[3])
solve(m)
argmin = getvalue(p)

using JuMP, Clp
using BenchmarkTools

# Lower Prob (5 vertices)
# p[1] = pq, p[2] = pc, p[3] = py
mlp = Model(solver = ClpSolver())
@variable(mlp, 0 <= p[i = 1:3] <= 1)
@objective(mlp, Min, 5 * p[1] + 10 * p[2] - 3 * p[3])
@constraint(mlp, sum(p) == 1)
@constraint(mlp, 0.15 <= p[1] <= 0.35)
@constraint(mlp, 0.55 <= p[2] <= 0.7)
@constraint(mlp, 0.05 <= p[3] <= 0.3)
solve(mlp)
argmin_mlp = getvalue(p)
obj_mlp = getobjectivevalue(mlp)

# P-box (4 vertices)
mpb = Model(solver = ClpSolver())
@variable(mpb, 0 <= p[i = 1:3] <= 1)
@objective(mpb, Min, 5 * p[1] + 10 * p[2] - 3 * p[3])
@constraint(mpb, sum(p) == 1)
@constraint(mpb, 0.15 <= p[1] <= 0.35)
@constraint(mpb, 0.05 <= p[3] <= 0.3)
solve(mpb)
argmin_mpb = getvalue(p)
obj_mpb = getobjectivevalue(mpb)

# Polytopic (3 vertices)
m1 = (0.05 - 0.3) / (0.6 - 0.55)
m2 = (0.1 - 0.3) / (0.7 - 0.55)
m3 = (0.05 - 0.1) / (0.6 - 0.7)
y1(c) = m1 * (c - 0.55) + 0.3
y2(c) = m2 * (c - 0.55) + 0.3
y3(c) = m3 * (c - 0.7) + 0.1
mp = Model(solver = ClpSolver())
@variable(mp, 0 <= p[i = 1:3] <= 1)
@objective(mp, Min, 5 * p[1] + 10 * p[2] - 3 * p[3])
@constraint(mp, sum(p) == 1)
@constraint(mp, p[3] >= y1(p[2]))
@constraint(mp, p[3] <= y2(p[2]))
@constraint(mp, p[3] >= y3(p[2]))
solve(mp)
argmin_mp = getvalue(p)
obj_mp = getobjectivevalue(mp)

(argmin_mlp, argmin_mpb, argmin_mp)
(obj_mlp, obj_mpb, obj_mp)

# benchmarks
@benchmark solve($mlp)
@benchmark solve($mpb)
@benchmark solve($mp)

# scratch

temp_array = Array{Float64}(2, 2)
temp_array[1,1] = 10.0

rpomdp = Baby3RPOMDP()
t_array = Array{Float64}(2, 2)
t_array[1, 1] = BoolDistribution(0.0)
t_array[1, 2] = BoolDistribution(1.0)
t_array[2, 1] = BoolDistribution(0.0)
t_array[2, 2] = BoolDistribution(0.1)
t_array[state_index(rpomdp, s), action_index(rpomdp, a)]
t_array = [BoolDistribution(0.0) BoolDistribution(1.0);
    BoolDistribution(0.0) BoolDistribution(0.1)]

prob = Baby3RPOMDP()
transition(prob, false, false)
observation(prob, false, true, "upper")
RPOMDPModels.observation_index(prob, :crying)

# DRO toy problem
# x ∈ R^3
# E(x) = [1,2,3]
# P(x ∈ [(0,5),(0,5),(0,5)]) ∈ [1,1]
# P(x ∈ [(0,3),(0,3),(0,3)]) ∈ [0.7, 0.9]

# Convert from v-rep to h-rep for JUMP model (polytopic ambiguity set)
using Polyhedra, CDDLib
using JuMP, Clp
polyv = vrep([0.1 0.1 0.8; 0.1 0.3 0.6; 0.3 0.3 0.4])
poly = polyhedron(polyv, CDDLibrary())
polyh = hrep(poly)
polyhv = polyhedron(polyh, CDDLibrary())
polyhvv = vrep(polyhv)
ph = MixedMatHRep(polyh)
pv = MixedMatVRep(polyhvv)

mpoly = Model(solver = ClpSolver())
@variable(mpoly, 0 <= p[i = 1:3] <= 1)
@objective(mpoly, Min, 5 * p[1] + 10 * p[2] - 3 * p[3])
@constraint(mpoly, sum(p) == 1)
@constraint(mpoly, ph.A * p .<= ph.b)
solve(mpoly)
argmin_mpoly = getvalue(p)
obj_mpoly = getobjectivevalue(mpoly)


text:
using PyPlot
using Glob
using JLD2

defaultcolors = plt.rcParams["axes.prop_cycle"].by_key()["color"]
removespine(side) = gca().spines[side].set_visible(false)
removespines(sides...) = [removespine(side) for side in sides]

include("ConvergenceTests/ConvergenceTests.jl")

using .ConvergenceTests
using .ConvergenceTests: compute_errors, extract_sizes

filenames = glob("taylor_green*.jld2", joinpath(@__DIR__, "data"))

errors = compute_errors(
            (x, y, z, t) -> ConvergenceTests.DoublyPeriodicTaylorGreen.u(x, y, t),
            filenames...)

sizes = extract_sizes(filenames...)

Nx = map(sz -> sz[1], sizes)
L₁ = map(err -> err.L₁, errors)
L∞ = map(err -> err.L∞, errors)

close("all")
fig, ax = subplots()

ax.loglog(Nx, L₁, basex=2, linestyle="None", marker="o", label="error, \$L_1\$-norm")
ax.loglog(Nx, L∞, basex=2, linestyle="None", marker="^", label="error, \$L_\\infty\$-norm")
ax.loglog(Nx, L₁[end] * (Nx[end] ./ Nx).^2, "k-", basex=2, linewidth=1, alpha=0.6,  label=L"\sim N_x^{-2}")

legend()

title("Convergence for freely-decaying Taylor-Green vortex")
removespines("top", "right")
ylabel("Norms of the absolute error, \$ | u_\\mathrm{simulation} - u_\\mathrm{analytical} | \$")
xlabel(L"N_x")

filepath = joinpath(@__DIR__, "figs", "taylor_green_convergence.png")
savefig(filepath, dpi=480)

test_rate_of_convergence(L₁, Nx, expected=-2.0, atol=0.001, name="Taylor-Green L₁")
test_rate_of_convergence(L∞, Nx, expected=-2.0, atol=0.050, name="Taylor-Green L∞")


text:
using DFTK
import DFTK: mpi_mean!
using MPI
using Test
using LinearAlgebra: norm

include("testcases.jl")

function test_chi0(testcase; symmetries=false, temperature=0,
                   spin_polarization=:none, eigensolver=lobpcg_hyper, Ecut=10,
                   kgrid=[3, 1, 1], fft_size=[15, 1, 15], compute_full_χ0=false)

    tol      = 1e-11
    ε        = 1e-6
    testtol  = 2e-6
    n_ep_extra = 3

    collinear = spin_polarization == :collinear
    is_metal = !isnothing(testcase.temperature)
    eigsol = eigensolver == lobpcg_hyper
    label = [
        is_metal        ? "    metal" : "insulator",
        eigsol          ? "   lobpcg" : "full diag",
        symmetries      ? "   symm" : "no symm",
        temperature > 0 ? "temp" : "  0K",
        collinear       ? "coll" : "none",
    ]
    @testset "Computing χ0 ($(join(label, ", ")))" begin
        spec = ElementPsp(testcase.atnum, psp=load_psp(testcase.psp))
        magnetic_moments = collinear ? [0.3, 0.7] : []
        model_kwargs = (; temperature, symmetries, magnetic_moments, spin_polarization)
        basis_kwargs = (; kgrid, fft_size, Ecut)

        model = model_LDA(testcase.lattice, testcase.atoms, testcase.positions;
                          model_kwargs...)
        basis = PlaneWaveBasis(model; basis_kwargs...)
        ρ0 = guess_density(basis, magnetic_moments)
        energies, ham0 = energy_hamiltonian(basis, nothing, nothing; ρ=ρ0)
        res = DFTK.next_density(ham0; tol, n_ep_extra, eigensolver)
        occ, εF = DFTK.compute_occupation(basis, res.eigenvalues)
        scfres = (ham=ham0, res..., n_ep_extra)

        # create external small perturbation εδV
        n_spin = model.n_spin_components
        δV = randn(eltype(basis), basis.fft_size..., n_spin)
        mpi_mean!(δV, MPI.COMM_WORLD)
        δV_sym = DFTK.symmetrize_ρ(basis, δV, symmetries=model.symmetries)
        if symmetries
            δV = δV_sym
        else
            @test δV_sym ≈ δV
        end

        function compute_ρ_FD(ε)
            term_builder = basis -> DFTK.TermExternal(ε * δV)
            model = model_LDA(testcase.lattice, testcase.atoms, testcase.positions;
                              model_kwargs..., extra_terms=[term_builder])
            basis = PlaneWaveBasis(model; basis_kwargs...)
            energies, ham = energy_hamiltonian(basis, nothing, nothing; ρ=ρ0)
            res = DFTK.next_density(ham; tol, n_ep_extra, eigensolver)
            res.ρout
        end

        # middle point finite difference for more precision
        ρ1 = compute_ρ_FD(-ε)
        ρ2 = compute_ρ_FD(ε)
        diff_findiff = (ρ2 - ρ1) / (2ε)

        # Test apply_χ0 and compare against finite differences
        diff_applied_χ0 = apply_χ0(scfres, δV)
        @test norm(diff_findiff - diff_applied_χ0) < testtol

        # just to cover it here
        if temperature > 0
            D = compute_dos(εF, basis, res.eigenvalues)
            LDOS = compute_ldos(εF, basis, res.eigenvalues, res.ψ)
        end

        if !symmetries
            #  Test compute_χ0 against finite differences
            #  (only works in reasonable time for small Ecut)
            if compute_full_χ0
                χ0 = compute_χ0(ham0)
                diff_computed_χ0 = reshape(χ0 * vec(δV), basis.fft_size..., n_spin)
                @test norm(diff_findiff - diff_computed_χ0) < testtol
            end

            # Test that apply_χ0 is self-adjoint
            δV1 = randn(eltype(basis), basis.fft_size..., n_spin)
            δV2 = randn(eltype(basis), basis.fft_size..., n_spin)
            mpi_mean!(δV1, MPI.COMM_WORLD)
            mpi_mean!(δV2, MPI.COMM_WORLD)

            χ0δV1 = apply_χ0(scfres, δV1)
            χ0δV2 = apply_χ0(scfres, δV2)
            @test abs(dot(δV1, χ0δV2) - dot(δV2, χ0δV1)) < testtol
        end
    end
end

@testset "Computing χ0" begin
    for (case, temperatures) in [(silicon, (0, 0.03)), (magnesium, (0.01, ))]
        for temperature in temperatures, spin_polarization in (:none, :collinear)
            for symmetries in (false, true)
                test_chi0(case; symmetries, temperature, spin_polarization)
            end
        end
    end

    # additional test for compute_χ0
    for spin_polarization in (:none, :collinear)
        test_chi0(silicon; symmetries=false, spin_polarization,
                  eigensolver=diag_full, Ecut=3, fft_size=[10, 1, 10],
                  compute_full_χ0=true)
    end
end


text:
"""
    Consistent(predictions::AbstractVector)

Create an object that can be used for resampling `predictions` and corresponding targets.

The set of predictions has to be provided as a vector consisting of
- probabilities,
- vectors of probabilities summing up to 1,
- or distributions that can be sampled from.

Consistency resampling can be performed with `rand` and `rand!`, as described in the
documentation of [`Random`](https://docs.julialang.org/en/v1/stdlib/Random/).

# Examples

```jldoctest
julia> predictions = rand(100);

julia> consistent = Consistent(predictions);

julia> rand(consistent) isa Tuple{Float64,Bool}
true

julia> first(rand(consistent)) in predictions
true
```

!!! note
    One can use `Random.Sampler` to create a sampler that is optimized for the
    generation of multiple samples, e.g., by building an alias table:
    ```jldoctest
    julia> using Random: Sampler, GLOBAL_RNG

    julia> sampler = Sampler(GLOBAL_RNG, Consistent(rand(100)));

    julia> rand(sampler) isa Tuple{Float64,Bool}
    true
    ```
    However, for small number of samples the setup cost can outweigh the
    improved sampling efficiency.

!!! note
    Multiple samples are returned as an array of tuples of predictions and corresponding
    targets. If you prefer a tuple of arrays of predictions and targets instead, you can use
    a package such as
    [StructArrays](https://github.com/JuliaArrays/StructArrays.jl) and perform in-place
    sampling:
    ```jldoctest
    julia> using StructArrays, Random

    julia> predictions = Vector{Float64}(undef, 100);

    julia> targets = Vector{Bool}(undef, 100);

    julia> rand!(StructVector((predictions, targets)), Consistent(rand(100)));
    ```

# References

Bröcker, J. and Smith, L.A., 2007. [Increasing the reliability of reliability diagrams](https://doi.org/10.1175/WAF993.1). Weather and forecasting, 22(3), pp. 651-661.
"""
struct Consistent{T,S,P<:AbstractVector{T}}
    predictions::P

    function Consistent{T,S}(predictions::AbstractVector{T}) where {T,S}
        return new{T,S,typeof(predictions)}(predictions)
    end
end

function Consistent(predictions::AbstractVector{<:Real})
    isempty(predictions) && error("at least one prediction is required")
    isprobabilities = all(predictions) do p
        zero(p) <= p <= one(p)
    end
    isprobabilities || error("predictions have to be probabilities")
    return Consistent{eltype(predictions),Bool}(predictions)
end

function Consistent(predictions::AbstractVector{<:AbstractVector{<:Real}})
    isempty(predictions) && error("at least one prediction is required")
    isnormalized = all(predictions) do p
        sum_p = sum(p)
        return sum_p ≈ one(sum_p)
    end
    isnormalized || error("predictions have to be normalized")
    return Consistent{eltype(predictions),Int}(predictions)
end

# fallback: sample one target to compute the type of targets
# `eltype` does not work reliably for determining e.g. the type of
# `rand(d::Distributions.Distribution)`
function Consistent(predictions::AbstractVector)
    isempty(predictions) && error("at least one prediction is required")
    S = typeof(rand(first(predictions)))
    return Consistent{eltype(predictions),S}(predictions)
end

Base.eltype(::Type{<:Consistent{T,S}}) where {T,S} = Tuple{T,S}


text:
@doc raw"""
`PSDE`: Stratonovich Partitioned Stochastic Differential Equation

Defines a partitioned stochastic differential initial value problem
```math
\begin{aligned}
dq (t) &= v(t, q(t)) \, dt + B(t, q(t)) \circ dW , & q(t_{0}) &= q_{0} , \\
dp (t) &= f(t, q(t)) \, dt + G(t, q(t)) \circ dW , & p(t_{0}) &= p_{0}
\end{aligned}
```
with the drift vector fields ``v`` and ``f``, diffusion matrices ``B`` and ``G``,
initial conditions ``q_{0}`` and ``p_{0}``, the dynamical variables ``(q,p)`` taking
values in ``\mathbb{R}^{d} \times \mathbb{R}^{d}``, and the m-dimensional Wiener process W

### Parameters

* `DT <: Number`: data type
* `TT <: Real`: time step type
* `AT <: AbstractArray{DT}`: array type
* `vType <: Function`: type of `v`
* `fType <: Function`: type of `f`
* `BType <: Function`: type of `B`
* `GType <: Function`: type of `G`
* `pType <: Union{NamedTuple,Nothing}`: parameters type

### Fields

* `d`:  dimension of dynamical variable ``q`` and the vector field ``v``
* `m`:  dimension of the Wiener process
* `ns`: number of sample paths
* `v`:  function computing the drift vector field for the position variable ``q``
* `f`:  function computing the drift vector field for the momentum variable ``p``
* `B`:  function computing the d x m diffusion matrix for the position variable ``q``
* `G`:  function computing the d x m diffusion matrix for the momentum variable ``p``
* `t₀`: initial time
* `q₀`: initial condition for dynamical variable ``q`` (may be a random variable itself)
* `p₀`: initial condition for dynamical variable ``p`` (may be a random variable itself)
* `parameters`: either a `NamedTuple` containing the equations parameters or `nothing`
* `periodicity`: determines the periodicity of the state vector `q` for cutting periodic solutions

The functions `v`, `f`, `B` and `G`, providing the drift vector fields and diffusion matrices, take four arguments,
`v(t, q, p, v)`, `f(t, q, p, f)`, `B(t, q, p,  B)` and `G(t, q, p, G)`, where `t` is the current time, `(q, p)` is the
current solution vector, and `v`, `f`, `B` and `G` are the variables which hold the result
of evaluating the vector fields ``v``, ``f`` and the matrices ``B``, ``G`` on `t` and `(q,p)`.

### Constructors

```julia
PSDE(m, ns, v, f, B, G, t₀, q₀, p₀; parameters=NullParameters(), periodicity=zero(q₀[begin]))
PSDE(m, ns, v, f, B, G, q₀::StateVector, p₀::StateVector; kwargs...) = PSDE(m, ns, v, f, B, G, 0.0, q₀, p₀; kwargs...)
PSDE(m, ns, v, f, B, G, t₀, q₀::State, p₀::State; kwargs...) = PSDE(m, ns, v, f, B, G, t₀, [q₀], [p₀]; kwargs...)
PSDE(m, ns, v, f, B, G, q₀::State, p₀::State; kwargs...) = PSDE(m, ns, v, f, B, G, 0.0, q₀, p₀; kwargs...)
```

### Example

```julia
    function v(λ, t, q, v)
        v[1] = λ*q[1]
        v[2] = λ*q[2]
    end

    function B(μ, t, q, B)
        B[1] = μ*q[1]
        B[2] = μ*q[2]
    end

    t₀ = 0.
    q₀ = [1., 1.]
    λ  = 2.
    μ  = 1.

    v_sde = (t, q, v) -> v(λ, t, q, v)
    B_sde = (t, q, B) -> B(μ, t, q, B)

    sde = SDE(v_sde, B_sde, t₀, q₀)
```
"""
struct PSDE{dType <: Number, tType <: Real, arrayType <: AbstractArray{dType},
            vType <: Function, fType <: Function,
            BType <: Function, GType <: Function,
            pType <: OptionalParameters} <: AbstractEquationPSDE{Nothing,Nothing,Nothing}

    d::Int
    m::Int
    ns::Int
    v::vType
    f::fType
    B::BType
    G::GType
    t₀::tType
    q₀::Vector{arrayType}
    p₀::Vector{arrayType}
    parameters::pType
    periodicity::arrayType

    function PSDE(m, ns, v::vType, f::fType, B::BType, G::GType,
                  t₀::tType, q₀::Vector{arrayType}, p₀::Vector{arrayType};
                  parameters=NullParameters(), periodicity=zero(q₀[begin])) where {
                        dType <: Number, tType <: Real, arrayType <: AbstractArray{dType},
                        vType <: Function, fType <: Function,
                        BType <: Function, GType <: Function}

        d  = length(q₀[begin])
        ni = length(q₀)

        @assert length(q₀) == length(p₀)
        @assert all(length(q) == d for q in q₀)
        @assert all(length(p) == d for p in p₀)

        @assert ns ≥ 1
        @assert ni == 1 || ns == 1
        # either multiple deterministic initial conditions and one sample path
        # or one deterministic initial condition and multiple sample paths

        new{dType, tType, arrayType, vType, fType, BType, GType, typeof(parameters)}(d, m, ns, v, f, B, G, t₀, q₀, p₀, parameters, periodicity)
    end
end

PSDE(m, ns, v, f, B, G, q₀::StateVector, p₀::StateVector; kwargs...) = PSDE(m, ns, v, f, B, G, 0.0, q₀, p₀; kwargs...)
PSDE(m, ns, v, f, B, G, t₀, q₀::State, p₀::State; kwargs...) = PSDE(m, ns, v, f, B, G, t₀, [q₀], [p₀]; kwargs...)
PSDE(m, ns, v, f, B, G, q₀::State, p₀::State; kwargs...) = PSDE(m, ns, v, f, B, G, 0.0, q₀, p₀; kwargs...)

# const PSDEHT{HT,DT,TT,AT,VT,FT,BT,GT,PT} = PSDE{DT,TT,AT,VT,FT,BT,GT,HT,PT} # type alias for dispatch on Hamiltonian type parameter
# const PSDEPT{PT,DT,TT,AT,VT,FT,BT,GT,HT} = PSDE{DT,TT,AT,VT,FT,BT,GT,HT,PT} # type alias for dispatch on parameters type parameter
const PSDEPT{PT,DT,TT,AT,VT,FT,BT,GT} = PSDE{DT,TT,AT,VT,FT,BT,GT,PT} # type alias for dispatch on parameters type parameter

Base.hash(sde::PSDE, h::UInt) = hash(sde.d, hash(sde.m, hash(sde.ns,
                                hash(sde.v, hash(sde.f, hash(sde.B, hash(sde.G,
                                hash(sde.t₀, hash(sde.q₀, hash(sde.p₀,
                                hash(sde.parameters, hash(sde.periodicity, h))))))))))))

Base.:(==)(sde1::PSDE, sde2::PSDE) = (
                                sde1.d == sde2.d
                             && sde1.m == sde2.m
                             && sde1.ns == sde2.ns
                             && sde1.v == sde2.v
                             && sde1.f == sde2.f
                             && sde1.B == sde2.B
                             && sde1.G == sde2.G
                             && sde1.t₀ == sde2.t₀
                             && sde1.q₀ == sde2.q₀
                             && sde1.p₀ == sde2.p₀
                             && sde1.parameters  == sde2.parameters
                             && sde1.periodicity == sde2.periodicity)

function Base.similar(equ::PSDE, t₀::Real, q₀::StateVector, p₀::StateVector, ns::Int=equ.ns)
    @assert all([length(q) == ndims(equ) for q in q₀])
    @assert all([length(p) == ndims(equ) for p in p₀])
    PSDE(equ.m, ns, equ.v, equ.f, equ.B, equ.G, t₀, q₀, p₀; parameters=equ.parameters, periodicity=equ.periodicity)
end

Base.similar(equ::PSDE, t₀::Real, q₀::State, p₀::State, ns::Int=equ.ns; kwargs...) = similar(equ, t₀, [q₀], [p₀], ns; kwargs...)
Base.similar(equ::PSDE, q₀::Union{State,StateVector}, p₀::Union{State,StateVector}, ns::Int=equ.ns; kwargs...) = similar(equ, equ.t₀, q₀, p₀, ns; kwargs...)

Base.ndims(equ::PSDE) = equ.d
Base.axes(equ::PSDE) = axes(equ.q₀[begin])
GeometricBase.nsamples(equ::PSDE) = length(equ.q₀)
GeometricBase.periodicity(equ::PSDE) = equ.periodicity

initial_conditions(equ::PSDE) = (equ.t₀, equ.q₀, equ.p₀)

# hashamiltonian(::PSDEHT{<:Nothing}) = false
# hashamiltonian(::PSDEHT{<:Function}) = true

hasparameters(::PSDEPT{<:NullParameters}) = false
hasparameters(::PSDEPT{<:NamedTuple}) = true

_get_v(equ::PSDE) = hasparameters(equ) ? (t,q,p,v) -> equ.v(t, q, p, v, equ.parameters) : equ.v
_get_f(equ::PSDE) = hasparameters(equ) ? (t,q,p,f) -> equ.f(t, q, p, f, equ.parameters) : equ.f
_get_B(equ::PSDE) = hasparameters(equ) ? (t,q,p,B) -> equ.B(t, q, p, B, equ.parameters) : equ.B
_get_G(equ::PSDE) = hasparameters(equ) ? (t,q,p,G) -> equ.G(t, q, p, G, equ.parameters) : equ.G
# _get_h(equ::PSDE) = hasparameters(equ) ? (t,q,p) -> equ.h(t, q, p, equ.parameters) : equ.h

function functions(equ::PSDE)
    names = (:v,:f,:B,:G)
    equs  = (_get_v(equ), _get_f(equ), _get_B(equ), _get_G(equ))

    # if hashamiltonian(equ)
    #     names = (names..., :h)
    #     equs  = (equs..., _get_h(equ))
    # end

    NamedTuple{names}(equs)
end