text:
# This file is a part of Julia. License is MIT: https://julialang.org/license

module REPL

using Base.Meta, Sockets
import InteractiveUtils

export
    AbstractREPL,
    BasicREPL,
    LineEditREPL,
    StreamREPL

import Base:
    AbstractDisplay,
    display,
    show,
    AnyDict,
    ==,
    catch_stack

include("Terminals.jl")
using .Terminals

include("LineEdit.jl")
using .LineEdit
import ..LineEdit:
    CompletionProvider,
    HistoryProvider,
    add_history,
    complete_line,
    history_next,
    history_next_prefix,
    history_prev,
    history_prev_prefix,
    history_first,
    history_last,
    history_search,
    accept_result,
    terminal,
    MIState

include("REPLCompletions.jl")
using .REPLCompletions

include("TerminalMenus/TerminalMenus.jl")
include("docview.jl")

@nospecialize # use only declared type signatures

function __init__()
    Base.REPL_MODULE_REF[] = REPL
end

abstract type AbstractREPL end

answer_color(::AbstractREPL) = ""

const JULIA_PROMPT = "julia> "

mutable struct REPLBackend
    "channel for AST"
    repl_channel::Channel
    "channel for results: (value, iserror)"
    response_channel::Channel
    "flag indicating the state of this backend"
    in_eval::Bool
    "current backend task"
    backend_task::Task

    REPLBackend(repl_channel, response_channel, in_eval) =
        new(repl_channel, response_channel, in_eval)
end

function eval_user_input(@nospecialize(ast), backend::REPLBackend)
    lasterr = nothing
    Base.sigatomic_begin()
    while true
        try
            Base.sigatomic_end()
            if lasterr !== nothing
                put!(backend.response_channel, (lasterr,true))
            else
                backend.in_eval = true
                value = Core.eval(Main, ast)
                backend.in_eval = false
                # note: use jl_set_global to make sure value isn't passed through `expand`
                ccall(:jl_set_global, Cvoid, (Any, Any, Any), Main, :ans, value)
                put!(backend.response_channel, (value,false))
            end
            break
        catch err
            if lasterr !== nothing
                println("SYSTEM ERROR: Failed to report error to REPL frontend")
                println(err)
            end
            lasterr = catch_stack()
        end
    end
    Base.sigatomic_end()
    nothing
end

function start_repl_backend(repl_channel::Channel, response_channel::Channel)
    backend = REPLBackend(repl_channel, response_channel, false)
    backend.backend_task = @async begin
        # include looks at this to determine the relative include path
        # nothing means cwd
        while true
            tls = task_local_storage()
            tls[:SOURCE_PATH] = nothing
            ast, show_value = take!(backend.repl_channel)
            if show_value == -1
                # exit flag
                break
            end
            eval_user_input(ast, backend)
        end
    end
    return backend
end
struct REPLDisplay{R<:AbstractREPL} <: AbstractDisplay
    repl::R
end

==(a::REPLDisplay, b::REPLDisplay) = a.repl === b.repl

function display(d::REPLDisplay, mime::MIME"text/plain", x)
    io = outstream(d.repl)
    get(io, :color, false) && write(io, answer_color(d.repl))
    show(IOContext(io, :limit => true, :module => Main), mime, x)
    println(io)
    nothing
end
display(d::REPLDisplay, x) = display(d, MIME("text/plain"), x)

function print_response(repl::AbstractREPL, @nospecialize(response), show_value::Bool, have_color::Bool)
    repl.waserror = response[2]
    io = IOContext(outstream(repl), :module => Main)
    print_response(io, response, show_value, have_color, specialdisplay(repl))
    nothing
end
function print_response(errio::IO, @nospecialize(response), show_value::Bool, have_color::Bool, specialdisplay=nothing)
    Base.sigatomic_begin()
    val, iserr = response
    while true
        try
            Base.sigatomic_end()
            if iserr
                Base.invokelatest(Base.display_error, errio, val)
            else
                if val !== nothing && show_value
                    try
                        if specialdisplay === nothing
                            Base.invokelatest(display, val)
                        else
                            Base.invokelatest(display, specialdisplay, val)
                        end
                    catch
                        println(errio, "Error showing value of type ", typeof(val), ":")
                        rethrow()
                    end
                end
            end
            break
        catch
            if iserr
                println(errio) # an error during printing is likely to leave us mid-line
                println(errio, "SYSTEM (REPL): showing an error caused an error")
                try
                    Base.invokelatest(Base.display_error, errio, catch_stack())
                catch e
                    # at this point, only print the name of the type as a Symbol to
                    # minimize the possibility of further errors.
                    println(errio)
                    println(errio, "SYSTEM (REPL): caught exception of type ", typeof(e).name.name,
                            " while trying to handle a nested exception; giving up")
                end
                break
            end
            val = catch_stack()
            iserr = true
        end
    end
    Base.sigatomic_end()
    nothing
end

# A reference to a backend
struct REPLBackendRef
    repl_channel::Channel
    response_channel::Channel
end

function run_repl(repl::AbstractREPL, @nospecialize(consumer = x -> nothing))
    repl_channel = Channel(1)
    response_channel = Channel(1)
    backend = start_repl_backend(repl_channel, response_channel)
    consumer(backend)
    run_frontend(repl, REPLBackendRef(repl_channel, response_channel))
    return backend
end

## BasicREPL ##

mutable struct BasicREPL <: AbstractREPL
    terminal::TextTerminal
    waserror::Bool
    BasicREPL(t) = new(t, false)
end

outstream(r::BasicREPL) = r.terminal

function run_frontend(repl::BasicREPL, backend::REPLBackendRef)
    d = REPLDisplay(repl)
    dopushdisplay = !in(d,Base.Multimedia.displays)
    dopushdisplay && pushdisplay(d)
    hit_eof = false
    while true
        Base.reseteof(repl.terminal)
        write(repl.terminal, JULIA_PROMPT)
        line = ""
        ast = nothing
        interrupted = false
        while true
            try
                line *= readline(repl.terminal, keep=true)
            catch e
                if isa(e,InterruptException)
                    try # raise the debugger if present
                        ccall(:jl_raise_debugger, Int, ())
                    catch
                    end
                    line = ""
                    interrupted = true
                    break
                elseif isa(e,EOFError)
                    hit_eof = true
                    break
                else
                    rethrow()
                end
            end
            ast = Base.parse_input_line(line)
            (isa(ast,Expr) && ast.head == :incomplete) || break
        end
        if !isempty(line)
            response = eval_with_backend(ast, backend)
            print_response(repl, response, !ends_with_semicolon(line), false)
        end
        write(repl.terminal, '\n')
        ((!interrupted && isempty(line)) || hit_eof) && break
    end
    # terminate backend
    put!(backend.repl_channel, (nothing, -1))
    dopushdisplay && popdisplay(d)
    nothing
end

## User Options

mutable struct Options
    hascolor::Bool
    extra_keymap::Union{Dict,Vector{<:Dict}}
    # controls the presumed tab width of code pasted into the REPL.
    # Must satisfy `0 < tabwidth <= 16`.
    tabwidth::Int
    # Maximum number of entries in the kill ring queue.
    # Beyond this number, oldest entries are discarded first.
    kill_ring_max::Int
    region_animation_duration::Float64
    beep_duration::Float64
    beep_blink::Float64
    beep_maxduration::Float64
    beep_colors::Vector{String}
    beep_use_current::Bool
    backspace_align::Bool
    backspace_adjust::Bool
    confirm_exit::Bool # ^D must be repeated to confirm exit
    auto_indent::Bool # indent a newline like line above
    auto_indent_tmp_off::Bool # switch auto_indent temporarily off if copy&paste
    auto_indent_bracketed_paste::Bool # set to true if terminal knows paste mode
    # cancel auto-indent when next character is entered within this time frame :
    auto_indent_time_threshold::Float64
end

Options(;
        hascolor = true,
        extra_keymap = AnyDict[],
        tabwidth = 8,
        kill_ring_max = 100,
        region_animation_duration = 0.2,
        beep_duration = 0.2, beep_blink = 0.2, beep_maxduration = 1.0,
        beep_colors = ["\e[90m"], # gray (text_colors not yet available)
        beep_use_current = true,
        backspace_align = true, backspace_adjust = backspace_align,
        confirm_exit = false,
        auto_indent = true,
        auto_indent_tmp_off = false,
        auto_indent_bracketed_paste = false,
        auto_indent_time_threshold = 0.005) =
            Options(hascolor, extra_keymap, tabwidth,
                    kill_ring_max, region_animation_duration,
                    beep_duration, beep_blink, beep_maxduration,
                    beep_colors, beep_use_current,
                    backspace_align, backspace_adjust, confirm_exit,
                    auto_indent, auto_indent_tmp_off, auto_indent_bracketed_paste, auto_indent_time_threshold)

# for use by REPLs not having an options field
const GlobalOptions = Options()

## LineEditREPL ##

mutable struct LineEditREPL <: AbstractREPL
    t::TextTerminal
    hascolor::Bool
    prompt_color::String
    input_color::String
    answer_color::String
    shell_color::String
    help_color::String
    history_file::Bool
    in_shell::Bool
    in_help::Bool
    envcolors::Bool
    waserror::Bool
    specialdisplay::Union{Nothing,AbstractDisplay}
    options::Options
    mistate::Union{MIState,Nothing}
    interface::ModalInterface
    backendref::REPLBackendRef
    LineEditREPL(t,hascolor,prompt_color,input_color,answer_color,shell_color,help_color,history_file,in_shell,in_help,envcolors) =
        new(t,true,prompt_color,input_color,answer_color,shell_color,help_color,history_file,in_shell,
            in_help,envcolors,false,nothing, Options(), nothing)
end
outstream(r::LineEditREPL) = r.t
specialdisplay(r::LineEditREPL) = r.specialdisplay
specialdisplay(r::AbstractREPL) = nothing
terminal(r::LineEditREPL) = r.t

LineEditREPL(t::TextTerminal, hascolor::Bool, envcolors::Bool=false) =
    LineEditREPL(t, hascolor,
        hascolor ? Base.text_colors[:green] : "",
        hascolor ? Base.input_color() : "",
        hascolor ? Base.answer_color() : "",
        hascolor ? Base.text_colors[:red] : "",
        hascolor ? Base.text_colors[:yellow] : "",
        false, false, false, envcolors
    )

mutable struct REPLCompletionProvider <: CompletionProvider end
mutable struct ShellCompletionProvider <: CompletionProvider end
struct LatexCompletions <: CompletionProvider end

beforecursor(buf::IOBuffer) = String(buf.data[1:buf.ptr-1])

function complete_line(c::REPLCompletionProvider, s)
    partial = beforecursor(s.input_buffer)
    full = LineEdit.input_string(s)
    ret, range, should_complete = completions(full, lastindex(partial))
    return unique!(map(completion_text, ret)), partial[range], should_complete
end

function complete_line(c::ShellCompletionProvider, s)
    # First parse everything up to the current position
    partial = beforecursor(s.input_buffer)
    full = LineEdit.input_string(s)
    ret, range, should_complete = shell_completions(full, lastindex(partial))
    return unique!(map(completion_text, ret)), partial[range], should_complete
end

function complete_line(c::LatexCompletions, s)
    partial = beforecursor(LineEdit.buffer(s))
    full = LineEdit.input_string(s)
    ret, range, should_complete = bslash_completions(full, lastindex(partial))[2]
    return unique!(map(completion_text, ret)), partial[range], should_complete
end

mutable struct REPLHistoryProvider <: HistoryProvider
    history::Array{String,1}
    history_file::Union{Nothing,IO}
    start_idx::Int
    cur_idx::Int
    last_idx::Int
    last_buffer::IOBuffer
    last_mode::Union{Nothing,Prompt}
    mode_mapping::Dict
    modes::Array{Symbol,1}
end
REPLHistoryProvider(mode_mapping) =
    REPLHistoryProvider(String[], nothing, 0, 0, -1, IOBuffer(),
                        nothing, mode_mapping, UInt8[])

