text:
#!/usr/bin/julia

using GLM, CSV, StatsModels, DataFrames, DataFramesMeta, Random, Statistics

mse(y1, y2) = mean((y1 - y2) .^ 2)

df = CSV.read("../csv/corei7_long.csv")

iterations = 100

cross_validation = DataFrame(MSE = [],
                             R2 = [],
                             PredictedParallel = [],
                             RealCompileTime = [],
                             PredictedCompileTime = [])

for i = 1:iterations
    println(i)
    train_ratio = 0.1
    train_size = round(Int, train_ratio * nrow(df))
    train_rows = shuffle(1:nrow(df))[1:train_size]
    df_train = df[train_rows, :]

    test_rows = [x for x in 1:nrow(df) if !(x in train_rows)]
    df_test = df[test_rows, :]

    regression = lm(@formula(CompileTime ~ ((Expected_insns +
                                             Expected_insns ^ 2) +
                                            (Functions +
                                             Functions ^ 2) +
                                            (Inlined_percentage +
                                             Inlined_percentage ^ 2) +
                                            (Num_partitions +
                                             Num_partitions ^ 2)) *
                             (Parallel + Parallel ^ 2)), df_train)

    prediction = predict(regression, select(df_test, Not(:CompileTime)))

    best_prediction = df_test[findmin(prediction)[2], :]

    push!(cross_validation, (MSE = mse(df_test[:, :CompileTime], prediction),
                             R2 = adjr2(regression),
                             PredictedCompileTime = findmin(prediction)[1],
                             RealCompileTime = best_prediction[:CompileTime],
                             PredictedParallel = best_prediction[:Parallel]))

end

CSV.write("../csv/cross_validation.csv", cross_validation)

cv_df = CSV.read("../csv/cross_validation.csv")
cv_df

println(mean(cv_df.MSE))
println((1.96 * std(cv_df.MSE)) / sqrt(length(cv_df.MSE)))

println(mean(cv_df.R2))
println((1.96 * std(cv_df.R2)) / sqrt(length(cv_df.R2)))

max_stars_repo_name: giulianobelinassi/LTO-Timing-Analysis

max_issues_repo_name: giulianobelinassi/LTO-Timing-Analysis

max_forks_repo_name: giulianobelinassi/LTO-Timing-Analysis


text:
module PkgServer

using Pkg
using HTTP
using Base.Threads: Event, @spawn
import Base: fetch
using Random
using LibGit2
using FilesystemDatastructures
using SimpleBufferStream
using JSON3, StructTypes
using Sockets
using Sockets: InetAddr
using Dates
using Tar
using Gzip_jll

include("task_utils.jl")
include("resource.jl")
include("meta.jl")
include("dynamic.jl")
mutable struct RegistryMeta
    # Upstream registry URL (e.g. "https://github.com/JuliaRegistries/General")
    upstream_url::String
    # The latest hash we know about for this registry
    latest_hash::Union{Nothing,String}

    function RegistryMeta(url::String)
        # Check to ensure this path actually exists
        if !url_exists(url)
            throw(ArgumentError("Invalid unreachable registry '$(url)'"))
        end

        # Auto-detect a repository that doesn't have `.git` at the end but could
        git_url = string(url, ".git")
        if !endswith(url, ".git") && url_exists(git_url)
            url = git_url
        end
        return new(url, nothing)
    end
end

struct ServerConfig
    root::String
    listen_addr::InetAddr
    listen_server::Ref{Any}
    cache::SizeConstrainedFileCache
    registries::Dict{String,RegistryMeta}
    storage_servers::Vector{String}

    # Default server config constructor
    function ServerConfig(; listen_addr = InetAddr(ip"127.0.0.1", 8000),
                            storage_root = "/tmp/pkgserver",
                            registries = Dict(
                                "23338594-aafe-5451-b93e-139f81909106" =>
                                RegistryMeta("https://github.com/JuliaRegistries/General")
                            ),
                            storage_servers = [
                                "https://us-east.storage.juliahub.com",
                                "https://kr.storage.juliahub.com",
                            ],
                            keep_free=3*1024^3)
        # Right now, the only thing we store in `static/` is `/registries`
        mkpath(joinpath(storage_root, "static"))
        # Downloads get stored into `temp`
        mkpath(joinpath(storage_root, "temp"))
        # Files get stored into `cache`
        mkpath(joinpath(storage_root, "cache"))
        return new(
            storage_root,
            listen_addr,
            Ref(nothing),
            SizeConstrainedFileCache(
                joinpath(storage_root, "cache"),
                TargetSizeKeepFree(keep_free),
                DiscardLRU()
            ),
            registries,
            sort!(storage_servers),
        )
    end
end

# Initialize a default config that will get overridden by `start()` below
global config = ServerConfig()

function __init__()
    # Set default HTTP useragent
    HTTP.setuseragent!("PkgServer (HTTP.jl)")

    # Record our starting time
    global time_start = now()
end

function start(;kwargs...)
    @info("Loading initial Server Config")
    # Create global config object
    global config = ServerConfig(;kwargs...)

    # Update registries first thing
    @info("Performing initial registry update")
    update_registries()
    global last_registry_update = now()

    # Experimental.@sync throws if _any_ of the tasks fail
    Base.Experimental.@sync begin
        global registry_update_task = @spawn begin
            while true
                sleep(1)
                @try_printerror begin
                    forget_failures()
                    update_registries()
                    last_registry_update = now()
                end
            end
        end

        # Registry watchdog; if `last_registry_update` doesn't change
        # for more than 20 minutes, we kill ourselves so that we can restart
        max_time_lag = Second(20 * 60)
        global registry_watchdog_task = @spawn begin
            while true
                time_lag = now() - last_registry_update
                if time_lag > max_time_lag
                    task_result = try fetch(registry_update_task); catch err; err end
                    @error "registry update watchdog timer tripped" time_lag max_time_lag task_result
                    exit(1)
                end
                sleep(max_time_lag.value)
            end
        end

        listen_server = Sockets.listen(config.listen_addr)
        config.listen_server[] = listen_server
        @info("server listening", config.listen_addr)
        num_requests = 0
        HTTP.listen(config.listen_addr.host, config.listen_addr.port; server=listen_server, max_connections=10) do http
            num_requests += 1
            resource = http.message.target
            request_id = HTTP.header(http, "X-Request-ID", "")
            # If the user is asking for `/meta`, generate the requisite JSON object and send it back
            if resource == "/meta"
                serve_meta(http)
                return
            end
            if resource == "/meta/stats"
                serve_meta_stats(http)
                return
            end
            if resource == "/meta/siblings"
                serve_siblings(http)
                return
            end
            if resource == "/meta/children"
                # PkgServers don't have children, but to be good neighbors,
                # we return an empty list here, to make it easier to recurse.
                serve_children(http)
                return
            end
            if resource == "/meta/parents"
                serve_parents(http)
                return
            end
            if resource == "/robots.txt"
                serve_robots_txt(http)
                return
            end

            if resource  == "/registries"
                open(joinpath(config.root, "static", "registries")) do io
                    serve_file(http, io, "text/plain")
                end
                return
            end

            # If the user asked for something that is an actual resource, send it directly
            if occursin(resource_re, resource)
                # If the resource already exists locally, yay!  Serve it and quit.
                resource_path = resource_filepath(resource)
                io = try_open(resource_path)
                if io !== nothing
                    hit!(config.cache, resource[2:end])
                    serve_file(http, io, "application/x-gzip")
                    close(io)
                    return
                end

                # If it doesn't exist locally, let's request a fetch on that resource.
                # This will return either `nothing` (e.g. resource does not exist) or
                # a `DownloadState` that represents a partial download.
                dl_state = fetch_resource(resource, request_id)
                if dl_state !== nothing
                    HTTP.setheader(http, "X-Cache-Miss" => "miss")
                    stream_path = temp_resource_filepath(resource)
                    # Wait until `stream_path` is created
                    while !isfile(stream_path) && dl_state.dl_task.state != :done
                        sleep(0.001)
                    end
                    # Try to serve `stream_path` file
                    stream_io = try_open(stream_path)
                    if stream_io !== nothing
                        serve_file(http, stream_io, "application/x-gzip";
                                   content_length=dl_state.content_length,
                                   dl_task=dl_state.dl_task)
                        close(stream_io)
                        return
                    end

                    # If we couldn't open `stream_path`, it may be because the file finished
                    # downloading since we last checked 20 lines ago.  Check again.
                    io = try_open(resource_path)
                    if io !== nothing
                        serve_file(http, io, "application/x-gzip")
                        close(io)
                        return
                    end
                end
            end

            # If a user asks for an `Artifacts.toml` for a particular artifact, we'll generate
            # one for them on the fly.
            m = match(artifact_toml_re, resource)
            if m !== nothing
                artifact_resource = m.captures[1]
                artifact_name = m.captures[2]

                # If we don't actually have the artifact locally, we need to fetch it first.
                artifact_path = resource_filepath(artifact_resource)
                if !isfile(artifact_path)
                    dl_state = fetch_resource(artifact_resource, request_id)
                    if dl_state !== nothing
                        # We actually need to wait for the download to finish so that we can
                        # calculate the SHA256 hash of the tarball.
                        wait(dl_state.dl_task)
                    end
                end

                # We check again, because it's possible the above download failed.
                if isfile(artifact_path)
                    serve_artifact_toml(http, artifact_path, artifact_name)
                    return
                end
            end

            if occursin(r"^/*$", resource)
                path = joinpath(dirname(@__DIR__), "static", "index.html")
                io = try_open(path)
                if io !== nothing
                    serve_file(http, io, "text/html")
                    close(io)
                    return
                end
            end

            global total_misses += 1
            HTTP.setstatus(http, 404)
            startwrite(http)
        end
    end
end

# precompilation
include(joinpath(dirname(@__DIR__), "deps", "precompile.jl"))
if get(ENV, "PKGSERVER_GENERATING_PRECOMPILE", nothing) === nothing
    _precompile_()
end

end # module

ext: jl


text:
using DotEnv, Pkg

"""
See README.md for build notes.
"""

env = joinpath(@__DIR__, ".env")
run(`touch $env`) # make sure .env file exists
DotEnv.config(path=env)
env_keys = ["gmshjl", "PYTHON"]

# configure gmsh
find_gmshjl() =
    try # search for gmsh.jl in /usr/local/
        path = read(pipeline(`find /usr/local/ -name "gmsh.jl"`, `head -n 1`), String)
        path[1:end-1]
    catch
        ""
    end
key = env_keys[1]
if !(key in keys(ENV))
    value = find_gmshjl()
    if value == ""
        # install gmsh
        install_gmsh_sh = joinpath(@__DIR__, "install_gmsh.sh")
        run(`bash $install_gmsh_sh`)
        value = find_gmshjl()
        @assert value != ""
    end
    ENV[key] = value
    println("$key=$value")
end

# configure python
find_conda() =
    try # use conda in $PATH
        path = read(`which conda`, String)
        path[1:end-1]
    catch
        try # check if miniconda3 is already in home directory
            path = joinpath(homedir(), "miniconda3", "bin", "conda")
            run(`test -f $path`) # check if file exists
            path
        catch
            ""
        end
    end
find_python(conda::String) =
    try
        path = read(pipeline(`$conda env list`, `grep DiscretePDEs`, `awk '{print $2}'`),
            String)
        joinpath(path[1:end-1], "bin", "python")
    catch
        ""
    end
key = env_keys[2]
if !(key in keys(ENV))
    conda = find_conda()
    if conda == ""
        # install conda
        install_conda_sh = joinpath(@__DIR__, "install_conda.sh")
        run(`bash $install_conda_sh`)
        conda = find_conda()
        @assert conda != ""
    end
    value = find_python(conda)
    if value == ""
        # create a virtual environment
        conda_env_yml = joinpath(@__DIR__, "conda_env.yml")
        run(`$conda env create -f $conda_env_yml`)
        value = find_python(conda)
        @assert value != ""
    end
    ENV[key] = value
    println("$key=$value")
end
Pkg.build("PyCall")

