text:
import Base: show

#
# Utility functions
#

function _showindent(io::IO)
	indent = get(io, :indent, 0)
	for _ in 1:indent
		write(io, "  ")
	end
end

function _xmlescape(t::AbstractString)
	# Technically these are not all always required, but to be safe
	t = replace(t, "&"=>"&")
	t = replace(t, "\""=>""")
	t = replace(t, "'"=>"'")
	t = replace(t, "<"=>"<")
	t = replace(t, ">"=>">")
	t
end
_xmlescape(t) = _xmlescape(string(t))

#
# Text Node
#

function Base.show(io::IO, mime::MIME"text/plain", el::TextNode)
	compact = get(io, :compact, false)
	compact || _showindent(io)
	show(io, mime, el.text)
	return nothing
end

function Base.show(io::IO, ::MIME"text/xml", el::TextNode)
	write(io, el.text |> string |> _xmlescape)
	return nothing
end

#
# Generic Element
#

function Base.show(io::IO, mime::MIME"text/plain", e::Element{ns, tag}) where{ns, tag}
	compact = get(io, :compact, false)
	compact || _showindent(io)
	write(io,"(")
	if ns !== :nothing
		write(io, ns)
		write(io, ":")
	end
	write(io, tag)
	for (prop, value) in e.props
		write(io, " ")
		write(io, prop)
		write(io,"=")
		show(io, mime, value)
	end
	if !isempty(e.children)
		depth = get(io, :depth, 0)
		child_io = IOContext(io,
			:indent => get(io, :indent, 0) + 1,
			:depth => depth + 1,
			:compact => compact || depth > 2
		)
		for child in e.children
			compact ? write(io, " ") : write(io, "\n")
			# if depth > 2 and !compact, then we should do our own indenting on the children, because they won't
			# Not sure if this makes sense
			compact || depth > 2 && _showindent(child_io)
			show(child_io, mime, child)
		end
		compact || write(io, "\n")
		compact || _showindent(io)
	end
	write(io, ")")
	return nothing
end

function Base.show(io::IO, mime::MIME"text/xml", e::Element{ns, tag}) where {ns, tag}
	write(io, "<")
	if ns !== :nothing
		write(io, _xmlescape(ns))
	end
	write(io, _xmlescape(string(tag)))
	for (prop, value) in e.props
		write(io, " ")
		write(io, _xmlescape(prop))
		write(io,"=\"")
		write(io, _xmlescape(value))
		write(io, "\"")
	end
	if !isempty(e.children)
		write(io, ">")
		for child in e.children
			show(io, mime, child)
		end
		write(io, "</")
		if ns !== :nothing
			write(io, _xmlescape(ns))
		end
		write(io, _xmlescape(string(tag)))
		write(io, ">")
	else
		write(io, "/>")
	end
	return nothing;
end
Base.show(io::IO, ::MIME"text/html", e::Element) = Base.show(io, MIME("text/xml"), e)
Base.show(io::IO, ::MIME"image/svg", e::Element) = Base.show(io, MIME("text/xml"), e)


text:
module common

    export file_parts
    function file_parts(a::Array{DT},maxsize::Integer) where {DT<:Number}
        nfiles = ceil(Int,sizeof(a)/maxsize)
        nelmts = length(a)
        parts = Vector{Int}(undef,nfiles)
        idxs = Vector{Int}(undef,nfiles+1)
        ids = Vector{Int}(undef,nfiles)
        ide = Vector{Int}(undef,nfiles)
        r::Int = rem(nelmts,nfiles)
        c::Int = ceil(Int,nelmts/nfiles)
        f::Int = floor(Int,nelmts/nfiles)
        parts[1:r]        .= c
        parts[r+1:nfiles] .= f
        @assert sum(parts)==nelmts "FATAL ERROR: failed to properly partition $nelmts to $nfiles files"
        for i=0:nfiles idxs[i+1]=sum(parts[1:i])+1 end
        for i=1:nfiles ids[i]=idxs[i] end
        for i=1:nfiles ide[i]=idxs[i+1]-1 end
        return (nfiles,nelmts,parts,ids,ide)
    end

end

ext: jl


text:
using PATH2
using ForwardDiff
using Base.Test

include("sparse_matrix.jl")
include("lcp.jl")
include("mcp.jl")


text:
using ModelingStructs
using Optim

rosen(u) = (1.0 - u.x)^2 + 100.0 * (u.y - u.x^2)^2

u₀ = MStruct(x=0, y=0)
lb = MStruct(x=-0.5, y=-0.5)
ub = MStruct(x=0.5, y=0.5)

df = TwiceDifferentiable(rosen, u₀; autodiff=:forward)
dfc = TwiceDifferentiableConstraints(lb, ub)

val = optimize(df, dfc, u₀, IPNewton())

max_stars_repo_name: jonniedie/NamedViewVectors


text:
"""
    edit_with(y::Dict{Any,Any}; kwargs...)
    edit_with(df::DataFrame, editor::T) where T<:Edit
    edit_with(df::DataFrame, lst::Array{T}) where T<:Edit
    edit_with(df::DataFrame, x::Describe, file::T) where T<:File
    edit_with(file::T, y::Dict{Any,Any}; kwargs...)
    edit_with(files::Array{T,N} where N, y::Dict{Any,Any}; kwargs...) where T<:File
This function edits the input DataFrame `df` and returns the resultant DataFrame.

# Arguments
- `df::DataFrame` on which to perform the edit.
- `editor::T where T<:Edit`: DataType containing information about which edit to perform.
    The following edit options are available and detailed below. If given a dictionary of
    edits, they will be made in this order:
    - [`SLiDE.Drop`](@ref): Remove information from the DataFrame -- either an entire column
        or rows containing specified values.
    - [`SLiDE.Rename`](@ref): Change column name `from` -> `to`.
    - [`SLiDE.Group`](@ref): Use to edit files containing data in successive dataframes with
        an identifying header cell or row.
    - [`SLiDE.Match`](@ref): Extract values from the specified column into a column or
        columns based on the specified regular expression.
    - [`SLiDE.Melt`](@ref): Normalize the dataframe by 'melting' columns into rows,
        lengthening the dataframe by duplicating values in the column `on` into new rows and
        defining 2 new columns:
        1. `var` with header names from the original dataframe.
        2. `val` with column values from the original dataframe.
    - [`SLiDE.Add`](@ref): Add new column `col` filled with `val`.
    - [`SLiDE.Map`](@ref): Define an `output` column containing values based on those in an
        `input` column. The mapping columns `from` -> `to` are contained in a .csv `file` in
        the coremaps directory. The columns `input` and `from` should contain the same
        values, as should `output` and `to`.
    - [`SLiDE.Replace`](@ref): Replace values in `col` `from` -> `to`.
    - [`SLiDE.Operate`](@ref): Perform an arithmetic operation across multiple DataFrame columns or rows.
    - [`SLiDE.Describe`](@ref): This DataType is required when multiple DataFrames will be
        appended into one output file (say, if multiple sheets from an XLSX file are
        included). Before the DataFrames are appended, a column `col` will be added and
        filled with the value in the file descriptor.
    - [`SLiDE.Order`](@ref): Rearranges columns in the order specified by `cols` and sets
        them to the specified type.
- `file::T where T <: File`: Data file containing information to read.
- `files::Array{T} where T <: File`: List of data files.
- `y::Dict{Any,Any}`: Dictionary containing all editing structures among other values read
    from the yaml file. Dictionary keys must correspond EXACTLY with SLiDE.Edit DataType
    names, or the edits will not be made.

# Returns
- `df::DataFrame` including edit(s)
"""
function edit_with(df::DataFrame, x::Add; file=nothing)
    # If adding the length of a string... (!!!! doc this option; used for naics code scaling)
    if typeof(x.val) == String && occursin("length", x.val)
        m = match(r"(?<col>\S*) length", x.val)

        # If this is not indicating a column length to add, add the value and exit.
        if (m === nothing || !(Symbol(m[:col]) in propertynames(df)))
            df[!, x.col] .= x.val
            return df
        end
        # If possible, return the length of characters  in the string.
        col_len = Symbol(m[:col])
        df[!, x.col] .= [ismissing(val_len) ? missing : length(convert_type(String, val_len))
            for val_len in df[:,col_len]]
    else
        df[!, x.col] .= x.val
    end
    return df
end

function edit_with(df::DataFrame, x::Combine; file=nothing)
    return combine_over(df, setdiff(findindex(df), x.output))
end

function edit_with(df::DataFrame, x::Deselect; file=nothing)
    if x.operation == "occursin"
        x.col = propertynames(df)[occursin.(x.col[1], propertynames(df))]
    end
    return select(df, setdiff(propertynames(df), x.col))
end

function edit_with(df::DataFrame, x::Drop; file=nothing)
    if x.val === "all" && x.operation == "occursin"
        df = edit_with(df, Drop.(propertynames(df)[occursin.(x.col, propertynames(df))], "all", "=="))
    end

    !(x.col in propertynames(df)) && (return df)
    if x.val === "all"  # Drop entire column to remove dead weight right away.
        df = df[:, setdiff(propertynames(df), [x.col])]
    else  # Drop rows using an operation or based on a value.
        if x.val === missing
            df = dropmissing(df, x.col)
        elseif x.val === "unique"
            df = unique(df, x.col)
        else
            df[!,x.col] .= convert_type.(typeof(x.val), df[:,x.col])
            df = if x.operation == "occursin"
                df[.!broadcast(datatype(x.operation), x.val, df[:,x.col]), :]
            else
                df[.!broadcast(datatype(x.operation), df[:,x.col], x.val), :]
            end
        end
    end
    return df
end

function edit_with(df::DataFrame, x::OrderedGroup)
    # !!!! can we merge this with Group? So we can group on either?
    # !!!! Can we group when we just need to fill columns with what's below them? (think: sctg codes)
    [df = _group_by(df, OrderedGroup(x.on, x.var, [val])) for val in x.val]
    return df
end

function edit_with(df::DataFrame, x::Group; file=nothing)
    # First, add a column to the original DataFrame indicating where the data set begins.
    cols = unique([propertynames(df); x.output])
    df[!,:start] = (1:size(df)[1]) .+ 1

    # Next, create a DataFrame describing where to "split" the input DataFrame.
    # Editing with a map will remove all rows that do not contain relevant information.
    # Add a column indicating where each data set STOPS, assuming all completely blank rows
    # were removed by read_file().
    df_split = edit_with(copy(df), convert_type(Map, x))
    df_split = sort(unique(df_split), :start)
    df_split[!, :stop] .= vcat(df_split[2:end, :start] .- 2, [size(df)[1]])

    # Add a new, blank output column to store identifying information about the data block.
    # Then, fill this column based on the identifying row numbers in df_split.
    for out in x.output
        df[!,out] .= ""
        [df[row[:start]:row[:stop], out] .= row[out] for row in eachrow(df_split)]
    end

    # Finally, remove header rows (these will be blank in the output column),
    # as well as the column describing where the sub-DataFrames begin.
    df = edit_with(df, Drop.(x.output, "", "=="))
    return df[:, cols]
end

edit_with(df::DataFrame, x::Map; file=nothing) = _map_with(df, x.file, x)

function edit_with(df::DataFrame, x::Match; file=nothing)
    if x.on == r"expand range"
        ROWS, COLS = size(df)
        cols = propertynames(df)
        df = [[DataFrame(Dict(cols[jj] =>
                cols[jj] == x.input ? _expand_range(df[ii,jj]) : df[ii,jj]
            for jj in 1:COLS)) for ii in 1:ROWS]...;]
    else
        # Ensure all row values are strings and can be matched with a Regex, and do so.
        # Temporarily remove missing values, just in case.
        df[:,x.input] .= convert_type.(String, df[:,x.input])
        col = edit_with(copy(df), Replace(x.input, missing, ""))[:,x.input]
        m = match.(x.on, col)

        # Add empty columns for all output columns not already in the DataFrame.
        # Where there is a match, fill empty cells. If values in the input column,
        # leave cells without a match unchanged.
        df = edit_with(df, Add.(setdiff(x.output, propertynames(df)), ""))
        [m[ii] !== nothing && ([df[ii,out] = m[ii][out] for out in x.output])
            for ii in 1:length(m)]
    end
    return df
end

function edit_with(df::DataFrame, x::Melt; file=nothing)
    on = intersect(x.on, propertynames(df))
    df = melt(df, on, variable_name=x.var, value_name=x.val)
    df[!, x.var] .= convert_type.(String, df[:, x.var])
    return df
end

function edit_with(df::DataFrame, x::Operate; file=nothing)
    # Append columns from before the operation with 0 if they might be replaced.
    # This is useful for debugging purposes.
    df_val = convert_type.(Float64, df[:,x.input])
    x.output in x.input && (df = edit_with(df, Rename(x.output, append(x.output, 0))))
    df[!,x.output] .= broadcast(datatype(x.operation), [col for col in eachcol(df_val)]...)

    # Adjust labeling columns: If both from/to descriptive columns are distinct and
    # in the DataFrame, Replace the column values from -> to.
    for (from, to) in zip(x.from, x.to)
        if length(intersect(propertynames(df), [from,to])) == 2
            df_comment = dropmissing(unique(df[:, [from; to]]))
            df[!, append(from, 0)] .= df[:,from]
            df = edit_with(df, Replace.(from, df_comment[:,from], df_comment[:,to]))
        end
    end
    return round!(df, x.output)
end

function edit_with(df::DataFrame, x::Order; file=nothing)
    # If not all columns are present, return the DataFrame as is. Such is the case when a
    # descriptor column must be added when appending multiple data sets in one DataFrame.
    if size(intersect(x.col, propertynames(df)))[1] < size(x.col)[1]
        return df
    # If all of the columns are present in the original DataFrame,
    # reorder the DataFrame columns and set them to the specified type.
    else
        df = df[!, x.col]
        [df[!, c] .= convert_type.(t, df[!, c]) for (c, t) in zip(x.col, x.type)]
        return df
    end
end

function edit_with(df::DataFrame, x::Rename; file=nothing)
    cols = propertynames(df)
    x.from in cols && (df = rename(df, x.from => x.to))
    x.to == :upper && (df = edit_with(df, Rename.(cols, uppercase.(cols))))
    x.to == :lower && (df = edit_with(df, Rename.(cols, lowercase.(cols))))
    return df
end

function edit_with(df::DataFrame, x::Replace; file=nothing)
    !(x.col in propertynames(df)) && (return df)

    # Check if we want to replace a value x.from in the column x.col with a value in the
    # from another column in the same row. !!!! add this option to the documentation.
    m = match(r"(?<col>\S*) value", string(x.to))
    if !isnothing(m) && (m[:col] in names(df))
        ii = df[:,x.col] .=== x.from
        df[ii,x.col] .= df[ii,Symbol(m[:col])]

    else
        if x.from === missing && Symbol(x.to) in propertynames(df)
            df[ismissing.(df[:,x.col]),x.col] .= df[ismissing.(df[:,x.col]), Symbol(x.to)]
            return df
        end

        # !!!! find out where we use this (i think check share_gsp, labor) and doc this option!
        if x.to === Not && eltype(df[:,x.col]) == Bool
            df[!,x.col] .= .!df[:,x.col]
        end

        df[!,x.col] .= if x.to === "lower"  lowercase.(df[:,x.col])
        elseif x.to === "upper"             uppercase.(df[:,x.col])
        elseif x.to === "uppercasefirst"    uppercasefirst.(lowercase.(df[:,x.col]))
        elseif x.to === "titlecase"         titlecase.(df[:,x.col])
        else
            replace(strip.(copy(df[:,x.col])), x.from => x.to)
        end
    end

    return df
end

function edit_with(df::DataFrame, x::Stack; file=nothing)
    df = [[edit_with(df[:, occursin.(indicator, propertynames(df))],
        [Rename.(propertynames(df)[occursin.(indicator, propertynames(df))], x.col);
            Add(x.var, replace(string(indicator), "_" => " "))]
    ) for indicator in x.on]...;]
    return dropmissing(df)
end

function edit_with(df::DataFrame, x::Describe, file; print_status::Bool=false)
    return edit_with(df, Add(x.col, file.descriptor))
end

# ----- SUPPORT FOR MULTIPLE EDITS ---------------------------------------------------------

# THIS is for those other functions
function edit_with(df::DataFrame, x::T, file; print_status::Bool=false) where T <: Edit
    print_status && _print_status(x)
    return edit_with(df, x)
end

function edit_with(
    df::DataFrame,
    lst::Array{T,1};
    print_status::Bool=false) where T <: Edit

    # [df = edit_with(df, x; print_status = print_status) for x in lst]
    [df = edit_with(df, x) for x in lst]
    return df
end

function edit_with(
    df::DataFrame,
    lst::Array{T,1},
    file;
    print_status::Bool=false) where T <: Edit

    [df = edit_with(df, x, file; print_status=print_status) for x in lst]
    return df
end

# ----- EDIT FROM FILE ---------------------------------------------------------------------

function edit_with(
    df::DataFrame,
    y::Dict{Any,Any},
    file::T;
    print_status::Bool=false) where T <: File

    # Specify the order in which edits must occur and which of these edits are included
    # in the yaml file of defined edits.
    EDITS = ["Deselect", "Rename", "OrderedGroup", "Group", "Stack", "Match", "Melt",
        "Add", "Map", "Replace", "Drop", "Operate", "Combine", "Describe", "Order"]
    KEYS = intersect(EDITS, collect(keys(y)))
    [df = edit_with(df, y[k], file; print_status=print_status) for k in KEYS]
    return df
end

function edit_with(
    file::T,
    y::Dict{Any,Any};
    print_status::Bool=false) where T <: File

    df = read_file(y["PathIn"], file; remove_notes=true)
    return edit_with(df, y, file)
end

function edit_with(
    files::Array{T},
    y::Dict{Any,Any};
    print_status::Bool=false) where T <: File

    df = vcat([edit_with(file, y; print_status=print_status) for file in files]...; cols=:union)

    if haskey(y,"Order")
        df = edit_with(df, y["Order"])
    end

    df = dropzero(df)

    df = _filter_datastream(df, y)
    df = _sort_datastream(df, y)
    return df
end

function edit_with(y::Dict{Any,Any}; print_status::Bool=false)
    # Find all dictionary keys corresponding to file names and save these in a list to
    # read, edit, and concattenate.
    files = ensurearray(values(find_oftype(y, File)))
    return edit_with(files, y; print_status=print_status)
end

"""
Allows for *basic* filtering over years. Will need to expand to include regions.
"""
function _filter_datastream(df::DataFrame, y::Dict)
    path = joinpath("data", "coresets")
    set = Dict()
    if "Filter" in keys(y)
        y["Filter"] in [true,"year"]  && push!(set, :yr => read_file(joinpath(path, "yr.csv"))[:,1])
        y["Filter"] in [true,"state"] && push!(set, :r => read_file(joinpath(path, "r", "state.csv"))[:,1])
    end

    !isempty(set) && (df = filter_with(df, set; extrapolate=false))
    return df
end

"""
    _sort_datastream(df::DataFrame)
Returns the edited DataFrame, stored in a nicely-sorted order. This is most helpful for
mapping and developing. Sorting isn't *necessary* and we could remove this function to save
some time for users.
"""
function _sort_datastream(df::DataFrame, y::Dict{Any,Any})
    sorting = "Sort" in keys(y) ? y["Sort"] : true

    sorting == false && (return df)

    df = if sorting == true
        sort(df, findindex(df))
    elseif occursin("unique", sorting)
        sort_unique(df, sorting)
    end

    return df
end

"This function prints an editing status message."
function _print_status(x::Add)
    println("\tAdding ", x.col, " = ", x.val)
end

function _print_status(x::Drop)
    if x.val === "all"
        println("\tDropping ", x.col)
    else
        println("\tDropping ", x.col, " = ", x.val)
    end
end

function _print_status(x::Map)
    println("\tMapping ", x.input, " -> ", x.output, " using ", x.file)
end

function _print_status(x::Rename)
    println("\tRenaming ", x.from, " -> ", x.to)
end

function _print_status(x::Replace)
    println("\tReplacing ", x.col, ": ", x.from, " -> ", x.to)
end

function _print_status(x::Any)
    println("\tEditing with ", typeof(x))
end

"""
    _expand_range()
"""
function _expand_range(x::T) where T <: AbstractString
    if occursin("-", x)
        x = string(strip(x))
        if all(x .!= ["31-33", "44-45", "48-49"])
            x = split(x, "-")
            x = ensurearray(convert_type(Int, x[1]):convert_type(Int, x[1][1:end - 1] * x[end][end]))
        end
    else
        x = convert_type(Int, x)
    end
    return x
end

function _expand_range(x::String)
    if match(r"\D", x) !== nothing
        m = String.([m.match for m in collect(eachmatch(r"\D.*?", x))])
        x = if length(setdiff(m, [",","-"," "])) .== 0
            [_expand_range.(split(x, ","))...;]
        else
            missing
        end
    else
        x = convert_type(Int, x)
    end
    return x
end

_expand_range(x::Missing) = x
_expand_range(x::Int) = x

"""
"""
function _group_by(df::DataFrame, x::OrderedGroup)
    id = x.val[1]
    colid = append.(convert_type(Symbol, id), x.on)
    cols = unique([propertynames(df); colid])

    # Find the line where each group begins and ends.
    if length(unique(df[:,x.var])) > 1
        N = size(df, 1)
        df[!,:line] .= (1:N)
        df_split = df[df[:,x.var] .== id, :]
        df_split = edit_with(df_split, Rename.(x.on, colid))
        df_split[!,:start] .= df_split[:,:line] .+ 1
        df_split[!,:stop] .= [df_split[:,:line][2:end] .- 1; N]

        # Can probably change a little to fill in values.
        [df[!,col] .= "" for col in colid]
        [df[row[:start]:row[:stop], col] .= row[col] for row in eachrow(df_split)
            for col in colid]

        # Drop "headers" that indicated where groups stopped and started.
        df = df[df[:,colid[1]] .!== "", cols]
    else
        df = edit_with(df, [Rename.(x.on, colid); Deselect([x.var], "==")])
    end

    return df
end

"""
Edits input DataFrame given either a DataFrame or file path.
"""
function _map_with(df::DataFrame, df_map::DataFrame, x::Map)
    cols = unique([propertynames(df); x.output])

    # Rename columns in the mapping DataFrame to temporary values in case any of these
    # columns were already present in the input DataFrame.
    from = _generate_id(x.from, :from)
    to = _generate_id(x.to, :to)

    df_map = unique(hcat(
        edit_with(df_map[:,x.from], Rename.(x.from, from)),
        edit_with(df_map[:,x.to], Rename.(x.to, to)),
    ))

    # Ensure the input and mapping DataFrames are consistent in type. Types from the mapping
    # DataFrame are used since all values in each column should be of the same type.
    if findtype(df[:,x.input]) !== findtype(df_map[:,from])
        for (ii, ff) in zip(x.input, from)
            try
                df[!,ii] .= convert_type.(findtype(df_map[:,ff]), df[:,ii])
            catch
                df[!,ii] .= convert_type.(String, df[:,ii])
                df_map[!,ff] .= convert_type.(String, df_map[:,ff])
            end
        end
    end

    join_cols = Pair.(x.input, from)

    x.kind == :inner && (df = innerjoin(df, df_map, on=join_cols; makeunique=true))
    x.kind == :outer && (df = outerjoin(df, df_map, on=join_cols; makeunique=true))
    x.kind == :left  && (df = leftjoin(df, df_map,  on=join_cols; makeunique=true))
    x.kind == :right && (df = rightjoin(df, df_map, on=join_cols; makeunique=true))
    x.kind == :semi  && (df = semijoin(df, df_map,  on=join_cols; makeunique=true))

    # Remove all output column propertynames that might already be in the DataFrame.
    # These will be overwritten by the columns from the mapping DataFrame. Finally,
    # remane mapping "to" columns from their temporary to output values.
    df = df[:, setdiff(propertynames(df), x.output)]
    df = edit_with(df, Rename.(to, x.output))
    return df[:,cols]
end

_map_with(df::DataFrame, file::String, x::Map) = _map_with(df, read_file(x), x)

"""
"""
function sort_unique(df::DataFrame, idx::Array{Symbol,1})
    idx = idx[sortperm(length.(unique.(skipmissing.(eachcol(df[:,idx])))))]
    return sort(df, idx)
end

function sort_unique(df::DataFrame, id::String)
    (id == "unique") && (return sort_unique(df))

    m = match.(r"unique\s+(\S+)", id)
    id = m !== nothing ? Symbol(m[1]) : Symbol(id)

    return sort_unique(df, id)
end

function sort_unique(df::DataFrame, id::Symbol)
    idx = propertynames(df)
    subidx = idx[occursin.(id, idx)]

    length(subidx) > 0 && (idx = subidx)

    return sort_unique(df, idx)
end

sort_unique(df::DataFrame) = sort_unique(df, findindex(df))

"""
"""
function DataFrames.select!(df::DataFrame, x::Parameter)
    return select!(df, intersect([x.index; :value], propertynames(df)))
end

max_stars_repo_name: NREL/SLiDE

max_issues_repo_name: NREL/SLiDE

max_forks_repo_name: NREL/SLiDE


text:
using Test
using Distributions
using PkgB
using Random

@testset "randchisq2" begin
  Random.seed!(0)
  x = [PkgB.randchisq2() for _ in 1:1_000_000]
  @test isapprox(mean(x), 2, rtol=0.01)
  @test isapprox(var(x), 4, rtol=0.01)
end


text:
using PaddedBlocks
using Documenter

makedocs(;
    modules=[PaddedBlocks],
    authors="Vilim Štih <vilim@neuro.mpg.de>",
    repo="https://github.com/portugueslab/PaddedBlocks.jl/blob/{commit}{path}#L{line}",
    sitename="PaddedBlocks.jl",
    format=Documenter.HTML(;
        prettyurls=get(ENV, "CI", "false") == "true",
        canonical="https://portugueslab.github.io/PaddedBlocks.jl",
        assets=String[],
    ),
    pages=[
        "Home" => "index.md",
    ],
)

deploydocs(;
    repo="github.com/portugueslab/PaddedBlocks.jl",
)


text:
struct TerminalCache
    v1::Vector{Vector{Int}}
    v2::Vector{Vector{Vector{Int}}}
    v3::Vector{Vector{Int}}
end

function TerminalCache(term::Matrix{Vector{Int}})
    v1 = vec(term)
    v2 = [term[i,:] for i in 1:size(term,1)]
    v3 = [[0,0]]

    return TerminalCache(v1,v2,v3)
end

Base.getindex(tc::TerminalCache) = tc.v1
Base.getindex(tc::TerminalCache, i::Int) = tc.v2[i]
function Base.getindex(tc::TerminalCache, i::Int, j::Int)
    tc.v3[1][1] = i
    tc.v3[1][2] = j
    return tc.v3
end

function Base.getindex(tc::TerminalCache, v)
    if length(v) === 0
        return tc[]
    elseif length(v) === 1
        return tc[v[1]]
    else
        return tc[v[1], v[2]]
    end
end

max_stars_repo_name: WhiffleFish/HelloCFR

max_issues_repo_name: WhiffleFish/HelloCFR

max_forks_repo_name: WhiffleFish/HelloCFR


text:
import FFTW

#
# Perform (i)FFTs.
#
# We perform two sets of (i)FFTs.

# For densities and potentials defined on the cubic basis set, r_to_G/G_to_r
# do a simple FFT/IFFT from the cubic basis set to the real-space grid.
# These functions do not take a k-point as input

# For orbitals, G_to_r converts the orbitals defined on a spherical
# basis set to the cubic basis set using zero padding, then performs
# an IFFT to get to the real-space grid. r_to_G performs an FFT, then
# restricts the output to the spherical basis set. These functions
# take a k-point as input.

"""
In-place version of `G_to_r`.
"""
function G_to_r!(f_real::AbstractArray3, basis::PlaneWaveBasis, f_fourier::AbstractArray3)
    mul!(f_real, basis.opBFFT, f_fourier)
    f_real .*= basis.G_to_r_normalization
end
function G_to_r!(f_real::AbstractArray3, basis::PlaneWaveBasis,
                 kpt::Kpoint, f_fourier::AbstractVector; normalize=true)
    @assert length(f_fourier) == length(kpt.mapping)
    @assert size(f_real) == basis.fft_size

    # Pad the input data
    fill!(f_real, 0)
    f_real[kpt.mapping] = f_fourier

    # Perform an IFFT
    mul!(f_real, basis.ipBFFT, f_real)
    normalize && (f_real .*= basis.G_to_r_normalization)
    f_real
end

"""
    G_to_r(basis::PlaneWaveBasis, [kpt::Kpoint, ] f_fourier)

Perform an iFFT to obtain the quantity defined by `f_fourier` defined
on the k-dependent spherical basis set (if `kpt` is given) or the
k-independent cubic (if it is not) on the real-space grid.
"""
function G_to_r(basis::PlaneWaveBasis, f_fourier::AbstractArray; assume_real=Val(true))
    # assume_real is true by default because this is the most common usage
    # (for densities & potentials). Val(true) to help const-prop;
    # see https://github.com/JuliaLang/julia/issues/44330
    f_real = similar(f_fourier)
    @assert length(size(f_fourier)) ∈ (3, 4)
    # this exploits trailing index convention
    for iσ = 1:size(f_fourier, 4)
        @views G_to_r!(f_real[:, :, :, iσ], basis, f_fourier[:, :, :, iσ])
    end
    (assume_real == Val(true)) ? real(f_real) : f_real
end
function G_to_r(basis::PlaneWaveBasis, kpt::Kpoint, f_fourier::AbstractVector; kwargs...)
    G_to_r!(similar(f_fourier, basis.fft_size...), basis, kpt, f_fourier; kwargs...)
end

@doc raw"""
In-place version of `r_to_G!`.
NOTE: If `kpt` is given, not only `f_fourier` but also `f_real` is overwritten.
"""
function r_to_G!(f_fourier::AbstractArray3, basis::PlaneWaveBasis, f_real::AbstractArray3)
    if eltype(f_real) <: Real
        f_real = complex.(f_real)
    end
    mul!(f_fourier, basis.opFFT, f_real)
    f_fourier .*= basis.r_to_G_normalization
end
function r_to_G!(f_fourier::AbstractVector, basis::PlaneWaveBasis,
                 kpt::Kpoint, f_real::AbstractArray3; normalize=true)
    @assert size(f_real) == basis.fft_size
    @assert length(f_fourier) == length(kpt.mapping)

    # FFT
    mul!(f_real, basis.ipFFT, f_real)

    # Truncate
    f_fourier .= view(f_real, kpt.mapping)
    normalize && (f_fourier .*= basis.r_to_G_normalization)
    f_fourier
end

"""
    r_to_G(basis::PlaneWaveBasis, [kpt::Kpoint, ] f_real)

Perform an FFT to obtain the Fourier representation of `f_real`. If
`kpt` is given, the coefficients are truncated to the k-dependent
spherical basis set.
"""
function r_to_G(basis::PlaneWaveBasis{T}, f_real::AbstractArray) where T
    f_fourier = similar(f_real, complex(promote_type(T, eltype(f_real))))
    @assert length(size(f_real)) ∈ (3, 4)
    # this exploits trailing index convention
    for iσ = 1:size(f_real, 4)
        @views r_to_G!(f_fourier[:, :, :, iσ], basis, f_real[:, :, :, iσ])
    end
    f_fourier
end
# TODO optimize this
function r_to_G(basis::PlaneWaveBasis, kpt::Kpoint, f_real::AbstractArray3; kwargs...)
    r_to_G!(similar(f_real, length(kpt.mapping)), basis, kpt, copy(f_real); kwargs...)
end

# returns matrix representations of the G_to_r and r_to_G matrices. For debug purposes.
function G_to_r_matrix(basis::PlaneWaveBasis{T}) where {T}
    ret = zeros(complex(T), prod(basis.fft_size), prod(basis.fft_size))
    for (iG, G) in enumerate(G_vectors(basis))
        for (ir, r) in enumerate(r_vectors(basis))
            ret[ir, iG] = cis2pi(dot(r, G)) / sqrt(basis.model.unit_cell_volume)
        end
    end
    ret
end
function r_to_G_matrix(basis::PlaneWaveBasis{T}) where {T}
    ret = zeros(complex(T), prod(basis.fft_size), prod(basis.fft_size))
    for (iG, G) in enumerate(G_vectors(basis))
        for (ir, r) in enumerate(r_vectors(basis))
            Ω = basis.model.unit_cell_volume
            ret[iG, ir] = cis2pi(-dot(r, G)) * sqrt(Ω) / prod(basis.fft_size)
        end
    end
    ret
end

@doc raw"""
Determine the minimal grid size for the cubic basis set to be able to
represent product of orbitals (with the default `supersampling=2`).

Optionally optimize the grid afterwards for the FFT procedure by
ensuring factorization into small primes.

The function will determine the smallest parallelepiped containing the wave vectors
 ``|G|^2/2 \leq E_\text{cut} ⋅ \text{supersampling}^2``.
For an exact representation of the density resulting from wave functions
represented in the spherical basis sets, `supersampling` should be at least `2`.

If `factors` is not empty, ensure that the resulting fft_size contains all the factors
"""
function compute_fft_size(model::Model{T}, Ecut, kcoords=nothing;
                          ensure_smallprimes=true, algorithm=:fast, factors=1, kwargs...) where T
    if algorithm == :fast
        Glims = compute_Glims_fast(model.lattice, Ecut; kwargs...)
    elseif algorithm == :precise
        @assert !isnothing(kcoords)

        # We build a temporary set of k-points here
        # We don't reuse this k-point construction for the pwbasis
        # because build_kpoints builds index mapping from the
        # k-point-specific basis to the global basis and thus the
        # fft_size needs to be final at k-point construction time
        Glims_temp    = compute_Glims_fast(model.lattice, Ecut; kwargs...)
        fft_size_temp = Tuple{Int, Int, Int}(2 .* Glims_temp .+ 1)
        kpoints_temp  = build_kpoints(model, fft_size_temp, kcoords, Ecut)

        Glims = compute_Glims_precise(model.lattice, Ecut, kpoints_temp; kwargs...)
    else
        error("Unknown fft_size_algorithm :$algorithm, try :fast or :precise")
    end

    # TODO Make default small primes type-dependent, since generic FFT is broken for some
    #      prime factors ... temporary workaround, see more details in workarounds/fft_generic.jl
    if ensure_smallprimes
        smallprimes = default_primes(T)  # Usually (2, 3 ,5)
    else
        smallprimes = ()
    end

    # Consider only sizes that are (a) a product of small primes and (b) contain the factors
    fft_size = Vec3(2 .* Glims .+ 1)
    fft_size = next_compatible_fft_size(fft_size; factors, smallprimes)
    Tuple{Int, Int, Int}(fft_size)
end

"""
Find the next compatible FFT size
Sizes must (a) be a product of small primes only and (b) contain the factors.
If smallprimes is empty (a) is skipped.
"""
function next_compatible_fft_size(size::Int; smallprimes=(2, 3, 5), factors=(1, ))
    # This could be optimized
    is_product_of_primes(n) = isempty(smallprimes) || (n == nextprod(smallprimes, n))
    @assert all(is_product_of_primes, factors) # ensure compatibility between (a) and (b)
    has_factors(n) = rem(n, prod(factors)) == 0

    while !(has_factors(size) && is_product_of_primes(size))
        size += 1
    end
    size
end
function next_compatible_fft_size(sizes::Union{Tuple, AbstractArray}; kwargs...)
    next_compatible_fft_size.(sizes; kwargs...)
end

# This uses a more precise and slower algorithm than the one above,
# simply enumerating all G vectors and seeing where their difference
# is. It needs the kpoints to do so.
@timing function compute_Glims_precise(lattice::AbstractMatrix{T}, Ecut, kpoints; supersampling=2) where T
    recip_lattice  = compute_recip_lattice(lattice)
    recip_diameter = diameter(recip_lattice)
    Glims = [0, 0, 0]
    # get the bounding rectangle that contains all G-G' vectors
    # (and therefore densities and potentials)
    # This handles the case `supersampling=2`
    for kpt in kpoints
        # TODO Hack: kpt.G_vectors is an internal detail, better use G_vectors(basis, kpt)
        for G in kpt.G_vectors
            if norm(recip_lattice * (G + kpt.coordinate)) ≤ sqrt(2Ecut) - recip_diameter
                # each of the 8 neighbors (in ∞-norm) also belongs to the grid
                # so we can safely skip the search knowing at least one of them
                # will have larger |G-G′|.
                # Savings with this trick are surprisingly small :
                # for silicon, 50% at Ecut 30, 70% at Ecut 100
                continue
            end
            # TODO Hack: kpt.G_vectors is an internal detail, better use G_vectors(basis, kpt)
            for G′ in kpt.G_vectors
                for α = 1:3
                    @inbounds Glims[α] = max(Glims[α], abs(G[α] - G′[α]))
                end
            end
        end
    end
    if supersampling != 2
        # no guarantees here, we just do our best to satisfy the
        # target supersampling ratio
        Glims = round.(Int, supersampling ./ 2 .* Glims)
    end
    Glims
end

# Fast implementation, but sometimes larger than necessary.
function compute_Glims_fast(lattice::AbstractMatrix{T}, Ecut; supersampling=2, tol=sqrt(eps(T))) where T
    Gmax = supersampling * sqrt(2 * Ecut)
    bounding_rectangle(lattice, Gmax; tol=tol)
end

# returns the lengths of the bounding rectangle in reciprocal space
# that encloses the sphere of radius Gmax
function bounding_rectangle(lattice::AbstractMatrix{T}, Gmax; tol=sqrt(eps(T))) where {T}
    # If |B G| ≤ Gmax, then
    # |Gi| = |e_i^T B^-1 B G| ≤ |B^-T e_i| Gmax = |A_i| Gmax
    # with B the reciprocal lattice matrix, e_i the i-th canonical
    # basis vector and A_i the i-th column of the lattice matrix
    Glims = [norm(lattice[:, i]) / 2T(π) * Gmax for i in 1:3]

    # Round up, unless exactly zero (in which case keep it zero in
    # order to just have one G vector for 1D or 2D systems)
    Glims = [Glim == 0 ? 0 : ceil(Int, Glim .- tol) for Glim in Glims]
    Glims
end

# For Float32 there are issues with aligned FFTW plans, so we
# fall back to unaligned FFTW plans (which are generally discouraged).
_fftw_flags(::Type{Float32}) = FFTW.MEASURE | FFTW.UNALIGNED
_fftw_flags(::Type{Float64}) = FFTW.MEASURE

"""
Plan a FFT of type `T` and size `fft_size`, spending some time on finding an
optimal algorithm. (Inplace, out-of-place) x (forward, backward) FFT plans are returned.
"""
function build_fft_plans(T::Union{Type{Float32}, Type{Float64}}, fft_size)
    tmp = Array{Complex{T}}(undef, fft_size...)
    ipFFT = FFTW.plan_fft!(tmp, flags=_fftw_flags(T))
    opFFT = FFTW.plan_fft(tmp, flags=_fftw_flags(T))
    # backward by inverting and stripping off normalizations
    ipFFT, opFFT, inv(ipFFT).p, inv(opFFT).p
end

# TODO Some grid sizes are broken in the generic FFT implementation
# in FourierTransforms, for more details see workarounds/fft_generic.jl
default_primes(::Type{Float32}) = (2, 3, 5)
default_primes(::Type{Float64}) = default_primes(Float32)
next_working_fft_size(::Type{Float32}, size::Int) = size
next_working_fft_size(::Type{Float64}, size::Int) = size
next_working_fft_size(T, sizes::Union{Tuple, AbstractArray}) = next_working_fft_size.(T, sizes)


text:
# custom extension of CuArray in CUDArt for sparse vectors/matrices
# using CSC format for interop with Julia's native sparse functionality

export CuSparseMatrixCSC, CuSparseMatrixCSR, CuSparseMatrixBSR, CuSparseMatrixCOO,
       CuSparseMatrix, AbstractCuSparseMatrix,
       CuSparseVector

using LinearAlgebra: BlasFloat
using SparseArrays: nonzeroinds, dimlub

abstract type AbstractCuSparseArray{Tv, Ti, N} <: AbstractSparseArray{Tv, Ti, N} end
const AbstractCuSparseVector{Tv, Ti} = AbstractCuSparseArray{Tv, Ti, 1}
const AbstractCuSparseMatrix{Tv, Ti} = AbstractCuSparseArray{Tv, Ti, 2}

Base.convert(T::Type{<:AbstractCuSparseArray}, m::AbstractArray) = m isa T ? m : T(m)

mutable struct CuSparseVector{Tv, Ti} <: AbstractCuSparseVector{Tv, Ti}
    iPtr::CuVector{Ti}
    nzVal::CuVector{Tv}
    dims::NTuple{2,Int}
    nnz::Ti

    function CuSparseVector{Tv, Ti}(iPtr::CuVector{<:Integer}, nzVal::CuVector,
                                dims::Integer) where {Tv, Ti <: Integer}
        new{Tv, Ti}(iPtr, nzVal, (dims,1), length(nzVal))
    end
end

function CUDA.unsafe_free!(xs::CuSparseVector)
    unsafe_free!(nonzeroinds(xs))
    unsafe_free!(nonzeros(xs))
    return
end

mutable struct CuSparseMatrixCSC{Tv, Ti} <: AbstractCuSparseMatrix{Tv, Ti}
    colPtr::CuVector{Ti}
    rowVal::CuVector{Ti}
    nzVal::CuVector{Tv}
    dims::NTuple{2,Int}
    nnz::Ti

    function CuSparseMatrixCSC{Tv, Ti}(colPtr::CuVector{<:Integer}, rowVal::CuVector{<:Integer},
                                   nzVal::CuVector, dims::NTuple{2,<:Integer}) where {Tv, Ti <: Integer}
        new{Tv, Ti}(colPtr, rowVal, nzVal, dims, length(nzVal))
    end
end

CuSparseMatrixCSC(A::CuSparseMatrixCSC) = A

function CUDA.unsafe_free!(xs::CuSparseMatrixCSC)
    unsafe_free!(xs.colPtr)
    unsafe_free!(rowvals(xs))
    unsafe_free!(nonzeros(xs))
    return
end

"""
    CuSparseMatrixCSR{Tv, Ti} <: AbstractCuSparseMatrix{Tv, Ti}

Container to hold sparse matrices in compressed sparse row (CSR) format on the
GPU.

!!! note
    Most CUSPARSE operations work with CSR formatted matrices, rather
    than CSC.

!!! compat "CUDA 11"
    Support of indices type rather than `Cint` (`Int32`) requires at least CUDA 11.
"""
mutable struct CuSparseMatrixCSR{Tv, Ti} <: AbstractCuSparseMatrix{Tv, Ti}
    rowPtr::CuVector{Ti}
    colVal::CuVector{Ti}
    nzVal::CuVector{Tv}
    dims::NTuple{2,Int}
    nnz::Ti

    function CuSparseMatrixCSR{Tv, Ti}(rowPtr::CuVector{<:Integer}, colVal::CuVector{<:Integer},
                                   nzVal::CuVector, dims::NTuple{2,Int}) where {Tv, Ti<:Integer}
        new{Tv, Ti}(rowPtr, colVal, nzVal, dims, length(nzVal))
    end
end

CuSparseMatrixCSR(A::CuSparseMatrixCSR) = A

function CUDA.unsafe_free!(xs::CuSparseMatrixCSR)
    unsafe_free!(xs.rowPtr)
    unsafe_free!(xs.colVal)
    unsafe_free!(nonzeros(xs))
    return
end

"""
Container to hold sparse matrices in block compressed sparse row (BSR) format on
the GPU. BSR format is also used in Intel MKL, and is suited to matrices that are
"block" sparse - rare blocks of non-sparse regions.
"""
mutable struct CuSparseMatrixBSR{Tv, Ti} <: AbstractCuSparseMatrix{Tv, Ti}
    rowPtr::CuVector{Ti}
    colVal::CuVector{Ti}
    nzVal::CuVector{Tv}
    dims::NTuple{2,Int}
    blockDim::Ti
    dir::SparseChar
    nnz::Ti

    function CuSparseMatrixBSR{Tv, Ti}(rowPtr::CuVector{<:Integer}, colVal::CuVector{<:Integer},
                                   nzVal::CuVector, dims::NTuple{2,<:Integer},
                                   blockDim::Integer, dir::SparseChar, nnz::Integer) where {Tv, Ti<:Integer}
        new{Tv, Ti}(rowPtr, colVal, nzVal, dims, blockDim, dir, nnz)
    end
end

CuSparseMatrixBSR(A::CuSparseMatrixBSR) = A

function CUDA.unsafe_free!(xs::CuSparseMatrixBSR)
    unsafe_free!(xs.rowPtr)
    unsafe_free!(xs.colVal)
    unsafe_free!(nonzeros(xs))
    return
end

"""
Container to hold sparse matrices in coordinate (COO) format on the GPU. COO
format is mainly useful to initially construct sparse matrices, afterwards
switch to [`CuSparseMatrixCSR`](@ref) for more functionality.
"""
mutable struct CuSparseMatrixCOO{Tv, Ti} <: AbstractCuSparseMatrix{Tv, Ti}
    rowInd::CuVector{Ti}
    colInd::CuVector{Ti}
    nzVal::CuVector{Tv}
    dims::NTuple{2,Int}
    nnz::Ti

    function CuSparseMatrixCOO{Tv, Ti}(rowInd::CuVector{<:Integer}, colInd::CuVector{<:Integer},
                                   nzVal::CuVector, dims::NTuple{2,Int}=(dimlub(rowInd),dimlub(colInd)),
                                   nnz::Integer=length(nzVal)) where {Tv, Ti}
        new{Tv, Ti}(rowInd,colInd,nzVal,dims,nnz)
    end
end

CuSparseMatrixCOO(A::CuSparseMatrixCOO) = A

"""
Utility union type of [`CuSparseMatrixCSC`](@ref), [`CuSparseMatrixCSR`](@ref),
[`CuSparseMatrixBSR`](@ref), [`CuSparseMatrixCOO`](@ref).
"""
const CuSparseMatrix{Tv, Ti} = Union{
    CuSparseMatrixCSC{Tv, Ti},
    CuSparseMatrixCSR{Tv, Ti},
    CuSparseMatrixBSR{Tv, Ti},
    CuSparseMatrixCOO{Tv, Ti}
}

# NOTE: we use Cint as default Ti on CUDA instead of Int to provide
# maximum compatiblity to old CUSPARSE APIs
function CuSparseVector{Tv}(iPtr::CuVector{<:Integer}, nzVal::CuVector, dims::Integer) where {Tv}
    CuSparseVector{Tv, Cint}(convert(CuVector{Cint}, iPtr), nzVal, dims)
end

function CuSparseMatrixCSC{Tv}(colPtr::CuVector{<:Integer}, rowVal::CuVector{<:Integer},
                                   nzVal::CuVector, dims::NTuple{2,<:Integer}) where {Tv}
    CuSparseMatrixCSC{Tv, Cint}(colPtr, rowVal, nzVal, dims)
end

function CuSparseMatrixCSR{Tv}(rowPtr::CuVector{<:Integer}, colVal::CuVector{<:Integer},
                                   nzVal::CuVector, dims::NTuple{2,Int}) where {Tv}
    CuSparseMatrixCSR{Tv, Cint}(rowPtr, colVal, nzVal, dims)
end

function CuSparseMatrixBSR{Tv}(rowPtr::CuVector{<:Integer}, colVal::CuVector{<:Integer},
                                   nzVal::CuVector, dims::NTuple{2,<:Integer},
                                   blockDim::Integer, dir::SparseChar, nnz::Integer) where {Tv}
    CuSparseMatrixBSR{Tv, Cint}(rowPtr, colVal, nzVal, dims, blockDim, dir, nnz)
end

function CuSparseMatrixCOO{Tv}(rowInd::CuVector{<:Integer}, colInd::CuVector{<:Integer},
                                   nzVal::CuVector, dims::NTuple{2,Int}=(dimlub(rowInd),dimlub(colInd)),
                                   nnz::Integer=length(nzVal)) where {Tv}
    CuSparseMatrixCOO{Tv, Cint}(rowInd,colInd,nzVal,dims,nnz)
end

## convenience constructors
CuSparseVector(iPtr::DenseCuArray{<:Integer}, nzVal::DenseCuArray{T}, dims::Int) where {T} =
    CuSparseVector{T}(iPtr, nzVal, dims)

CuSparseMatrixCSC(colPtr::DenseCuArray{<:Integer}, rowVal::DenseCuArray{<:Integer},
                  nzVal::DenseCuArray{T}, dims::NTuple{2,Int}) where {T} =
    CuSparseMatrixCSC{T}(colPtr, rowVal, nzVal, dims)

CuSparseMatrixCSR(rowPtr::DenseCuArray, colVal::DenseCuArray, nzVal::DenseCuArray{T}, dims::NTuple{2,Int}) where T =
    CuSparseMatrixCSR{T}(rowPtr, colVal, nzVal, dims)

CuSparseMatrixBSR(rowPtr::DenseCuArray, colVal::DenseCuArray, nzVal::DenseCuArray{T}, blockDim, dir, nnz,
                  dims::NTuple{2,Int}) where T =
    CuSparseMatrixBSR{T}(rowPtr, colVal, nzVal, dims, blockDim, dir, nnz)

Base.similar(Vec::CuSparseVector) = CuSparseVector(copy(nonzeroinds(Vec)), similar(nonzeros(Vec)), Vec.dims[1])
Base.similar(Mat::CuSparseMatrixCSC) = CuSparseMatrixCSC(copy(Mat.colPtr), copy(rowvals(Mat)), similar(nonzeros(Mat)), Mat.dims)
Base.similar(Mat::CuSparseMatrixCSR) = CuSparseMatrixCSR(copy(Mat.rowPtr), copy(Mat.colVal), similar(nonzeros(Mat)), Mat.dims)
Base.similar(Mat::CuSparseMatrixBSR) = CuSparseMatrixBSR(copy(Mat.rowPtr), copy(Mat.colVal), similar(nonzeros(Mat)), Mat.blockDim, Mat.dir, nnz(Mat), Mat.dims)

## array interface

Base.length(g::CuSparseVector) = prod(g.dims)
Base.size(g::CuSparseVector) = g.dims
Base.ndims(g::CuSparseVector) = 1

Base.length(g::CuSparseMatrix) = prod(g.dims)
Base.size(g::CuSparseMatrix) = g.dims
Base.ndims(g::CuSparseMatrix) = 2

function Base.size(g::CuSparseVector, d::Integer)
    if d == 1
        return g.dims[d]
    elseif d > 1
        return 1
    else
        throw(ArgumentError("dimension must be ≥ 1, got $d"))
    end
end

function Base.size(g::CuSparseMatrix, d::Integer)
    if d in [1, 2]
        return g.dims[d]
    elseif d > 1
        return 1
    else
        throw(ArgumentError("dimension must be ≥ 1, got $d"))
    end
end

Base.eltype(g::CuSparseMatrix{T}) where T = T

## sparse array interface

SparseArrays.nnz(g::AbstractCuSparseArray) = g.nnz
SparseArrays.nonzeros(g::AbstractCuSparseArray) = g.nzVal

SparseArrays.nonzeroinds(g::AbstractCuSparseVector) = g.iPtr

SparseArrays.rowvals(g::CuSparseMatrixCSC) = g.rowVal

LinearAlgebra.issymmetric(M::Union{CuSparseMatrixCSC,CuSparseMatrixCSR}) = false
LinearAlgebra.ishermitian(M::Union{CuSparseMatrixCSC,CuSparseMatrixCSR}) = false
LinearAlgebra.issymmetric(M::Symmetric{CuSparseMatrixCSC}) = true
LinearAlgebra.ishermitian(M::Hermitian{CuSparseMatrixCSC}) = true

LinearAlgebra.istriu(M::UpperTriangular{T,S}) where {T<:BlasFloat, S<:AbstractCuSparseMatrix} = true
LinearAlgebra.istril(M::UpperTriangular{T,S}) where {T<:BlasFloat, S<:AbstractCuSparseMatrix} = false
LinearAlgebra.istriu(M::LowerTriangular{T,S}) where {T<:BlasFloat, S<:AbstractCuSparseMatrix} = false
LinearAlgebra.istril(M::LowerTriangular{T,S}) where {T<:BlasFloat, S<:AbstractCuSparseMatrix} = true

Hermitian{T}(Mat::CuSparseMatrix{T}) where T = Hermitian{T,typeof(Mat)}(Mat,'U')

## indexing

# translations
Base.getindex(A::AbstractCuSparseVector, ::Colon)          = copy(A)
Base.getindex(A::AbstractCuSparseMatrix, ::Colon, ::Colon) = copy(A)
Base.getindex(A::AbstractCuSparseMatrix, i, ::Colon)       = getindex(A, i, 1:size(A, 2))
Base.getindex(A::AbstractCuSparseMatrix, ::Colon, i)       = getindex(A, 1:size(A, 1), i)
Base.getindex(A::AbstractCuSparseMatrix, I::Tuple{Integer,Integer}) = getindex(A, I[1], I[2])

# column slices
function Base.getindex(x::CuSparseMatrixCSC, ::Colon, j::Integer)
    checkbounds(x, :, j)
    r1 = convert(Int, x.colPtr[j])
    r2 = convert(Int, x.colPtr[j+1]) - 1
    CuSparseVector(rowvals(x)[r1:r2], nonzeros(x)[r1:r2], size(x, 1))
end

function Base.getindex(x::CuSparseMatrixCSR, i::Integer, ::Colon)
    checkbounds(x, i, :)
    c1 = convert(Int, x.rowPtr[i])
    c2 = convert(Int, x.rowPtr[i+1]) - 1
    CuSparseVector(x.colVal[c1:c2], nonzeros(x)[c1:c2], size(x, 2))
end

# row slices
Base.getindex(A::CuSparseMatrixCSC, i::Integer, ::Colon) = CuSparseVector(sparse(A[i, 1:end]))  # TODO: optimize
Base.getindex(A::CuSparseMatrixCSR, ::Colon, j::Integer) = CuSparseVector(sparse(A[1:end, j]))  # TODO: optimize

function Base.getindex(A::CuSparseMatrixCSC{T}, i0::Integer, i1::Integer) where T
    m, n = size(A)
    if !(1 <= i0 <= m && 1 <= i1 <= n)
        throw(BoundsError())
    end
    r1 = Int(A.colPtr[i1])
    r2 = Int(A.colPtr[i1+1]-1)
    (r1 > r2) && return zero(T)
    r1 = searchsortedfirst(rowvals(A), i0, r1, r2, Base.Order.Forward)
    ((r1 > r2) || (rowvals(A)[r1] != i0)) ? zero(T) : nonzeros(A)[r1]
end

function Base.getindex(A::CuSparseMatrixCSR{T}, i0::Integer, i1::Integer) where T
    m, n = size(A)
    if !(1 <= i0 <= m && 1 <= i1 <= n)
        throw(BoundsError())
    end
    c1 = Int(A.rowPtr[i0])
    c2 = Int(A.rowPtr[i0+1]-1)
    (c1 > c2) && return zero(T)
    c1 = searchsortedfirst(A.colVal, i1, c1, c2, Base.Order.Forward)
    ((c1 > c2) || (A.colVal[c1] != i1)) ? zero(T) : nonzeros(A)[c1]
end

function SparseArrays._spgetindex(m::Integer, nzind::CuVector{Ti}, nzval::CuVector{Tv},
                                  i::Integer) where {Tv,Ti}
    ii = searchsortedfirst(nzind, convert(Ti, i))
    (ii <= m && nzind[ii] == i) ? nzval[ii] : zero(Tv)
end

## interop with sparse CPU arrays

# cpu to gpu
# NOTE: we eagerly convert the indices to Cint here to avoid additional conversion later on
CuSparseVector{T}(Vec::SparseVector) where {T} =
    CuSparseVector(CuVector{Cint}(Vec.nzind), CuVector{T}(Vec.nzval), length(Vec))
CuSparseVector{T}(Mat::SparseMatrixCSC) where {T} =
    size(Mat,2) == 1 ?
        CuSparseVector(CuVector{Cint}(Mat.rowval), CuVector{T}(Mat.nzval), size(Mat)[1]) :
        throw(ArgumentError("The input argument must have a single column"))
CuSparseMatrixCSC{T}(Vec::SparseVector) where {T} =
    CuSparseMatrixCSC{T}(CuVector{Cint}([1]), CuVector{Cint}(Vec.nzind),
                         CuVector{T}(Vec.nzval), size(Vec))
CuSparseMatrixCSC{T}(Mat::SparseMatrixCSC) where {T} =
    CuSparseMatrixCSC{T}(CuVector{Cint}(Mat.colptr), CuVector{Cint}(Mat.rowval),
                         CuVector{T}(Mat.nzval), size(Mat))
CuSparseMatrixCSR{T}(Mat::Transpose{Tv, <:SparseMatrixCSC}) where {T, Tv} =
    CuSparseMatrixCSR{T}(CuVector{Cint}(parent(Mat).colptr), CuVector{Cint}(parent(Mat).rowval),
                         CuVector{T}(parent(Mat).nzval), size(Mat))
CuSparseMatrixCSR{T}(Mat::Adjoint{Tv, <:SparseMatrixCSC}) where {T, Tv} =
    CuSparseMatrixCSR{T}(CuVector{Cint}(parent(Mat).colptr), CuVector{Cint}(parent(Mat).rowval),
                         CuVector{T}(conj.(parent(Mat).nzval)), size(Mat))
CuSparseMatrixCSR{T}(Mat::SparseMatrixCSC) where {T} = CuSparseMatrixCSR(CuSparseMatrixCSC{T}(Mat))
CuSparseMatrixBSR{T}(Mat::SparseMatrixCSC, blockdim) where {T} = CuSparseMatrixBSR(CuSparseMatrixCSR{T}(Mat), blockdim)
CuSparseMatrixCOO{T}(Mat::SparseMatrixCSC) where {T} = CuSparseMatrixCOO(CuSparseMatrixCSR{T}(Mat))

# untyped variants
CuSparseVector(x::AbstractSparseArray{T}) where {T} = CuSparseVector{T}(x)
CuSparseMatrixCSC(x::AbstractSparseArray{T}) where {T} = CuSparseMatrixCSC{T}(x)
CuSparseMatrixCSR(x::AbstractSparseArray{T}) where {T} = CuSparseMatrixCSR{T}(x)
CuSparseMatrixBSR(x::AbstractSparseArray{T}, blockdim) where {T} = CuSparseMatrixBSR{T}(x, blockdim)
CuSparseMatrixCOO(x::AbstractSparseArray{T}) where {T} = CuSparseMatrixCOO{T}(x)
CuSparseMatrixCSR(x::Transpose{T}) where {T} = CuSparseMatrixCSR{T}(x)
CuSparseMatrixCSR(x::Adjoint{T}) where {T} = CuSparseMatrixCSR{T}(x)
CuSparseMatrixCSC(x::Transpose{T}) where {T} = CuSparseMatrixCSC{T}(x)
CuSparseMatrixCSC(x::Adjoint{T}) where {T} = CuSparseMatrixCSC{T}(x)

# gpu to cpu
SparseVector(x::CuSparseVector) = SparseVector(length(x), Array(nonzeroinds(x)), Array(nonzeros(x)))
SparseMatrixCSC(x::CuSparseMatrixCSC) = SparseMatrixCSC(size(x)..., Array(x.colPtr), Array(rowvals(x)), Array(nonzeros(x)))
SparseMatrixCSC(x::CuSparseMatrixCSR) = SparseMatrixCSC(CuSparseMatrixCSC(x))  # no direct conversion
SparseMatrixCSC(x::CuSparseMatrixBSR) = SparseMatrixCSC(CuSparseMatrixCSR(x))  # no direct conversion
SparseMatrixCSC(x::CuSparseMatrixCOO) = SparseMatrixCSC(CuSparseMatrixCSR(x))  # no direct conversion

# collect to Array
Base.collect(x::CuSparseVector) = collect(SparseVector(x))
Base.collect(x::CuSparseMatrixCSC) = collect(SparseMatrixCSC(x))
Base.collect(x::CuSparseMatrixCSR) = collect(SparseMatrixCSC(x))
Base.collect(x::CuSparseMatrixBSR) = collect(CuSparseMatrixCSR(x))  # no direct conversion
Base.collect(x::CuSparseMatrixCOO) = collect(CuSparseMatrixCSR(x))  # no direct conversion

Adapt.adapt_storage(::Type{CuArray}, xs::SparseVector) = CuSparseVector(xs)
Adapt.adapt_storage(::Type{CuArray}, xs::SparseMatrixCSC) = CuSparseMatrixCSC(xs)
Adapt.adapt_storage(::Type{CuArray{T}}, xs::SparseVector) where {T} = CuSparseVector{T}(xs)
Adapt.adapt_storage(::Type{CuArray{T}}, xs::SparseMatrixCSC) where {T} = CuSparseMatrixCSC{T}(xs)

Adapt.adapt_storage(::CUDA.CuArrayAdaptor, xs::AbstractSparseArray) =
  adapt(CuArray, xs)
Adapt.adapt_storage(::CUDA.CuArrayAdaptor, xs::AbstractSparseArray{<:AbstractFloat}) =
  adapt(CuArray{Float32}, xs)

Adapt.adapt_storage(::Type{Array}, xs::CuSparseVector) = SparseVector(xs)
Adapt.adapt_storage(::Type{Array}, xs::CuSparseMatrixCSC) = SparseMatrixCSC(xs)

## copying between sparse GPU arrays

function Base.copyto!(dst::CuSparseVector, src::CuSparseVector)
    if dst.dims != src.dims
        throw(ArgumentError("Inconsistent Sparse Vector size"))
    end
    copyto!(nonzeroinds(dst), nonzeroinds(src))
    copyto!(nonzeros(dst), nonzeros(src))
    dst.nnz = nnz(src)
    dst
end

function Base.copyto!(dst::CuSparseMatrixCSC, src::CuSparseMatrixCSC)
    if dst.dims != src.dims
        throw(ArgumentError("Inconsistent Sparse Matrix size"))
    end
    copyto!(dst.colPtr, src.colPtr)
    copyto!(rowvals(dst), rowvals(src))
    copyto!(nonzeros(dst), nonzeros(src))
    dst.nnz = nnz(src)
    dst
end

function Base.copyto!(dst::CuSparseMatrixCSR, src::CuSparseMatrixCSR)
    if dst.dims != src.dims
        throw(ArgumentError("Inconsistent Sparse Matrix size"))
    end
    copyto!(dst.rowPtr, src.rowPtr)
    copyto!(dst.colVal, src.colVal)
    copyto!(nonzeros(dst), nonzeros(src))
    dst.nnz = nnz(src)
    dst
end

function Base.copyto!(dst::CuSparseMatrixBSR, src::CuSparseMatrixBSR)
    if dst.dims != src.dims
        throw(ArgumentError("Inconsistent Sparse Matrix size"))
    end
    copyto!(dst.rowPtr, src.rowPtr)
    copyto!(dst.colVal, src.colVal)
    copyto!(nonzeros(dst), nonzeros(src))
    dst.dir = src.dir
    dst.nnz = nnz(src)
    dst
end

function Base.copyto!(dst::CuSparseMatrixCOO, src::CuSparseMatrixCOO)
    if dst.dims != src.dims
        throw(ArgumentError("Inconsistent Sparse Matrix size"))
    end
    copyto!(dst.rowInd, src.rowInd)
    copyto!(dst.colInd, src.colInd)
    copyto!(nonzeros(dst), nonzeros(src))
    dst.nnz = nnz(src)
    dst
end

Base.copy(Vec::CuSparseVector) = copyto!(similar(Vec), Vec)
Base.copy(Mat::CuSparseMatrixCSC) = copyto!(similar(Mat), Mat)
Base.copy(Mat::CuSparseMatrixCSR) = copyto!(similar(Mat), Mat)
Base.copy(Mat::CuSparseMatrixBSR) = copyto!(similar(Mat), Mat)
Base.copy(Mat::CuSparseMatrixCOO) = copyto!(similar(Mat), Mat)

# input/output

for (gpu, cpu) in [CuSparseVector => SparseVector,
                   CuSparseMatrixCSC => SparseMatrixCSC,
                   CuSparseMatrixCSR => SparseMatrixCSC,
                   CuSparseMatrixBSR => SparseMatrixCSC,
                   CuSparseMatrixCOO => SparseMatrixCSC]
    @eval Base.show(io::IOContext, x::$gpu) =
        show(io, $cpu(x))

    @eval function Base.show(io::IO, mime::MIME"text/plain", S::$gpu)
        xnnz = nnz(S)
        m, n = size(S)
        print(io, m, "×", n, " ", typeof(S), " with ", xnnz, " stored ",
                  xnnz == 1 ? "entry" : "entries")
        if !(m == 0 || n == 0)
            println(io, ":")
            io = IOContext(io, :typeinfo => eltype(S))
            if ndims(S) == 1
                show(io, $cpu(S))
            else
                # so that we get the nice Braille pattern
                Base.print_array(io, $cpu(S))
            end
        end
    end
end

# interop with device arrays

function Adapt.adapt_structure(to::CUDA.Adaptor, x::CuSparseVector)
    return CuSparseDeviceVector(
        adapt(to, x.iPtr),
        adapt(to, x.nzVal),
        x.dims, x.nnz
    )
end

function Adapt.adapt_structure(to::CUDA.Adaptor, x::CuSparseMatrixCSR)
    return CuSparseDeviceMatrixCSR(
        adapt(to, x.rowPtr),
        adapt(to, x.colVal),
        adapt(to, x.nzVal),
        x.dims, x.nnz
    )
end

function Adapt.adapt_structure(to::CUDA.Adaptor, x::CuSparseMatrixCSC)
    return CuSparseDeviceMatrixCSC(
        adapt(to, x.colPtr),
        adapt(to, x.rowVal),
        adapt(to, x.nzVal),
        x.dims, x.nnz
    )
end

function Adapt.adapt_structure(to::CUDA.Adaptor, x::CuSparseMatrixBSR)
    return CuSparseDeviceMatrixBSR(
        adapt(to, x.rowPtr),
        adapt(to, x.colVal),
        adapt(to, x.nzVal),
        x.dims, x.blockDim,
        x.dir, x.nnz
    )
end

function Adapt.adapt_structure(to::CUDA.Adaptor, x::CuSparseMatrixCOO)
    return CuSparseDeviceMatrixCOO(
        adapt(to, x.rowInd),
        adapt(to, x.colInd),
        adapt(to, x.nzVal),
        x.dims, x.nnz
    )
end


text:
function laplace_matrix(::Type{T}, n, dims) where T
    D = second_order_central_diff(T, n);
    A = copy(D);

    for idx = 2 : dims
        A = kron(A, speye(n)) + kron(speye(size(A, 1)), D);
    end

    A
end

second_order_central_diff(::Type{T}, dim) where {T} = convert(SparseMatrixCSC{T, Int}, SymTridiagonal(fill(2 * one(T), dim), fill(-one(T), dim - 1)))


text:
using .ForwardDiff: gradient!
using .DiffResults: GradientResult, value, gradient

struct MALA{D} <: MHSampler
    proposal::D
end

# Create a RandomWalkProposal if we weren't given one already.
MALA(d) = MALA(RandomWalkProposal(d))

# If we were given a RandomWalkProposal, just use that instead.
MALA(d::RandomWalkProposal) = MALA{typeof(d)}(d)

struct GradientTransition{T<:Union{Vector, Real, NamedTuple}, L<:Real, G<:Union{Vector, Real, NamedTuple}} <: AbstractTransition
    params::T
    lp::L
    gradient::G
end

transition(::MALA, model, params) = GradientTransition(model, params)

# Store the new draw, its log density and its gradient
GradientTransition(model::DensityModel, params) = GradientTransition(params, logdensity_and_gradient(model, params)...)

propose(rng::Random.AbstractRNG, ::MALA, model) = error("please specify initial parameters")

function propose(
    rng::Random.AbstractRNG,
    spl::MALA{<:Proposal},
    model::DensityModel,
    params_prev::GradientTransition
)
    proposal = propose(rng, spl.proposal(params_prev.gradient), model, params_prev.params)
    return GradientTransition(model, proposal)
end

function q(
    spl::MALA{<:Proposal},
    t::GradientTransition,
    t_cond::GradientTransition
)
    return q(spl.proposal(-t_cond.gradient), t.params, t_cond.params)
end

"""
    logdensity_and_gradient(model::DensityModel, params)

Efficiently returns the value and gradient of the model
"""
function logdensity_and_gradient(model::DensityModel, params)
    res = GradientResult(params)
    gradient!(res, model.logdensity, params)
    return (value(res), gradient(res))
end

logdensity(model::DensityModel, t::GradientTransition) = t.lp


text:
function GlobalPressureTemperatureModel(arg0::jdouble, arg1::jdouble, arg2::Frame)
    return GlobalPressureTemperatureModel((jdouble, jdouble, Frame), arg0, arg1, arg2)
end

function GlobalPressureTemperatureModel(arg0::jdouble, arg1::jdouble, arg2::Frame, arg3::DataContext)
    return GlobalPressureTemperatureModel((jdouble, jdouble, Frame, DataContext), arg0, arg1, arg2, arg3)
end

function get_pressure(obj::GlobalPressureTemperatureModel)
    return jcall(obj, "getPressure", jdouble, ())
end

function get_temperature(obj::GlobalPressureTemperatureModel)
    return jcall(obj, "getTemperature", jdouble, ())
end

function weather_parameters(obj::GlobalPressureTemperatureModel, arg0::jdouble, arg1::AbsoluteDate)
    return jcall(obj, "weatherParameters", void, (jdouble, AbsoluteDate), arg0, arg1)
end


text:
function _transform!(d::IdDict, trans::Vector{Symbol}, gejac::GEJacobian)
    varvals = gejac.jacs.varvals
    for irfs in values(d)
        for v in trans
            irf = get(irfs, v, nothing)
            if irf !== nothing
                ss = varvals[v]
                irf .= 100.0.*irf./ss
            end
        end
    end
end

function _transform!(d::IdDict, trans::Bool, gejac::GEJacobian)
    if trans
        varvals = gejac.jacs.varvals
        for irfs in values(d)
            for (v, irf) in irfs
                ss = varvals[v]
                irf .= 100.0.*irf./ss
            end
        end
    end
end

function linirf(gejac::GEJacobian{TF}, dshocks::ValidPathInput, endovars=nothing;
        transform=false) where TF
    jacs = gejac.jacs
    dshocks isa Pair && (dshocks = (dshocks,))
    endovars isa Symbol && (endovars = (endovars,))
    endovars === nothing && (endovars = setdiff(jacs.vars, gejac.exovars, jacs.tars))
    isempty(endovars) && throw(ArgumentError("endovars cannot be empty"))
    nT = jacs.nT
    out = IdDict{Symbol,IdDict{Symbol,VecOrMat{TF}}}()
    for (exo, dZ) in dshocks
        size(dZ,1) == nT || throw(ArgumentError("the length of shock $exo is not $nT"))
        d = IdDict{Symbol,VecOrMat{TF}}()
        out[exo] = d
        for endo in endovars
            G = getG!(gejac, exo, endo)
            d[endo] = G * dZ
        end
    end
    _transform!(out, transform, gejac)
    return out
end

function linirf(jacs::TotalJacobian, dshocks::ValidPathInput, endovars=nothing;
        transform=false, keepH_U::Bool=false)
    dshocks isa Pair && (dshocks = (dshocks,))
    exovars = (exo for (exo, _) in dshocks)
    gejac = GEJacobian(jacs, exovars; keepH_U=keepH_U)
    irfs = linirf(gejac, dshocks isa Tuple ? dshocks[1] : dshocks, endovars; transform=transform)
    return irfs, gejac
end

function _transform!(d::IdDict, trans::Vector{Symbol}, jacs::TotalJacobian)
    varvals = jacs.varvals
    for v in trans
        irf = get(d, v, nothing)
        if irf !== nothing
            ss = varvals[v]
            irf .= 100.0.*irf./ss .- 100.0
        end
    end
end

function _transform!(d::IdDict, trans::Bool, jacs::TotalJacobian)
    if trans
        varvals = jacs.varvals
        for (v, irf) in d
            ss = varvals[v]
            irf .= 100.0.*irf./ss .- 100.0
        end
    end
end

function nlirf(tr::Transition{TF}, endovars=nothing; transform=false) where TF
    jacs = tr.jacs
    endovars isa Symbol && (endovars = (endovars,))
    endovars === nothing && (endovars = setdiff(jacs.vars, tr.exovars, jacs.tars))
    isempty(endovars) && throw(ArgumentError("endovars cannot be empty"))
    out = IdDict{Symbol,VecOrMat{TF}}()
    for endo in endovars
        path = get(tr.varpaths, endo, nothing)
        path === nothing && throw(ArgumentError("$endo is not an endogenous variable"))
        out[endo] = path[2]
    end
    _transform!(out, transform, jacs)
    return out
end

function nlirf(jacs::TotalJacobian, shocks::ValidPathInput, endovars=nothing;
        transform=false, initials=nothing, H_U=nothing, kwargs...)
    tr = Transition(jacs, shocks, initials; H_U=H_U)
    solve!(tr, kwargs...)
    return nlirf(tr, endovars; transform=transform), tr
end


text:
module TrackerCompat

using ..Tracker: Tracker,
                 TrackedReal,
                 TrackedVector,
                 TrackedMatrix,
                 TrackedArray,
                 TrackedVecOrMat,
                 @grad,
                 track,
                 data,
                 param

import ..Bijectors
using ..Bijectors: Log, SimplexBijector, ADBijector,
    TrackerAD, Inverse, Stacked, Exp

import ChainRulesCore
import LogExpFunctions

using Compat: eachcol
using LinearAlgebra
using Distributions: LocationScale

Bijectors.maporbroadcast(f, x::TrackedArray...) = f.(x...)
function Bijectors.maporbroadcast(
    f,
    x1::TrackedArray{T, N},
    x::AbstractArray{<:TrackedReal}...,
) where {T, N}
    return f.(convert(Array{TrackedReal{T}, N}, x1), x...)
end

Bijectors._eps(::Type{<:TrackedReal{T}}) where {T} = Bijectors._eps(T)
function Base.minimum(d::LocationScale{<:TrackedReal})
    m = minimum(d.ρ)
    if isfinite(m)
        return d.μ + d.σ * m
    else
        return m
    end
end
function Base.maximum(d::LocationScale{<:TrackedReal})
    m = maximum(d.ρ)
    if isfinite(m)
        return d.μ + d.σ * m
    else
        return m
    end
end

# AD implementations
function Bijectors.jacobian(
    b::Union{<:ADBijector{<:TrackerAD}, Inverse{<:ADBijector{<:TrackerAD}}},
    x::Real
)
    return data(Tracker.gradient(b, x)[1])
end
function Bijectors.jacobian(
    b::Union{<:ADBijector{<:TrackerAD}, Inverse{<:ADBijector{<:TrackerAD}}},
    x::AbstractVector{<:Real}
)
    # We extract `data` so that we don't return a `Tracked` type
    return data(Tracker.jacobian(b, x))
end

# implementations for Shift bijector
function Bijectors._logabsdetjac_shift(a::TrackedReal, x::Real, ::Val{0})
    return tracker_shift_logabsdetjac(a, x, Val(0))
end
function Bijectors._logabsdetjac_shift(a::TrackedReal, x::AbstractVector{<:Real}, ::Val{0})
    return tracker_shift_logabsdetjac(a, x, Val(0))
end
function Bijectors._logabsdetjac_shift(
    a::Union{TrackedReal, TrackedVector{<:Real}},
    x::AbstractVector{<:Real},
    ::Val{1}
)
    return tracker_shift_logabsdetjac(a, x, Val(1))
end
function Bijectors._logabsdetjac_shift(
    a::Union{TrackedReal, TrackedVector{<:Real}},
    x::AbstractMatrix{<:Real},
    ::Val{1}
)
    return tracker_shift_logabsdetjac(a, x, Val(1))
end
function tracker_shift_logabsdetjac(a, x, ::Val{N}) where {N}
    return param(Bijectors._logabsdetjac_shift(data(a), data(x), Val(N)))
end

# Log bijector

@grad function Bijectors.logabsdetjac(b::Log{1}, x::AbstractVector)
    return -sum(log, data(x)), Δ -> (nothing, -Δ ./ data(x))
end
@grad function Bijectors.logabsdetjac(b::Log{1}, x::AbstractMatrix)
    return -vec(sum(log, data(x); dims = 1)), Δ -> (nothing, .- Δ' ./ data(x))
end
@grad function Bijectors.logabsdetjac(b::Log{2}, x::AbstractMatrix)
    return -sum(log, data(x)), Δ -> (nothing, -Δ ./ data(x))
end

# implementations for Scale bijector
# Adjoints for 0-dim and 1-dim `Scale` using `Real`
function Bijectors._logabsdetjac_scale(a::TrackedReal, x::Real, ::Val{0})
    return track(Bijectors._logabsdetjac_scale, a, data(x), Val(0))
end
@grad function Bijectors._logabsdetjac_scale(a::Real, x::Real, ::Val{0})
    return Bijectors._logabsdetjac_scale(data(a), data(x), Val(0)), Δ -> (inv(data(a)) .* Δ, nothing, nothing)
end
# Need to treat `AbstractVector` and `AbstractMatrix` separately due to ambiguity errors
function Bijectors._logabsdetjac_scale(a::TrackedReal, x::AbstractVector, ::Val{0})
    return track(Bijectors._logabsdetjac_scale, a, data(x), Val(0))
end
@grad function Bijectors._logabsdetjac_scale(a::Real, x::AbstractVector, ::Val{0})
    da = data(a)
    J = fill(inv.(da), length(x))
    return Bijectors._logabsdetjac_scale(da, data(x), Val(0)), Δ -> (transpose(J) * Δ, nothing, nothing)
end
function Bijectors._logabsdetjac_scale(a::TrackedReal, x::AbstractMatrix, ::Val{0})
    return track(Bijectors._logabsdetjac_scale, a, data(x), Val(0))
end
@grad function Bijectors._logabsdetjac_scale(a::Real, x::AbstractMatrix, ::Val{0})
    da = data(a)
    J = fill(size(x, 1) / da, size(x, 2))
    return Bijectors._logabsdetjac_scale(da, data(x), Val(0)), Δ -> (transpose(J) * Δ, nothing, nothing)
end
# adjoints for 1-dim and 2-dim `Scale` using `AbstractVector`
function Bijectors._logabsdetjac_scale(a::TrackedVector, x::AbstractVector, ::Val{1})
    return track(Bijectors._logabsdetjac_scale, a, data(x), Val(1))
end
@grad function Bijectors._logabsdetjac_scale(a::TrackedVector, x::AbstractVector, ::Val{1})
    # ∂ᵢ (∑ⱼ log|aⱼ|) = ∑ⱼ δᵢⱼ ∂ᵢ log|aⱼ|
    #                 = ∂ᵢ log |aᵢ|
    #                 = (1 / aᵢ) ∂ᵢ aᵢ
    #                 = (1 / aᵢ)
    da = data(a)
    J = inv.(da)
    return Bijectors._logabsdetjac_scale(da, data(x), Val(1)), Δ -> (J .* Δ, nothing, nothing)
end
function Bijectors._logabsdetjac_scale(a::TrackedVector, x::AbstractMatrix, ::Val{1})
    return track(Bijectors._logabsdetjac_scale, a, data(x), Val(1))
end
@grad function Bijectors._logabsdetjac_scale(a::TrackedVector, x::AbstractMatrix, ::Val{1})
    da = data(a)
    Jᵀ = repeat(inv.(da), 1, size(x, 2))
    return Bijectors._logabsdetjac_scale(da, data(x), Val(1)), Δ -> (Jᵀ * Δ, nothing, nothing)
end
# TODO: implement analytical gradient for scaling a vector using a matrix
# function _logabsdetjac_scale(a::TrackedMatrix, x::AbstractVector, ::Val{1})
#     track(_logabsdetjac_scale, a, data(x), Val{1})
# end
# @grad function _logabsdetjac_scale(a::TrackedMatrix, x::AbstractVector, ::Val{1})
#     throw
# end
# implementations for Stacked bijector
function Bijectors.logabsdetjac(b::Stacked, x::TrackedMatrix{<:Real})
    return map(eachcol(x)) do c
        Bijectors.logabsdetjac(b, c)
    end
end
# TODO: implement custom adjoint since we can exploit block-diagonal nature of `Stacked`
function (sb::Stacked)(x::TrackedMatrix{<:Real})
    return Bijectors.eachcolmaphcat(sb, x)
end
# Simplex adjoints
function Bijectors._simplex_bijector(X::TrackedVecOrMat, b::SimplexBijector{1})
    return track(Bijectors._simplex_bijector, X, b)
end
function Bijectors._simplex_inv_bijector(Y::TrackedVecOrMat, b::SimplexBijector{1})
    return track(Bijectors._simplex_inv_bijector, Y, b)
end
@grad function Bijectors._simplex_bijector(X::AbstractVector, b::SimplexBijector{1})
    Xd = data(X)
    return Bijectors._simplex_bijector(Xd, b), Δ -> (Bijectors.simplex_link_jacobian(Xd)' * Δ, nothing)
end
@grad function Bijectors._simplex_inv_bijector(Y::AbstractVector, b::SimplexBijector{1})
    Yd = data(Y)
    return Bijectors._simplex_inv_bijector(Yd, b), Δ -> (Bijectors.simplex_invlink_jacobian(Yd)' * Δ, nothing)
end

@grad function Bijectors._simplex_bijector(X::AbstractMatrix, b::SimplexBijector{1})
    Xd = data(X)
    return Bijectors._simplex_bijector(Xd, b), Δ -> begin
        Bijectors.maphcat(eachcol(Xd), eachcol(Δ)) do c1, c2
            Bijectors.simplex_link_jacobian(c1)' * c2
        end, nothing
    end
end
@grad function Bijectors._simplex_inv_bijector(Y::AbstractMatrix, b::SimplexBijector{1})
    Yd = data(Y)
    return Bijectors._simplex_inv_bijector(Yd, b), Δ -> begin
        Bijectors.maphcat(eachcol(Yd), eachcol(Δ)) do c1, c2
            Bijectors.simplex_invlink_jacobian(c1)' * c2
        end, nothing
    end
end

Bijectors.replace_diag(::typeof(log), X::TrackedMatrix) = track(Bijectors.replace_diag, log, X)
@grad function Bijectors.replace_diag(::typeof(log), X)
    Xd = data(X)
    f(i, j) = i == j ? log(Xd[i, j]) : Xd[i, j]
    out = f.(1:size(Xd, 1), (1:size(Xd, 2))')
    out, ∇ -> begin
        g(i, j) = i == j ? ∇[i, j]/Xd[i, j] : ∇[i, j]
        return (nothing, g.(1:size(Xd, 1), (1:size(Xd, 2))'))
    end
end

Bijectors.replace_diag(::typeof(exp), X::TrackedMatrix) = track(Bijectors.replace_diag, exp, X)
@grad function Bijectors.replace_diag(::typeof(exp), X)
    Xd = data(X)
    f(i, j) = ifelse(i == j, exp(Xd[i, j]), Xd[i, j])
    out = f.(1:size(Xd, 1), (1:size(Xd, 2))')
    out, ∇ -> begin
        g(i, j) = ifelse(i == j, ∇[i, j]*exp(Xd[i, j]), ∇[i, j])
        return (nothing, g.(1:size(Xd, 1), (1:size(Xd, 2))'))
    end
end

Bijectors.logabsdetjac(b::SimplexBijector{1}, x::TrackedVecOrMat) = track(Bijectors.logabsdetjac, b, x)
@grad function Bijectors.logabsdetjac(b::SimplexBijector{1}, x::AbstractVector)
    xd = data(x)
    return Bijectors.logabsdetjac(b, xd), Δ -> begin
        (nothing, Bijectors.simplex_logabsdetjac_gradient(xd) * Δ)
    end
end
@grad function Bijectors.logabsdetjac(b::SimplexBijector{1}, x::AbstractMatrix)
    xd = data(x)
    return Bijectors.logabsdetjac(b, xd), Δ -> begin
        (nothing, Bijectors.maphcat(eachcol(xd), Δ) do c, g
            Bijectors.simplex_logabsdetjac_gradient(c) * g
        end)
    end
end

for header in [
    (:(α_::TrackedReal), :β, :z_0, :(z::AbstractVector)),
    (:α_, :(β::TrackedReal), :z_0, :(z::AbstractVector)),
    (:α_, :β, :(z_0::TrackedVector), :(z::AbstractVector)),
    (:α_, :β, :z_0, :(z::TrackedVector)),
    (:(α_::TrackedReal), :(β::TrackedReal), :z_0, :(z::AbstractVector)),
    (:(α_::TrackedReal), :β, :(z_0::TrackedVector), :(z::AbstractVector)),
    (:(α_::TrackedReal), :β, :z_0, :(z::TrackedVecOrMat)),
    (:(α_::TrackedReal), :(β::TrackedReal), :(z_0::TrackedVector), :(z::AbstractVector)),
    (:(α_::TrackedReal), :(β::TrackedReal), :z_0, :(z::TrackedVector)),
    (:(α_::TrackedReal), :β, :(z_0::TrackedVector), :(z::TrackedVector)),
    (:α_, :(β::TrackedReal), :(z_0::TrackedVector), :(z::TrackedVector)),
    (:(α_::TrackedReal), :(β::TrackedReal), :(z_0::TrackedVector), :(z::TrackedVector)),
]
    @eval begin
        function Bijectors._radial_transform($(header...))
            α = LogExpFunctions.log1pexp(α_)            # from A.2
            β_hat = -α + LogExpFunctions.log1pexp(β)    # from A.2
            if β_hat isa TrackedReal
                TV = vectorof(typeof(β_hat))
                T = vectorof(typeof(β_hat))
            elseif z_0 isa TrackedVector
                TV = typeof(z_0)
                T = typeof(z_0)
            else
                T = TV = typeof(z)
            end
            Tr = promote_type(eltype(z), eltype(z_0))
            r::Tr = norm((z .- z_0)::TV)
            transformed::T = z .+ β_hat ./ (α .+ r') .* (z .- z_0)   # from eq(14)
            return (transformed = transformed, α = α, β_hat = β_hat, r = r)
        end
    end
end

for header in [
    (:(α_::TrackedReal), :β, :z_0, :(z::AbstractMatrix)),
    (:α_, :(β::TrackedReal), :z_0, :(z::AbstractMatrix)),
    (:α_, :β, :(z_0::TrackedVector), :(z::AbstractMatrix)),
    (:α_, :β, :z_0, :(z::TrackedMatrix)),
    (:(α_::TrackedReal), :(β::TrackedReal), :z_0, :(z::AbstractMatrix)),
    (:(α_::TrackedReal), :β, :(z_0::TrackedVector), :(z::AbstractMatrix)),
    (:(α_::TrackedReal), :β, :z_0, :(z::TrackedMatrix)),
    (:(α_::TrackedReal), :(β::TrackedReal), :(z_0::TrackedVector), :(z::AbstractMatrix)),
    (:(α_::TrackedReal), :(β::TrackedReal), :z_0, :(z::TrackedMatrix)),
    (:(α_::TrackedReal), :β, :(z_0::TrackedVector), :(z::TrackedMatrix)),
    (:α_, :(β::TrackedReal), :(z_0::TrackedVector), :(z::TrackedMatrix)),
    (:(α_::TrackedReal), :(β::TrackedReal), :(z_0::TrackedVector), :(z::TrackedMatrix)),
]
    @eval begin
        function Bijectors._radial_transform($(header...))
            α = LogExpFunctions.log1pexp(α_)            # from A.2
            β_hat = -α + LogExpFunctions.log1pexp(β)    # from A.2
            if β_hat isa TrackedReal
                TV = vectorof(typeof(β_hat))
                T = matrixof(TV)
            elseif z_0 isa TrackedVector
                TV = typeof(z_0)
                T = matrixof(TV)
            else
                T = typeof(z)
                TV = vectorof(T)
            end
            r::TV = eachcolnorm(z .- z_0)
            transformed::T = z .+ β_hat ./ (α .+ r') .* (z .- z_0)   # from eq(14)
            return (transformed = transformed, α = α, β_hat = β_hat, r = r)
        end
    end
end
eachcolnorm(X) = map(norm, eachcol(X))
eachcolnorm(X::TrackedMatrix) = track(eachcolnorm, X)
@grad function eachcolnorm(X)
    Xd = data(X)
    y = map(norm, eachcol(Xd))
    y, Δ -> begin
        (Xd .* (Δ ./ y)',)
    end
end

function matrixof(::Type{<:TrackedArray{T,1,Vector{T}}}) where {T<:Real}
    return TrackedArray{T,2,Matrix{T}}
end
function matrixof(::Type{TrackedReal{T}}) where {T<:Real}
    return TrackedArray{T,2,Matrix{T}}
end
function vectorof(::Type{<:TrackedArray{T,2,Matrix{T}}}) where {T<:Real}
    return TrackedArray{T,1,Vector{T}}
end
function vectorof(::Type{TrackedReal{T}}) where {T<:Real}
    return TrackedArray{T,1,Vector{T}}
end

(b::Exp{0})(x::TrackedVector) = exp.(x)::vectorof(float(eltype(x)))
(b::Exp{1})(x::TrackedVector) = exp.(x)::vectorof(float(eltype(x)))
(b::Exp{1})(x::TrackedMatrix) = exp.(x)::matrixof(float(eltype(x)))
(b::Exp{2})(x::TrackedMatrix) = exp.(x)::matrixof(float(eltype(x)))

(b::Log{0})(x::TrackedVector) = log.(x)::vectorof(float(eltype(x)))
(b::Log{1})(x::TrackedVector) = log.(x)::vectorof(float(eltype(x)))
(b::Log{1})(x::TrackedMatrix) = log.(x)::matrixof(float(eltype(x)))
(b::Log{2})(x::TrackedMatrix) = log.(x)::matrixof(float(eltype(x)))

Bijectors.logabsdetjac(b::Log{0}, x::TrackedVector) = .-log.(x)::vectorof(float(eltype(x)))
Bijectors.logabsdetjac(b::Log{1}, x::TrackedMatrix) = - vec(sum(log.(x); dims = 1))

Bijectors.getpd(X::TrackedMatrix) = track(Bijectors.getpd, X)
@grad function Bijectors.getpd(X::AbstractMatrix)
    Xd = data(X)
    return Bijectors.LowerTriangular(Xd) * Bijectors.LowerTriangular(Xd)', Δ -> begin
        Xl = Bijectors.LowerTriangular(Xd)
        return (Bijectors.LowerTriangular(Δ' * Xl + Δ * Xl),)
    end
end

Bijectors.lower(A::TrackedMatrix) = track(Bijectors.lower, A)
@grad function Bijectors.lower(A::AbstractMatrix)
    Ad = data(A)
    return Bijectors.lower(Ad), Δ -> (Bijectors.lower(Δ),)
end

Bijectors._inv_link_chol_lkj(y::TrackedMatrix) = track(Bijectors._inv_link_chol_lkj, y)
@grad function Bijectors._inv_link_chol_lkj(y_tracked)
    y = data(y_tracked)

    K = LinearAlgebra.checksquare(y)

    w = similar(y)

    z_mat = similar(y) # cache for adjoint
    tmp_mat = similar(y)

    @inbounds for j in 1:K
        w[1, j] = 1
        for i in 2:j
            z = tanh(y[i-1, j])
            tmp = w[i-1, j]

            z_mat[i, j] = z
            tmp_mat[i, j] = tmp

            w[i-1, j] = z * tmp
            w[i, j] = tmp * sqrt(1 - z^2)
        end
        for i in (j+1):K
            w[i, j] = 0
        end
    end

    function pullback_inv_link_chol_lkj(Δw)
        LinearAlgebra.checksquare(Δw)

        Δy = zero(y)

        @inbounds for j in 1:K
            Δtmp = Δw[j,j]
            for i in j:-1:2
                Δz = Δw[i-1, j] * tmp_mat[i, j] - Δtmp * tmp_mat[i, j] / sqrt(1 - z_mat[i, j]^2) * z_mat[i, j]
                Δy[i-1, j] = Δz / cosh(y[i-1, j])^2
                Δtmp = Δw[i-1, j] * z_mat[i, j] + Δtmp * sqrt(1 - z_mat[i, j]^2)
            end
        end

        return (Δy,)
    end

    return w, pullback_inv_link_chol_lkj
end

Bijectors._link_chol_lkj(w::TrackedMatrix) = track(Bijectors._link_chol_lkj, w)
@grad function Bijectors._link_chol_lkj(w_tracked)
    w = data(w_tracked)

    K = LinearAlgebra.checksquare(w)

    z = similar(w)

    @inbounds z[1, 1] = 0

    tmp_mat = similar(w) # cache for pullback.

    @inbounds for j=2:K
        z[1, j] = atanh(w[1, j])
        tmp = sqrt(1 - w[1, j]^2)
        tmp_mat[1, j] = tmp
        for i in 2:(j - 1)
            p = w[i, j] / tmp
            tmp *= sqrt(1 - p^2)
            tmp_mat[i, j] = tmp
            z[i, j] = atanh(p)
        end
        z[j, j] = 0
    end

    function pullback_link_chol_lkj(Δz)
        LinearAlgebra.checksquare(Δz)

        Δw = similar(w)

        @inbounds Δw[1,1] = zero(eltype(Δz))

        @inbounds for j=2:K
            Δw[j, j] = 0
            Δtmp = zero(eltype(Δz)) # Δtmp_mat[j-1,j]
            for i in (j-1):-1:2
                p = w[i, j] / tmp_mat[i-1, j]
                ftmp = sqrt(1 - p^2)
                d_ftmp_p = -p / ftmp
                d_p_tmp = -w[i,j] / tmp_mat[i-1, j]^2

                Δp = Δz[i,j] / (1-p^2) + Δtmp * tmp_mat[i-1, j] * d_ftmp_p
                Δw[i, j] = Δp / tmp_mat[i-1, j]
                Δtmp = Δp * d_p_tmp + Δtmp * ftmp # update to "previous" Δtmp
            end
            Δw[1, j] = Δz[1, j] / (1-w[1,j]^2) - Δtmp / sqrt(1 - w[1,j]^2) * w[1,j]
        end

        return (Δw,)
    end

    return z, pullback_link_chol_lkj
end

function Bijectors.find_alpha(wt_y::T, wt_u_hat::T, b::T) where {T<:TrackedReal}
    return track(Bijectors.find_alpha, wt_y, wt_u_hat, b)
end
@grad function Bijectors.find_alpha(wt_y::TrackedReal, wt_u_hat::TrackedReal, b::TrackedReal)
    α = Bijectors.find_alpha(data(wt_y), data(wt_u_hat), data(b))

    ∂wt_y = inv(1 + wt_u_hat * sech(α + b)^2)
    ∂wt_u_hat = - tanh(α + b) * ∂wt_y
    ∂b = ∂wt_y - 1
    find_alpha_pullback(Δ::Real) = (Δ * ∂wt_y, Δ * ∂wt_u_hat, Δ * ∂b)

    return α, find_alpha_pullback
end

# `OrderedBijector`
Bijectors._transform_ordered(y::Union{TrackedVector,TrackedMatrix}) = track(Bijectors._transform_ordered, y)
@grad function Bijectors._transform_ordered(y::AbstractVecOrMat)
    x, dx = ChainRulesCore.rrule(Bijectors._transform_ordered, data(y))
    return x, (wrap_chainrules_output ∘ Base.tail ∘ dx)
end

function Bijectors._transform_inverse_ordered(x::Union{TrackedVector,TrackedMatrix})
    return track(Bijectors._transform_inverse_ordered, x)
end
@grad function Bijectors._transform_inverse_ordered(x::AbstractVecOrMat)
    y, dy = ChainRulesCore.rrule(Bijectors._transform_inverse_ordered, data(x))
    return y, (wrap_chainrules_output ∘ Base.tail ∘ dy)
end

# NOTE: Probably doesn't work in complete generality.
wrap_chainrules_output(x) = x
wrap_chainrules_output(x::ChainRulesCore.AbstractZero) = nothing
wrap_chainrules_output(x::Tuple) = map(wrap_chainrules_output, x)

end


text:
"""
    Combined(n) = [ x-> x^k for k = 0:n-1 ]
"""
struct Combined{B} <: AbstractBasis
    bases::B
end
Combined(bases...) = Combined{typeof(bases)}(bases)

Base.length(B::Combined) = mapreduce(length,+,B.bases)

Base.:|(B::Combined,x) = Iterators.flatten(IterTools.imap(Bi->Bi|x,B.bases))

function Base.getindex(B::Combined,i)
    for Bi in B.bases
        i <= length(Bi) && return Bi[i]
        i -= length(Bi)
    end
    throw(BoundsError(B,i))
end

function evaltransform(B::Combined, x, c)
    T = typeof(one(eltype(c)) * one(eltype(B,x[1])))
    f = zeros(T, length(x), Base.tail(size(c))...)
    i = 0
    for Bi in B.bases
        fidx = ( 1:length(x), Base.OneTo.(Base.tail(size(c)))... )
        cidx = ( i.+(1:length(Bi)), Base.OneTo.(Base.tail(size(c)))... )
        if all(length.(cidx) .!= 0)
            @views f[fidx...] .+= evaltransform(Bi,x,c[cidx...])
            i += length(Bi)
        end
    end
    return f
end


text:
## ExoplanetsSysSim/src/limb_darkening.jl
## (c) 2018 Eric B. Ford

@compat abstract type LimbDarkeningParamAbstract end

struct LimbDarkeningParamLinear <: LimbDarkeningParamAbstract
  coeff::Tuple{Float64}
  # TODO SCI DETAIL: Replace with sensible limits on LD params
  LimbDarkeningParamLinear(_u1::Real ) = (!( (-2.5<=_u1<=2.5) )) ? error(string("Invalid quadratic limb darkening parameters: ",_u1," & ", _u2)) : new((_u1,))
end

function depth_at_midpoint(radius_ratio::Float64, ld::LimbDarkeningParamLinear)
    c0 = 1-sum(ld.coeff)
    omega = c0/4+ld.coeff[1]/6
    ksq = 1-radius_ratio^2
    tmp0 = c0/4*ksq
    tmp2 = ld.coeff[1]/6*ksq^(3//2)
    return  1-(tmp0+tmp2)/omega
end

function ratio_from_depth(depth::Float64, ld::LimbDarkeningParamLinear)
    c0 = 1-sum(ld.coeff)
    omega = c0/4+ld.coeff[1]/6
    xn = 1-depth
    y0 = (1-depth)*omega
    for i in 1:20
        tmp0 = c0/4*xn
        tmp2 = ld.coeff[1]/6*xn^(3//2)
        fn = (tmp0+tmp2)-y0

        tmp_0 = c0
        tmp_2 = ld.coeff[1]*xn^(1//2)
        f_n = 0.25*(tmp_0+tmp_2)

        xn -= (fn/f_n)
        if (fn/f_n)/xn < 1e-8
           break
        end
    end
    return sqrt(1-xn)
end

struct LimbDarkeningParamQuadratic <: LimbDarkeningParamAbstract
  coeff::Tuple{Float64,Float64}
  # TODO SCI DETAIL: Replace with sensible limits on LD params
  LimbDarkeningParamQuadratic(_u1::Real, _u2::Real ) = (!( (-2.5<=_u1<=2.5) && (-2.5<=_u2<=2.5) )) ? error(string("Invalid quadratic limb darkening parameters: ",_u1," & ", _u2)) : new((_u1,_u2))
end

function depth_at_midpoint(radius_ratio::Float64, ld::LimbDarkeningParamQuadratic)
    c0 = 1-sum(ld.coeff)
    omega = c0/4+(ld.coeff[1]+2*ld.coeff[2])/6-ld.coeff[2]/8
    ksq = 1-radius_ratio^2
    tmp0 = c0/4*ksq
    tmp2 = (ld.coeff[1]+2*ld.coeff[2])/6*ksq^(3//2)
    tmp4 = -ld.coeff[2]/8*ksq^2
    return  1-(tmp0+tmp2+tmp4)/omega
end

function ratio_from_depth(depth::Float64, ld::LimbDarkeningParamQuadratic)
    c0 = 1-sum(ld.coeff)
    omega = c0/4+(ld.coeff[1]+2*ld.coeff[2])/6-ld.coeff[2]/8
    xn = 1-depth
    y0 = (1-depth)*omega
    for i in 1:20
        tmp0 = c0/4*xn
        tmp2 = (ld.coeff[1]+2*ld.coeff[2])/6*xn^(3//2)
        tmp4 = -ld.coeff[2]/8*xn^2
        fn = (tmp0+tmp2+tmp4)-y0

        tmp_0 = c0
        tmp_2 = (ld.coeff[1]+2*ld.coeff[2])*xn^(1//2)
        tmp_4 = -ld.coeff[2]*xn
        f_n = 0.25*(tmp_0+tmp_2+tmp_4)

        xn -= (fn/f_n)
        if (fn/f_n)/xn < 1e-8
           break
        end
    end
    return sqrt(1-xn)
end

struct LimbDarkeningParam4thOrder <: LimbDarkeningParamAbstract
  coeff::Tuple{Float64,Float64,Float64,Float64}
  # TODO SCI DETAIL: Replace with sensible limits on LD params
  LimbDarkeningParam4thOrder(_c1::Real, _c2::Real, _c3::Real, _c4::Real ) = (!( (-2.5<=_c1<=2.5) && (-2.5<=_c2<=2.5)  && (-2.5<=_c3<=2.5) && (-2.5<=_c4<=2.5) )) ? error(string("Invalid limb darkening parameters: ",_c1," , ", _c2, ", ",_c3," , ",_c4)) : new((_c1,_c2,_c3,_c4))
end

function depth_at_midpoint(radius_ratio::Float64, ld::LimbDarkeningParam4thOrder)
    c0 = 1-sum(ld.coeff)
    omega = c0/4+sum(ld.coeff./(5:8))
    ksq = max(0.0, 1-radius_ratio^2)
    # likely false positive? not sure why this is happening
    tmp0 = c0/4*ksq
    tmp1 = ld.coeff[1]/5*ksq^(5//4)
    tmp2 = ld.coeff[2]/6*ksq^(3//2)
    tmp3 = ld.coeff[3]/7*ksq^(7//4)
    tmp4 = ld.coeff[4]/8*ksq^2
    return 1-(tmp0+tmp1+tmp2+tmp3+tmp4)/omega
end

# Cite Mandel & Agol 2002
# See https://arxiv.org/pdf/astro-ph/0210099.pdf and https://faculty.washington.edu/agol/mandel_agol_errata.pdf
# https://github.com/ericagol/ExoJulia/blob/master/ExoJulia/Transit/occultsmall.jl
function depth_small_planet(z::T,p::T,ld::LimbDarkeningParam4thOrder) where T<:Real
  @assert p<=0.2
  @assert sum(c)<=1
  c = ld.coeff

  one_over_Omega_times_4 = 1/(1-c[1]/5-c[2]/3-3*c[3]/7-c[4]/2)
  if z == zero(z)
    F = p*p*one_over_Omega_times_4
  elseif z <= 1-p
    sig1=sqrt(sqrt(1-(z-p)^2))
    sig2=sqrt(sqrt(1-(z+p)^2))
    ave_surface_intensity=(1-c[1]*(1+(sig2^5-sig1^5)/(5*p*z))-c[2]*(1+(sig2^6-sig1^6)/(6*p*z))-c[3]*(1+(sig2^7-sig1^7)/(7*p*z))-c[4]*(p^2+z^2))
    F = p*p*ave_surface_intensity*one_over_Omega_times_4
  elseif z<1+p
    x=1-(z-p)^2
    tmp=1-c[1]*(1-0.8*x^(1//4))-c[2]*(1-2//3*sqrt(x))-c[3]*(1-4//7*x^(3//4))-c[4]*(1-0.5*x)
    F = tmp*(p^2*acos((z-1)/p)-(z-1)*sqrt(p^2-(z-1)^2))*one_over_Omega_times_4/pi
  else
    F = zero(T)
  end
  return 1-F
end

function ratio_from_depth(depth::Float64, ld::LimbDarkeningParam4thOrder)
    c0 = 1-sum(ld.coeff)
    omega = c0/4+sum(ld.coeff./(5:8))
    xn = 1-depth
    y0 = (1-depth)*omega
    for i in 1:20
        tmp0 = c0/4*xn
        tmp1 = ld.coeff[1]/5*xn^(5//4)
        tmp2 = ld.coeff[2]/6*xn^(3//2)
        tmp3 = ld.coeff[3]/7*xn^(7//4)
        tmp4 = ld.coeff[4]/8*xn^2
        fn = (tmp0+tmp1+tmp2+tmp3+tmp4)-y0

        tmp_0 = c0
        tmp_1 = ld.coeff[1]*xn^(1//4)
        tmp_2 = ld.coeff[2]*xn^(1//2)
        tmp_3 = ld.coeff[3]*xn^(3//4)
        tmp_4 = ld.coeff[4]*xn
        f_n = 0.25*(tmp_0+tmp_1+tmp_2+tmp_3+tmp_4)

        xn -= (fn/f_n)
        if (fn/f_n)/xn < 1e-8
           break
        end
    end
    return sqrt(1-xn)
end


text:
# ---
# title: Block matching
# author: Johnny Chen
# date: 2020-10-15
# ---

# Block matching algorithm is one key stage in patch-based image processing algorithm. This example
# shows how block matching algorithm can be implemented efficiently and easily with the help of
# LazyDistances.jl.

using Images
using TestImages
using LazyDistances

img = imresize(float.(testimage("cameraman")), (64, 64))

# For given patch radius `rₚ`, we can get a patch at each pixel `p` in the image.

p = CartesianIndex(20, 30)
rₚ = CartesianIndex(3, 3) # patch size (7, 7)
patch_p = img[p-rₚ:p+rₚ]

# For each pixel `p`, what block matching does is to find a group of patches that are similar to
# `patch_p`.

function block_matching(f, img, p, rₚ; num_patches=10)
    R = CartesianIndices(img)
    ## to simplify the implementation and not consider border case
    candidates = first(R)+rₚ:last(R)-rₚ

    ## Measure the similarity of two patches with `f`. Typically, `f` is `sqeuclidean`
    patch_p = @view img[p-rₚ:p+rₚ]
    dist = map(candidates) do q
        patch_q = @view img[q-rₚ:q+rₚ]
        f(patch_p, patch_q)
    end |> vec

    ## sort from small(the most similar) to large(the least similar)
    ## TODO: use partialsortperm! to reduce allocation
    order = partialsortperm(dist, 1:num_patches) # 34.394 μs (4 allocations: 26.44 KiB)
    qs = @view candidates[order]
end

# Here we get a group of patches that are similar to `patch_p` in the sense of F-norm.

matched_patches = block_matching(SqEuclidean(), img, p, rₚ) # 212.266 μs (7 allocations: 39.83 KiB)
mosaicview(map(q->img[q-rₚ:q+rₚ], matched_patches); npad=2, nrow=2)

# This implementation looks good, but there's one key challenge here, that we need to loop over the
# whole image and do block matching for each pixel. This means we are doing a massive computation
# and there are a lot of unnecessary redundant calculation if we are doing a simple loop.
#
# For example, this is what a naive implementation of patch-based algorithm might looks like

function patched_mean(img, rₚ; num_patches=10)
    out = fill(zero(eltype(img)), axes(img))

    R = CartesianIndices(img)
    for p in first(R)+rₚ:last(R)-rₚ
        matched_patches = block_matching(SqEuclidean(), img, p, rₚ; num_patches=num_patches)
        out[p] = mapreduce(q->img[q], +, matched_patches)/length(matched_patches)
    end
    out
end

## @btime patched_mean($img, $rₚ) # 799.883 ms (23549 allocations: 130.86 MiB)
patched_mean(img, rₚ)

# What's wrong here? There are a lot of repeated calculation in `f(patch_q, patch_q)`. For example,
# when `f = SqEuclidean()`
#
# ```math
# d = \sum_{p, q} (I[p] - I[q])^2
# ```
#
# Please do note that `(I[p] -  I[q])^2` is repeatedly calculated in the whole iteration. One single
# calculation of such is not expensive, but it is a siginficant amount of work when it's in a nested
# for-loops.
#
# Among all the existing MATLAB and Python implementations, there are two things done to work around
# this key challenge. The first workaround is to pre-calculate the pixelwise distances and then
# reuse it in the inner loop. It works quite well but it also introduces another key challenge, that
# the memeory will complain for large image size; to store the result for image with size `(256,
# 256)` we need an array of size `(256, 256, 256, 256)`, which is 32 GB memory and
# unaccptable. The second workaround, is then to specify a search window, that to only search
# similar patches in a larger neighborhood of pixel `p`. For search window size `(17, 17)`, the
# memeory needed to store the result is 0.14 GB, which is more acceptable.
#
# We are not going to explore how search window and pre-calculation are implemented because this is
# quite a dirty work; you almost have to write the whole implementation in a long long for-loop.
# Instead, we are going to see how this can be efficiently and easily implemented with
# LazyDistances.jl.

# A [`PairwiseDistance`](@ref) is a lazy array that mimics the output of pairwise distance.

eval_op(x, y) = abs2(x - y)
pointwise_dist = PairwiseDistance(eval_op, (img, img)); # 1.982 ns (0 allocations: 0 bytes)

## pointwise_dist[I, J] is defined as f(img[I], img[J])
pq1 = pointwise_dist[CartesianIndex(1, 1), CartesianIndex(2, 2)]
pq2 = eval_op(img[CartesianIndex(1, 1)], img[CartesianIndex(2, 2)])
pq1 == pq2

# Generating this array does not doing any actual computation; the computation does not happen until
# you ask for the data. We can also build a patchwise distance with this.

R = CartesianIndices(img)
valid_R = first(R)+rₚ:last(R)-rₚ

## For simplicity, we didn't deal with boundary condition here, so it will error
## when we index with `patchwise_dist[1, 1, 1, 1]`.
patchwise_dist = let rₚ = rₚ, img = img
    PairwiseDistance(SqEuclidean(), (img, img)) do i
        i-rₚ:i+rₚ
    end
end; # 4.903 ns (0 allocations: 0 bytes)

p = CartesianIndex(4, 4)
q = CartesianIndex(5, 5)
## @btime getindex($patchwise_dist, $p, $q) # 53.874 ns (0 allocations: 0 bytes)
patchwise_dist[p, q] == sqeuclidean(img[p-rₚ:p+rₚ], img[q-rₚ:q+rₚ])

# This way we have generated the patchwise distances, although the actual computation doesn't happen
# until we need it.

function patched_mean_lazy(img, rₚ; num_patches=10)
    out = fill(zero(eltype(img)), axes(img))

    patchwise_dist = PairwiseDistance(SqEuclidean(), (img, img)) do i
        i-rₚ:i+rₚ
    end

    R = CartesianIndices(img)
    R0 = first(R)+rₚ:last(R)-rₚ
    for p in R0
        dist = vec(patchwise_dist[p, R0])
        matched_patches = R0[partialsortperm(dist, 1:num_patches)]
        out[p] = mapreduce(q->img[q], +, matched_patches)/length(matched_patches)
    end
    out
end

## @btime patched_mean_lazy($img, $rₚ);
##  643.413 ms (37015 allocations: 131.99 MiB)
## @btime patched_mean($img, $rₚ);
##  824.653 ms (23549 allocations: 130.86 MiB)
patched_mean_lazy(img, rₚ) == patched_mean(img, rₚ)

# Great! We haven't incoporate the pre-calculation trick yet, we still get a bit faster by making
# things lazy.

# More generally speaking, pre-calculation is just a cache. `PairwiseDistance` supports a so-called
# [`LocalWindowCache`](@ref) cache which stores the results in neighborhood `(p, q-r:q+r)` for each pixel
# pair `(p, q)`.

eval_op(x, y) = abs2(x - y)
pointwise_dist = PairwiseDistance(eval_op, (img, img), LocalWindowCache((7, 7))); # 32.575 μs (3 allocations: 980.12 KiB)

pq1 = pointwise_dist[CartesianIndex(1, 1), CartesianIndex(2, 2)]
pq2 = eval_op(img[CartesianIndex(1, 1)], img[CartesianIndex(2, 2)])
pq1 == pq2

# Instead of caching the result of pixel distances, we choose to cache the result of patch distances:

patchwise_dist = let img=img, rₚ=rₚ
    PairwiseDistance(SqEuclidean(), (img, img), LocalWindowCache(size(img))) do i
        i-rₚ:i+rₚ
    end;
end;

p = CartesianIndex(4, 4)
q = CartesianIndex(5, 5)
## @btime getindex($patchwise_dist, $p, $q) # 7.946 ns (0 allocations: 0 bytes)
## @btime sqeuclidean($(img[p-rₚ:p+rₚ]), $(img[q-rₚ:q+rₚ])) # 19.858 ns (0 allocations: 0 bytes)
patchwise_dist[p, q] == sqeuclidean(img[p-rₚ:p+rₚ], img[q-rₚ:q+rₚ])

# The time 7.946 ns is the cache read overhead. This says as long as our calculate takes more than
# this amount of time, it worths caching the results. The patch distance takes about 20 ns, so caching
# it is expected to make a performance improvment.

function patched_mean_cache(img, rₚ; num_patches=10)
    out = fill(zero(eltype(img)), axes(img))

    patchwise_dist = let rₚ=rₚ
        PairwiseDistance(SqEuclidean(), (img, img), LocalWindowCache(size(img))) do i
            i-rₚ:i+rₚ
        end;
    end;

    R = CartesianIndices(img)
    R0 = first(R)+rₚ:last(R)-rₚ
    for p in R0
        dist = vec(patchwise_dist[p, R0])
        matched_patches = R0[partialsortperm(dist, 1:num_patches)]
        out[p] = mapreduce(q->img[q], +, matched_patches)/length(matched_patches)
    end
    out
end

## @btime patched_mean_cache($img, $rₚ);
##  752.450 ms (26920 allocations: 214.15 MiB)
## @btime patched_mean_lazy($img, $rₚ);
##  532.812 ms (26916 allocations: 131.63 MiB)
## @btime patched_mean($img, $rₚ);
##  824.653 ms (23549 allocations: 130.86 MiB)
patched_mean_cache(img, rₚ) == patched_mean(img, rₚ)

# Oops! It's worse than our non-cache version. TODO: "fix" it


text:
module Poincare

using GPUBenchmarks, BenchmarkTools

description = """
Mapreduce, e.g. sum!.
Interestingly, for the sum benchmark the arrayfire opencl backend is the fastest, while GPUArrays OpenCL backend is the slowest.
This means we should be able to remove the slowdown for GPUArrays + OpenCL and maybe also for all the CUDA backends.
"""

function makeresult(name, bench, N, device, hardware, mdiff)
    BenchResult(
        name,
        bench,
        N,
        Float32,
        string(device),
        hardware,
        @__FILE__,
        mdiff
    )
end
relerr(x,y) = abs(x - y) * 2 / (abs(x) + abs(y))
function execute(device)
    results = BenchResult[]
    srand(2)
    for i = 1:7
        N = 10^i
        hardware, array_type = init(device)
        A = rand(Float32, N)
        Agpu = array_type(A)
        bench = @benchmark begin
            sum($Agpu)
            synchronize($(Agpu)) # synchronize for the benchmark
        end
        jl_results = sum(A)
        result = sum(Agpu)
        mdiff = relerr(jl_results, result)
        @show mdiff
        push!(results, makeresult("sum", bench, N, device, hardware, mdiff))
        free(Agpu)
    end
    results
end

end


text:
function load(filename::AbstractString)
        return load_object(filename)
end

function pad1D(X::AbstractArray, n::Int)
        return [X zeros(eltype(X), (size(X, 1), n))]
end

function pad2D(X::AbstractArray, n::Int)
        return [X zeros(eltype(X), (size(X, 1), n)); zeros(eltype(X), (n, n + size(X, 2)))]
end

function regularActivity(N::Int, time::Int, frequency::Real; delay=0)
        period = repeat([1; repeat([0], 1000 ÷ frequency - 1)...], 1, N)
        activity = repeat(period, Int(ceil(time * frequency / 1000)))
        delayed_activity = [zeros(delay, N); activity][1:time, :]
        return delayed_activity
end

function randomActivity(N::Int, time::Int, frequency::Real; delay=0, seed=nothing)
        if !isnothing(seed)
                Random.seed!(seed)
        end
        activity = rand(time, N) .< (frequency / 1000)
        delayed_activity = [zeros(delay, N); activity][1:time, :]
        return delayed_activity
end


text:
include("../src/RandomField.jl")
using .RandomField

# choose parameters:
## remember, we need that 4|n_i for i=1,...,d
N = [400,400]		# number pixels per dimension
L = [20.0,80.0]		# real space box is [-l,l]^d
f(V::Vector{T} where T<:Real) = gaus_rf(V)
					# spectral density
n_real = 1000		# number of realizations

# derived parameters
d = length(N) 		# dimension
H = (pi/2) .* N ./L
var_y(V::Vector{T} where T<:Real) = 2*(2pi)^(d/2)*prod(L)/pi^d*f(V)

# allocate memory
V = Cloud(d,N)
X = Cloud(d,N)
Y = Cscalar(V)
Z = Cscalar(X)
M_est = Cscalar(X)	# estimated mean

# initialize grids
init!(X,L)
init!(V,H)

# estimate mean
@. M_est.values = 0
for i=1:n_real
	draw_realization!(Y,var_y)
	ift!(Z,Y)	# Y ↦ Z
	@. M_est.values += Z.values
end
@. M_est.values /= n_real

# print:
tovtk(M_est, "../data/x.mu.vtk")
tocsv(slice(Z,1), "../data/x.z1.csv")
tocsv(slice(M_est,1), "../data/x.mu1.csv")