invalid_history_message(path::String) = """
Invalid history file ($path) format:
If you have a history file left over from an older version of Julia,
try renaming or deleting it.
Invalid character: """

munged_history_message(path::String) = """
Invalid history file ($path) format:
An editor may have converted tabs to spaces at line """

function hist_getline(file)
    while !eof(file)
        line = readline(file, keep=true)
        isempty(line) && return line
        line[1] in "\r\n" || return line
    end
    return ""
end

function hist_from_file(hp, file, path)
    hp.history_file = file
    seek(file, 0)
    countlines = 0
    while true
        mode = :julia
        line = hist_getline(file)
        isempty(line) && break
        countlines += 1
        line[1] != '#' &&
            error(invalid_history_message(path), repr(line[1]), " at line ", countlines)
        while !isempty(line)
            m = match(r"^#\s*(\w+)\s*:\s*(.*?)\s*$", line)
            m === nothing && break
            if m.captures[1] == "mode"
                mode = Symbol(m.captures[2])
            end
            line = hist_getline(file)
            countlines += 1
        end
        isempty(line) && break
        # Make sure starts with tab
        line[1] == ' '  &&
            error(munged_history_message(path), countlines)
        line[1] != '\t' &&
            error(invalid_history_message(path), repr(line[1]), " at line ", countlines)
        lines = String[]
        while !isempty(line)
            push!(lines, chomp(line[2:end]))
            eof(file) && break
            ch = Char(Base.peek(file))
            ch == ' '  && error(munged_history_message(path), countlines)
            ch != '\t' && break
            line = hist_getline(file)
            countlines += 1
        end
        push!(hp.modes, mode)
        push!(hp.history, join(lines, '\n'))
    end
    seekend(file)
    hp.start_idx = length(hp.history)
    return hp
end

function mode_idx(hist::REPLHistoryProvider, mode)
    c = :julia
    for (k,v) in hist.mode_mapping
        isequal(v, mode) && (c = k)
    end
    return c
end

function add_history(hist::REPLHistoryProvider, s)
    str = rstrip(String(take!(copy(s.input_buffer))))
    isempty(strip(str)) && return
    mode = mode_idx(hist, LineEdit.mode(s))
    !isempty(hist.history) &&
        isequal(mode, hist.modes[end]) && str == hist.history[end] && return
    push!(hist.modes, mode)
    push!(hist.history, str)
    hist.history_file === nothing && return
    entry = """
    # time: $(Libc.strftime("%Y-%m-%d %H:%M:%S %Z", time()))
    # mode: $mode
    $(replace(str, r"^"ms => "\t"))
    """
    # TODO: write-lock history file
    seekend(hist.history_file)
    print(hist.history_file, entry)
    flush(hist.history_file)
    nothing
end

function history_move(s::Union{LineEdit.MIState,LineEdit.PrefixSearchState}, hist::REPLHistoryProvider, idx::Int, save_idx::Int = hist.cur_idx)
    max_idx = length(hist.history) + 1
    @assert 1 <= hist.cur_idx <= max_idx
    (1 <= idx <= max_idx) || return :none
    idx != hist.cur_idx || return :none

    # save the current line
    if save_idx == max_idx
        hist.last_mode = LineEdit.mode(s)
        hist.last_buffer = copy(LineEdit.buffer(s))
    else
        hist.history[save_idx] = LineEdit.input_string(s)
        hist.modes[save_idx] = mode_idx(hist, LineEdit.mode(s))
    end

    # load the saved line
    if idx == max_idx
        last_buffer = hist.last_buffer
        LineEdit.transition(s, hist.last_mode) do
            LineEdit.replace_line(s, last_buffer)
        end
        hist.last_mode = nothing
        hist.last_buffer = IOBuffer()
    else
        if haskey(hist.mode_mapping, hist.modes[idx])
            LineEdit.transition(s, hist.mode_mapping[hist.modes[idx]]) do
                LineEdit.replace_line(s, hist.history[idx])
            end
        else
            return :skip
        end
    end
    hist.cur_idx = idx

    return :ok
end

# REPL History can also transitions modes
function LineEdit.accept_result_newmode(hist::REPLHistoryProvider)
    if 1 <= hist.cur_idx <= length(hist.modes)
        return hist.mode_mapping[hist.modes[hist.cur_idx]]
    end
    return nothing
end

function history_prev(s::LineEdit.MIState, hist::REPLHistoryProvider,
                      num::Int=1, save_idx::Int = hist.cur_idx)
    num <= 0 && return history_next(s, hist, -num, save_idx)
    hist.last_idx = -1
    m = history_move(s, hist, hist.cur_idx-num, save_idx)
    if m === :ok
        LineEdit.move_input_start(s)
        LineEdit.reset_key_repeats(s) do
            LineEdit.move_line_end(s)
        end
        return LineEdit.refresh_line(s)
    elseif m === :skip
        return history_prev(s, hist, num+1, save_idx)
    else
        return Terminals.beep(s)
    end
end

function history_next(s::LineEdit.MIState, hist::REPLHistoryProvider,
                      num::Int=1, save_idx::Int = hist.cur_idx)
    if num == 0
        Terminals.beep(s)
        return
    end
    num < 0 && return history_prev(s, hist, -num, save_idx)
    cur_idx = hist.cur_idx
    max_idx = length(hist.history) + 1
    if cur_idx == max_idx && 0 < hist.last_idx
        # issue #6312
        cur_idx = hist.last_idx
        hist.last_idx = -1
    end
    m = history_move(s, hist, cur_idx+num, save_idx)
    if m === :ok
        LineEdit.move_input_end(s)
        return LineEdit.refresh_line(s)
    elseif m === :skip
        return history_next(s, hist, num+1, save_idx)
    else
        return Terminals.beep(s)
    end
end

history_first(s::LineEdit.MIState, hist::REPLHistoryProvider) =
    history_prev(s, hist, hist.cur_idx - 1 -
                 (hist.cur_idx > hist.start_idx+1 ? hist.start_idx : 0))

history_last(s::LineEdit.MIState, hist::REPLHistoryProvider) =
    history_next(s, hist, length(hist.history) - hist.cur_idx + 1)

function history_move_prefix(s::LineEdit.PrefixSearchState,
                             hist::REPLHistoryProvider,
                             prefix::AbstractString,
                             backwards::Bool,
                             cur_idx = hist.cur_idx)
    cur_response = String(take!(copy(LineEdit.buffer(s))))
    # when searching forward, start at last_idx
    if !backwards && hist.last_idx > 0
        cur_idx = hist.last_idx
    end
    hist.last_idx = -1
    max_idx = length(hist.history)+1
    idxs = backwards ? ((cur_idx-1):-1:1) : ((cur_idx+1):max_idx)
    for idx in idxs
        if (idx == max_idx) || (startswith(hist.history[idx], prefix) && (hist.history[idx] != cur_response || hist.modes[idx] != LineEdit.mode(s)))
            m = history_move(s, hist, idx)
            if m === :ok
                if idx == max_idx
                    # on resuming the in-progress edit, leave the cursor where the user last had it
                elseif isempty(prefix)
                    # on empty prefix search, move cursor to the end
                    LineEdit.move_input_end(s)
                else
                    # otherwise, keep cursor at the prefix position as a visual cue
                    seek(LineEdit.buffer(s), sizeof(prefix))
                end
                LineEdit.refresh_line(s)
                return :ok
            elseif m === :skip
                return history_move_prefix(s,hist,prefix,backwards,idx)
            end
        end
    end
    Terminals.beep(s)
    nothing
end
history_next_prefix(s::LineEdit.PrefixSearchState, hist::REPLHistoryProvider, prefix::AbstractString) =
    history_move_prefix(s, hist, prefix, false)
history_prev_prefix(s::LineEdit.PrefixSearchState, hist::REPLHistoryProvider, prefix::AbstractString) =
    history_move_prefix(s, hist, prefix, true)

function history_search(hist::REPLHistoryProvider, query_buffer::IOBuffer, response_buffer::IOBuffer,
                        backwards::Bool=false, skip_current::Bool=false)

    qpos = position(query_buffer)
    qpos > 0 || return true
    searchdata = beforecursor(query_buffer)
    response_str = String(take!(copy(response_buffer)))

    # Alright, first try to see if the current match still works
    a = position(response_buffer) + 1 # position is zero-indexed
    # FIXME: I'm pretty sure this is broken since it uses an index
    # into the search data to index into the response string
    b = a + sizeof(searchdata)
    b = b ≤ ncodeunits(response_str) ? prevind(response_str, b) : b-1
    b = min(lastindex(response_str), b) # ensure that b is valid

    searchfunc1, searchfunc2, searchstart, skipfunc = backwards ?
                                                      (findlast, findprev, b, prevind) :
                                                      (findfirst, findnext, a, nextind)

    if searchdata == response_str[a:b]
        if skip_current
            searchstart = skipfunc(response_str, searchstart)
        else
            return true
        end
    end

    # Start searching
    # First the current response buffer
    if 1 <= searchstart <= lastindex(response_str)
        match = searchfunc2(searchdata, response_str, searchstart)
        if match !== nothing
            seek(response_buffer, first(match) - 1)
            return true
        end
    end

    # Now search all the other buffers
    idxs = backwards ? ((hist.cur_idx-1):-1:1) : ((hist.cur_idx+1):length(hist.history))
    for idx in idxs
        h = hist.history[idx]
        match = searchfunc1(searchdata, h)
        if match !== nothing && h != response_str && haskey(hist.mode_mapping, hist.modes[idx])
            truncate(response_buffer, 0)
            write(response_buffer, h)
            seek(response_buffer, first(match) - 1)
            hist.cur_idx = idx
            return true
        end
    end

    return false
end

function history_reset_state(hist::REPLHistoryProvider)
    if hist.cur_idx != length(hist.history) + 1
        hist.last_idx = hist.cur_idx
        hist.cur_idx = length(hist.history) + 1
    end
    nothing
end
LineEdit.reset_state(hist::REPLHistoryProvider) = history_reset_state(hist)

function return_callback(s)
    ast = Base.parse_input_line(String(take!(copy(LineEdit.buffer(s)))), depwarn=false)
    return !(isa(ast, Expr) && ast.head === :incomplete)
end

find_hist_file() = get(ENV, "JULIA_HISTORY",
    joinpath(homedir(), ".julia", "logs", "repl_history.jl"))

backend(r::AbstractREPL) = r.backendref

function eval_with_backend(ast, backend::REPLBackendRef)
    put!(backend.repl_channel, (ast, 1))
    take!(backend.response_channel) # (val, iserr)
end

function respond(f, repl, main; pass_empty = false)
    return function do_respond(s, buf, ok)
        if !ok
            return transition(s, :abort)
        end
        line = String(take!(buf))
        if !isempty(line) || pass_empty
            reset(repl)
            local response
            try
                ast = Base.invokelatest(f, line)
                response = eval_with_backend(ast, backend(repl))
            catch
                response = (catch_stack(), true)
            end
            print_response(repl, response, !ends_with_semicolon(line), Base.have_color)
        end
        prepare_next(repl)
        reset_state(s)
        return s.current_mode.sticky ? true : transition(s, main)
    end
end

function reset(repl::LineEditREPL)
    raw!(repl.t, false)
    print(repl.t, Base.text_colors[:normal])
end

function prepare_next(repl::LineEditREPL)
    println(terminal(repl))
end

function mode_keymap(julia_prompt::Prompt)
    AnyDict(
    '\b' => function (s,o...)
        if isempty(s) || position(LineEdit.buffer(s)) == 0
            buf = copy(LineEdit.buffer(s))
            transition(s, julia_prompt) do
                LineEdit.state(s, julia_prompt).input_buffer = buf
            end
        else
            LineEdit.edit_backspace(s)
        end
    end,
    "^C" => function (s,o...)
        LineEdit.move_input_end(s)
        LineEdit.refresh_line(s)
        print(LineEdit.terminal(s), "^C\n\n")
        transition(s, julia_prompt)
        transition(s, :reset)
        LineEdit.refresh_line(s)
    end)
end