# save to .env
env_contents = reduce(*, ["$k=$(ENV[k])\n" for k in env_keys])
open(env, "w") do io
    write(io, env_contents)
end


text:
""" Serialization of raw and semantic flow graphs.
"""
module Serialization
export read_raw_graphml, read_raw_graph_json,
  read_semantic_graphml, read_semantic_graph_json

using Nullables

using Catlab, Catlab.WiringDiagrams
import Catlab.WiringDiagrams: convert_from_graph_data, convert_to_graph_data
using ..Doctrine
using ..RawFlowGraphs

# Raw flow graphs
#################

read_raw_graphml(xml) = read_graphml(RawNode, RawPort, Nothing, xml)
read_raw_graph_json(json) = read_json_graph(RawNode, RawPort, Nothing, json)

function convert_from_graph_data(::Type{RawNode}, data::AbstractDict)
  annotation = to_nullable(String, pop!(data, "annotation", nothing))
  annotation_index = to_nullable(Int, pop!(data, "annotation_index", nothing))
  annotation_kind_str = to_nullable(String, pop!(data, "annotation_kind", nothing))
  annotation_kind = isnull(annotation_kind_str) ? FunctionAnnotation :
    convert(RawNodeAnnotationKind, get(annotation_kind_str))
  RawNode(data, annotation, annotation_index, annotation_kind)
end

function convert_from_graph_data(::Type{RawPort}, data::AbstractDict)
  annotation = to_nullable(String, pop!(data, "annotation", nothing))
  annotation_index = to_nullable(Int, pop!(data, "annotation_index", nothing))
  value = to_nullable(Any, pop!(data, "value", nothing))
  RawPort(data, annotation, annotation_index, value)
end

to_nullable(T::Type, x) = isnothing(x) ? Nullable{T}() : Nullable{T}(x)

# Semantic flow graphs
######################

read_semantic_graphml(xml) = read_graphml(
  Union{Monocl.Hom,Nothing}, Union{Monocl.Ob,Nothing}, Nothing, xml)
read_semantic_graph_json(json) = read_json_graph(
  Union{Monocl.Hom,Nothing}, Union{Monocl.Ob,Nothing}, Nothing, json)

function convert_from_graph_data(::Type{Monocl.Ob}, data::AbstractDict)
  parse_json_sexpr(Monocl, data["ob"]; symbols=false)
end
function convert_from_graph_data(::Type{Monocl.Hom}, data::AbstractDict)
  parse_json_sexpr(Monocl, data["hom"]; symbols=false)
end
function convert_from_graph_data(::Type{Union{Monocl.Ob,Nothing}}, data::AbstractDict)
  isempty(data) ? nothing : convert_from_graph_data(Monocl.Ob, data)
end
function convert_from_graph_data(::Type{Union{Monocl.Hom,Nothing}}, data::AbstractDict)
  isempty(data) ? nothing : convert_from_graph_data(Monocl.Hom, data)
end

function convert_to_graph_data(expr::Monocl.Ob)
  Dict("ob" => to_json_sexpr(expr))
end
function convert_to_graph_data(expr::Monocl.Hom)
  Dict("hom" => to_json_sexpr(expr))
end

function convert_from_graph_data(::Type{MonoclElem}, data::AbstractDict)
  ob = haskey(data, "ob") ?
    parse_json_sexpr(Monocl, data["ob"]; symbols=false) : nothing
  value = get(data, "value", Nullable())
  MonoclElem(ob, value)
end

function convert_to_graph_data(elem::MonoclElem)
  data = Dict{String,Any}()
  if (elem.ob != nothing) data["ob"] = to_json_sexpr(elem.ob) end
  if (!isnull(elem.value)) data["value"] = get(elem.value) end
  return data
end

end

max_stars_repo_name: cclauss/semanticflowgraph

max_issues_repo_name: cclauss/semanticflowgraph

max_forks_repo_name: cclauss/semanticflowgraph


text:
module PlotlyJS

using Reexport
@reexport using PlotlyBase
using JSON
using REPL, Pkg

# need to import some functions because methods are meta-generated
import PlotlyBase:
    restyle!, relayout!, update!, addtraces!, deletetraces!, movetraces!,
    redraw!, extendtraces!, prependtraces!, purge!, to_image, download_image,
    restyle, relayout, update, addtraces, deletetraces, movetraces, redraw,
    extendtraces, prependtraces, prep_kwargs, sizes, savefig, _tovec, html_body,
    react, react!

using WebIO
using JSExpr
using JSExpr: @var, @new
using Blink
using Requires

export plot

# globals for this package
const _pkg_root = dirname(dirname(@__FILE__))
const _js_path = joinpath(_pkg_root, "assets", "plotly-latest.min.js")
const _js_cdn_path = "https://cdn.plot.ly/plotly-latest.min.js"
const _mathjax_cdn_path =
    "https://cdnjs.cloudflare.com/ajax/libs/mathjax/2.7.1/MathJax.js?config=TeX-AMS-MML_SVG"

struct PlotlyJSDisplay <: AbstractDisplay end

# include the rest of the core parts of the package
include("display.jl")
include("util.jl")

function docs()
    schema_path = joinpath(dirname(dirname(@__FILE__)), "deps", "schema.html")
    if !isfile(schema_path)
        msg = "schema docs not built. Run `Pkg.build(\"PlotlyJS\")` to generate"
        error(msg)
    end
    w = Blink.Window()
    wait(w.content)
    Blink.content!(w, "html", open(f->read(f, String), schema_path), fade=false, async=false)
end

PlotlyBase.savefig(p::SyncPlot, a...; k...) = savefig(p.plot, a...; k...)
PlotlyBase.savefig(io::IO, p::SyncPlot, a...; k...) = savefig(io, p.plot, a...; k...)

for (mime, fmt) in PlotlyBase._KALEIDO_MIMES
    @eval function Base.show(
        io::IO, ::MIME{Symbol($mime)}, plt::SyncPlot,
        width::Union{Nothing,Int}=nothing,
        height::Union{Nothing,Int}=nothing,
        scale::Union{Nothing,Real}=nothing,
    )
        savefig(io, plt.plot, format = $fmt)
    end
end

function __init__()
    _build_log = joinpath(_pkg_root, "deps", "build.log")
    if isfile(_build_log) && occursin("Warning:", read(_build_log, String))
        @warn("Warnings were generated during the last build of PlotlyJS:  please check the build log at $_build_log")
    end

    if !isfile(_js_path)
        @info("plotly.js javascript libary not found -- downloading now")
        include(joinpath(_pkg_root, "deps", "build.jl"))
    end

    insert!(Base.Multimedia.displays, findlast(x -> x isa Base.TextDisplay || x isa REPL.REPLDisplay, Base.Multimedia.displays) + 1, PlotlyJSDisplay())

    atreplinit(i -> begin
        while PlotlyJSDisplay() in Base.Multimedia.displays
            popdisplay(PlotlyJSDisplay())
        end
        insert!(Base.Multimedia.displays, findlast(x -> x isa REPL.REPLDisplay, Base.Multimedia.displays) + 1, PlotlyJSDisplay())
    end)
end

end # module


text:
# See https://udel.edu/~arthij/MD.pdf for information on forces
# See https://arxiv.org/pdf/1401.1181.pdf for applying forces to atoms
# See OpenMM documentation and Gromacs manual for other aspects of forces

export
    mass,
    ustripvec,
    force,
    accelerations,
    LennardJones,
    SoftSphere,
    Mie,
    Coulomb,
    CoulombReactionField,
    Gravity,
    HarmonicBond,
    HarmonicAngle,
    PeriodicTorsion,
    RBTorsion

"""
    mass(atom)

The mass of an atom.
"""
mass(atom::Atom) = atom.mass

"""
    ustripvec(x)

Broadcasted form of `ustrip` from Unitful.jl, allowing e.g. `ustripvec.(coords)`.
"""
ustripvec(x) = ustrip.(x)

function checkforcetype(fdr, force_type)
    if unit(first(fdr)) != force_type
        error("Simulation force type is ", force_type,
                " but encountered force type ", unit(first(fdr)))
    end
end

"""
    force(inter, coord_i, coord_j, atom_i, atom_j, box_size)

Calculate the force between a pair of atoms due to a given interation type.
Custom interaction types should implement this function.
"""
function force end

# Allow 2D broadcasting whilst eliminating the diagonal corresponding to self interaction
@inline @inbounds function force(inter, coord_i, coord_j, atom_i, atom_j, box_size, self_interaction, force_type)
    if isone(self_interaction)
        return ustrip.(zero(coord_i))
    else
        fdr = force(inter, coord_i, coord_j, atom_i, atom_j, box_size)
        checkforcetype(fdr, force_type)
        return ustrip.(fdr)
    end
end

@inline @inbounds function force!(forces, inter, s::Simulation, i::Integer, j::Integer, force_type, weight_14::Bool=false)
    fdr = force(inter, s.coords[i], s.coords[j], s.atoms[i], s.atoms[j], s.box_size)
    checkforcetype(fdr, force_type)
    if weight_14
        fdr_ustrip = ustrip.(fdr) * inter.weight_14
    else
        fdr_ustrip = ustrip.(fdr)
    end
    forces[i] -= fdr_ustrip
    forces[j] += fdr_ustrip
    return nothing
end

"""
    accelerations(simulation; parallel=true)

Calculate the accelerations of all atoms using the general and specific
interactions and Newton's second law.
"""
function accelerations(s::Simulation; parallel::Bool=true)
    n_atoms = length(s.coords)

    if parallel && nthreads() > 1 && n_atoms >= 100
        forces_threads = [ustripvec.(zero(s.coords)) for i in 1:nthreads()]

        # Loop over interactions and calculate the acceleration due to each
        for inter in values(s.general_inters)
            if inter.nl_only
                neighbors = s.neighbors
                @threads for ni in 1:length(neighbors)
                    i, j, w = neighbors[ni]
                    force!(forces_threads[threadid()], inter, s, i, j, s.force_unit, w)
                end
            else
                @threads for i in 1:n_atoms
                    for j in 1:(i - 1)
                        force!(forces_threads[threadid()], inter, s, i, j, s.force_unit)
                    end
                end
            end
        end

        forces = sum(forces_threads)
    else
        forces = ustripvec.(zero(s.coords))

        for inter in values(s.general_inters)
            if inter.nl_only
                neighbors = s.neighbors
                for ni in 1:length(neighbors)
                    i, j, w = neighbors[ni]
                    force!(forces, inter, s, i, j, s.force_unit, w)
                end
            else
                for i in 1:n_atoms
                    for j in 1:(i - 1)
                        force!(forces, inter, s, i, j, s.force_unit)
                    end
                end
            end
        end
    end

    for inter_list in values(s.specific_inter_lists)
        sparse_forces = force.(inter_list, (s.coords,), (s,))
        ge1, ge2 = getindex.(sparse_forces, 1), getindex.(sparse_forces, 2)
        checkforcetype(first(first(ge2)), s.force_unit)
        sparse_vec = sparsevec(reduce(vcat, ge1), reduce(vcat, ge2), n_atoms)
        forces += ustripvec.(Array(sparse_vec))
    end

    return (forces * s.force_unit) ./ mass.(s.atoms)
end

function accelerations(s::Simulation, coords, coords_is, coords_js, atoms_is, atoms_js, self_interactions)
    n_atoms = length(coords)
    forces = ustripvec.(zero(coords))

    for inter in values(s.general_inters)
        # Currently the neighbor list is not used for this implementation
        forces -= reshape(sum(force.((inter,), coords_is, coords_js, atoms_is, atoms_js,
                                        (s.box_size,), self_interactions, s.force_unit);
                                        dims=2), n_atoms)
    end

    for inter_list in values(s.specific_inter_lists)
        # Take coords off the GPU if they are on there
        coords_cpu = Array(coords)
        sparse_forces = force.(inter_list, (coords_cpu,), (s,))
        ge1, ge2 = getindex.(sparse_forces, 1), getindex.(sparse_forces, 2)
        checkforcetype(first(first(ge2)), s.force_unit)
        sparse_vec = sparsevec(reduce(vcat, ge1), reduce(vcat, ge2), n_atoms)
        # Move back to GPU if required
        forces += convert(typeof(forces), ustripvec.(Array(sparse_vec)))
    end

    return (forces * s.force_unit) ./ mass.(s.atoms)
end

include("interactions/lennard_jones.jl")
include("interactions/soft_sphere.jl")
include("interactions/mie.jl")
include("interactions/coulomb.jl")
include("interactions/coulomb_reaction_field.jl")
include("interactions/gravity.jl")
include("interactions/harmonic_bond.jl")
include("interactions/harmonic_angle.jl")
include("interactions/periodic_torsion.jl")
include("interactions/rb_torsion.jl")


text:
# global state management
#
# The functionality in this file serves to create a coherent environment to perform
# computations in, with support for Julia constructs like tasks (and executing those on
# multiple threads), using a GPU of your choice (with the ability to reset that device, or
# use different devices on different threads).
#
# this is complicated by the fact that CUDA uses thread-bound contexts, while we want to be
# able to switch between tasks executing on the same thread. the current approach involves
# prefixing every CUDA API call with code that ensures the CUDA thread-bound state is
# identical to Julia's task-local one.

export context, context!, device, device!, device_reset!, deviceid, stream, stream!

## task-local state

@enum MathMode begin
    # use prescribed precision and standardized arithmetic for all calculations.
    # this may serialize operations, and reduce performance.
    PEDANTIC_MATH

    # use at least the required precision, and allow reordering operations for performance.
    DEFAULT_MATH

    # additionally allow downcasting operations for better use of hardware resources.
    # whenever possible the `precision` flag passed to `math_mode!` will be used
    # to constrain those downcasts.
    FAST_MATH
end

# math mode and precision are sticky (once set on a task, inherit to newly created tasks)
const default_math_mode = Ref{Union{Nothing,MathMode}}(nothing)
const default_math_precision = Ref{Union{Nothing,Symbol}}(nothing)

# the default device unitialized tasks will use, set when switching devices.
# this behavior differs from the CUDA Runtime, where device 0 is always used.
# this setting won't be used when switching tasks on a pre-initialized thread.
const default_device = Ref{Union{Nothing,CuDevice}}(nothing)

mutable struct TaskLocalState
    device::CuDevice
    context::CuContext
    streams::Vector{Union{Nothing,CuStream}}
    math_mode::MathMode
    math_precision::Symbol

    function TaskLocalState(dev::CuDevice=something(default_device[], CuDevice(0)),
                            ctx::CuContext = context(dev))
        math_mode = something(default_math_mode[],
                              Base.JLOptions().fast_math==1 ? FAST_MATH : DEFAULT_MATH)
        math_precision = something(default_math_precision[], :TensorFloat32)
        new(dev, ctx, Base.fill(nothing, ndevices()), math_mode, math_precision)
    end
end

function validate_task_local_state(state::TaskLocalState)
    # NOTE: the context may be invalid if another task reset it (which we detect here
    #       since we can't touch other tasks' local state from `device_reset!`)
    if !isvalid(state.context)
        device!(state.device)
        @inbounds state.streams[deviceid(state.device)+1] = nothing
    end
    return state
end

# get or create the task local state, and make sure it's valid
function task_local_state!(args...)
    tls = task_local_storage()
    if haskey(tls, :CUDA)
        validate_task_local_state(@inbounds(tls[:CUDA]))
    else
        tls[:CUDA] = TaskLocalState(args...)
    end::TaskLocalState
end

# only get the task local state (it may be invalid!), or return nothing if unitialized
function task_local_state()
    tls = task_local_storage()
    if haskey(tls, :CUDA)
        @inbounds(tls[:CUDA])
    else
        nothing
    end::Union{TaskLocalState,Nothing}
end

@inline function prepare_cuda_state()
    state = task_local_state!()

    # NOTE: current_context() is too slow to use here (taking a lock, accessing a dict)
    #       so we use the raw handle. is that safe though, when we reset the device?
    #ctx = current_context()
    ctx = Ref{CUcontext}()
    cuCtxGetCurrent(ctx)
    if ctx[] != state.context.handle
        activate(state.context)
    end

    return
end

# convenience function to get all relevant state
# without querying task local storage multiple times
@inline function active_state()
    # inline to remove unused state properties
    state = task_local_state!()
    return (device=state.device, context=state.context, stream=stream(state),
            math_mode=state.math_mode, math_precision=state.math_precision)
end

## context-based API

"""
    context()::CuContext

Get or create a CUDA context for the current thread (as opposed to
`current_context` which may return `nothing` if there is no context bound to the
current thread).
"""
function context()
    task_local_state!().context
end

"""
    context!(ctx::CuContext)
    context!(ctx::CuContext) do ... end

Bind the current host thread to the context `ctx`. Returns the previously-bound context. If
used with do-block syntax, the change is only temporary.

Note that the contexts used with this call should be previously acquired by calling
[`context`](@ref), and not arbitrary contexts created by calling the `CuContext`
constructor.
"""
function context!(ctx::CuContext)
    # switch contexts
    # NOTE: if we actually need to switch contexts, we eagerly activate it so that we can
    #       query its device (we normally only do so lazily in `prepare_cuda_state`)
    state = task_local_state()
    if state === nothing
        old_ctx = nothing
        activate(ctx)
        dev = current_device()
        task_local_state!(dev, ctx)
    else
        old_ctx = state.context
        if old_ctx != ctx
            activate(ctx)
            dev = current_device()
            state.device = dev
            state.context = ctx
        end
    end

    return old_ctx
end

macro context!(ex...)
    body = ex[end]
    ctx = ex[end-1]
    kwargs = ex[1:end-2]

    skip_destroyed = false
    for kwarg in kwargs
        Meta.isexpr(kwarg, :(=)) || throw(ArgumentError("non-keyword argument like option '$kwarg'"))
        key, val = kwarg.args
        isa(key, Symbol) || throw(ArgumentError("non-symbolic keyword '$key'"))

        if key == :skip_destroyed
            skip_destroyed = val
        else
            throw(ArgumentError("unrecognized keyword argument '$kwarg'"))
        end
    end

    quote
        ctx = $(esc(ctx))
        if isvalid(ctx)
            old_ctx = context!(ctx)
            try
                $(esc(body))
            finally
                if old_ctx !== nothing && old_ctx != ctx && isvalid(old_ctx)
                    context!(old_ctx)
                end
            end
        elseif !$(esc(skip_destroyed))
            error("Cannot switch to an invalidated context.")
        end
    end
end

@inline function context!(f::Function, ctx::CuContext; skip_destroyed::Bool=false)
    # @inline so that the kwarg method is inlined too and we can const-prop skip_destroyed
    @context! skip_destroyed=skip_destroyed ctx f()
end

## device-based API

"""
    device()::CuDevice

Get the CUDA device for the current thread, similar to how [`context()`](@ref) works
compared to [`current_context()`](@ref).
"""
function device()
    task_local_state!().device
end

const __device_contexts = LazyInitialized{Vector{Union{Nothing,CuContext}}}()
device_contexts() = get!(__device_contexts) do
    [nothing for _ in 1:ndevices()]
end
function device_context(i::Int)
    contexts = device_contexts()
    assume(isassigned(contexts, i))
    @inbounds contexts[i]
end
function device_context!(i::Int, ctx)
    contexts = device_contexts()
    @inbounds contexts[i] = ctx
    return
end
device_context(dev::CuDevice) = device_context(deviceid(dev)+1)
device_context!(dev::CuDevice, ctx) = device_context!(deviceid(dev)+1, ctx)

function context(dev::CuDevice)
    devidx = deviceid(dev)+1

    # querying the primary context for a device is expensive (~100ns), so cache it
    ctx = device_context(devidx)
    if ctx !== nothing
        return ctx
    end

    if capability(dev) < v"3.5"
        @warn("""Your $(name(dev)) GPU does not meet the minimal required compute capability ($(capability(dev)) < 3.5).
                 Some functionality might be unavailable.""",
              maxlog=1, _id=devidx)
    end

    # configure the primary context
    pctx = CuPrimaryContext(dev)
    ctx = CuContext(pctx)
    device_context!(devidx, ctx)
    return ctx
end

"""
    device!(dev::Integer)
    device!(dev::CuDevice)
    device!(dev) do ... end

Sets `dev` as the current active device for the calling host thread. Devices can be
specified by integer id, or as a `CuDevice` (slightly faster). Both functions can be used
with do-block syntax, in which case the device is only changed temporarily, without changing
the default device used to initialize new threads or tasks.
"""
device!

function device!(dev::CuDevice, flags=nothing)
    # configure the primary context flags
    if flags !== nothing
        devidx = deviceid(dev)+1
        if device_context(devidx) !== nothing
            error("Cannot set flags for an active device. Do so before calling any CUDA function, or reset the device first.")
        end
        cuDevicePrimaryCtxSetFlags(dev, flags)
    end

    # make this device the new default
    default_device[] = dev

    # switch contexts
    state = task_local_state()
    if state === nothing
        task_local_state!(dev)
    else
        state.device = dev
        state.context = context(dev)
    end

    dev
end

macro device!(dev, body)
    quote
        ctx = context($(esc(dev)))
        @context! ctx $(esc(body))
    end
end

function device!(f::Function, dev::CuDevice)
    @device! dev f()
end

# NVIDIA bug #3240770
can_reset_device() = !(release() == v"11.2" && any(dev->stream_ordered(dev), devices()))

"""
    device_reset!(dev::CuDevice=device())

Reset the CUDA state associated with a device. This call with release the underlying
context, at which point any objects allocated in that context will be invalidated.
"""
function device_reset!(dev::CuDevice=device())
    if !can_reset_device()
        @error "Due to a bug in CUDA, resetting the device is not possible on CUDA 11.2 when using the stream-ordered memory allocator."
        return
    end

    # unconditionally reset the primary context (don't just release it),
    # as there might be users outside of CUDA.jl
    pctx = CuPrimaryContext(dev)
    unsafe_reset!(pctx)

    # wipe the device-specific state
    devidx = deviceid(dev)+1
    device_context!(devidx, nothing)

    return
end

## integer device-based API

device!(dev::Integer, flags=nothing) = device!(CuDevice(dev), flags)
device!(f::Function, dev::Integer) = device!(f, CuDevice(dev))

"""
    deviceid()::Int
    deviceid(dev::CuDevice)::Int

Get the ID number of the current device of execution. This is a 0-indexed number,
corresponding to the device ID as known to CUDA.
"""
deviceid(dev::CuDevice=device()) = Int(convert(CUdevice, dev))

## math mode

function math_mode!(mode::MathMode; precision=nothing)
    state = task_local_state!()

    state.math_mode = mode
    default_math_mode[] = mode

    if precision !== nothing
        state.math_precision = math_precision
        default_math_precision[] = precision
    end

    return
end

math_mode() = task_local_state!().math_mode
math_precision() = task_local_state!().math_precision

## streams

"""
    stream()

Get the CUDA stream that should be used as the default one for the currently executing task.
"""
@inline function stream(state=task_local_state!())
    # @inline so that it can be DCE'd when unused from active_state
    devidx = deviceid(state.device)+1
    if state.streams[devidx] === nothing
        stream = CuStream()

        # register the name of this task
        t = current_task()
        tptr = pointer_from_objref(current_task())
        tptrstr = string(convert(UInt, tptr), base=16, pad=Sys.WORD_SIZE>>2)
        NVTX.nvtxNameCuStreamA(stream, "Task(0x$tptrstr)")

        state.streams[devidx] = stream
    else
        state.streams[devidx]::CuStream
    end
end

function stream!(stream::CuStream)
    state = task_local_state!()
    devidx = deviceid(state.device)+1
    state.streams[devidx] = stream
    return
end

function stream!(f::Function, stream::CuStream)
    state = task_local_state!()
    devidx = deviceid(state.device)+1
    old_stream = state.streams[devidx]
    state.streams[devidx] = stream
    try
        f()
    finally
        state.streams[devidx] = old_stream
    end