repl_filename(repl, hp::REPLHistoryProvider) = "REPL[$(max(length(hp.history)-hp.start_idx, 1))]"
repl_filename(repl, hp) = "REPL"

const JL_PROMPT_PASTE = Ref(true)
enable_promptpaste(v::Bool) = JL_PROMPT_PASTE[] = v

setup_interface(
    repl::LineEditREPL;
    # those keyword arguments may be deprecated eventually in favor of the Options mechanism
    hascolor::Bool = repl.options.hascolor,
    extra_repl_keymap::Union{Dict,Vector{<:Dict}} = repl.options.extra_keymap
) = setup_interface(repl, hascolor, extra_repl_keymap)

# This non keyword method can be precompiled which is important
function setup_interface(
    repl::LineEditREPL,
    hascolor::Bool,
    extra_repl_keymap::Union{Dict,Vector{<:Dict}},
)
    ###
    #
    # This function returns the main interface that describes the REPL
    # functionality, it is called internally by functions that setup a
    # Terminal-based REPL frontend, but if you want to customize your REPL
    # or embed the REPL in another interface, you may call this function
    # directly and append it to your interface.
    #
    # Usage:
    #
    # repl_channel,response_channel = Channel(),Channel()
    # start_repl_backend(repl_channel, response_channel)
    # setup_interface(REPLDisplay(t),repl_channel,response_channel)
    #
    ###

    ###
    # We setup the interface in two stages.
    # First, we set up all components (prompt,rsearch,shell,help)
    # Second, we create keymaps with appropriate transitions between them
    #   and assign them to the components
    #
    ###

    ############################### Stage I ################################

    # This will provide completions for REPL and help mode
    replc = REPLCompletionProvider()

    # Set up the main Julia prompt
    julia_prompt = Prompt(JULIA_PROMPT;
        # Copy colors from the prompt object
        prompt_prefix = hascolor ? repl.prompt_color : "",
        prompt_suffix = hascolor ?
            (repl.envcolors ? Base.input_color : repl.input_color) : "",
        repl = repl,
        complete = replc,
        on_enter = return_callback)

    # Setup help mode
    help_mode = Prompt("help?> ",
        prompt_prefix = hascolor ? repl.help_color : "",
        prompt_suffix = hascolor ?
            (repl.envcolors ? Base.input_color : repl.input_color) : "",
        repl = repl,
        complete = replc,
        # When we're done transform the entered line into a call to help("$line")
        on_done = respond(helpmode, repl, julia_prompt, pass_empty=true))

    # Set up shell mode
    shell_mode = Prompt("shell> ";
        prompt_prefix = hascolor ? repl.shell_color : "",
        prompt_suffix = hascolor ?
            (repl.envcolors ? Base.input_color : repl.input_color) : "",
        repl = repl,
        complete = ShellCompletionProvider(),
        # Transform "foo bar baz" into `foo bar baz` (shell quoting)
        # and pass into Base.repl_cmd for processing (handles `ls` and `cd`
        # special)
        on_done = respond(repl, julia_prompt) do line
            Expr(:call, :(Base.repl_cmd),
                :(Base.cmd_gen($(Base.shell_parse(line)[1]))),
                outstream(repl))
        end)

    ################################# Stage II #############################

    # Setup history
    # We will have a unified history for all REPL modes
    hp = REPLHistoryProvider(Dict{Symbol,Any}(:julia => julia_prompt,
                                              :shell => shell_mode,
                                              :help  => help_mode))
    if repl.history_file
        try
            hist_path = find_hist_file()
            mkpath(dirname(hist_path))
            f = open(hist_path, read=true, write=true, create=true)
            finalizer(replc) do replc
                close(f)
            end
            hist_from_file(hp, f, hist_path)
        catch
            print_response(repl, (catch_stack(),true), true, Base.have_color)
            println(outstream(repl))
            @info "Disabling history file for this session"
            repl.history_file = false
        end
    end
    history_reset_state(hp)
    julia_prompt.hist = hp
    shell_mode.hist = hp
    help_mode.hist = hp

    julia_prompt.on_done = respond(x->Base.parse_input_line(x,filename=repl_filename(repl,hp)), repl, julia_prompt)

    search_prompt, skeymap = LineEdit.setup_search_keymap(hp)
    search_prompt.complete = LatexCompletions()

    # Canonicalize user keymap input
    if isa(extra_repl_keymap, Dict)
        extra_repl_keymap = [extra_repl_keymap]
    end

    repl_keymap = AnyDict(
        ';' => function (s,o...)
            if isempty(s) || position(LineEdit.buffer(s)) == 0
                buf = copy(LineEdit.buffer(s))
                transition(s, shell_mode) do
                    LineEdit.state(s, shell_mode).input_buffer = buf
                end
            else
                edit_insert(s, ';')
            end
        end,
        '?' => function (s,o...)
            if isempty(s) || position(LineEdit.buffer(s)) == 0
                buf = copy(LineEdit.buffer(s))
                transition(s, help_mode) do
                    LineEdit.state(s, help_mode).input_buffer = buf
                end
            else
                edit_insert(s, '?')
            end
        end,

        # Bracketed Paste Mode
        "\e[200~" => (s,o...)->begin
            input = LineEdit.bracketed_paste(s) # read directly from s until reaching the end-bracketed-paste marker
            sbuffer = LineEdit.buffer(s)
            curspos = position(sbuffer)
            seek(sbuffer, 0)
            shouldeval = (bytesavailable(sbuffer) == curspos && !occursin(UInt8('\n'), sbuffer))
            seek(sbuffer, curspos)
            if curspos == 0
                # if pasting at the beginning, strip leading whitespace
                input = lstrip(input)
            end
            if !shouldeval
                # when pasting in the middle of input, just paste in place
                # don't try to execute all the WIP, since that's rather confusing
                # and is often ill-defined how it should behave
                edit_insert(s, input)
                return
            end
            LineEdit.push_undo(s)
            edit_insert(sbuffer, input)
            input = String(take!(sbuffer))
            oldpos = firstindex(input)
            firstline = true
            isprompt_paste = false
            while oldpos <= lastindex(input) # loop until all lines have been executed
                if JL_PROMPT_PASTE[]
                    # Check if the next statement starts with "julia> ", in that case
                    # skip it. But first skip whitespace
                    while input[oldpos] in ('\n', ' ', '\t')
                        oldpos = nextind(input, oldpos)
                        oldpos >= sizeof(input) && return
                    end
                    # Check if input line starts with "julia> ", remove it if we are in prompt paste mode
                    jl_prompt_len = 7
                    if (firstline || isprompt_paste) && startswith(SubString(input, oldpos), JULIA_PROMPT)
                        isprompt_paste = true
                        oldpos += jl_prompt_len
                    # If we are prompt pasting and current statement does not begin with julia> , skip to next line
                    elseif isprompt_paste
                        while input[oldpos] != '\n'
                            oldpos = nextind(input, oldpos)
                            oldpos >= sizeof(input) && return
                        end
                        continue
                    end
                end
                ast, pos = Meta.parse(input, oldpos, raise=false, depwarn=false)
                if (isa(ast, Expr) && (ast.head == :error || ast.head == :incomplete)) ||
                        (pos > ncodeunits(input) && !endswith(input, '\n'))
                    # remaining text is incomplete (an error, or parser ran to the end but didn't stop with a newline):
                    # Insert all the remaining text as one line (might be empty)
                    tail = input[oldpos:end]
                    if !firstline
                        # strip leading whitespace, but only if it was the result of executing something
                        # (avoids modifying the user's current leading wip line)
                        tail = lstrip(tail)
                    end
                    if isprompt_paste # remove indentation spaces corresponding to the prompt
                        tail = replace(tail, r"^ {7}"m => "") # 7: jl_prompt_len
                    end
                    LineEdit.replace_line(s, tail, true)
                    LineEdit.refresh_line(s)
                    break
                end
                # get the line and strip leading and trailing whitespace
                line = strip(input[oldpos:prevind(input, pos)])
                if !isempty(line)
                    if isprompt_paste # remove indentation spaces corresponding to the prompt
                        line = replace(line, r"^ {7}"m => "") # 7: jl_prompt_len
                    end
                    # put the line on the screen and history
                    LineEdit.replace_line(s, line)
                    LineEdit.commit_line(s)
                    # execute the statement
                    terminal = LineEdit.terminal(s) # This is slightly ugly but ok for now
                    raw!(terminal, false) && disable_bracketed_paste(terminal)
                    LineEdit.mode(s).on_done(s, LineEdit.buffer(s), true)
                    raw!(terminal, true) && enable_bracketed_paste(terminal)
                    LineEdit.push_undo(s) # when the last line is incomplete
                end
                oldpos = pos
                firstline = false
            end
        end,

        # Open the editor at the location of a stackframe or method
        # This is accessing a global variable that gets set in
        # the show_backtrace and show_method_table functions.
        "^Q" => (s, o...) -> begin
            linfos = Base.LAST_SHOWN_LINE_INFOS
            str = String(take!(LineEdit.buffer(s)))
            n = tryparse(Int, str)
            n === nothing && @goto writeback
            if n <= 0 || n > length(linfos) || startswith(linfos[n][1], "./REPL")
                @goto writeback
            end
            InteractiveUtils.edit(linfos[n][1], linfos[n][2])
            LineEdit.refresh_line(s)
            return
            @label writeback
            write(LineEdit.buffer(s), str)
            return
        end,
    )

    prefix_prompt, prefix_keymap = LineEdit.setup_prefix_keymap(hp, julia_prompt)

    a = Dict{Any,Any}[skeymap, repl_keymap, prefix_keymap, LineEdit.history_keymap, LineEdit.default_keymap, LineEdit.escape_defaults]
    prepend!(a, extra_repl_keymap)

    julia_prompt.keymap_dict = LineEdit.keymap(a)

    mk = mode_keymap(julia_prompt)

    b = Dict{Any,Any}[skeymap, mk, prefix_keymap, LineEdit.history_keymap, LineEdit.default_keymap, LineEdit.escape_defaults]
    prepend!(b, extra_repl_keymap)

    shell_mode.keymap_dict = help_mode.keymap_dict = LineEdit.keymap(b)

    allprompts = [julia_prompt, shell_mode, help_mode, search_prompt, prefix_prompt]
    return ModalInterface(allprompts)
end

function run_frontend(repl::LineEditREPL, backend::REPLBackendRef)
    d = REPLDisplay(repl)
    dopushdisplay = repl.specialdisplay === nothing && !in(d,Base.Multimedia.displays)
    dopushdisplay && pushdisplay(d)
    if !isdefined(repl,:interface)
        interface = repl.interface = setup_interface(repl)
    else
        interface = repl.interface
    end
    repl.backendref = backend
    repl.mistate = LineEdit.init_state(terminal(repl), interface)
    run_interface(terminal(repl), interface, repl.mistate)
    dopushdisplay && popdisplay(d)
    nothing
end

## StreamREPL ##

mutable struct StreamREPL <: AbstractREPL
    stream::IO
    prompt_color::String
    input_color::String
    answer_color::String
    waserror::Bool
    StreamREPL(stream,pc,ic,ac) = new(stream,pc,ic,ac,false)
end
StreamREPL(stream::IO) = StreamREPL(stream, Base.text_colors[:green], Base.input_color(), Base.answer_color())
run_repl(stream::IO) = run_repl(StreamREPL(stream))

outstream(s::StreamREPL) = s.stream

answer_color(r::LineEditREPL) = r.envcolors ? Base.answer_color() : r.answer_color
answer_color(r::StreamREPL) = r.answer_color
input_color(r::LineEditREPL) = r.envcolors ? Base.input_color() : r.input_color
input_color(r::StreamREPL) = r.input_color