end

"""
    stream!(::CuStream)
    stream!(::CuStream) do ... end

Change the default CUDA stream for the currently executing task, temporarily if using the
do-block version of this function.
"""
stream!

## helpers

"""
    PerDevice{T}()

A helper struct for maintaining per-device state that's lazily initialized and automatically
invalidated when the device is reset. Use `get!(per_device, dev) do ... end` to initialize
and fetch a value.

Mutating or deleting state is not supported. If this is required, use a boxed value, like
a `Ref` or a `Threads.Atomic`.

Furthermore, even though the initialization of this helper, fetching its value for a
given device, and clearing it when the device is reset are all performed in a thread-safe
manner, you should still take care about thread-safety when using the contained value.
For example, if you need to update the value, use atomics; if it's a complex structure like
an array or a dictionary, use additional locks.
"""
struct PerDevice{T}
    lock::ReentrantLock
    values::LazyInitialized{Vector{Union{Nothing,Tuple{CuContext,T}}}}
end

function PerDevice{T}() where {T}
    values = LazyInitialized{Vector{Union{Nothing,Tuple{CuContext,T}}}}()
    PerDevice{T}(ReentrantLock(), values)
end

get_values(x::PerDevice) = get!(x.values) do
    Base.fill(nothing, ndevices())
end

function Base.get(x::PerDevice, dev::CuDevice, val)
    y = get_values(x)
    id = deviceid(dev)+1
    ctx = device_context(id)    # may be nothing
    @inbounds begin
        # test-lock-test
        if y[id] === nothing || y[id][1] !== ctx
            val
        else
            y[id][2]
        end
    end
end

function Base.get!(constructor::F, x::PerDevice, dev::CuDevice) where {F}
    y = get_values(x)
    id = deviceid(dev)+1
    ctx = device_context(id)    # may be nothing
    @inbounds begin
        # test-lock-test
        if y[id] === nothing || y[id][1] !== ctx
            Base.@lock x.lock begin
                if y[id] === nothing || y[id][1] !== ctx
                    y[id] = (context(), constructor())
                end
            end
        end
        y[id][2]
    end
end

Base.length(x::PerDevice) = length(get_values(x))
Base.size(x::PerDevice) = size(get_values(x))
Base.keys(x::PerDevice) = keys(get_values(x))

function Base.show(io::IO, mime::MIME"text/plain", x::PerDevice{T}) where {T}
    print(io, "PerDevice{$T}:")
    for dev in devices()
        print(io, "\n $(deviceid(dev)): ", get(x, dev, "#undef"))
    end
end


text:
# Copyright (c) 2019 Matthew Help and contributors                               #src
#                                                                                #src
# Permission is hereby granted, free of charge, to any person obtaining a copy   #src
# of this software and associated documentation files (the "Software"), to deal  #src
# in the Software without restriction, including without limitation the rights   #src
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell      #src
# copies of the Software, and to permit persons to whom the Software is          #src
# furnished to do so, subject to the following conditions:                       #src
#                                                                                #src
# The above copyright notice and this permission notice shall be included in all #src
# copies or substantial portions of the Software.                                #src
#                                                                                #src
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR     #src
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,       #src
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE    #src
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER         #src
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,  #src
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE  #src
# SOFTWARE.                                                                      #src

# # N-Queens

# **Originally Contributed by**: Matthew Helm ([with help from @mtanneau on Julia Discourse](https://discourse.julialang.org/t/which-jump-jl-solver-for-this-problem/43350/17?u=mthelm85))

# The N-Queens problem involves placing N queens on an N x N chessboard such
# that none of the queens attacks another. In chess, a queen can move
# vertically, horizontally, and diagonally so there cannot be more than one
# queen on any given row, column, or diagonal.

# ![Four Queens](../../assets/n_queens4.png)

# *Note that none of the queens above are able to attack any other as a result
# of their careful placement.*

using JuMP
import GLPK
import LinearAlgebra

# N-Queens
N = 8

model = Model(GLPK.Optimizer)

# Next, let's create an N x N chessboard of binary values. 0 will represent an
# empty space on the board and 1 will represent a space occupied by one of our
# queens:

@variable(model, x[1:N, 1:N], Bin)

# Now we can add our constraints:

# There must be exactly one queen in a given row/column
for i = 1:N
    @constraint(model, sum(x[i, :]) == 1)
    @constraint(model, sum(x[:, i]) == 1)
end

# There can only be one queen on any given diagonal
for i in -(N-1):(N-1)
    @constraint(model, sum(LinearAlgebra.diag(x, i)) <= 1)
    @constraint(model, sum(LinearAlgebra.diag(reverse(x, dims=1), i)) <= 1)
end

# That's it! We are ready to put our model to work and see if it is able to find
# a feasible solution:

optimize!(model)

# We can now review the solution that our model found:

solution = convert.(Int, value.(x))

# ![Four Queens](../../assets/n_queens.png)

max_stars_repo_name: ethansaxenian/RosettaDecode

max_issues_repo_name: ethansaxenian/RosettaDecode

max_forks_repo_name: ethansaxenian/RosettaDecode


text:
using Test

include("batch_ts.jl")

@testset "multi batch - single element y" begin
    for _ in 1:100
        dim1 = rand(1:20)
        dim2 = rand(1:20)
        num = rand(2:100) # single item batch not used with single element

        tups = [(rand(dim1, dim2), rand()) for i in 1:num]
        b_tups = batch_ts(tups)
        @test b_tups |> length == 2
        @test size(b_tups[1]) == (dim1, dim2, num)
        @test size(b_tups[2]) == (1, 1, num)
    end
end

@testset "multi batch - multi element y" begin
    for _ in 1:100
        dim1 = rand(1:20)
        dim2 = rand(1:20)
        dim3 = rand(1:20)
        dim4 = rand(1:20)
        num = rand(2:100) # single item batch

        tups = [(rand(dim1, dim2), rand(dim3, dim4)) for i in 1:num]
        b_tups = batch_ts(tups)
        @test b_tups |> length == 2
        @test size(b_tups[1]) == (dim1, dim2, num)
        @test size(b_tups[2]) == (dim3, dim4, num)
    end
end

@testset "single item batch" begin
    for _ in 1:100
        dim1 = rand(1:20)
        dim2 = rand(1:20)

        tups = (rand(dim1, dim2), rand(dim1, 1))
        b_tups = batch_ts(tups)
        @test ndims(b_tups[1]) == ndims(b_tups[2]) == 3
    end
end

max_stars_repo_name: adinhobl/model-zoo

max_issues_repo_name: adinhobl/model-zoo

max_forks_repo_name: adinhobl/model-zoo


text:
const RNEXTLINE = r" # .*, line (\d*):"
typealias MaybeInt Nullable{Int}

lines(s::String) = split(s, "\n")
function has_linum(s::AbstractString)
    ismatch(RNEXTLINE, s)
end

function linum(s::AbstractString)
    if has_linum(s)
        i = parse(Int, match(RNEXTLINE, s)[1])
        return MaybeInt(i)
    else
        return MaybeInt()
    end
end

function purge_linum(s::AbstractString)
    has_linum(s) && return split(s, RNEXTLINE)[1]
    s
end

split_linum(l) = purge_linum(l), linum(l)
split_limums(s::AbstractString) = map(split_linum, lines(s))

maybe_linums(lines::Vector) = map(linum, lines)

function min_linums(maybe_nums::Vector{MaybeInt})
    min_numbers = Int[]
    min_linum = 1
    for maybe_i in maybe_nums
        push!(min_numbers, min_linum)
        if !(isnull(maybe_i))
            min_linum = get(maybe_i)
        end
    end
    min_numbers
end

const INF = typemax(Int64)
function max_linums(maybe_nums::Vector{MaybeInt})
    max_num = INF
    max_numbers = Int[]
    for maybe_i in reverse(maybe_nums)
        if !(isnull(maybe_i))
            max_num = get(maybe_i)
        end
        unshift!(max_numbers, max_num)
    end
    max_numbers
end

immutable LinePosition
    guess::Int
    lower::Int
    upper::Int
    exact::Bool
    function LinePosition(guess, lower, upper, exact)
        @assert 1 <= lower <= guess <= upper (lower, guess, upper)
        new(guess, lower, upper, exact)
    end
end

LinePosition(g::Int,l::Int,u::Int) = LinePosition(g,l,u, l == u)
LinePosition(i::Int) = LinePosition(i,i,i,true)

function linepositions{S <:AbstractString}(lines::Vector{S})
    mays = maybe_linums(lines)
    minis = min_linums(mays)
    maxis = max_linums(mays)
    positions = LinePosition[]

    last_mini = 0
    last_guess = 0
    N = length(minis)
    for i in 1:N
        mini = minis[i]
        maxi = maxis[i]
        if mini > last_mini
            pos = LinePosition(mini)
        else
            guess = clamp(last_guess+1, mini, maxi)
            pos = LinePosition(guess, mini, maxi)
        end
        last_guess = pos.guess
        last_mini = mini
        push!(positions, pos)
    end
    positions
end

hascomment(s::AbstractString) = length(split(s, '#')) > 1

function get_comment(s)
    @assert hascomment(s)
    split(s, '#', limit=2)[2]
end

function comment_dict{S <: AbstractString}(lines::Vector{S})
    ret = Dict{Int, String}()
    for i in 1:length(lines)
        l = lines[i]
        if hascomment(l)
            ret[i] = get_comment(l)
        end
    end
    ret
end

comment_dict(s::String) = s |> lines |> comment_dict

function extract!(d::Dict, key)
    val = d[key]
    delete!(d, key)
    val
end

function add_comments(lines, positions, com_dict::Dict)
    com_dict = deepcopy(com_dict)  # we eat it up
    ret = String[]
    for (line, pos) in zip(lines, positions)
        guess = pos.guess
        if guess in keys(com_dict)
            com = extract!(com_dict, guess)
            line = string(line, " #", com)
            push!(ret, line)
        else
            push!(ret, line)
        end
    end
    last_guess = positions[end].guess
    # after lines
    for key in com_dict |> keys |> collect |> sort
        @assert key > last_guess
        com = extract!(com_dict, key)
        val = string("#", com_dict[key])
        push!(ret, val)
    end
    @assert isempty(com_dict)
    ret
end

function fill_missing_lines_positions(ls, poss::Vector{LinePosition})
    N = length(ls)

    @assert N == length(poss)
    ls_ret = String[]
    poss_ret = LinePosition[]
    N == 0 && return ls_ret, poss_ret

    last_pos = first(poss)
    for i in 1:N
        next_line = ls[i]
        next_pos = poss[i]
        if next_pos.guess > last_pos.guess + 1 # missing lines
            lower = min(last_pos.lower, next_pos.lower)
            upper = max(last_pos.upper, next_pos.upper)
            l_missing = match(r"(\s*)", next_line)[1]  # current indentation level

            for missing in (last_pos.guess + 1) : (next_pos.guess -1)
                p_missing = LinePosition(missing, lower, upper)
                push!(poss_ret, p_missing)
                push!(ls_ret, l_missing)
            end
        end
        push!(poss_ret, next_pos)
        push!(ls_ret, next_line)
        last_pos = next_pos
    end
    ls_ret, poss_ret
end

tabs2spaces(s, n=4) = join(split(s, '\t'), " "^n)
unindent(s) = try return match(r"\s\s\s\s(.*)", tabs2spaces(s))[1] catch e error("cannot unindent $s") end

wrap_begin_end(s::String) = """begin
        $s
    end"""
function unwrap_begin_end(ls::Vector)
    @assert strip(ls[1]) == "begin" ls[1]
    @assert strip(ls[end]) == "end" ls[end]
    ret = ls[2:end-1]
    map(unindent, ret)
end