# heuristic function to decide if the presence of a semicolon
# at the end of the expression was intended for suppressing output
function ends_with_semicolon(line::AbstractString)
    match = findlast(isequal(';'), line)
    if match !== nothing
        # state for comment parser, assuming that the `;` isn't in a string or comment
        # so input like ";#" will still thwart this to give the wrong (anti-conservative) answer
        comment = false
        comment_start = false
        comment_close = false
        comment_multi = 0
        for c in line[(match + 1):end]
            if comment_multi > 0
                # handle nested multi-line comments
                if comment_close && c == '#'
                    comment_close = false
                    comment_multi -= 1
                elseif comment_start && c == '='
                    comment_start = false
                    comment_multi += 1
                else
                    comment_start = (c == '#')
                    comment_close = (c == '=')
                end
            elseif comment
                # handle line comments
                if c == '\r' || c == '\n'
                    comment = false
                end
            elseif comment_start
                # see what kind of comment this is
                comment_start = false
                if c == '='
                    comment_multi = 1
                else
                    comment = true
                end
            elseif c == '#'
                # start handling for a comment
                comment_start = true
            else
                # outside of a comment, encountering anything but whitespace
                # means the semi-colon was internal to the expression
                isspace(c) || return false
            end
        end
        return true
    end
    return false
end

function run_frontend(repl::StreamREPL, backend::REPLBackendRef)
    have_color = Base.have_color
    Base.banner(repl.stream)
    d = REPLDisplay(repl)
    dopushdisplay = !in(d,Base.Multimedia.displays)
    dopushdisplay && pushdisplay(d)
    while !eof(repl.stream)
        if have_color
            print(repl.stream,repl.prompt_color)
        end
        print(repl.stream, "julia> ")
        if have_color
            print(repl.stream, input_color(repl))
        end
        line = readline(repl.stream, keep=true)
        if !isempty(line)
            ast = Base.parse_input_line(line)
            if have_color
                print(repl.stream, Base.color_normal)
            end
            response = eval_with_backend(ast, backend)
            print_response(repl, response, !ends_with_semicolon(line), have_color)
        end
    end
    # Terminate Backend
    put!(backend.repl_channel, (nothing, -1))
    dopushdisplay && popdisplay(d)
    nothing
end

function start_repl_server(port::Int)
    return listen(port) do server, status
        client = accept(server)
        run_repl(client)
        nothing
    end
end

end # module

max_stars_repo_name: nstiurca/julia

max_issues_repo_name: nstiurca/julia

max_forks_repo_name: nstiurca/julia


text:
import Base.min

export get_primobj, get_primslacks, get_primfeaserr
export get_dualobj, get_dualslacks, get_dualfeaserr
export evaluate
export min, mu, delta
export issol

function evaluate(A::Vector{PointE{T, U}}, x::PointE{T}) where {T<:Number, U<:AbstractArray}
    res = zeros(T, length(A))

    for (i, ai) in enumerate(A)
        res[i] = dot(ai, x)
    end

    return res
end

function evaluate(A::Vector{PointE{T, U}}, y::Vector{T}) where {T<:Number, U<:AbstractArray}
    @assert length(A) == length(y)

    dims = first(A).dims
    res = PointE(dims, length(first(A).vec), T, Dense{T})

    for (i, ai) in enumerate(A)
        add!(res, densify(product(A[i], y[i])))
    end

    return res
end

function get_primobj(pb::SDCOContext{T}, x::PointE{T}) where T<:Number
    @assert pb.c.dims == x.dims
    return dot(pb.c, x)
end

get_primobj(pb::SDCOContext{T}, z::PointPrimalDual{T}) where T<:Number = get_primobj(pb, z.x)

function get_dualobj(pb::SDCOContext{T}, y::Vector{T}) where T<:Number
    @assert length(y) == length(pb.b)
    return dot(pb.b, y)
end

get_dualobj(pb::SDCOContext{T}, z::PointPrimalDual{T}) where T<:Number = get_dualobj(pb, z.y)

function get_primslacks(pb::SDCOContext{T}, x::PointE{T}) where T<:Number
    return evaluate(pb.A, x) .- pb.b
end

get_primslacks(pb::SDCOContext{T}, z::PointPrimalDual{T}) where T<:Number = get_primslacks(pb, z.x)

function get_dualslacks(pb::SDCOContext{T}, y::Vector{T}, s::PointE{T}) where T<:Number
    slacks = copy(s)
    add!(slacks, evaluate(pb.A, y))
    add!(slacks, -pb.c)
    return slacks
end

get_dualslacks(pb::SDCOContext{T}, z::PointPrimalDual{T}) where T<:Number = get_dualslacks(pb, z.y, z.s)

function get_primfeaserr(pb::SDCOContext{T}, x::PointE{T}) where T<:Number
    return norm(get_primslacks(pb, x))
end

get_primfeaserr(pb::SDCOContext{T}, z::PointPrimalDual{T}) where T<:Number = get_primfeaserr(pb, z.x)

function get_dualfeaserr(pb::SDCOContext{T}, y::Vector{T}, s::PointE{T}) where T<:Number
    return norm(get_dualslacks(pb, y, s))
end

get_dualfeaserr(pb::SDCOContext{T}, z::PointPrimalDual{T}) where T<:Number = get_dualfeaserr(pb, z.y, z.s)

function min(pb::SDCOContext{T}, x::PointE{T, Dense{T}}) where {T<:Number, U<:Dense}
    minx = Inf
    if length(x.vec) > 0
        minx = minimum(x.vec)
    end

    for (i, mat) in enumerate(x.mats)
        if norm(mat - transpose(mat)) == 0
            eigmin = LinearAlgebra.eigmin(mat)
        elseif norm(mat - transpose(mat)) < 1e-14
            eigmin = LinearAlgebra.eigmin(mat + transpose(mat))
        else
            @show norm(mat - transpose(mat))
            @error("Matrix s=$i is non symmetric.\nnorm(mat-transpose(mat)) = ", norm(mat-transpose(mat)))
        end

        minx = min(minx, eigmin)
    end

    return minx
end

ext: jl


text:
using LinearAlgebra
using Oceananigans: array_type

function ∇²!(grid, f, ∇²f)
    @loop for k in (1:grid.Nz; (blockIdx().z - 1) * blockDim().z + threadIdx().z)
        @loop for j in (1:grid.Ny; (blockIdx().y - 1) * blockDim().y + threadIdx().y)
            @loop for i in (1:grid.Nx; (blockIdx().x - 1) * blockDim().x + threadIdx().x)
                @inbounds ∇²f[i, j, k] = ∇²(i, j, k, grid, f)
            end
        end
    end
end

function fftw_planner_works(FT, Nx, Ny, Nz, planner_flag)
    grid = RegularCartesianGrid(FT; size=(Nx, Ny, Nz), length=(100, 100, 100))
    solver = PoissonSolver(CPU(), PPN(), grid)
    true  # Just making sure the PoissonSolver does not error/crash.
end

function poisson_ppn_planned_div_free_cpu(FT, Nx, Ny, Nz, planner_flag)
    arch = CPU()
    grid = RegularCartesianGrid(FT; size=(Nx, Ny, Nz), length=(1.0, 2.5, 3.6))
    solver = PoissonSolver(arch, PPN(), grid)
    fbcs = HorizontallyPeriodicBCs()

    RHS = CellField(FT, arch, grid)
    interior(RHS) .= rand(Nx, Ny, Nz)
    interior(RHS) .= interior(RHS) .- mean(interior(RHS))

    RHS_orig = deepcopy(RHS)
    solver.storage .= interior(RHS)
    solve_poisson_3d!(solver, grid)

    ϕ   = CellField(FT, arch, grid)
    ∇²ϕ = CellField(FT, arch, grid)

    interior(ϕ) .= real.(solver.storage)

    fill_halo_regions!(ϕ.data, fbcs, arch, grid)
    ∇²!(grid, ϕ, ∇²ϕ)

    fill_halo_regions!(∇²ϕ.data, fbcs, arch, grid)

    interior(∇²ϕ) ≈ interior(RHS_orig)
end

function poisson_pnn_planned_div_free_cpu(FT, Nx, Ny, Nz, planner_flag)
    arch = CPU()
    grid = RegularCartesianGrid(FT; size=(Nx, Ny, Nz), length=(1.0, 2.5, 3.6))
    solver = PoissonSolver(arch, PNN(), grid)
    fbcs = ChannelBCs()

    RHS = CellField(FT, arch, grid)
    interior(RHS) .= rand(Nx, Ny, Nz)
    interior(RHS) .= interior(RHS) .- mean(interior(RHS))

    RHS_orig = deepcopy(RHS)

    solver.storage .= interior(RHS)

    solve_poisson_3d!(solver, grid)

    ϕ   = CellField(FT, arch, grid)
    ∇²ϕ = CellField(FT, arch, grid)

    interior(ϕ) .= real.(solver.storage)

    fill_halo_regions!(ϕ.data, fbcs, arch, grid)
    ∇²!(grid, ϕ, ∇²ϕ)

    fill_halo_regions!(∇²ϕ.data, fbcs, arch, grid)

    interior(∇²ϕ) ≈ interior(RHS_orig)
end

function poisson_ppn_planned_div_free_gpu(FT, Nx, Ny, Nz)
    arch = GPU()
    grid = RegularCartesianGrid(FT; size=(Nx, Ny, Nz), length=(1.0, 2.5, 3.6))
    solver = PoissonSolver(arch, PPN(), grid)
    fbcs = HorizontallyPeriodicBCs()

    RHS = rand(Nx, Ny, Nz)
    RHS .= RHS .- mean(RHS)
    RHS = CuArray(RHS)

    RHS_orig = copy(RHS)

    solver.storage .= RHS

    # Performing the permutation [a, b, c, d, e, f] -> [a, c, e, f, d, b]
    # in the z-direction in preparation to calculate the DCT in the Poisson
    # solver.
    solver.storage .= cat(solver.storage[:, :, 1:2:Nz], solver.storage[:, :, Nz:-2:2]; dims=3)

    solve_poisson_3d!(solver, grid)

    # Undoing the permutation made above to complete the IDCT.
    solver.storage .= CuArray(reshape(permutedims(cat(solver.storage[:, :, 1:Int(Nz/2)],
                                                      solver.storage[:, :, end:-1:Int(Nz/2)+1]; dims=4), (1, 2, 4, 3)), Nx, Ny, Nz))

    ϕ   = CellField(FT, arch, grid)
    ∇²ϕ = CellField(FT, arch, grid)

    interior(ϕ) .= real.(solver.storage)

    fill_halo_regions!(ϕ.data, fbcs, arch, grid)
    ∇²!(grid, ϕ.data, ∇²ϕ.data)

    fill_halo_regions!(∇²ϕ.data, fbcs, arch, grid)
    interior(∇²ϕ) ≈ RHS_orig
end

function poisson_pnn_planned_div_free_gpu(FT, Nx, Ny, Nz)
    arch = GPU()
    grid = RegularCartesianGrid(FT; size=(Nx, Ny, Nz), length=(1.0, 2.5, 3.6))
    solver = PoissonSolver(arch, PNN(), grid)
    fbcs = ChannelBCs()

    RHS = rand(Nx, Ny, Nz)
    RHS .= RHS .- mean(RHS)
    RHS = CuArray(RHS)

    RHS_orig = copy(RHS)

    solver.storage .= RHS

    solver.storage .= cat(solver.storage[:, :, 1:2:Nz], solver.storage[:, :, Nz:-2:2]; dims=3)
    solver.storage .= cat(solver.storage[:, 1:2:Ny, :], solver.storage[:, Ny:-2:2, :]; dims=2)

    solve_poisson_3d!(solver, grid)

    ϕ   = CellField(FT, arch, grid)
    ∇²ϕ = CellField(FT, arch, grid)

    ϕ_p = view(interior(ϕ), 1:Nx, solver.p_y_inds, solver.p_z_inds)

    @. ϕ_p = real(solver.storage)

    fill_halo_regions!(ϕ.data, fbcs, arch, grid)
    ∇²!(grid, ϕ.data, ∇²ϕ.data)

    fill_halo_regions!(∇²ϕ.data, fbcs, arch, grid)
    interior(∇²ϕ) ≈ RHS_orig
end

"""
    poisson_ppn_recover_sine_cosine_solution(FT, Nx, Ny, Nz, Lx, Ly, Lz, mx, my, mz)

Test that the Poisson solver can recover an analytic solution. In this test, we
are trying to see if the solver can recover the solution

    ``\\Psi(x, y, z) = cos(\\pi m_z z / L_z) sin(2\\pi m_y y / L_y) sin(2\\pi m_x x / L_x)``

by giving it the source term or right hand side (RHS), which is

    ``f(x, y, z) = \\nabla^2 \\Psi(x, y, z) =
    -((\\pi m_z / L_z)^2 + (2\\pi m_y / L_y)^2 + (2\\pi m_x/L_x)^2) \\Psi(x, y, z)``.
"""
function poisson_ppn_recover_sine_cosine_solution(FT, Nx, Ny, Nz, Lx, Ly, Lz, mx, my, mz)
    grid = RegularCartesianGrid(FT; size=(Nx, Ny, Nz), length=(Lx, Ly, Lz))
    solver = PoissonSolver(CPU(), PPN(), grid)

    xC, yC, zC = grid.xC, grid.yC, grid.zC
    xC = reshape(xC, (Nx, 1, 1))
    yC = reshape(yC, (1, Ny, 1))
    zC = reshape(zC, (1, 1, Nz))

    Ψ(x, y, z) = cos(π*mz*z/Lz) * sin(2π*my*y/Ly) * sin(2π*mx*x/Lx)
    f(x, y, z) = -((mz*π/Lz)^2 + (2π*my/Ly)^2 + (2π*mx/Lx)^2) * Ψ(x, y, z)

    @. solver.storage = f(xC, yC, zC)
    solve_poisson_3d!(solver, grid)
    ϕ = real.(solver.storage)

    error = norm(ϕ - Ψ.(xC, yC, zC)) / √(Nx*Ny*Nz)

    @info "Error (ℓ²-norm), $FT, N=($Nx, $Ny, $Nz), m=($mx, $my, $mz): $error"

    isapprox(ϕ, Ψ.(xC, yC, zC); rtol=5e-2)
end

function can_solve_single_tridiagonal_system(arch, N)
    ArrayType = array_type(arch)

    a = rand(N-1)
    b = 3 .+ rand(N) # +3 to ensure diagonal dominance.
    c = rand(N-1)
    f = rand(N)

    # Solve the system with backslash on the CPU to avoid scalar operations on the GPU.
    M = Tridiagonal(a, b, c)
    ϕ_correct = M \ f

    # Convert to CuArray if needed.
    a, b, c, f = ArrayType.([a, b, c, f])

    ϕ = reshape(zeros(N), (1, 1, N)) |> ArrayType

    grid = RegularCartesianGrid(size=(1, 1, N), length=(1, 1, 1))
    btsolver = BatchedTridiagonalSolver(arch; dl=a, d=b, du=c, f=f, grid=grid)

    solve_batched_tridiagonal_system!(ϕ, arch, btsolver)

    return Array(ϕ[:]) ≈ ϕ_correct
end

function can_solve_single_tridiagonal_system_with_functions(arch, N)
    ArrayType = array_type(arch)

    grid = RegularCartesianGrid(size=(1, 1, N), length=(1, 1, 1))

    a = rand(N-1)
    c = rand(N-1)

    @inline b(i, j, k, grid, p) = 3 .+ cos(2π*grid.zC[k])  # +3 to ensure diagonal dominance.
    @inline f(i, j, k, grid, p) = sin(2π*grid.zC[k])

    bₐ = [b(1, 1, k, grid, nothing) for k in 1:N]
    fₐ = [f(1, 1, k, grid, nothing) for k in 1:N]

    # Solve the system with backslash on the CPU to avoid scalar operations on the GPU.
    M = Tridiagonal(a, bₐ, c)
    ϕ_correct = M \ fₐ

    # Convert to CuArray if needed.
    a, c = ArrayType.([a, c])

    ϕ = reshape(zeros(N), (1, 1, N)) |> ArrayType

    btsolver = BatchedTridiagonalSolver(arch; dl=a, d=b, du=c, f=f, grid=grid)

    solve_batched_tridiagonal_system!(ϕ, arch, btsolver)

    return Array(ϕ[:]) ≈ ϕ_correct
end

function can_solve_batched_tridiagonal_system_with_3D_RHS(arch, Nx, Ny, Nz)
    ArrayType = array_type(arch)

    a = rand(Nz-1)
    b = 3 .+ rand(Nz) # +3 to ensure diagonal dominance.
    c = rand(Nz-1)
    f = rand(Nx, Ny, Nz)

    M = Tridiagonal(a, b, c)
    ϕ_correct = zeros(Nx, Ny, Nz)

    # Solve the systems with backslash on the CPU to avoid scalar operations on the GPU.
    for i = 1:Nx, j = 1:Ny
        ϕ_correct[i, j, :] .= M \ f[i, j, :]
    end

    # Convert to CuArray if needed.
    a, b, c, f = ArrayType.([a, b, c, f])

    grid = RegularCartesianGrid(size=(Nx, Ny, Nz), length=(1, 1, 1))
    btsolver = BatchedTridiagonalSolver(arch; dl=a, d=b, du=c, f=f, grid=grid)

    ϕ = zeros(Nx, Ny, Nz) |> ArrayType

    solve_batched_tridiagonal_system!(ϕ, arch, btsolver)

    return Array(ϕ) ≈ ϕ_correct
end

function can_solve_batched_tridiagonal_system_with_3D_functions(arch, Nx, Ny, Nz)
    ArrayType = array_type(arch)

    grid = RegularCartesianGrid(size=(Nx, Ny, Nz), length=(1, 1, 1))

    a = rand(Nz-1)
    c = rand(Nz-1)

    @inline b(i, j, k, grid, p) = 3 + grid.xC[i]*grid.yC[j] * cos(2π*grid.zC[k])
    @inline f(i, j, k, grid, p) = (grid.xC[i] + grid.yC[j]) * sin(2π*grid.zC[k])

    ϕ_correct = zeros(Nx, Ny, Nz)

    # Solve the system with backslash on the CPU to avoid scalar operations on the GPU.
    for i = 1:Nx, j = 1:Ny
        bₐ = [b(i, j, k, grid, nothing) for k in 1:Nz]
        M = Tridiagonal(a, bₐ, c)

        fₐ = [f(i, j, k, grid, nothing) for k in 1:Nz]
        ϕ_correct[i, j, :] .= M \ fₐ
    end

    # Convert to CuArray if needed.
    a, c = ArrayType.([a, c])

    btsolver = BatchedTridiagonalSolver(arch; dl=a, d=b, du=c, f=f, grid=grid)

    ϕ = zeros(Nx, Ny, Nz) |> ArrayType
    solve_batched_tridiagonal_system!(ϕ, arch, btsolver)

    return Array(ϕ) ≈ ϕ_correct
end

@testset "Solvers" begin
    println("Testing Solvers...")

    # @testset "FFTW plans" begin
    #     println("  Testing FFTW planning...")
    #
    #     for FT in float_types
    #         @test fftw_planner_works(FT, 32, 32, 32, FFTW.ESTIMATE)
    #         @test fftw_planner_works(FT, 1,  32, 32, FFTW.ESTIMATE)
    #         @test fftw_planner_works(FT, 32,  1, 32, FFTW.ESTIMATE)
    #         @test fftw_planner_works(FT,  1,  1, 32, FFTW.ESTIMATE)
    #     end
    # end

    @testset "Divergence-free solution [CPU]" begin
        println("  Testing divergence-free solution [CPU]...")

        for N in [7, 10, 16, 20]
            for FT in float_types
                for planner_flag in (FFTW.ESTIMATE, FFTW.MEASURE)
                    @test poisson_ppn_planned_div_free_cpu(FT, N, N, N, planner_flag)
                    @test poisson_ppn_planned_div_free_cpu(FT, 1, N, N, planner_flag)
                    @test poisson_ppn_planned_div_free_cpu(FT, N, 1, N, planner_flag)
                    @test poisson_ppn_planned_div_free_cpu(FT, 1, 1, N, planner_flag)
                end
            end
        end

        Ns = [5, 11, 20, 32]
        for Nx in Ns, Ny in Ns, Nz in Ns, FT in float_types
            @test poisson_ppn_planned_div_free_cpu(FT, Nx, Ny, Nz, FFTW.ESTIMATE)
            @test poisson_pnn_planned_div_free_cpu(FT, Nx, Ny, Nz, FFTW.ESTIMATE)
        end
    end

    @testset "Divergence-free solution [GPU]" begin
        println("  Testing divergence-free solution [GPU]...")
        @hascuda begin
            for FT in [Float64]
                @test poisson_ppn_planned_div_free_gpu(FT, 16, 16, 16)
                @test poisson_ppn_planned_div_free_gpu(FT, 32, 32, 32)
                @test poisson_ppn_planned_div_free_gpu(FT, 32, 32, 16)
                @test poisson_ppn_planned_div_free_gpu(FT, 16, 32, 24)

                @test poisson_pnn_planned_div_free_gpu(FT, 16, 16, 16)
                @test poisson_pnn_planned_div_free_gpu(FT, 32, 32, 32)
                @test poisson_pnn_planned_div_free_gpu(FT, 32, 32, 16)
                @test poisson_pnn_planned_div_free_gpu(FT, 16, 32, 24)
            end
        end
    end

    @testset "Analytic solution reconstruction" begin
        println("  Testing analytic solution reconstruction...")
        for N in [32, 48, 64], m in [1, 2, 3]
            @test poisson_ppn_recover_sine_cosine_solution(Float64, N, N, N, 100, 100, 100, m, m, m)
        end
    end

    for arch in archs
        @testset "Batched tridiagonal solver [$arch]" begin
            for Nz in [8, 11, 18]
                @test can_solve_single_tridiagonal_system(arch, Nz)
                @test can_solve_single_tridiagonal_system_with_functions(arch, Nz)
            end

            for Nx in [3, 8], Ny in [5, 16], Nz in [8, 11]
                @test can_solve_batched_tridiagonal_system_with_3D_RHS(arch, Nx, Ny, Nz)
                @test can_solve_batched_tridiagonal_system_with_3D_functions(arch, Nx, Ny, Nz)
            end
        end
    end
end


text:
################################################################################
# 3-D model