function transform_ast_preserving_comments(f,s::String)  # better name?
    s = wrap_begin_end(s)
    com_dict = comment_dict(s)
    ls_ret = s |> parse |> f |> string |> lines
    poss = linepositions(ls_ret)
    ls_ret = map(purge_linum, ls_ret)
    ls_ret, poss = fill_missing_lines_positions(ls_ret, poss)
    ls_ret = add_comments(ls_ret, poss, com_dict)
    ls_ret = unwrap_begin_end(ls_ret)
    map!(rstrip, ls_ret)
    join(ls_ret, "\n")
end


text:
export LinearSubspace,
    ExtrinsicDescription,
    IntrinsicDescription,
    Coordinates,
    Intrinsic,
    Extrinsic,
    intrinsic,
    is_linear,
    extrinsic,
    dim,
    codim,
    ambient_dim,
    rand_subspace,
    coord_change,
    geodesic_distance,
    geodesic,
    translate

# References
const _LKK19 = """Lim, Lek-Heng, Ken Sze-Wai Wong, and Ke Ye. "Numerical algorithms on the affine Grassmannian." SIAM Journal on Matrix Analysis and Applications 40.2 (2019): 371-393"""

abstract type AbstractSubspace{T} end

"""
    Coordinates

A type used for encoding the used coordinates and for performing coordinate changes.

Currently supported coordinates are:
- [`Intrinsic`](@ref)
- [`Extrinsic`](@ref)
"""
struct Coordinates{T} end
Base.broadcastable(C::Coordinates) = Ref(C)

"""
    Intrinsic <: Coordinates

Indicates the use of the intrinsic description of an (affine) linear subspace. See also [`IntrinsicDescription`](@ref).
"""
const Intrinsic = Coordinates{:Intrinsic}()

"""
    Extrinsic <: Coordinates

Indicates the use of the extrinsic description of an (affine) linear subspace. See also [`ExtrinsicDescription`](@ref).
"""
const Extrinsic = Coordinates{:Extrinsic}()
# These are not part of the public API
const IntrinsicStiefel = Coordinates{:IntrinsicStiefel}()
const ExtrinsicStiefel = Coordinates{:ExtrinsicStiefel}()