"""
    rotstressvec!(::Type{DeforModelRed3D},  outstress::FVec{T},  instress::FVec{T}, Rm::_RotationMatrix) where {T}

Rotate the stress vector by the supplied rotation matrix.

Calculate the rotation of the stress vector to the 'bar' coordinate system given
by the columns of the rotation matrix `Rm`.

- `outstress` = output stress vector, overwritten inside
- `instress` = input stress vector
- `Rm` = columns are components of 'bar' basis vectors on the 'plain'
     basis vectors
"""
function  rotstressvec!(::Type{DeforModelRed3D},  outstress::FVec{T},  instress::FVec{T}, Rm::_RotationMatrix) where {T}
      # # Derivation of the transformation matrix [T]
    # #This is from Barbero''s  book Finite element analysis of composite
    # #materials  using Abaqus.  Note that his matrix "a"  is the transpose of
    # #the FinEALE matrix "Rm".
    # #  We also use the FinEALE numbering of the strains.
    #
    # syms T alpha R
    # syms a a11 a12 a13 a21 a22 a23 a31 a32 a33 real
    # a = [a11,a12,a13;
    #     a21,a22,a23;
    #     a31,a32,a33];
    # a = a';#his matrix "a"  is the transpose of the FinEALE matrix "Rm".
    # # # it can be done in terms of l,m,n's as well
    # # syms a l1 m1 n1 l2 m2 n2 l3 m3 n3
    # # a = [l1,m1,n1;l2,m2,n2;l3,m3,n3]
    # #  We also use the FinEALE numbering of the strains.
    # Numbering =[1,4,5;
    #     4,2,6;
    #     5,6,3];
    # T(1:6,1:6) = 0;
    # for i=1:1:3
    #     for j=1:1:3
    #         #if i==j; alpha = j; else alpha = 9-i-j; end
    #         alpha = Numbering(i,j);
    #         for p=1:1:3
    #             for q=1:1:3
    #                 #   if p==q beta = p; else beta = 9-p-q; end
    #                 beta = Numbering(p,q);
    #                 T(alpha,beta) = 0;
    #                 if alpha<=3 & beta<= 3; T(alpha,beta)=a(i,p)*a(i,p); end
    #                 if alpha> 3 & beta<= 3; T(alpha,beta)=a(i,p)*a(j,p); end
    #                 if alpha<=3 & beta>3; T(alpha,beta)=a(i,q)*a(i,p)+a(i,p)*a(i,q);end
    #                 if alpha>3 & beta>3; T(alpha,beta)=a(i,p)*a(j,q)+a(i,q)*a(j,p);end
    #             end
    #         end
    #     end
    # end
    # T
    # R = eye(6,6); R(4,4)=2; R(5,5)=2; R(6,6)=2; # Reuter matrix
    # Tbar = R*T*R^(-1)
    a11=Rm[1,1]; a12=Rm[1,2]; a13=Rm[1,3];
    a21=Rm[2,1]; a22=Rm[2,2]; a23=Rm[2,3];
    a31=Rm[3,1]; a32=Rm[3,2]; a33=Rm[3,3];

    outstress[1] =  (a11^2)*instress[1] + (a21^2)*instress[2] +
                    (a31^2)*instress[3] + (2*a11*a21)*instress[4] +
                    (2*a11*a31)*instress[5] + (2*a21*a31)*instress[6]
    outstress[2] =  (a12^2)*instress[1] +  (a22^2)*instress[2] +
                    (a32^2)*instress[3] + (2*a12*a22)*instress[4] +
                    (2*a12*a32)*instress[5] + (2*a22*a32)*instress[6]
    outstress[3] =  (a13^2)*instress[1] + (a23^2)*instress[2] +
                    (a33^2)*instress[3] + (2*a13*a23)*instress[4] +
                    (2*a13*a33)*instress[5] + (2*a23*a33)*instress[6]
    outstress[4] =  a11*a12*instress[1] +  a21*a22*instress[2] +
                    a31*a32*instress[3] + (a11*a22 + a12*a21)*instress[4] +
                    (a11*a32 + a12*a31)*instress[5] + (a21*a32 + a22*a31)*instress[6]
    outstress[5] =  (a11*a13)*instress[1] + (a21*a23)*instress[2] +
                    (a31*a33)*instress[3] + (a11*a23 + a13*a21)*instress[4] +
                    (a11*a33 + a13*a31)*instress[5] + (a21*a33 + a23*a31)*instress[6]
    outstress[6] =  (a12*a13)*instress[1] + (a22*a23)*instress[2] +
                    (a32*a33)*instress[3] + (a12*a23 + a13*a22)*instress[4] +
                    (a12*a33 + a13*a32)*instress[5] + (a22*a33 + a23*a32)*instress[6]
    return outstress
end

################################################################################
# 2-D plane strain model

"""
    rotstressvec!(::Type{DeforModelRed2DStrain},  outstress::FVec{T},  instress::FVec{T},  Rm::_RotationMatrix) where {T}

Rotate the stress vector by the supplied rotation matrix.

Calculate the rotation of the stress vector to the 'bar' coordinate system given
by the columns of the rotation matrix `Rm`.

- `outstress` = output stress vector, overwritten inside
- `instress` = input stress vector
- `Rm` = columns are components of 'bar' basis vectors on the 'plain'
     basis vectors
"""
function  rotstressvec!(::Type{DeforModelRed2DStrain},  outstress::FVec{T},  instress::FVec{T},  Rm::_RotationMatrix) where {T}
    a11=Rm[1,1]; a12=Rm[1,2]; a13=0.0;
    a21=Rm[2,1]; a22=Rm[2,2]; a23=0.0;
    a31=0.0; a32=0.0; a33=1.0;
    # Note the special arrangement  of the components for plane strain
    outstress[1] =  (a11^2)*instress[1] + (a21^2)*instress[2] + (a31^2)*instress[4] + (2*a11*a21)*instress[3]
    outstress[2] =  (a12^2)*instress[1] +  (a22^2)*instress[2] + (a32^2)*instress[4] + (2*a12*a22)*instress[3]
    outstress[4] =  (a13^2)*instress[1] + (a23^2)*instress[2] + (a33^2)*instress[4] + (2*a13*a23)*instress[3]
    outstress[3] =  a11*a12*instress[1] +  a21*a22*instress[2] + a31*a32*instress[4] + (a11*a22 + a12*a21)*instress[3]
    return outstress
end

################################################################################
# 2-D plane stress model

"""
    rotstressvec!(::Type{DeforModelRed2DStress},  outstress::FVec{T},  instress::FVec{T},  Rm::_RotationMatrix) where {T}

Rotate the stress vector by the supplied rotation matrix.

Calculate the rotation of the stress vector to the 'bar' coordinate system given
by the columns of the rotation matrix `Rm`.

- `outstress` = output stress vector, overwritten inside
- `instress` = input stress vector
- `Rm` = columns are components of 'bar' basis vectors on the 'plain'
     basis vectors
"""
function  rotstressvec!(::Type{DeforModelRed2DStress},  outstress::FVec{T},  instress::FVec{T},  Rm::_RotationMatrix) where {T}
    a11=Rm[1,1]; a12=Rm[1,2];
    a21=Rm[2,1]; a22=Rm[2,2];
    outstress[1] =  (a11^2)*instress[1] + (a21^2)*instress[2] + (2*a11*a21)*instress[3]
    outstress[2] =  (a12^2)*instress[1] +  (a22^2)*instress[2] + (2*a12*a22)*instress[3]
    outstress[3] =  (a11*a12)*instress[1] + (a21*a22)*instress[2] + (a11*a22+a12*a21)*instress[3]
    return outstress
end

################################################################################
# 2-D axially symmetric stress model

"""
    rotstressvec!(::Type{DeforModelRed2DAxisymm},  outstress::FVec{T},  instress::FVec{T},  Rm::_RotationMatrix) where {T}

Rotate the stress vector by the supplied rotation matrix.

Calculate the rotation of the stress vector to the 'bar' coordinate system given
by the columns of the rotation matrix `Rm`.

- `outstress` = output stress vector, overwritten inside
- `instress` = input stress vector
- `Rm` = columns are components of 'bar' basis vectors on the 'plain'
     basis vectors
"""
function  rotstressvec!(::Type{DeforModelRed2DAxisymm},  outstress::FVec{T},  instress::FVec{T},  Rm::_RotationMatrix) where {T}
    a11=Rm[1,1]; a12=Rm[1,2];
    a21=Rm[2,1]; a22=Rm[2,2];
    outstress[1] =  (a11^2)*instress[1] + (a21^2)*instress[2] + (0.0)*instress[3] + (2*a11*a21)*instress[4]
    outstress[2] =  (a12^2)*instress[1] + (a22^2)*instress[2] + (0.0)*instress[3] + (2*a12*a22)*instress[4]
    outstress[3] =  (0.0)*instress[1] + (0.0)*instress[2] + (1.0)*instress[3] + (0.0)*instress[4]
    outstress[4] =  (a11*a12)*instress[1] + (a21*a22)*instress[2] + (0.0)*instress[3] + (a11*a22 + a12*a21)*instress[4]
    return outstress
end

################################################################################
# 1-D stress model

"""
    rotstressvec!(::Type{DeforModelRed1D},  outstress::FVec{T},  instress::FVec{T},  Rm::_RotationMatrix) where {T}

Rotate the stress vector by the supplied rotation matrix.

Calculate the rotation of the stress vector to the 'bar' coordinate system given
by the columns of the rotation matrix `Rm`.

- `outstress` = output stress vector, overwritten inside
- `instress` = input stress vector
- `Rm` = columns are components of 'bar' basis vectors on the 'plain'
     basis vectors
"""
function  rotstressvec!(::Type{DeforModelRed1D},  outstress::FVec{T},  instress::FVec{T},  Rm::_RotationMatrix) where {T}
    copyto!(outstress, instress)
    return outstress
end


text:
module GenericRegex

module IR
include("ir/graph.jl")
include("ir/opcode.jl")
include("ir/graph2opcode.jl")
end

module Parse
include("parse/unicode.jl")
end

end


text:
using StatsBase

filename() = "day10.txt"
read_data() = parse.(Int, readlines(filename()))

function part1(data)
    sort!(data)
    v = vcat(0, data, data[end]+3)
    return prod(values(countmap(diff(v))))
end

function part2(data)
    sort!(data)
    v = vcat(0, data, data[end]+3)
    len = length(v)
    dct = Dict{Int,Int}()
    function helper(v, i)
        haskey(dct, i) && return dct[i]
        i == len && return 1
        n1 =               v[i+1] - v[i] <= 3 ? helper(v, i+1) : 0
        n2 = i+2 <= len && v[i+2] - v[i] <= 3 ? helper(v, i+2) : 0
        n3 = i+3 <= len && v[i+3] - v[i] <= 3 ? helper(v, i+3) : 0
        val = n1 + n2 + n3
        dct[i] = val
        return val
    end
    helper(v, 1)
end

main() = let v = read_data()
    @show part1(v)
    @show part2(v)
    nothing
end

# Performance

#=
julia> @btime part1($d)
  2.171 μs (28 allocations: 4.09 KiB)
1914

julia> @btime part2($d)
  13.625 μs (158 allocations: 10.22 KiB)
9256148959232
=#

#-- community solutions--

# James Doss-Gollin
function james()
    parse_file(fname) = parse.(Int, readlines(fname))

    function solve1(input::Vector{Int})
        push!(input, 0, maximum(input) + 3)
        input = sort(unique(input))
        diffs = diff(input)
        return sum(diffs .== 1) * sum(diffs .== 3)
    end

    function count_continuous_ones(Δ::Vector{Int})
        counts = Vector{Int}[]
        return string.(Δ) |> join |> x -> split(x, "3") |> x -> filter(y -> length(y) > 0, x) .|> length
    end

    function n_possible_steps(n::Int)
        if n == 1
            return 1
        elseif n == 2
            return 2
        elseif n == 3
            return 4
        else
            return n_possible_steps(n - 1) + n_possible_steps(n - 2) + n_possible_steps(n - 3)
        end
    end

    function solve2(input::Vector{Int})
        push!(input, 0, maximum(input) + 3)
        input = sort(unique(input))
        diffs = diff(input)
        return mapreduce(n_possible_steps, *, count_continuous_ones(diffs))
    end

    function main()
        input = parse_file(filename())
        sol1 = solve1(input)
        sol2 = solve2(input)
        @show sol1, sol2
    end

    main()
end

# Pablo Zubieta
function pablo()
    input = parse.(Int, eachline(filename()))
    sort!(input)
    prepend!(input, 0)
    append!(input, last(input) + 3)

    # Part 1
    d = diff(input)
    @show count(==(1), d) * count(==(3), d)

    #Part 2
    # This is only because maximum(length, s) == 4 for my example and the examples given, so it's not general enough
    @show s = split(join(d), '3')
    @show 2^count(==("11"), s) * 4^count(==("111"), s) * 7^count(==("1111"), s)
end

# Sukera
function sukera()
    function solve2(input)
        diff(input) |>
            x -> mapreduce(string, *, x) |>
            x -> findall(r"1+", x) .|>
            length .|>
            tribonacci |>
            x -> reduce(*, x)
    end

    function tribonacci(n)
        n <= 1 && return 1  # 1x => 1, x => 1s
        n == 2 && return 2  # 11x => 1.(x+1), .2x
        return tribonacci(n-1) + tribonacci(n-2) + tribonacci(n-3)
    end

    data = parse.(Int, readlines(filename())) |> sort
    data = vcat(0, data, data[end]+3)
    @show solve2(data)

    diff(data) |> x -> mapreduce(string, *, x) |> x -> findall(r"1+", x) .|> length .|> tribonacci
end


text:
using SubstitutionTilings
using SubstitutionTilings.Heisenberg
using Luxor

width = 1280
height = 1280
sc = 10
@draw begin
    colors = ["#21ABC2", "#43848F", "#F6504F", "#C2216C"]
    tiling = He(0,0,0)*substitute(heis(), dilate(-1//7 ,Heisenberg.He(0,10,1))*hn(0),9);
    println(sizeof(tiling))

    setline(0.5)
    for tile in tiling
        draw(tile, sc, colors[Heisenberg.color(tile)], :fill)

        draw(tile, sc, "black", :stroke)
    end
end width height


text:
# by lazarusA # HIDE
# using GLMakie # HIDE
using CairoMakie
CairoMakie.activate!() # HIDE
let
    x = -2π:0.1:2π
    fig = Figure(resolution = (700, 450), fontsize = 18, font = "sans")
    ax = Axis(fig, xlabel = "x", ylabel = "f(x)", xgridcolor = :red,
                xgridstyle=:dash, xgridwidth=0.85, xtickalign=1, xticksize=20)
    line1 = lines!(x, sin.(x), color = "#56B4E9", linewidth = 2)
    line2 = lines!(x, cos.(x), color = :black, linewidth = 2, linestyle = :dash)
    leg = Legend(fig, [line1, line2], ["sin(x)", "cos(x)"], tellheight = false,
    tellwidth = false, halign = :left, valign = :bottom, labelsize = 14,
    linewidth = 2, margin = (10, 10, 10, 10), framevisible = true)
    limits!(ax, -2π, 2π, -1,1)
    fig[1, 1] = ax
    fig[1, 1] = leg
    fig
    save(joinpath(@__DIR__, "output", "Fig2Lines.png"), fig, px_per_unit = 2) # HIDE
end

using Pkg # HIDE
Pkg.status("CairoMakie") # HIDE

max_stars_repo_name: MisterBiggs/BeautifulMakie

max_issues_repo_name: MisterBiggs/BeautifulMakie

max_forks_repo_name: MisterBiggs/BeautifulMakie


text:
type FastMMSE
    rho

    v
    bias

    FastMMSE() = new()
end

function FastMMSE(snr,c)
    est = FastMMSE()
    est.rho = 10^(0.1*snr)

    w = c./(c .+ (1/est.rho));
    est.v = fft(w)
    est.bias = sum(log.(1 .- w))

    return est
end

function estimate(est::FastMMSE, y)
    nAntennas  = size(y,1)
    nCoherence = size(y,2)
    nBatches   = size(y,3)

    y = fft(y,1)./sqrt(nAntennas)

    cest = mean(abs.(y).^2,2)
    cest = reshape(cest,nAntennas,nBatches)

    exps = real(ifft(conj(est.v).*fft(cest,1),1))*est.rho*nCoherence .+ est.bias*nCoherence
    exps = exps .- maximum(exps,1)

    weights = exp.(exps)
    weights = weights./sum(weights,1)

    F = real(ifft(est.v .* fft(weights,1),1));
    F = reshape(F,nAntennas,1,nBatches)

    x = y .* F;

    x = ifft(x,1).*sqrt(nAntennas);
end

max_stars_repo_name: tum-msv/learning-mmse-est


text:
# S. M. Katz modified code from E. R. Mueller
# Most code is kept in its original form
# Data was reprocessed using BayesNets.jl since I did not have the original final bayes net files
# Final Bayes Nets are stored in iBN_data.jld2 and tBN_data.jld2
# See the following citation and Eric's PhD thesis for more information about this model
# E. R. Mueller and M. J. Kochenderfer,
# “Simulation comparison of collision avoidance algorithms for small multi-rotor aircraft,”
# in AIAA Modeling and Simulation Technologies Conference, 2016, p. 3674

################################################################################################
############## THIS MIGHT HAVE A UNITS ISSUE SO CHECK HERE WHEN THINGS GO WRONG ################
mutable struct IntruderState
	altAGL::Float64
	altdot::Float64
	heading::Float64
	rrHeading::Float64
	rrHeaddot::Float64
	range::Float64
	vel::Float64   		# Note that this variable corresponds to dxdt in the BNs and the feature name
	vdot::Float64  		# Note that this variable corresponds to d2xdt2 in the BNs and the feature name
	x::Float64
	y::Float64
	x0::Float64
	y0::Float64
	xdd::Float64
	ydd::Float64
	t::Float64

	IntruderState(altAGL, altdot, heading, rrHeading, rrHeaddot, range, vel, vdot, x, y, x0, y0, xdd, ydd, t) =
			  new(altAGL, altdot, heading, rrHeading, rrHeaddot, range, vel, vdot, x, y, x0, y0, xdd, ydd, t)
end

function IntruderState(;
	altAGL::Float64=0.,
	altdot::Float64=0.,
	heading::Float64=0.,
	rrHeading::Float64=0.,
	rrHeaddot::Float64=0.,
	range::Float64=0.,
	vel::Float64=0.,
	vdot::Float64=0.,
	x::Float64=0.,
	y::Float64=0.,
	x0::Float64=0.,
	y0::Float64=0.,
	xdd::Float64=0.,
	ydd::Float64=0.,
	t::Float64=0.)

	IntruderState(altAGL, altdot, heading, rrHeading, rrHeaddot, range, vel, vdot, x, y, x0, y0, xdd, ydd, t)
end

mutable struct BNParams
	# Holds the discretized bin sizes for the Bayes net.  It would be nice to be able to read this in directly from the xdsl file,
	# but it doesn't appear to be stored there.

	binDisc::Dict

	BNParams(binDisc) = new(binDisc)

end

function BNParams()
# This discretization serves two functions: allows us to bin the continuous data before using Kyle's program (if we use that), and
# do continuous sampling within a bin once a set of discrete bins has been selected from the BN.  To do the latter, we need outer limits
# on the bin sizes, which adds two values to each discretization.  During the former binning, the first and last elements of each of
# these arrays is ignored and -Inf and +Inf are used instead (so no data is left out).

# A limitation of GeNIe is that discretization has to happen from zero, it cannot "wrap around".  This is a problem for states that
# represent angles because, for example, we'd like to bin everything from -22.5 to 22.5 in one bin, 22.5 to 67.5 in the second, etc.
# I should be able to add this capability myself if I do the binning.

	binDisc = Dict(
		:range     => [0., 11., 29., 101., 822., 10000.],
		:rrHeading => [0., 45., 90., 135., 180., 225., 270., 315., 360],
		:rrHeaddot =>  [-50., -8.698, -0.0324, 0.032, 8.079, 50.],
		:altAGL    => [0., 10., 50., 100., 400., 5000.],
		:altdot    => [-5, -0.3921, -0.060, 0.061, 0.41447, 5],
		:dxdt      =>  [0., 0.03, 0.10, 0.30, 1.75, 40.],
		:d2xdt2    => [-5., -0.46, -5.11e-6, 4.906e-5, 0.464, 5.],  # This may need to be narrower for the 2nd and 5th bins.
		:d2xdt2_tp1    => [-5., -0.46, -5.11e-6, 4.906e-5, 0.464, 5.],
		:altdot_tp1    => [-5, -0.3921, -0.060, 0.061, 0.41447, 5],
		:rrHeaddot_tp1 =>  [-50., -8.698, -0.0324, 0.032, 8.079, 50.])

		return BNParams(binDisc)
end
################################################################################################
################################################################################################

# Sample initial state distribution to obtain an initial state
function sampleBayesInitial(iBN::BayesNet; x0::Float64=0., y0::Float64=0., evidenceVal::Assignment=Assignment(), bnp::BNParams=BNParams(),
							maxAttempts::Int64=1000, bearingInit::Float64=0., VERBOSE::Bool=false)
	# The initial heading is not part of the bayes net, and shouldn't affect the dynamics or interactions with other aircraft, so a
	# default value of 0 will be used.
	# Function returns a continuous sample of the bayes net, optionally accepting set values for specific parameters in evidenceVal.
	# All variables coming out of this model will be in ft, ft/s, or ft/s^2.

	# x0, y0, x, y, xd, yd, xdd, ydd, h, hd = sampleBayesInitial(iBN)
	# x0, y0, x, y, xd, yd, xdd, ydd, h, hd = sampleBayesInitial(iBN, x0=0., y0=0., evidenceVal=Assignment(), bnp=BNParams(), maxAttempts=1000, bearingInit=0.)

	# If there is evidence, set that in the bayes net.  First have to get bins:
	evidenceBins = convertVals2Bins(iBN, evidenceVal, bnp=bnp)

	# Sample the network, getting bins.
	initSamp = rand(iBN, LikelihoodWeightedSampler(evidenceBins), 1) # rand_table_weighted(iBN, numSamples=1, consistentWith=evidenceBins)
	ind = 1
	while (size(initSamp,1)==0) && (ind<maxAttempts)
		initSamp = rand_table(iBN, numSamples=1, consistentWith=evidenceBins)
		ind+=1
	end
	if size(initSamp,1)==0
		if VERBOSE
			display("Could not find initial sample consistent with evidence.  Returning random sample.")
		end
		ind = 1
		while (size(initSamp,1)==0) && (ind<maxAttempts)
			initSamp = rand_table(iBN, numSamples=1)
			ind+=1
		end
	end

	# Convert the bin assignments to a table of boundaries for the bins
	binSampleBounds = convertBins2Bounds(iBN, initSamp, bnp=bnp)

	# Pull a continuous sample from the bins.
	contSample = disc2contSample(binSampleBounds)

	# Need to do the unit conversion right after sampling from the BN so that our state conversion equations are more
	# straightforward.  This should be cleaner becaues only the BN will be in aviation units, everything else will
	# be done in ft, ft/s, (check on altitude rate, make sure it's not in ft/min), radians.
	contSampleBasic = convertAvUnits2Basic(contSample)

	# Convert the sample from the existing variables to the starter variables
	x, y, xd, yd, xdd, ydd, h, hd = convertRRVars2EuclVars(contSampleBasic, x0=x0, y0=y0, bearingInit=bearingInit)

	# To test with a specific initial condition:
	# x=0.
	# y=0.
	# xd=0.
	# yd=0.
	# xdd=0.
	# ydd=0.
	# h=10.
	# hd=0.

	return x0, y0, x, y, xd, yd, xdd, ydd, h, hd

end

# Convert values to discretized bins
function convertVals2Bins(bn::BayesNet, evidenceVal::Assignment; bnp::BNParams=BNParams())

	evidenceBin = Assignment()

	for key in keys(evidenceVal)
		binInd=1
		while (binInd<length(bnp.binDisc[key])) && (evidenceVal[key]>bnp.binDisc[key][binInd+1])
			binInd+=1
		end
		# If the value is larger than the last boundary in the discretization, the value will be binned in the last bin.
		binInd = clamp(binInd,1,length(bnp.binDisc[key])-1) #clamp(binInd,1, length(domain(bn,key).elements))
		# evidenceBin[key] = bn.domains[bn.index[key]].elements[binInd]    # This line was for when indexing into an integer
		#****** I think the way I implemented the discretization, this will work ********
		evidenceBin[key] = binInd #bn.nodes[bn.name_to_index[key]].domain.elements[binInd]

		# evidenceBin[key] = binInd
	end

	return evidenceBin

end

function convertBins2Bounds(bn::BayesNet, samp::DataFrame; bnp::BNParams=BNParams())
	# Given a dataframe with a sample, look up the bin boundaries that correspond to the bin names (strings) that
	# comprise the domains of the BayesNet.  The nodes of the BayesNet, bn, will be labeled with the same symbol
	# as the dataframe columns.  The row of samp with data will have strings that are in the domains of the
	# appropriate node in the BN.

	j,n = size(samp)
	if j == 0
		display("Zero samples in bin->bound conversion, exiting.")
		return -1
	end

	if j>1
		display("More than one sample in bin->bound conversion, will only return the bins of the first: $(samp)")
		samp = samp[1,:]
	end

	binBounds = DataFrame()
	for key in names(samp)
		key = Symbol(key)
		if key!=:p  	# This is a probability key generated by the likelihood-weighted sampling process.  Shouldn't need it.
			bins = bnp.binDisc[key]   # This is an array of floats with the bin boundaries for key (one of the bin labels)
			binVal = samp[!, key][1]  	  # This is the text string label for the bin in the sample corresponding to the above bin discretization
			#indBin = find(bn.domains[bn.index[key]].elements.==binVal)[1]   # This is the bin index corresponding to the string label

			# SMK: I just need to find the bin index that is the bottom of the bin (maybe I already have this from the way I discretized?)
			# - I think I do ******
			binInd = binVal #find(bn.nodes[bn.name_to_index[key]].domain.elements.==binVal)[1]

			#binVal = samp[key][1]
			#binInd = find(domain(bn,key).elements.==binVal)[1]
			binBounds[!, key] = bins[binInd:binInd+1]
		end
	end

	return binBounds

end

function disc2contSample(binBounds::DataFrame)

	contSamp = DataFrame()
	for key in names(binBounds)
		l,h = binBounds[!, key]
		contSamp[!, key] = [(h-l)*rand() + l]
	end

	return contSamp

end

################### Will be interesting to see if there are issues here ################
function convertAvUnits2Basic(samp::DataFrame)

	sampConv = DataFrame()
	for key in names(samp)

		if key == :rrHeaddot
			sampConv[!, key] = pi/180 * samp[!, key]
		elseif key == :rrHeading
			sampConv[!, key] = pi/180 * samp[!, key]
		elseif key == :rrHeaddot_tp1
			sampConv[!, key] = pi/180 * samp[!, key]
		else
			sampConv[!, key] = samp[!, key]
		end
	end

	return sampConv
end

# Converts bayes net sampled params to actual variables needed to propagate the trajectory
function convertRRVars2EuclVars(contSample::DataFrame; x0::Float64=0., y0::Float64=0., bearingInit::Float64=0.)
	# bearingInit in degrees

	R = contSample[!, :range][1]
	rrHeading = contSample[!, :rrHeading][1]
	rrHeaddot = contSample[!, :rrHeaddot][1]
	v = contSample[!, :dxdt][1]
	vd = contSample[!, :d2xdt2][1]

	x = x0 + R*sin(bearingInit)
	y = y0 + R*cos(bearingInit)

	xd = v*sin(bearingInit+rrHeading)
	yd = v*cos(bearingInit+rrHeading)

	bearingdot = (xd*(y-y0)-yd*(x-x0))/R^2

	xdd = vd*sin(bearingInit+rrHeading) + v*(rrHeaddot+bearingdot)*cos(bearingInit+rrHeading)
	ydd = vd*cos(bearingInit+rrHeading) - v*(rrHeaddot+bearingdot)*sin(bearingInit+rrHeading)

	h = contSample[!, :altAGL][1]
	hd = contSample[!, :altdot][1]

	return x, y, xd, yd, xdd, ydd, h, hd

end

# Sets intruder state
function setInitState(x0::Float64, y0::Float64, x::Float64, y::Float64, xd::Float64, yd::Float64, xdd::Float64, ydd::Float64,
					  h::Float64, hd::Float64; t::Float64=0.)
# Will probably just want to specify a few states and then calculate all the others.  This function will do that.

	dx = x-x0
	dy = y-y0
	R = sqrt(dx^2+dy^2)

	psi_t = atan(dx, dy)
	psi_t_dot = 0.
	if R!=0
		psi_t_dot = (xd*dy - yd*dx)/R^2
	end

	psi_h = atan(xd, yd)
	psi_h_dot = 0.
	if (xd^2+yd^2)!=0.
		psi_h_dot = (xdd*yd - ydd*xd)/(xd^2+yd^2)
	end

	psi_v = mod(psi_h - psi_t, 2*pi)
	psi_v_dot = psi_h_dot - psi_t_dot

	v = sqrt(xd^2+yd^2)
	v_dot = sqrt(xdd^2 + ydd^2)

	IntruderState(h, hd, psi_h, psi_v, psi_v_dot, R, v, v_dot, x, y, x0, y0, xdd, ydd, t)
end

function trajComplete(iState::IntruderState; tMax::Float64=120)
	# For now, will just simulate to a max time

	if iState.t>= tMax
		return true
	end

	return false
end

function sampleBayesTransition(tBN::BayesNet, iState::IntruderState; bnp::BNParams=BNParams(), maxAttempts::Int64=1000, VERBOSE::Bool=false)

	# Set the given evidence (pulls out the relevant state from iState, sets bins)
	#display("$iState")
	evidenceBins = setTransitionEvidence(tBN, iState, bnp=bnp)

	# Pull a sample from network, getting bins
	tSamp = rand(tBN, LikelihoodWeightedSampler(evidenceBins), 1) #rand_table_weighted(tBN, numSamples=1, consistentWith=evidenceBins)
	ind=1
	while (size(tSamp,1)==0) && (ind<maxAttempts)
		tSamp = #rand_table_weighted(tBN, numSamples=1, consistentWith=evidenceBins)
		ind+=1
	end
	if size(tSamp,1)==0
		if VERBOSE
			display("Could not find transition sample consistent with evidence.  Returning random sample.")
		end
		ind = 1
		while (size(tSamp,1)==0) && (ind<maxAttempts)
			tSamp = rand_table_weighted(tBN, numSamples=1)
			ind+=1
		end
	end

	# Go through each bin, comparing whether it is a different bin from the last bin.  If it's the same, resample a continuous point
	# with probability 1-epsilon (epsilon ~ 0.02-0.08, or learn that value from the data). If its a different bin, just sample agoin.
	# As a first implementation I could just resample every time.

	# Convert the bin assignments to a table of boundaries for the bins
	binSampleBounds = convertBins2Bounds(tBN, tSamp, bnp=bnp)

	# Pull a continuous sample from the bins.
	contSample = disc2contSample(binSampleBounds)

	# Check whether the bin contains zero. If it does, return exactly zero.  Note that this will operate on all fields, while
	# we only care about zeroing out the t+1 fields (altdot, rrHeaddot and vdot).  Implementation will be cleaner if we
	# zero out every field and then just return the appropriate ones.
	zeroOutSpanningBins!(contSample, binSampleBounds)

	# Need to do the unit conversion right after sampling from the BN so that our state conversion equations are more
	# straightforward.  This should be cleaner becaues only the BN will be in aviation units, everything else will
	# be done in ft, ft/s, (check on altitude rate, make sure it's not in ft/min), radians.
	contSampleBasic = convertAvUnits2Basic(contSample)

	hdot_t2 = contSampleBasic[!, :altdot_tp1][1]
	vdot_t2 = contSampleBasic[!, :d2xdt2_tp1][1]
	psivdot_t2 = contSampleBasic[!, :rrHeaddot_tp1][1]

	return hdot_t2, vdot_t2, psivdot_t2

end

function setTransitionEvidence(tBN::BayesNet, iState::IntruderState; bnp::BNParams=BNParams())
# this function sets the seven relevant feature values in iState into a dictionary, returning a mapping from
# features to feature values (not bins).  These feature

	evidenceVal = Assignment()

	evidenceVal[:altAGL] = iState.altAGL
	evidenceVal[:altdot] = iState.altdot
	evidenceVal[:d2xdt2] = iState.vdot
	evidenceVal[:rrHeaddot] = iState.rrHeaddot
	evidenceVal[:range] = iState.range
	evidenceVal[:dxdt] = iState.vel
	evidenceVal[:rrHeading] = iState.rrHeading

	convertVals2Bins(tBN, evidenceVal, bnp=bnp)

end

function zeroOutSpanningBins!(sample, binSampleBounds)

	for key in names(sample)
		if (maximum(binSampleBounds[!, key])>0.) & (minimum(binSampleBounds[!, key])<0.)
			sample[!, key] = [0.]
		end
	end

end

function getNextState(intState::IntruderState, hdot_t2::Float64, vdot_t2::Float64, psivdot_t2::Float64, dt::Float64;
					  psiTol::Float64=1e-5, psidotTol::Float64=1e-6, maxIters::Int64=20, VERBOSE::Bool=false)
# Given a current state and the values of the d/dt terms (h, v, psi) at the next time step, return
# a state that represents the next state.  Note that psivdot is the range-relative heading rate, a pretty wierd
# state, but one that hopefully captures the necessary behavior from the simulated trajectories.

	# Basic state updates:
	h_t1 = intState.altAGL
	hdot_t1 = intState.altdot
	v_t1 = intState.vel
	vdot_t1 = intState.vdot
	psiv_t1 = intState.rrHeading
	psivdot_t1 = intState.rrHeaddot

	h_t2 = h_t1 + 0.5*(hdot_t1 + hdot_t2)*dt
	v_t2 = v_t1 + 0.5*(vdot_t1 + vdot_t2)*dt
	psiv_t2 = mod(psiv_t1 + 0.5*(psivdot_t1 + psivdot_t2)*dt, 2*pi)

	# If v_t2 < 0 then strange things are going to happen.  My assumption is that velocity is always positive, but the
	# Bayes net could certainly sample a "deceleration" (negative vdot_t2) when velocity is low, which would indeed
	# cause a negative velocity.  I could either just set the velocity to zero in that case, or take the absolute value
	# of the velocity and rotate the heading and rrheading vectors 180 degrees.  Would have ot figure out what to do
	# with psivdot_t2
	if v_t2 < 0.
		v_t2=0.
		vdot_t2=0.
	end
	t1 = intState.t

	# Derived variable state updates:
	x_t1 = intState.x
	y_t1 = intState.y
	x0 = intState.x0
	y0 = intState.y0
	psih_t1 = intState.heading
	xd_t1 = v_t1*sin(psih_t1)
	yd_t1 = v_t1*cos(psih_t1)
	xdd_t1 = intState.xdd
	ydd_t1 = intState.ydd
	v_t1==0. ? psihdot_t1=0. : psihdot_t1 = (xdd_t1*yd_t1 - ydd_t1*xd_t1)/(v_t1^2)

	# Initialize variables for iterative psihdot search (plus other variables)
	psih_t2 = psih_t1
	psihdot_t2 = psihdot_t1
	psih_old = Inf
	psihdot_old = Inf
	iters = 0
	xdd_t2 = xdd_t1
	ydd_t2 = ydd_t1
	xd_t2 = xd_t1
	yd_t2 = yd_t1
	x_t2 = x_t1
	y_t2 = y_t1
	R_t2 = sqrt(x_t1^2+y_t1^2)
	while ((abs(psih_t2-psih_old) > psiTol) | (abs(psihdot_t2-psihdot_old) > psidotTol)) & (iters<maxIters)
		xdd_t2 = vdot_t2*sin(psih_t2) + v_t2*psihdot_t2*cos(psih_t2)
		ydd_t2 = vdot_t2*cos(psih_t2) - v_t2*psihdot_t2*sin(psih_t2)
		xd_t2 = xd_t1 + 0.5*(xdd_t1+xdd_t2)*dt
		yd_t2 = yd_t1 + 0.5*(ydd_t1+ydd_t2)*dt
		x_t2 = x_t1 + 0.5*(xd_t1+xd_t2)*dt
		y_t2 = y_t1 + 0.5*(yd_t1+yd_t2)*dt
		R_t2 = sqrt(x_t2^2+y_t2^2)

		psih_old = psih_t2
		psihdot_old = psihdot_t2
		psih_t2 = atan(xd_t2,yd_t2)
		#	psihdot_t2 = (xdd_t2*yd_t2 - ydd_t2*xd_t2)/(xd_t2^2+yd_t2^2)  # Try the following instead?:
		R_t2==0. ? psitdot_t2 = 0. : psitdot_t2 = (xd_t2*(y_t2-y0) - yd_t2*(x_t2-x0))/R_t2^2
		psihdot_t2 = psivdot_t2-psitdot_t2

		# if VERBOSE
		# 	display("After iteration $(iters+1), error on heading is $(psih_t2-psih_old) and error on heading rate is $(psihdot_t2-psihdot_old)")
		# end

		iters+=1
	end
	t2 = t1+dt

	IntruderState(h_t2, hdot_t2, psih_t2, psiv_t2, psivdot_t2, R_t2, v_t2, vdot_t2, x_t2, y_t2, intState.x0, intState.y0, xdd_t2, ydd_t2, t2)

end

max_issues_repo_name: sisl/em-model-hobbydrone