"""
    ExtrinsicDescription(A, b)

Extrinsic description of an ``m``-dimensional (affine) linear subspace ``L`` in ``n``-dimensional space.
That is ``L = \\{ x | A x = b \\}``.
Note that internally `A` and `b` will be stored such that the rows of `A` are orthonormal.
"""
struct ExtrinsicDescription{T}
    A::Matrix{T}
    b::Vector{T}

    function ExtrinsicDescription(
        A::Matrix{T},
        b::Vector{T};
        orthonormal::Bool = false,
    ) where {T}
        if orthonormal
            new{T}(A, b)
        else
            svd = LA.svd(A)
            new{T}(svd.Vt, (inv.(svd.S) .* (svd.U' * b)))
        end
    end
end
(A::ExtrinsicDescription)(x::AbstractVector) = A.A * x - A.b

function Base.convert(::Type{ExtrinsicDescription{T}}, A::ExtrinsicDescription) where {T}
    ExtrinsicDescription(
        convert(Matrix{T}, A.A),
        convert(Vector{T}, A.b);
        orthonormal = true,
    )
end

"""
    dim(A::ExtrinsicDescription)

Dimension of the (affine) linear subspace `A`.
"""
dim(A::ExtrinsicDescription) = size(A.A, 2) - size(A.A, 1)

"""
    codim(A::ExtrinsicDescription)

Codimension of the (affine) linear subspace `A`.
"""
codim(A::ExtrinsicDescription) = size(A.A, 1)

Base.:(==)(A::ExtrinsicDescription, B::ExtrinsicDescription) = A.A == B.A && A.b == B.b
function Base.copy!(A::ExtrinsicDescription, B::ExtrinsicDescription)
    copy!(A.A, B.A)
    copy!(A.b, B.b)
    A
end
Base.copy(A::ExtrinsicDescription) =
    ExtrinsicDescription(copy(A.A), copy(A.b); orthonormal = true)

function Base.show(io::IO, A::ExtrinsicDescription{T}) where {T}
    println(io, "ExtrinsicDescription{$T}:")
    println(io, "A:")
    show(io, A.A)
    println(io, "\nb:")
    show(io, A.b)
end

Base.broadcastable(A::ExtrinsicDescription) = Ref(A)

"""
    IntrinsicDescription(A, b₀)

Intrinsic description of an ``m``-dimensional (affine) linear subspace ``L`` in ``n``-dimensional space.
That is ``L = \\{ u | A u + b₀ \\}``. Here, ``A`` and ``b₀`` are in orthogonal coordinates.
That is, the columns of ``A`` are orthonormal and ``A' b₀ = 0``.
"""
struct IntrinsicDescription{T}
    # orthogonal coordinates
    A::Matrix{T}
    b₀::Vector{T}
    # stiefel coordinates
    Y::Matrix{T}
end
IntrinsicDescription(A::Matrix, b₀::Vector) =
    IntrinsicDescription(A, b₀, stiefel_coordinates(A, b₀))

function Base.convert(::Type{IntrinsicDescription{T}}, A::IntrinsicDescription) where {T}
    IntrinsicDescription(
        convert(Matrix{T}, A.A),
        convert(Vector{T}, A.b₀),
        convert(Matrix{T}, A.Y),
    )
end

(A::IntrinsicDescription)(u::AbstractVector, ::Coordinates{:Intrinsic}) = A.A * u + A.b₀
(A::IntrinsicDescription)(u::AbstractVector, ::Coordinates{:IntrinsicStiefel}) = A.Y * u

function stiefel_coordinates(A::AbstractMatrix, b::AbstractVector)
    n, k = size(A)
    Y = zeros(eltype(A), n + 1, k + 1)
    stiefel_coordinates!(Y, A, b)
    Y
end
function stiefel_coordinates!(Y, A::AbstractMatrix, b::AbstractVector)
    γ = sqrt(1 + sum(abs2, b))
    n, k = size(A)
    Y[1:n, 1:k] .= A
    Y[1:n, k+1] .= b ./ γ
    Y[n+1, k+1] = 1 / γ
    Y
end

"""
    dim(A::IntrinsicDescription)

Dimension of the (affine) linear subspace `A`.
"""
dim(I::IntrinsicDescription) = size(I.A, 2)

"""
    codim(A::IntrinsicDescription)

Codimension of the (affine) linear subspace `A`.
"""
codim(I::IntrinsicDescription) = size(I.A, 1) - size(I.A, 2)

function Base.:(==)(A::IntrinsicDescription, B::IntrinsicDescription)
    A.A == B.A && A.b₀ == B.b₀ && A.Y == B.Y
end

function Base.show(io::IO, A::IntrinsicDescription{T}) where {T}
    println(io, "IntrinsicDescription{$T}:")
    println(io, "A:")
    show(io, A.A)
    println(io, "\nb₀:")
    show(io, A.b₀)
end

function Base.copy!(A::IntrinsicDescription, B::IntrinsicDescription)
    copy!(A.A, B.A)
    copy!(A.b₀, B.b₀)
    copy!(A.Y, B.Y)
    A
end
Base.copy(A::IntrinsicDescription) = IntrinsicDescription(copy(A.A), copy(A.b₀), copy(A.Y))

Base.broadcastable(A::IntrinsicDescription) = Ref(A)

function IntrinsicDescription(E::ExtrinsicDescription)
    svd = LA.svd(E.A; full = true)
    m, n = size(E.A)
    A = Matrix((@view svd.Vt[m+1:end, :])')
    if iszero(E.b)
        b₀ = zeros(eltype(E.b), size(A, 1))
    else
        b₀ = svd \ E.b
    end
    IntrinsicDescription(A, b₀)
end

function ExtrinsicDescription(I::IntrinsicDescription)
    svd = LA.svd(I.A; full = true)
    m, n = size(I.A)
    A = Matrix((@view svd.U[:, n+1:end])')
    if iszero(I.b₀)
        b = zeros(eltype(A), size(A, 1))
    else
        b = A * I.b₀
    end
    ExtrinsicDescription(A, b; orthonormal = true)
end

"""
    LinearSubspace(A, b)

An ``m``-dimensional (affine) linear subspace ``L`` in ``n``-dimensional space given by
the extrinsic description ``L = \\{ x | A x = b \\}``.

```julia-repl
julia> A = LinearSubspace([1 0 3; 2 1 3], [5, -2])
1-dim. (affine) linear subspace {x|Ax=b} with eltype Float64:
A:
2×3 Array{Float64,2}:
 1.0  0.0  3.0
 2.0  1.0  3.0
b:
2-element Array{Float64,1}:
  5.0
 -2.0

julia> dim(A)
1

julia> codim(A)
2

julia> ambient_dim(A)
3
```

A `LinearSubspace` holds always its [`extrinsic`](@ref) description,
see also [`ExtrinsicDescription`](@ref), as well as its [`intrinsic`](@ref) description,
see [`IntrinsicDescription`](@ref).

```julia-repl
julia> intrinsic(A)
IntrinsicDescription{Float64}:
A:
3×1 Array{Float64,2}:
 -0.6882472016116853
  0.6882472016116853
  0.22941573387056186
b₀:
3-element Array{Float64,1}:
 -3.0526315789473677
 -3.947368421052632
  2.684210526315789
```

A `LinearSubspace` can be evaluated with either using [`Intrinsic`](@ref)
or [`Extrinsic`](@ref) coordinates.

```julia-repl
julia> u = [0.5]
1-element Array{Float64,1}:
 0.5

julia> x = A(u, Intrinsic)
3-element Array{Float64,1}:
 -3.3967551797532103
 -3.6032448202467893
  2.79891839325107

julia> A(x, Extrinsic)
  2-element Array{Float64,1}:
   0.0
   0.0
```

To change the used coordinates you can use [`coord_change`](@ref).
```julia-repl
julia> coord_change(A, Extrinsic, Intrinsic, x)
1-element Array{Float64,1}:
 0.49999999999999994
```
"""
struct LinearSubspace{T} <: AbstractSubspace{T}
    extrinsic::ExtrinsicDescription{T}
    intrinsic::IntrinsicDescription{T}
end

LinearSubspace(I::IntrinsicDescription) = LinearSubspace(ExtrinsicDescription(I), I)
LinearSubspace(E::ExtrinsicDescription) = LinearSubspace(E, IntrinsicDescription(E))

function LinearSubspace(
    A::AbstractMatrix{T},
    b::AbstractVector{T} = zeros(eltype(A), size(A, 1)),
) where {T}
    size(A, 1) == length(b) || throw(ArgumentError("Size of A and b not compatible."))
    0 < size(A, 1) ≤ size(A, 2) || throw(
        ArgumentError(
            "Affine subspace has to be given in extrinsic coordinates, i.e., by A x = b.",
        ),
    )

    LinearSubspace(ExtrinsicDescription(Matrix(float.(A)), Vector(float.(b))))
end

function Base.convert(::Type{LinearSubspace{T}}, A::LinearSubspace) where {T}
    LinearSubspace(
        convert(ExtrinsicDescription{T}, A.extrinsic),
        convert(IntrinsicDescription{T}, A.intrinsic),
    )
end

Base.broadcastable(A::LinearSubspace) = Ref(A)

"""
    dim(A::LinearSubspace)

Dimension of the (affine) linear subspace `A`.
"""
dim(A::LinearSubspace) = dim(A.intrinsic)

"""
    codim(A::LinearSubspace)

Codimension of the (affine) linear subspace `A`.
"""
codim(A::LinearSubspace) = codim(A.intrinsic)

"""
    ambient_dim(A::LinearSubspace)

Dimension of ambient space of the (affine) linear subspace `A`.
"""
ambient_dim(A::LinearSubspace) = dim(A) + codim(A)

"""
    is_linear(L::LinearSubspace)

Returns `true` if the space is proper linear subspace, i.e., described by
`L = \\{ x | Ax = 0 \\}.
"""
is_linear(A::LinearSubspace) = iszero(extrinsic(A).b)

function Base.show(io::IO, A::LinearSubspace{T}) where {T}
    if is_linear(A)
        println(io, "$(dim(A))-dim. linear subspace {x | Ax=0} with eltype $T:")
        println(io, "A:")
        show(io, A.extrinsic.A)
    else
        println(io, "$(dim(A))-dim. affine linear subspace {x | Ax=b} with eltype $T:")
        println(io, "A:")
        show(io, A.extrinsic.A)
        println(io, "\nb:")
        show(io, A.extrinsic.b)
    end
end

"""
    intrinsic(A::LinearSubspace)

Obtain the intrinsic description of `A`, see also [`IntrinsicDescription`](@ref).
"""
intrinsic(A::LinearSubspace) = A.intrinsic

"""
    extrinsic(A::LinearSubspace)

Obtain the extrinsic description of `A`, see also [`ExtrinsicDescription`](@ref).
"""
extrinsic(A::LinearSubspace) = A.extrinsic

function Base.copy!(A::LinearSubspace, B::LinearSubspace)
    copy!(A.intrinsic, B.intrinsic)
    copy!(A.extrinsic, B.extrinsic)
    A
end
Base.copy(A::LinearSubspace) = LinearSubspace(copy(A.extrinsic), copy(A.intrinsic))

function Base.:(==)(A::LinearSubspace, B::LinearSubspace)
    intrinsic(A) == intrinsic(B) && extrinsic(A) == extrinsic(B)
end
Base.isequal(A::LinearSubspace, B::LinearSubspace) = A === B

function (A::LinearSubspace)(x::AbstractVector, ::Coordinates{:Intrinsic})
    intrinsic(A)(x, Intrinsic)
end
function (A::LinearSubspace)(x::AbstractVector, ::Coordinates{:IntrinsicStiefel})
    intrinsic(A)(x, IntrinsicStiefel)
end
function (A::LinearSubspace)(x::AbstractVector, ::Coordinates{:Extrinsic} = Extrinsic)
    extrinsic(A)(x)
end

"""
    rand_subspace(n::Integer; dim | codim, affine = true, real = false)

Generate a random [`LinearSubspace`](@ref) with given dimension `dim` or codimension `codim`
(one of them has to be provided) in ambient space of dimension `n`.
If `real` is `true`, then the extrinsic description is real.
If `affine` then an affine linear subspace is generated.
The subspace is generated by drawing each entry of the extrinsic description indepdently
from a normal distribuation using `randn`.

    rand_subspace(x::AbstractVector; dim | codim, affine = true)

Generate a random [`LinearSubspace`](@ref) with given dimension `dim` or codimension `codim`
(one of them has to be provided) in ambient space of dimension `length(x)` going through
the given point `x`.

## Example

Construction of a general random subspace:
```julia-repl
julia> rand_subspace(3; dim = 1)
1-dim. (affine) linear subspace {x|Ax=b} with eltype Complex{Float64}:
A:
2×3 Array{Complex{Float64},2}:
  -1.73825+1.27987im   -0.0871343+0.840408im  -0.551957+0.106397im
 -0.597132-0.343965im   -0.122543-0.172715im   -1.04949+0.370917im
b:
2-element Array{Complex{Float64},1}:
  0.47083334430689394 + 0.8099804422599071im
 -0.12018696822943896 + 0.11723026326952792im

julia> rand_subspace(4; codim = 1)
3-dim. (affine) linear subspace {x|Ax=b} with eltype Complex{Float64}:
A:
1×4 Array{Complex{Float64},2}:
 0.345705+0.0893881im  -0.430867-0.663249im  0.979969-0.569378im  -0.29722-0.192493im
b:
1-element Array{Complex{Float64},1}:
 0.7749708228192062 + 0.9762873764567546im
```
"""
function rand_subspace(
    n::Integer;
    dim::Union{Nothing,Integer} = nothing,
    codim::Union{Nothing,Integer} = nothing,
    real::Bool = false,
    affine::Bool = true,
)
    !isnothing(dim) ||
        !isnothing(codim) ||
        throw(ArgumentError("Neither `dim` nor `codim` specified."))

    if !isnothing(dim)
        0 < dim < n || throw(ArgumentError("`dim` has to be between 0 and `n`."))
        k = dim
    else
        0 < codim < n || throw(ArgumentError("`codim` has to be between 0 and `n`."))
        k = n - codim
    end
    T = real ? Float64 : ComplexF64
    A = randn(T, n - k, n)
    if affine
        LinearSubspace(A, randn(T, n - k))
    else
        LinearSubspace(A)
    end

end
rand_subspace(x::AbstractVector{Variable}; kwargs...) = rand_subspace(length(x); kwargs...)
function rand_subspace(
    x::AbstractVector;
    dim::Union{Nothing,Integer} = nothing,
    codim::Union{Nothing,Integer} = nothing,
    affine::Bool = true,
)
    n = length(x)
    !isnothing(dim) ||
        !isnothing(codim) ||
        throw(ArgumentError("Neither `dim` nor `codim` specified."))

    if !isnothing(dim)
        0 < dim < n || throw(ArgumentError("`dim` has to be between 0 and `n`."))
        k = dim
    else
        0 < codim < n || throw(ArgumentError("`codim` has to be between 0 and `n`."))
        k = n - codim
    end

    if affine
        A = randn(eltype(x), n - k, n)
        b = A * x
        LinearSubspace(A, b)
    else
        N = LA.nullspace(Matrix(x'))'
        A = randn(eltype(N), n - k, size(N, 1)) * N
        LinearSubspace(A)
    end
end

# Coordinate changes
"""
    coord_change(A::LinearSubspace, C₁::Coordinates, C₂::Coordinates, p)

Given an (affine) linear subspace `A` and a point `p` in coordinates `C₁` compute the point
`x` describing p in coordinates `C₂`.

## Example
```julia-repl
julia> A = LinearSubspace([1 0 3; 2 1 3], [5, -2]);

julia> u = [1.25];

julia> x = coord_change(A, Intrinsic, Extrinsic, u)
3-element Array{Float64,1}:
 -3.9129405809619744
 -3.087059419038025
  2.9709801936539915

julia> A(x, Extrinsic)
2-element Array{Float64,1}:
 0.0
 0.0

julia> x - A(u, Intrinsic)
3-element Array{Float64,1}:
 0.0
 0.0
 0.0
```
"""
coord_change(A::LinearSubspace, ::C, ::C, x) where {C<:Coordinates} = x
coord_change(A::LinearSubspace, ::Coordinates{:Intrinsic}, ::Coordinates{:Extrinsic}, u) =
    A(u, Intrinsic)
function coord_change(
    A::LinearSubspace,
    ::Coordinates{:Intrinsic},
    ::Coordinates{:IntrinsicStiefel},
    u,
)
    [
        u
        1 / A.intrinsic.Y[end, end]
    ]
end
function coord_change(
    A::LinearSubspace,
    ::Coordinates{:Intrinsic},
    ::Coordinates{:ExtrinsicStiefel},
    u,
)
    v = coord_change(A, Intrinsic, IntrinsicStiefel, u)
    coord_change(A, IntrinsicStiefel, ExtrinsicStiefel, v)
end

coord_change(A::LinearSubspace, ::Coordinates{:Extrinsic}, ::Coordinates{:Intrinsic}, x) =
    A.intrinsic.A' * (x - A.intrinsic.b₀)
function coord_change(
    A::LinearSubspace,
    ::Coordinates{:Extrinsic},
    ::Coordinates{:IntrinsicStiefel},
    x,
)
    u = coord_change(A, Extrinsic, Intrinsic, x)
    coord_change(A, Intrinsic, IntrinsicStiefel, u)
end
function coord_change(
    A::LinearSubspace,
    ::Coordinates{:Extrinsic},
    ::Coordinates{:ExtrinsicStiefel},
    x,
)
    [x; 1]
end
function coord_change(
    A::LinearSubspace,
    ::Coordinates{:IntrinsicStiefel},
    ::Coordinates{:Intrinsic},
    u,
)
    γ = A.intrinsic.Y[end, end]
    [u[i] / (γ * u[end]) for i = 1:(length(u)-1)]
end
function coord_change(
    A::LinearSubspace,
    ::Coordinates{:IntrinsicStiefel},
    ::Coordinates{:Extrinsic},
    u,
)
    x = coord_change(A, IntrinsicStiefel, ExtrinsicStiefel, u)
    coord_change(A, ExtrinsicStiefel, Extrinsic, x)
end
function coord_change(
    A::LinearSubspace,
    ::Coordinates{:IntrinsicStiefel},
    ::Coordinates{:ExtrinsicStiefel},
    u,
)
    A(u, IntrinsicStiefel)
end
function coord_change(
    A::LinearSubspace,
    ::Coordinates{:ExtrinsicStiefel},
    ::Coordinates{:Intrinsic},
    x,
)
    u = coord_change(A, ExtrinsicStiefel, IntrinsicStiefel, x)
    coord_change(A, IntrinsicStiefel, Intrinsic, u)
end
function coord_change(
    A::LinearSubspace,
    ::Coordinates{:ExtrinsicStiefel},
    ::Coordinates{:Extrinsic},
    x,
)
    x[1:end-1] ./ x[end]
end
function coord_change(
    A::LinearSubspace,
    ::Coordinates{:ExtrinsicStiefel},
    ::Coordinates{:IntrinsicStiefel},
    x,
)
    A.intrinsic.Y' * x
end

"""
    geodesic_distance(A::LinearSubspace, B::LinearSubspace)

Compute the geodesic distance between `A` and `B` in the affine Grassmanian `Graff(k, n)`
where `k = dim(A)` and `n` is the amebient dimension.
This follows the derivation in [^LKK19].

[^LKK19]: $_LKK19.
"""
geodesic_distance(A::LinearSubspace, B::LinearSubspace) =
    geodesic_distance(A.intrinsic, B.intrinsic)
function geodesic_distance(A::IntrinsicDescription, B::IntrinsicDescription)
    sqrt(sum(σᵢ -> acos(min(σᵢ, 1.0))^2, LA.svdvals(A.Y' * B.Y)))
end

# geodesic
"""
    geodesic_svd(A::LinearSubspace, B::LinearSubspace)

Computes the factors ``Q``, ``Θ`` and ``U`` from Corollary 4.3 in [^LKK19].
These values are necessary to construct the geodesic between `A` and `B`.

[^LKK19]: $_LKK19
"""
geodesic_svd(A::LinearSubspace, B::LinearSubspace) = geodesic_svd(A.intrinsic, B.intrinsic)
function geodesic_svd(A::IntrinsicDescription, B::IntrinsicDescription)
    U, Σ, V = LA.svd!(A.Y' * B.Y)
    # Θ = acos.(min.(Σ, 1.0))
    # We have acos(1-eps()) = 2.1073424255447017e-8
    # So this is numerically super unstable if the singular value is off by only eps()
    # We try to correct this fact by treating σ >= 1 - 2eps() as 1.0
    Θ = map(σ -> σ + 2 * eps() > 1.0 ? 0.0 : acos(σ), Σ)

    n, k = size(A.Y)
    # M = (LA.I - A.Y * A.Y') * B.Y * inv(A.Y' * B.Y)
    # Have to compute an SVD of M s.t. M = Q tanΘ U'
    # Equivalently M * U = Q tan(Θ)
    # We can achieve this by using a *pivoted* QR
    # since then R will be a diagonal matrix s.t. the absolute value of R_ii is θ_{k-i}

    # inv(A.Y' * B.Y) * U = V * LA.diagm(inv.(Σ))
    MU = (LA.I - A.Y * A.Y') * B.Y * V * LA.diagm(inv.(Σ))
    Q, R = LA.qr!(MU, Val(true))
    # correct signs and ordering of Q
    Q′ = Q[:, k:-1:1]
    for j = 1:k
        # Look if we need to flip signs
        real(R[k-j+1, k-j+1]) < 0 || continue
        for i = 1:n
            Q′[i, j] = -Q′[i, j]
        end
    end

    Q′, Θ, U
end

"""
    geodesic(A::LinearSubspace, B::LinearSubspace)

Returns the geodesic ``γ(t)`` connecting `A` and `B` in the Grassmanian ``Gr(k+1,n+1)``
where ``k`` is the dimension of ``A`` and ``n`` is the ambient dimension.
See also Corollary 4.3 in [^LKK19].

[^LKK19]: $_LKK19
"""
geodesic(A::LinearSubspace, B::LinearSubspace) = geodesic(A.intrinsic, B.intrinsic)
function geodesic(A::IntrinsicDescription, B::IntrinsicDescription)
    Q, Θ, U = geodesic_svd(A, B)
    t -> A.Y * U * LA.diagm(cos.(t .* Θ,)) * U' + Q * LA.diagm(sin.(t .* Θ,)) * U'
end
#
"""
    translate(L::LinearSubspace, δb, ::Coordinates = Extrinsic)

Translate the (affine) linear subspace `L` by `δb`.
"""
function translate(L::LinearSubspace, δb, ::Coordinates{:Extrinsic} = Extrinsic)
    translate!(copy(L), δb, Extrinsic)
end
function translate!(L::LinearSubspace, δb, ::Coordinates{:Extrinsic} = Extrinsic)
    ext = extrinsic(L)
    ext.b .+= δb
    int = intrinsic(L)
    LA.mul!(int.b₀, ext.A', δb, true, true)
    stiefel_coordinates!(int.Y, int.A, int.b₀)
    L
end

"""
    Base.intersect(L₁::LinearSubspace, L₂::LinearSubspace)

Intersect the two given linear subspaces.
Throws an `ErrorException` if the intersection is the sum of the
codimensions is larger than the ambient dimension.
"""
function Base.intersect(L₁::LinearSubspace, L₂::LinearSubspace)
    ambient_dim(L₁) == ambient_dim(L₂) || error("Ambient dimensions don't match.")
    codim(L₁) + codim(L₂) ≤ ambient_dim(L₁) ||
        error("Sum of codimensions larger than ambient dimension.")
    ext₁ = extrinsic(L₁)
    ext₂ = extrinsic(L₂)
    LinearSubspace([ext₁.A; ext₂.A], [ext₁.b; ext₂.b])
end


text:
function PolynomialParser(arg0::jchar, arg1::PolynomialParser_Unit)
    return PolynomialParser((jchar, PolynomialParser_Unit), arg0, arg1)
end

function parse(obj::PolynomialParser, arg0::JString)
    return jcall(obj, "parse", Vector{jdouble}, (JString,), arg0)
end


text:
__precompile__()

module PkgBenchmark

using BenchmarkTools
using JSON
using Pkg
using LibGit2
using Dates
using InteractiveUtils
using Printf
using Logging: with_logger
using TerminalLoggers: TerminalLogger
using UUIDs: UUID

export benchmarkpkg, judge, writeresults, readresults, export_markdown, memory
export BenchmarkConfig, BenchmarkResults, BenchmarkJudgement

include("benchmarkconfig.jl")
include("benchmarkresults.jl")
include("benchmarkjudgement.jl")
include("runbenchmark.jl")
include("judge.jl")
include("util.jl")

end # module


text:
function idx(I)
    if I == ALL
        return I, 0 #ni doesn't matter if I=ALL
    elseif I isa UnitRange
        return [I.start, I.stop], libgb.GxB_RANGE #Simple ranges
    elseif I isa StepRange
        #The step must survive tozerobased(I), so we add 1 to it.
        if I.step > 0
            return [I.start, I.stop, I.step + 1], libgb.GxB_STRIDE #Normal stried ranges
        elseif I.step < 0
            #Strided ranges with negative increment
            return [I.start, I.stop, -I.step + 1], libgb.GxB_BACKWARDS
        end
    elseif I isa Vector
        return Vector{libgb.GrB_Index}(I), length(I) #Assume ni = length(I) otherwise
    elseif I isa Integer
        return UInt64(I), 0
    end
end


text:
using luteos
using SymPy

# let # limit scope

Ps = [P1(), P2(), P3()] # Range of polynomial order
Ns = [9, 17, 33] # Range of grid size

dim = 2

# Material
mat = Material( ν=0.33, E=1.0, dim = dim )

## Set up problem
#   Get functions for exact solution
function ExactSol( Cstiff )

  dim = 2

  @syms x1 x2

  u1 = (x2 - x2.^2) .* cos.(pi*x1) .* (1/2 - x1) .* (1 - x2/2)   #10 * (x2 - x2.^2) .* sin.(pi*x1) .* (1 - x1) .* (1 - x2/2)
  u2 = 2 * (x1 - x1.^2) .* sin.(pi*x2) .* (1 - x2) .* (1 - x1/2) # 2 * (x1 - x1.^2) .* sin.(pi*x2) .* (1 - x2) .* (1 - x1/2)

  ∂u1∂x1 = diff( u1, x1 )
  ∂u1∂x2 = diff( u1, x2 )
  ∂u2∂x1 = diff( u2, x1 )
  ∂u2∂x2 = diff( u2, x2 )

  ∇u = [ ∂u1∂x1 ∂u2∂x1; ∂u1∂x2 ∂u2∂x2 ]

  ϵ = 1/2 * ( ∇u + ∇u.' )

  σ = fill( 0*x1, dim, dim )

  for ii in 1:dim, jj in 1:dim, kk in 1:dim, ll in 1:dim
    σ[ii,jj] += (Cstiff[ii,jj,kk,ll] * ϵ[kk,ll])[1,1,1,1]
  end

  F  = fill( 0*x1, dim )

  ∇σ = fill( 0*x1, dim, dim, dim )

  for ii in 1:dim, jj in 1:dim
    ∇σ[ii,jj,1] = diff( σ[ii,jj], x1 )
    ∇σ[ii,jj,2] = diff( σ[ii,jj], x2 )
  end

  for ii in 1:dim, jj in 1:dim
    F[ii] -= ∇σ[ii,jj,jj]
  end

  u1func_org = lambdify( u1 )
  u2func_org = lambdify( u2 )

  F1func_org = lambdify( F[1] )
  F2func_org = lambdify( F[2] )
  σ1func_org = lambdify( σ[1] )
  σ2func_org = lambdify( σ[2] )
  σ4func_org = lambdify( σ[4] )
  ϵ1func_org = lambdify( ϵ[1] )
  ϵ2func_org = lambdify( ϵ[2] )
  ϵ4func_org = lambdify( ϵ[4] )

  # make sure you can call it with matrices
  u1func = x -> u1func_org.(x[:,1], x[:,2])
  u2func = x -> u2func_org.(x[:,1], x[:,2])

  F1func = x -> F1func_org.(x[:,1], x[:,2])
  F2func = x -> F2func_org.(x[:,1], x[:,2])
  σ1func = x -> σ1func_org.(x[:,1], x[:,2])
  σ2func = x -> σ2func_org.(x[:,1], x[:,2])
  σ4func = x -> σ4func_org.(x[:,1], x[:,2])
  ϵ1func = x -> ϵ1func_org.(x[:,1], x[:,2])
  ϵ2func = x -> ϵ2func_org.(x[:,1], x[:,2])
  ϵ4func = x -> ϵ4func_org.(x[:,1], x[:,2])

  return( u1func, u2func, F1func, F2func, σ1func, σ2func, σ4func, ϵ1func, ϵ2func, ϵ4func )

end

( u1func, u2func, F1func, F2func, σ1func, σ2func, σ4func, ϵ1func, ϵ2func, ϵ4func ) = ExactSol( mat.Cstiff )

#   Setup boundary conditions
function funcB( p::Array{Float64} )
  return fill(0.0, size(p,1), 2)
end
funcNE = (p) -> [ σ1func(p)  σ2func(p)]
funcNW = (p) -> [-σ1func(p) -σ2func(p)]
bctype = [1,2,1,2] # All Dirichlet
#   Setup source function
source = (p) -> [F1func(p) F2func(p)]

## Compute error for different polynomial orders and grid sizes
#   Initialize arrays
Err_uh1 = fill( 0.0, length(Ps), length(Ns) )
Err_uh2 = fill( 0.0, length(Ps), length(Ns) )
Err_σh1 = fill( 0.0, length(Ps), length(Ns) )
Err_σh2 = fill( 0.0, length(Ps), length(Ns) )
Err_σh4 = fill( 0.0, length(Ps), length(Ns) )
# Err_ϵh1 = fill( 0.0, length(Ps), length(Ns) )
# Err_ϵh2 = fill( 0.0, length(Ps), length(Ns) )
# Err_ϵh4 = fill( 0.0, length(Ps), length(Ns) )

#   Loop over polynomial order and grid size
for ii in 1:length(Ps), jj in 1:length(Ns)

  P = Ps[ii]; N = Ns[jj]
  println("P ", P.p)

  mesh   = Mesh2D( "square", P, N = N)
  master = Master2D( P )

  prob = Elas( @sprintf("Reg %i %i", P.p, N), mat, source, bctype, false, [funcB, funcNE, funcB, funcNW] )

  (uhathTri, uh, σh) = hdgSolve( master, mesh, prob )

  #   Initialize arrays
  err_uh = fill( 0.0, dim )
  err_σh = fill( 0.0, dim^2 )
  # err_ϵh = fill( 0.0, dim^2 )

  # preallocate
  jcwd = fill( 0.0, size(master.∇ϕ,2), size(master.∇ϕ,2) )
  ∂ξ∂x = fill( 0.0, size(master.∇ϕ,2), dim^2 )
  ∂x∂ξ = fill( 0.0, size(master.∇ϕ,2), dim^2 )

  for kk in 1:size(mesh.t,1)

    # Compute Jacobians
    luteos.compJacob!( master, mesh.nodes[:,:,kk], ∂ξ∂x, jcwd, ∂x∂ξ )

    # u
    Δuh1       = master.ϕ' * ( uh[:,1,kk] - u1func( mesh.nodes[:,:,kk] ) )
    Δuh2       = master.ϕ' * ( uh[:,2,kk] - u2func( mesh.nodes[:,:,kk] ) )
    err_uh[1] += Δuh1' * jcwd * Δuh1
    err_uh[2] += Δuh2' * jcwd * Δuh2

    # σ
    Δσh1       = master.ϕ' * ( σh[:,1,kk] - σ1func( mesh.nodes[:,:,kk] ) )
    Δσh2       = master.ϕ' * ( σh[:,2,kk] - σ2func( mesh.nodes[:,:,kk] ) )
    Δσh4       = master.ϕ' * ( σh[:,4,kk] - σ4func( mesh.nodes[:,:,kk] ) )
    err_σh[1] += Δσh1' * jcwd * Δσh1
    err_σh[2] += Δσh2' * jcwd * Δσh2
    err_σh[4] += Δσh4' * jcwd * Δσh4

    # # ϵ
    # Δϵh1       = master.ϕ' * ( ϵh[:,1,kk] - ϵ1func( mesh.nodes[:,:,kk] ) )
    # Δϵh2       = master.ϕ' * ( ϵh[:,2,kk] - ϵ2func( mesh.nodes[:,:,kk] ) )
    # Δϵh4       = master.ϕ' * ( ϵh[:,4,kk] - ϵ4func( mesh.nodes[:,:,kk] ) )
    # err_ϵh[1] += Δϵh1' * jcwd * Δϵh1
    # err_ϵh[2] += Δϵh2' * jcwd * Δϵh2
    # err_ϵh[4] += Δϵh4' * jcwd * Δϵh4

  end

  Err_uh1[ii,jj] = sqrt(err_uh[1])
  Err_uh2[ii,jj] = sqrt(err_uh[2])
  Err_σh1[ii,jj] = sqrt(err_σh[1])
  Err_σh2[ii,jj] = sqrt(err_σh[2])
  Err_σh4[ii,jj] = sqrt(err_σh[4])
  # Err_ϵh1[ii,jj] = sqrt(err_ϵh[1])
  # Err_ϵh2[ii,jj] = sqrt(err_ϵh[2])
  # Err_ϵh4[ii,jj] = sqrt(err_ϵh[4])

end

# Compute convergence rates
h = 1 ./ ( Ns - 1 )
conv_uh1 = (log.( Err_uh1[:,end-2]) - log.( Err_uh1[:,end] ) ) / (log.( h[end-2]) - log.( h[end] ));
conv_uh2 = (log.( Err_uh2[:,end-2]) - log.( Err_uh2[:,end] ) ) / (log.( h[end-2]) - log.( h[end] ));

conv_σh1 = (log.( Err_σh1[:,end-2]) - log.( Err_σh1[:,end] ) ) / (log.( h[end-2]) - log.( h[end] ));
conv_σh2 = (log.( Err_σh2[:,end-2]) - log.( Err_σh2[:,end] ) ) / (log.( h[end-2]) - log.( h[end] ));
conv_σh4 = (log.( Err_σh4[:,end-2]) - log.( Err_σh4[:,end] ) ) / (log.( h[end-2]) - log.( h[end] ));

# conv_ϵh1 = (log.( Err_ϵh1[:,end-2]) - log.( Err_ϵh1[:,end] ) ) / (log.( h[end-2]) - log.( h[end] ));
# conv_ϵh2 = (log.( Err_ϵh2[:,end-2]) - log.( Err_ϵh2[:,end] ) ) / (log.( h[end-2]) - log.( h[end] ));
# conv_ϵh4 = (log.( Err_ϵh4[:,end-2]) - log.( Err_ϵh4[:,end] ) ) / (log.( h[end-2]) - log.( h[end] ));

# Output to terminal
@printf("\n")
@printf("   Convergence rates for 2D Dirichlet/Neumann problem\n\n")
@printf("   --------------------------------------------------\n\n")
@printf( "P   ")
for jj in 1:size(Ps,1)
  @printf( " %6i", Ps[jj].p )
end
@printf( "\n" )

#   u
@printf( "u₁  ")
for jj in 1:size(Ps,1)
  @printf( " %6.4f", conv_uh1[jj] )
end
@printf( "\n" )
@printf( "u₂  ")
for jj in 1:size(Ps,1)
  @printf( " %6.4f", conv_uh2[jj] )
end
@printf( "\n" )

#   σ
@printf( "σ₁  ")
for jj in 1:length(Ps)
  @printf( " %6.4f", conv_σh1[jj] )
end
@printf( "\n" )
@printf( "σ₂  ")
for jj in 1:length(Ps)
  @printf( " %6.4f", conv_σh2[jj] )
end
@printf( "\n" )
@printf( "σ₄  ")
for jj in 1:length(Ps)
  @printf( " %6.4f", conv_σh4[jj] )
end
@printf( "\n" )
#
# #   ϵ
# @printf( "ϵ₁  ")
# for jj in 1:length(Ps)
#   @printf( " %6.4f", conv_ϵh1[jj] )
# end
# @printf( "\n" )
# @printf( "ϵ₂  ")
# for jj in 1:length(Ps)
#   @printf( " %6.4f", conv_ϵh2[jj] )
# end
# @printf( "\n" )
# @printf( "ϵ₄  ")
# for jj in 1:length(Ps)
#   @printf( " %6.4f", conv_ϵh4[jj] )
# end
# @printf( "\n" )

open("errors_Elas_Neumann2D.dat", "w") do f
  @printf(f, "P \t N \t E_uh1 \t E_uh2 \t E_σh1 \t E_σh2 \t E_σh4 \n")#\t E_ϵh1 \t E_ϵh2 \t E_ϵh4 \n")#\t E_J\n")
  for ii in 1:length(Ps), jj in 1:length(Ns)
    @printf(f, "%i \t %i \t %16.15e \t %16.15e \t %16.15e \t %16.15e \t %16.15e \n",#\t %16.15e \t %16.15e \t %16.15e\n",
      Ps[ii].p, Ns[jj], Err_uh1[ii,jj], Err_uh2[ii,jj],
      Err_σh1[ii,jj], Err_σh2[ii,jj], Err_σh4[ii,jj] )#,
      #Err_ϵh1[ii,jj], Err_ϵh2[ii,jj], Err_ϵh4[ii,jj] )
  end
end

# end # limit scope


text:
# 2D LS-RTM gradient test with 1 source
# The receiver positions and the source wavelets are the same for each of the four experiments.
# Author: Philipp Witte, pwitte@eos.ubc.ca
# Date: January 2017
#
# Mathias Louboutin, mlouboutin3@gatech.edu
# Updated July 2020
#
# Ziyi Yin, ziyi.yin@gatech.edu
# Updated July 2021

### Model
model, model0, dm = setup_model(tti, viscoacoustic, 4)
q, srcGeometry, recGeometry, f0 = setup_geom(model)
dt = srcGeometry.dt[1]

opt = Options(sum_padding=true, free_surface=fs, f0=f0)
F = judiModeling(model, srcGeometry, recGeometry; options=opt)
F0 = judiModeling(model0, srcGeometry, recGeometry; options=opt)
J = judiJacobian(F0, q)

# Observed data
dobs = F*q
dobs0 = F0*q

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

@testset "LSRTM gradient test with $(nlayer) layers, tti $(tti), viscoacoustic $(viscoacoustic). freesurface $(fs), nlind $(nlind)" for nlind=[true, false]
	@timeit TIMEROUTPUT "LSRTM gradient test, nlind=$(nlind)" begin
		# Gradient test
		ftol = (tti && fs) ? 1f-1 : 5f-2
		h = 5f-2
		maxiter = 5
		err1 = zeros(maxiter)
		err2 = zeros(maxiter)
		h_all = zeros(maxiter)

		# LS-RTM gradient and function value for m0
		Jm0, grad = lsrtm_objective(model0, q, dobs, dm; options=opt, nlind=nlind)
		dD = nlind ? (dobs - dobs0) : dobs

		# Perturbation
		dmp = 2f0*circshift(dm, 10)
		dJ = dot(grad, dmp)

		for j=1:maxiter
			dmloc = dm + h*dmp
			# LS-RTM gradient and function falue for m0 + h*dm
			Jm = lsrtm_objective(model0, q, dobs, dmloc; options=opt, nlind=nlind)[1]
			@printf("h = %2.2e, J0 = %2.2e, Jm = %2.2e \n", h, Jm0, Jm)
			# Check convergence
			err1[j] = abs(Jm - Jm0)
			err2[j] = abs(Jm - Jm0 - h*dJ)

			j == 1 ? prev = 1 : prev = j - 1
			@printf("h = %2.2e, e1 = %2.2e, rate = %2.2e", h, err1[j], err1[prev]/err1[j])
			@printf(", e2  = %2.2e, rate = %2.2e \n", err2[j], err2[prev]/err2[j])
			h_all[j] = h
			h = h * .5f0
		end

		# Check convergence rates
		rate_1 = sum(err1[1:end-1]./err1[2:end])/(maxiter - 1)
		rate_2 = sum(err2[1:end-1]./err2[2:end])/(maxiter - 1)

		# This is a linearized problem, so the whole expansion is O(dm) and
		# "second order error" should be first order
		@test isapprox(rate_1, 2f0; rtol=ftol)
		@test isapprox(rate_2, 4f0; rtol=ftol)

		# test that with zero dm we get the same as fwi_objective for residual
		if nlind
			ENV["OMP_NUM_THREADS"]=1
			Jls, gradls = lsrtm_objective(model0, q, dobs, 0f0.*dm; options=opt, nlind=nlind)
			Jfwi, gradfwi = fwi_objective(model0, q, dobs; options=opt)
			@test isapprox(Jls, Jfwi; rtol=0f0, atol=0f0)
			@test isapprox(gradls, gradfwi; rtol=0f0, atol=0f0)
		end
	end
end

# Test if lsrtm_objective produces the same value/gradient as is done by the correct algebra
@testset "LSRTM gradient linear algebra test with $(nlayer) layers, tti $(tti), viscoacoustic $(viscoacoustic), freesurface $(fs)" begin
	# Draw a random case to avoid long CI.
	isic, dft, optchk = rand([true, false], 3)
	optchk = optchk && !dft
    @timeit TIMEROUTPUT "LSRTM validity (isic=$(isic), checkpointing=$(optchk), dft=$(dft))" begin
		ftol = fs ? 1f-3 : 5f-4
		freq = dft ? [[2.5, 4.5],[3.5, 5.5],[10.0, 15.0], [30.0, 32.0]] : []
		J.options.free_surface = fs
		J.options.isic = isic
		J.options.optimal_checkpointing = optchk
		J.options.frequencies = freq

		dm1 = 2f0*circshift(dm, 10)
		d_res = dobs0 + J*dm1 - dobs
		Jm0_1 = 0.5f0 * norm(d_res)^2f0
		grad_1 = J'*d_res

		opt = J.options
		Jm0, grad = lsrtm_objective(model0, q, dobs, dm1; options=opt, nlind=true)
		Jm01, grad1 = lsrtm_objective(model0, q, dobs-dobs0, dm1; options=opt, nlind=false)

		@test isapprox(grad, grad_1; rtol=ftol)
		@test isapprox(Jm0, Jm0_1; rtol=ftol)
		@test isapprox(grad, grad1; rtol=ftol)
		@test isapprox(Jm0, Jm01; rtol=ftol)
	end
end


text:
module Insertion

export add!

import ..Common
using ..Common
using ..Hulls: ConicHull
using ..Hulls: antidominates_replaced

function find_dominated_facet(hull::ConicHull, generator::Generator)
    NF = nfacet(hull)

    # Make sure that the primary is not in the plane of the facet;
    # if the generator also is it will be unable to escape.
    # todo: more efficient way to find a primary generator and starting facet
    # that span a volume
    facet = first(hull.facets)
    primary = generatorof(facet, 1) # get a primary that is extreme
    for f in hull.facets
        if Common.dot(f, primary) != 0
            facet = f
            break
        end
    end
    @assert Common.dot(facet, primary) != 0

    while true
        if dominates(generator, facet); return facet; end
        found = false
        for face in facesof(facet)
            if antidominates_replaced(generator, face, primary)
                facet = nbof(face)
                found = true
                break
            end
        end
        if !found; return nothing; end
    end
end

function mark_dominated!{F}(newfacets::Vector{F}, removedfacets::Vector{F}, hull::ConicHull{F},
                            generator::Generator, facet::F)
    if ismarked(hull, facet); return true; end
    if antidominates(generator, facet); return false; end

    mark_facet!(hull, facet)
    push!(removedfacets, facet)
    for (face, nb) in zip(facesof(facet), nbsof(facet))
        if !mark_dominated!(newfacets, removedfacets, hull, generator, nb)
            # Found border: facet is (weakly) dominated but not nb
            newfacet = add_facet!(hull, face, generator)
            mark_facet!(hull, newfacet)
            set_nb!(newfacet, nb, generator)
            set_nb!(nb, newfacet, facet)
            push!(newfacets, newfacet)
        end
    end
    return true
end

function create_links!(hull::ConicHull, newfacet::Facet, generator::Generator)
    facet0 = nbof(newfacet, generator)
    for (face, gen) in zip(facesof(newfacet), generatorsof(newfacet))
        if gen == generator; continue; end
        if !(nbof(face) === nothing); continue; end

        lastfacet = newfacet
        gfrom = generator
        gto = generatorof(face)
        facet = facet0
        while !ismarked(hull, facet)
            gnew = generatorof(facet, lastfacet)
            gto, gfrom = gnew, gto
            fnew = nbof(facet, gfrom)
            facet, lastfacet = fnew, facet
        end

        set_nb!(face, facet)
        set_nb!(facet, newfacet, gto)
    end
end

add!{F,G}(hull::ConicHull{F,G}, x) = add!(hull, G(x))
function add!{F,G}(hull::ConicHull{F,G}, generator::G)
    @assert nconic(generator) == nconic(hull)
    dominated = find_dominated_facet(hull, generator)
    if dominated === nothing; return nothing; end

    push!(hull.generators, generator)

    newfacets, removedfacets = F[], F[]
    mark_dominated!(newfacets, removedfacets, hull, generator, dominated)

    for facet in removedfacets; del_facet!(hull, facet); end
    for facet in newfacets; create_links!(hull, facet, generator); end
    for facet in newfacets; mark_facet!(hull, facet, false); end

    return generator
end

end # module