text:
export lex_xml

function lex_begin_tag(l::Lexer)
    if !accept_char(l, "<")
        return error(l, "XML tag must begin  with '<'")
    end

    if accept_char(l, "/")
        if scan_identifier(l) && accept_char(l, ">")
            emit_token(l, CLOSE_TAG)
            return lex_xml
        else
           return error(l, "XML closing tag must have alphanumeric identifier and end with '>'")
        end
    elseif scan_identifier(l)
        emit_token(l, BEGIN_TAG)
        if accept_char(l, ">")
            emit_token(l, END_TAG)
            return lex_xml
        else
           return lex_inside_tag
        end
    else
        return error(l, "Encountered neither a open or closing tag")
    end
end

"lex an identifier such as a variable or function name"
function lex_attrib_ident(l::Lexer)
    scan_identifier(l)
    emit_token(l, IDENT)
    return lex_inside_tag
end

"Lex a string enclosed in \" quotes"
function lex_attrib_string(l::Lexer)
    if scan_string(l)
        emit_token(l, STRING)
        lex_inside_tag
    else
        error(l, "Not a valid quoted string")
    end
end

"Lex the text content between two tags. Includes whitespace"
function lex_text(l::Lexer)
    ch = peek_char(l)
    while ch ∉ ['<', '>', EOFChar]
        next_char(l)
        ch = peek_char(l)
    end
    emit_token(l, TEXT)
    return lex_xml
end

function lex_inside_tag(l::Lexer)
    while true
        ignore_whitespace(l)
    	ch = peek_char(l)

        if ch == '>'
            next_char(l)
            emit_token(l, END_TAG)
            return lex_xml
    	elseif ch == '/'
            if accept_char(l, "/") && accept_char(l, ">")
                emit_token(l, END_AND_CLOSE_TAG)
                return lex_xml
            else
                return error(l, "Not properly formed end and close tag")
            end
    	elseif ch == '"'
    		return lex_attrib_string
    	elseif isletter(ch)
    		return lex_attrib_ident
        elseif ch == '='
            next_char(l)
            emit_token(l, EQUAL)
            return lex_inside_tag
        else
            return error(l, "Don't know how to handle '$ch'")
    	end
    end
end

function lex_xml(l::Lexer)
    while true
        ignore_whitespace(l)
        ch = peek_char(l)

        if ch == EOFChar
            emit_token(l, EOF)
            return lex_end
        elseif ch == '<'
            return lex_begin_tag
        else
           return lex_text
        end
    end
end

function lex_xml(input::String)
    lex(input, lex_xml)
end


text:
using SparseArrays
import SparseArrays: AbstractSparseVector, nonzeros, nonzeroinds

#=
Sparse vector
=#

"""
Space complexity: O(m)
"""
mutable struct SimpleSparseVector{T} <: AbstractSparseVector{T, Int}
    n::Int
    indices::Vector{Int}
    values::Vector{T}

    function SimpleSparseVector(n::Integer,
                                indices::Vector{<:Integer},
                                values::Vector{T}) where T
        n ≥ 0 || throw(ArgumentError("`n` must be ≥ 0"))
        length(indices) == length(values) ||
            throw(ArgumentError("`indices` and `values` must be the same length"))
        new{T}(n, indices, values)
    end
end

#=
Constructors
=#

SimpleSparseVector{T}(n::Integer) where T = SimpleSparseVector(n, Int[], T[])

#=
Core methods
=#

"""
Time complexity: O(m)
"""
function Base.getindex(ssv::SimpleSparseVector{T}, i::Integer) where T
    checkbounds(ssv, i)
    indices = nonzeroinds(ssv)
    values = nonzeros(ssv)
    j = searchsortedfirst(indices, i)
    if j ≤ length(indices) && indices[j] == i
        return values[j]
    end
    return zero(T)
end

"""
Time complexity: O(m)
"""
function Base.setindex!(ssv::SimpleSparseVector{T}, v::T, i::Integer) where T
    checkbounds(ssv, i)
    indices = nonzeroinds(ssv)
    values = nonzeros(ssv)
    j = searchsortedfirst(indices, i)
    if 1 ≤ j ≤ length(indices) && indices[j] == i
        values[j] = v
    elseif !iszero(v)
        insert!(indices, j, i)
        insert!(values, j, v)
    end
    return ssv
end

"""
Time complexity: O(m)
"""
function Base.insert!(ssv::SimpleSparseVector{T}, index::Integer, item::T) where T
    checkbounds(ssv, index)
    indices = nonzeroinds(ssv)
    ssv.n += 1
    for i in 1:length(indices)
        if indices[i] ≥ index
            indices[i] += 1
        end
    end
    setindex!(ssv, item, index)
end

"""
Time complexity: O(m)
"""
function Base.deleteat!(ssv::SimpleSparseVector, index::Integer)
    checkbounds(ssv, index)
    indices = nonzeroinds(ssv)
    values = nonzeros(ssv)
    j = findfirst(==(index), indices)
    if !isnothing(j)
        deleteat!(indices, j)
        deleteat!(values, j)
    end
    for i in 1:length(indices)
        if indices[i] ≥ index
            indices[i] -= 1
        end
    end
    ssv.n -= 1
    return ssv
end

#=
Interface
=#

SparseArrays.nonzeros(ssv::SimpleSparseVector) = ssv.values

SparseArrays.nonzeroinds(ssv::SimpleSparseVector) = ssv.indices

Base.size(ssv::SimpleSparseVector) = (ssv.n,)

IndexStyle(::SimpleSparseVector) = IndexLinear

ext: jl


text:
@Generated 1:32 immutable Table{Names, Types <: Tuple}
    if !isa(Names, Tuple) || eltype(Names) != Symbol || length(Names) != length(unique(Names)) # Also rules out length-0
        str = "Table parameter 1 (Names) is expected to be a tuple of unique symbols, got $Names" # TODO: reinsert colons in symbols?
        error(str)
    end
    N = length(Names)
    if length(Types.parameters) != N || reduce((a,b) -> a | !isa(b, DataType), false, Types.parameters)
        str = "Table parameter 2 (Types) is expected to be a Tuple{} of $N DataTypes, got $Types"
        error(str)
    end

    # Column containers should have a valid eltype
    for i = 1:N
        if eltype(Types.parameters[i]) == Types.parameters[i] # Probably should be more refined here...
            warn("Column $i storage type $(Types.parameters[i]) appears not to be a container")
        end
    end

    exprs = [:( $(Names[i])::$(Types.parameters[i]) ) for i = 1:N]
    return Expr(:block, exprs...)
end

Base.call{Names,T <: Tuple}(::Type{Table{Names,T}}, data::T) = Table{Names,T}(data...) # Constructed from a tuple (signature must match exactly...)

@generated function Base.call{Names}(::Type{Table{Names}}, data...)
    if !isa(Names, Tuple) || eltype(Names) != Symbol || length(Names) != length(unique(Names))
        str = "Table parameter 1 (Names) is expected to be a tuple of unique symbols, got $Names" # TODO: reinsert colons in symbols?
        return :(error($str))
    end

    if length(Names) == length(data)
        return :(Table{Names, $(Expr(:curly, :Tuple, data...))}(data...))
    elseif length(data) == 1 && data[1] <: Tuple && length(Names) == length(data[1].parameters)  # Constructed from a tuple (with correct length)
        return :(Table{Names, $(data[1])}(data[1]...))
    else
        return :(error("Can't construct Table with $(length(Names)) columns with input $data"))
    end
end

Base.call{Names, Types, Types_new}(::Type{Table{Names,Types_new}}, t::Table{Names,Types}) = convert(Table{Names,Types_new}, t)
@generated function Base.convert{Names, Names_new, Types, Types_new <: Tuple}(::Type{Table{Names_new,Types_new}}, t::Table{Names,Types})
    if !isa(Names_new, Tuple) || eltype(Names_new) != Symbol || length(Names_new) != length(Names) || length(Names_new) != length(unique(Names_new))
        str = "Cannot convert $(length(Names)) columns to new names $(Names_new)."
        return :(error($str))
    end
    if length(Types_new.parameters) != length(Types.parameters)
        str = "Cannot convert $(length(Types.pareters)) columns to $(length(Types_new.parameters)) new types."
        return :(error($str))
    end
    exprs = [:(convert(Types_new.parameters[$j], t.($j))) for j = 1:length(Names)]
    return Expr(:call, Table{Names_new,Types_new}, exprs...)
end

@inline colnames{Names}(::Table{Names}) = Names
@inline colnames{Names, Types <: Tuple}(::Type{Table{Names,Types}}) = Names
@inline colnames{Names}(::Type{Table{Names}}) = Names
@inline colnames{T <: Table}(::Type{T}) = colnames(super(T))
@inline eltypes{Names, Types <: Tuple}(::Table{Names,Types}) = Types
@inline eltypes{Names, Types <: Tuple}(::Type{Table{Names,Types}}) = Types
@inline eltypes{T <: Table}(::Type{T}) = eltypes(super(T))

@inline nrow(t::Table) = length(t.(1))
@inline ncol{Names}(::Table{Names}) = length(Names)

# reordering
@generated function permutecols{Names1,Names2}(t::Table{Names1}, ::Type{Val{Names2}})
    if Names1 == Names2
        return :(t)
    else
        if !(isa(Names2, Tuple)) || eltype(Names2) != Symbol || length(Names2) != length(Names1) || length(Names2) != length(unique(Names2))
            str = "New column names $Names2 do not match existing names $Names1"
            return :(error($str))
        end

        order = permutator(Names1, Names2)

        exprs = [:(t.($(order[j]))) for j = 1:length(Names1)]
        return Expr(:call, Table{Names2}, exprs...)
    end
end

function permutator{N}(names1::NTuple{N,Symbol}, names2::NTuple{N,Symbol})
    order = zeros(Int, N)
    for i = 1:N
        isfound = false
        for j = 1:N
            if names1[i] == names2[j]
                isfound = true
                order[j] = i
                break
            end
        end

        if !isfound
            str = "New column names $names2 do not match existing names $names1"
            return :(error($str))
        end
    end

    return order
end

# iterating

# TODO don't assume all containers are compatible with each other... (e.g. different implementations of dictionaries).
# Probably should define some kind of table key for this

# TODO very strong assumption about columns being of same length. Probably
# should check in start? And hope the user doesn't change the column sizes
# differently? Or do the safe thing and get the user to wrap it in @inbounds??

Base.start(t::Table) = 1
@generated function Base.next{Names}(t::Table{Names}, state)
    exprs = [:(Base.unsafe_getindex(t.($i),state)) for i = 1:length(Names)]
    return Expr(:tuple, Expr(:call, Row{Names}, exprs...), :(state + 1))
end
Base.done(t::Table, state) = state > nrow(t)

# indexing
Base.endof(t::Table) = nrow(t)
Base.length(t::Table) = nrow(t)
Base.size(t::Table) = (nrow(t),)
Base.size(t::Table, d) = size(t.(1), d)

@generated function getindex{Names}(t::Table{Names}, i::Integer)
    exprs = [:(Base.getindex(t.($c), i)) for c = 1:length(Names)]
    return Expr(:call, Row{Names}, exprs...)
end
@generated function getindex{Names}(t::Table{Names}, inds)
    exprs = [:(getindex(t.($c), inds)) for c = 1:length(Names)]
    return Expr(:call, Table{Names}, exprs...)
end

@generated function setindex!{Names1,Names2}(t::Table{Names1}, v::Row{Names2}, i::Integer)
    if Names1 == Names2
        exprs = [:(setindex!(t.($c), v.$(c), i)) for c = 1:length(Names1)]
        return Expr(:block, exprs...)
    else
        if length(Names1) != length(Names2)
            str = "Cannot assign $(length(v.parameters)) columns to $(length(Names)) columns"
            return :(error($str))
        end

        order = permutator(Names1, Names2)
        exprs = [:(setindex!(t.($(order[c])), v.$(c), i)) for c = 1:length(Names1)]
        return Expr(:block, exprs...)
    end
end
@generated function setindex!{Names}(t::Table{Names}, v::Tuple, i)
    if length(v.parameters) != length(Names)
        str = "Cannot assign a $(length(v.parameters))-tuple to $(length(Names)) columns"
        return :(error($str))
    end
    exprs = [:(setindex!(t.($c), v.$(c), i)) for c = 1:length(Names)]
    return Expr(:block, exprs...)
end
@generated function setindex!{Names1,Names2}(t::Table{Names1}, v::Table{Names2}, inds)
    if Names1 == Names2
        exprs = [:(setindex!(t.($c), v.$(c), inds)) for c = 1:length(Names1)]
        return Expr(:block, exprs...)
    else
        if length(Names1) != length(Names2)
            str = "Cannot assign $(length(v.parameters)) columns to $(length(Names)) columns"
            return :(error($str))
        end

        order = permutator(Names1, Names2)
        exprs = [:(setindex!(t.($(order[c])), v.$(c), inds)) for c = 1:length(Names1)]
        return Expr(:block, exprs...)
    end
end

@generated function unsafe_getindex{Names}(t::Table{Names}, i::Integer)
    exprs = [:(Base.unsafe_getindex(t.($c), i)) for c = 1:length(Names)]
    return Expr(:call, Row{Names}, exprs...)
end
@generated function unsafe_getindex{Names}(t::Table{Names}, inds)
    exprs = [:(unsafe_getindex(t.($c), inds)) for c = 1:length(Names)]
    return Expr(:call, Table{Names}, exprs...)
end

@generated function unsafe_setindex!{Names1,Names2}(t::Table{Names1}, v::Row{Names2}, i::Integer)
    if Names1 == Names2
        exprs = [:(unsafe_setindex!(t.($c), v.$(c), i)) for c = 1:length(Names1)]
        return Expr(:block, exprs...)
    else
        if length(Names1) != length(Names2)
            str = "Cannot assign $(length(v.parameters)) columns to $(length(Names)) columns"
            return :(error($str))
        end

        order = permutator(Names1, Names2)
        exprs = [:(unsafe_setindex!(t.($(order[c])), v.$(c), i)) for c = 1:length(Names1)]
        return Expr(:block, exprs...)
    end
end
@generated function unsafe_setindex!{Names}(t::Table{Names}, v::Tuple, i)
    if length(v.parameters) != length(Names)
        str = "Cannot assign a $(length(v.parameters))-tuple to $(length(Names)) columns"
        return :(error($str))
    end
    exprs = [:(unsafe_setindex!(t.($c), v.$(c), i)) for c = 1:length(Names)]
    return Expr(:block, exprs...)
end
@generated function unsafe_setindex!{Names1,Names2}(t::Table{Names1}, v::Table{Names2}, inds)
    if Names1 == Names2
        exprs = [:(unsafe_setindex!(t.($c), v.$(c), inds)) for c = 1:length(Names1)]
        return Expr(:block, exprs...)
    else
        if length(Names1) != length(Names2)
            str = "Cannot assign $(length(v.parameters)) columns to $(length(Names)) columns"
            return :(error($str))
        end

        order = permutator(Names1, Names2)
        exprs = [:(unsafe_setindex!(t.($(order[c])), v.$(c), inds)) for c = 1:length(Names1)]
        return Expr(:block, exprs...)
    end
end

# push, pop, etc
@generated function Base.pop!{Names}(t::Table{Names})
    exprs = [:(pop!(t.($c))) for c = 1:length(Names)]
    return Expr(:call, Row{Names}, exprs...)
end

@generated function Base.shift!{Names}(t::Table{Names})
    exprs = [:(shift!(t.($c))) for c = 1:length(Names)]
    return Expr(:call, Row{Names}, exprs...)
end

@generated function Base.push!{Names}(t::Table{Names}, v::Tuple)
    if length(Names) != length(v.parameters)
        str = "Cannot assign a $(length(v.parameters))-tuple to $(length(Names)) columns"
        return :(error($str))
    end
    exprs = [:(push!(t.($c), v.($c))) for c = 1:length(Names)]
    return Expr(:block, exprs..., :t)
end
@generated function Base.push!{Names1,Names2}(t::Table{Names1}, r::Row{Names2})
    order = permutator(Names1, Names2)
    exprs = [:(push!(t.($(order[c])), r.($c))) for c = 1:length(Names1)]
    return Expr(:block, exprs..., :t)
end

@generated function Base.unshift!{Names}(t::Table{Names}, v::Tuple)
    if length(Names) != length(v.parameters)
        str = "Cannot assign a $(length(v.parameters))-tuple to $(length(Names)) columns"
        return :(error($str))
    end
    exprs = [:(unshift!(t.($c), v.($c))) for c = 1:length(Names)]
    return Expr(:block, exprs..., :t)
end
@generated function Base.unshift!{Names1,Names2}(t::Table{Names1}, r::Row{Names2})
    order = permutator(Names1, Names2)
    exprs = [:(unshift!(t.($(order[c])), r.($c))) for c = 1:length(Names1)]
    return Expr(:block, exprs..., :t)
end

@generated function Base.append!{Names}(t::Table{Names}, v::Tuple)
    if length(Names) != length(v.parameters)
        str = "Cannot assign a $(length(v.parameters))-tuple to $(length(Names)) columns"
        return :(error($str))
    end
    exprs = [:(append!(t.($c), v.($c))) for c = 1:length(Names)]
    return Expr(:block, exprs..., :t)
end
@generated function Base.append!{Names1,Names2}(t::Table{Names1}, t2::Table{Names2})
    order = permutator(Names1, Names2)
    exprs = [:(append!(t.($(order[c])), t2.($c))) for c = 1:length(Names1)]
    return Expr(:block, exprs..., :t)
end

@generated function Base.prepend!{Names}(t::Table{Names}, v::Tuple)
    if length(Names) != length(v.parameters)
        str = "Cannot assign a $(length(v.parameters))-tuple to $(length(Names)) columns"
        return :(error($str))
    end
    exprs = [:(prepend!(t.($c), v.($c))) for c = 1:length(Names)]
    return Expr(:block, exprs..., :t)
end
@generated function Base.prepend!{Names1,Names2}(t::Table{Names1}, t2::Table{Names2})
    order = permutator(Names1, Names2)
    exprs = [:(prepend!(t.($(order[c])), t2.($c))) for c = 1:length(Names1)]
    return Expr(:block, exprs..., :t)
end

@generated function Base.insert!{Names}(t::Table{Names}, i::Integer, v::Tuple)
    if length(Names) != length(v.parameters)
        str = "Cannot assign a $(length(v.parameters))-tuple to $(length(Names)) columns"
        return :(error($str))
    end
    exprs = [:(insert!(t.($c), i, v.($c))) for c = 1:length(Names)]
    return Expr(:block, exprs..., :t)
end
@generated function Base.insert!{Names1,Names2}(t::Table{Names1}, i::Integer, r::Row{Names2})
    order = permutator(Names1, Names2)
    exprs = [:(insert!(t.($(order[c])), i, r.($c))) for c = 1:length(Names1)]
    return Expr(:block, exprs..., :t)
end

@generated function Base.deleteat!{Names}(t::Table{Names}, i)
    exprs = [:(deleteat!(t.($c), i)) for c = 1:length(Names)]
    return Expr(:block, exprs..., :t)
end

@generated function Base.splice!{Names}(t::Table{Names}, i::Integer)
    exprs = [:(splice!(t.($c), i)) for c = 1:length(Names)]
    return Expr(:call, Row{Names}, exprs...)
end
@generated function Base.splice!{Names}(t::Table{Names}, i::Integer, v::Tuple)
    exprs = [:(splice!(t.($c), i, v.($c))) for c = 1:length(Names)]
    return Expr(:call, Row{Names}, exprs...)
end
@generated function Base.splice!{Names1,Names2}(t::Table{Names1}, i::Integer, v::Union{Row{Names2},Table{Names2}})
    order = permutator(Names2, Names1)
    exprs = [:(splice!(t.($c), i, v.($(order[c])))) for c = 1:length(Names1)]
    return Expr(:call, Row{Names1}, exprs...)
end
@generated function Base.splice!{Names}(t::Table{Names}, i)
    exprs = [:(splice!(t.($c), i)) for c = 1:length(Names)]
    return Expr(:call, Table{Names}, exprs...)
end
@generated function Base.splice!{Names}(t::Table{Names}, i, v::Tuple)
    exprs = [:(splice!(t.($c), i, v.($c))) for c = 1:length(Names)]
    return Expr(:call, Table{Names}, exprs...)
end
@generated function Base.splice!{Names1,Names2}(t::Table{Names1}, i, v::Union{Row{Names2},Table{Names2}})
    order = permutator(Names2, Names1)
    exprs = [:(splice!(t.($c), i, v.($(order[c])))) for c = 1:length(Names1)]
    return Expr(:call, Table{Names1}, exprs...)
end

# Vertically concatenate rows and tables into tables
@generated function Base.vcat{Names}(t1::Union{Row{Names}, Table{Names}})
    exprs = [:(vcat(t1.($c))) for c = 1:length(Names)]
    return Expr(:call, Table{Names}, exprs...)
end

@generated function Base.vcat{Names1, Names2}(t1::Union{Row{Names1}, Table{Names1}}, t2::Union{Row{Names2}, Table{Names2}})
    if Names1 == Names2
        exprs = [:(vcat(t1.($c), t2.($c))) for c = 1:length(Names1)]
        return Expr(:call, Table{Names1}, exprs...)
    else
        if length(Names1) != length(Names2)
            str = "Cannot match $(length(v.parameters)) columns to $(length(Names)) columns"
        end

        order = permutator(Names1, Names2)
        exprs = [:(vcat(t1.($c), t2.($(order[c])))) for c = 1:length(Names1)]
        return Expr(:call, Table{Names1}, exprs...)
    end
end

Base.vcat{Names1, Names2}(t1::Union{Row{Names1}, Table{Names1}}, t2::Union{Row{Names2}, Table{Names2}}, ts::Union{Row, Table}...) = vcat(vcat(t1, t2), ts...)

# Horizontally concatenate columns and tables into tables
@generated Base.hcat{Name}(c::Column{Name}) = :(Table{$((Name,))}(c.(1)))
Base.hcat(t::Table) = t

@generated function Base.hcat(r1::Union{Column,Table}, r2::Union{Column,Table})
    names1 = (r1 <: Column ? (colname(r1),) : colnames(r1))
    names2 = (r2 <: Column ? (colname(r2),) : colnames(r2))

    if length(intersect(names1, names2)) != 0
        str = "Column names are not distinct. Got $names1 and $names2"
        return :(error($str))
    end

    newnames = (names1..., names2...)
    exprs = vcat([:(r1.($j)) for j = 1:length(names1)], [:(r2.($j)) for j = 1:length(names2)])

    return Expr(:call, Table{newnames}, exprs...)
end

Base.hcat(t1::Union{Column,Table}, t2::Union{Column,Table}, ts::Union{Column,Table}...) = hcat(hcat(t1, t2), ts...)

# copy
@generated function Base.copy{Names}(t::Table{Names})
    exprs = [:(copy(t.($j))) for j = 1:length(Names)]
    return Expr(:call, Table{Names}, exprs...)
end


text:
# Parts of this code were taken / derived from NetworkX. See LICENSE for
# licensing details.

"""
    pagerank(g, α=0.85, n=100, ϵ=1.0e-6)

Calculate the [PageRank](https://en.wikipedia.org/wiki/PageRank) of the
directed graph `g` parameterized by damping factor `α`, number of
iterations `n`, and convergence threshold `ϵ`. Return a vector representing
the centrality calculated for each node in `g`, or an error if convergence
is not reached within `n` iterations.
"""
function pagerank end

@traitfn function pagerank(g::::IsDirected, α=0.85, n=100::Integer, ϵ=1.0e-6)
    A = adjacency_matrix(g, :in, Float64)
    S = vec(sum(A, 1))
    S = 1 ./ S
    S[find(S .== Inf)] = 0.0
    M = A'  # need a separate line due to bug #17456 in julia
    # scaling the adjmat to stochastic adjacency matrix
    M = (Diagonal(S) * M)'
    N = Int(nv(g))
    # solution vector
    x = fill(1.0 / N, N)
    # personalization vector
    p = fill(1.0 / N, N)
    # temporary to hold the results of SpMV
    y = zeros(Float64, N)
    # adjustment for leaf nodes in digraph
    dangling_weights = p
    is_dangling = find(S .== 0)
    # save some flops by precomputing this
    pscaled = (1 .- α) .* p
    for _ in 1:n
        xlast = x
        # in place SpMV to conserve memory
        A_mul_B!(y, M, x)
        # using broadcast to avoid temporaries
        x = α .* (y .+ sum(x[is_dangling]) .* dangling_weights) .+ pscaled
        # l1 change in solution convergence criterion
        err = sum(abs, (x .- xlast))
        if (err < N * ϵ)
            return x
        end
    end
    error("Pagerank did not converge after $n iterations.")
end


text:
#
# Copyright (c) 2021 Tobias Thummerer, Lars Mikelsons, Josef Kircher
# Licensed under the MIT license. See LICENSE file in the project root for details.
#

"""
IN PROGRESS:
The mutable struct representing an FMU in the FMI 3.0 Standard.
Also contains the paths to the FMU and ZIP folder as well als all the FMI 3.0 function pointers
"""
mutable struct FMU3

    # Constructor
    fmu3() = new()
end


text:
using IntervalRootFinding
using Test

include("branch_and_bound.jl")
include("search_interface.jl")
include("roots.jl")
include("test_smiley.jl")
include("newton1d.jl")
include("quadratic.jl")
include("linear_eq.jl")
include("slopes.jl")


text:
using DataFrames, CSV, Test, FixedEffectModels

df = CSV.read(joinpath(dirname(@__FILE__), "..", "dataset/Cigar.csv"))
df[:pState] = categorical(df[:State])
df[:pYear] = categorical(df[:Year])

if Base.USE_GPL_LIBS
	method_s = [:cholesky, :qr, :lsmr, :lsmr_parallel, :lsmr_threads]
else
	method_s = [:lsmr, :lsmr_parallel, :lsmr_threads]
end

for method in method_s
	global model = @model Sales ~ Price fe = pYear save = true method = $(method)
	global result = reg(df, model)
	@test fes(result)[1, :pYear] ≈ 164.77833189721005

	global model = @model Sales ~ Price fe = pYear + pState save = true method = $(method)
	global result = reg(df, model)
	@test fes(result)[1, :pYear] + fes(result)[1, :pState] ≈ 140.6852 atol = 1e-3

	global model = @model Sales ~ Price fe = Year&pState save = true method = $(method)
	global result = reg(df, model)
	@test fes(result)[1, :YearxpState] ≈ 1.742779  atol = 1e-3

	global model = @model Sales ~ Price fe = pState + Year&pState save = true method = $(method)
	global result = reg(df, model)
	@test fes(result)[1, :pState] ≈ -91.690635 atol = 1e-1

	global model = @model Sales ~ Price fe = pState save = true subset = (State .<= 30) method = $(method)
	global result = reg(df, model)
	@test fes(result)[1,:pState] ≈  124.913976 atol = 1e-1
	@test ismissing(fes(result)[1380,:pState])
end
	# add test with IV & weight

df[:Price2] = df[:Price]
model = @model Sales ~ Price + Price2 fe = pYear save = true
result = reg(df, model)
@test fes(result)[1, :pYear] ≈ 164.77833189721005

model = @model Sales ~ (State ~ Price) fe = pYear save = true
result = reg(df, model)
@test fes(result)[1, :pYear] ≈ -167.48093490413623

model = @model Sales ~ (State ~ Price + Price2) fe = pYear save = true
result = reg(df, model)
@test fes(result)[1, :pYear] ≈ -167.48093490413623

# check does not change r2
model1 = @model Sales ~ Price weights = Pop save = true
result1 = reg(df, model1)
model2 = @model Sales ~ Price weights = Pop
result2 = reg(df, model2)
@test r2(result1) ≈ r2(result2)

# check with weights
model = @model Sales ~ Price weights = Pop fe = pYear save = true
result = reg(df, model)
@test fes(result)[2, :pYear] -  fes(result)[1, :pYear] ≈ -3.0347149502496222

# I should check with IV and cehck with Weights an check with both


text:
using DataFrames, CSV

const CCLE = CSV.read("../gene_expression/tpm_values.csv", DataFrame)

function get_tpm(gene::String, cell_line::String)::Float64
    return CCLE[CCLE[!,Symbol("Column1")].==gene,Symbol(cell_line)][1]
end

function init_E1(p::Vector{Float64}, cell_line::String)::Float64
    return p[C.scale_E1] * get_tpm("EGFR", cell_line)
end

function init_E2(p::Vector{Float64}, cell_line::String)::Float64
    return p[C.scale_E2] * get_tpm("ERBB2", cell_line)
end

function init_E3(p::Vector{Float64}, cell_line::String)::Float64
    return p[C.scale_E3] * get_tpm("ERBB3", cell_line)
end

function init_E4(p::Vector{Float64}, cell_line::String)::Float64
    return p[C.scale_E4] * get_tpm("ERBB4", cell_line)
end

function init_G(p::Vector{Float64}, cell_line::String)::Float64
    return p[C.scale_G] * get_tpm("GRB2", cell_line)
end

function init_S(p::Vector{Float64}, cell_line::String)::Float64
    return (p[C.a_SHC1] * get_tpm("SHC1", cell_line) +
            p[C.a_SHC2] * get_tpm("SHC2", cell_line) +
            p[C.a_SHC3] * get_tpm("SHC3", cell_line) +
            p[C.a_SHC4] * get_tpm("SHC4", cell_line))
end

function init_I(p::Vector{Float64}, cell_line::String)::Float64
    return (p[C.a_PIK3CA] * get_tpm("PIK3CA", cell_line) +
            p[C.a_PIK3CB] * get_tpm("PIK3CB", cell_line) +
            p[C.a_PIK3CD] * get_tpm("PIK3CD", cell_line) +
            p[C.a_PIK3CG] * get_tpm("PIK3CG", cell_line))
end

function init_PTEN(p::Vector{Float64}, cell_line::String)::Float64
    return p[C.scale_PTEN] * get_tpm("PTEN", cell_line)
end

function init_R(p::Vector{Float64}, cell_line::String)::Float64
    return (p[C.a_RASA1] * get_tpm("RASA1", cell_line) +
            p[C.a_RASA2] * get_tpm("RASA2", cell_line) +
            p[C.a_RASA3] * get_tpm("RASA3", cell_line))
end

function init_O(p::Vector{Float64}, cell_line::String)::Float64
    return (p[C.a_SOS1] * get_tpm("SOS1", cell_line) +
            p[C.a_SOS2] * get_tpm("SOS2", cell_line))
end

function init_A(p::Vector{Float64}, cell_line::String)::Float64
    return p[C.scale_A] * get_tpm("GAB1", cell_line)
end

function init_Akt(p::Vector{Float64}, cell_line::String)::Float64
    return (p[C.a_AKT1] * get_tpm("AKT1", cell_line) +
            p[C.a_AKT2] * get_tpm("AKT2", cell_line))
end

function init_RsD(p::Vector{Float64}, cell_line::String)::Float64
    return (p[C.a_HRAS] * get_tpm("HRAS", cell_line) +
            p[C.a_KRAS] * get_tpm("KRAS", cell_line) +
            p[C.a_NRAS] * get_tpm("NRAS", cell_line))
end

function init_Raf(p::Vector{Float64}, cell_line::String)::Float64
    return (p[C.a_ARAF] * get_tpm("ARAF", cell_line) +
            p[C.a_BRAF] * get_tpm("BRAF", cell_line) +
            p[C.a_RAF1] * get_tpm("RAF1", cell_line))
end

function init_MEK(p::Vector{Float64}, cell_line::String)::Float64
    return (p[C.a_MAP2K1] * get_tpm("MAP2K1", cell_line) +
            p[C.a_MAP2K2] * get_tpm("MAP2K2", cell_line))
end

function init_T(p::Vector{Float64}, cell_line::String)::Float64
    return p[C.scale_T] * get_tpm("PTPN1", cell_line)
end

function init_CREBn(p::Vector{Float64}, cell_line::String)::Float64
    return p[C.scale_CREBn] * get_tpm("CREB1", cell_line)
end

function init_ERKc(p::Vector{Float64}, cell_line::String)::Float64
    return (p[C.a_MAPK1] * get_tpm("MAPK1", cell_line) +
            p[C.a_MAPK3] * get_tpm("MAPK3", cell_line))
end

function init_Elk1n(p::Vector{Float64}, cell_line::String)::Float64
    return p[C.scale_Elk1n] * get_tpm("ELK1", cell_line)
end

function init_RSKc(p::Vector{Float64}, cell_line::String)::Float64
    return (p[C.a_RPS6KA1] * get_tpm("RPS6KA1", cell_line) +
            p[C.a_RPS6KA2] * get_tpm("RPS6KA2", cell_line) +
            p[C.a_RPS6KA3] * get_tpm("RPS6KA3", cell_line) +
            p[C.a_RPS6KA6] * get_tpm("RPS6KA6", cell_line))
end

function tpm2ival!(p::Vector{Float64}, u0::Vector{Float64}, cell_line::String)
    p[C.PTEN] *= init_PTEN(p, cell_line)

    u0[V.E1] *= init_E1(p, cell_line)
    u0[V.E2] *= init_E2(p, cell_line)
    u0[V.E3] *= init_E3(p, cell_line)
    u0[V.E4] *= init_E4(p, cell_line)
    u0[V.G] *= init_G(p, cell_line)
    u0[V.S] *= init_S(p, cell_line)
    u0[V.I] *= init_I(p, cell_line)
    u0[V.R] *= init_R(p, cell_line)
    u0[V.O] *= init_O(p, cell_line)
    u0[V.A] *= init_A(p, cell_line)
    u0[V.Akt] *= init_Akt(p, cell_line)
    u0[V.RsD] *= init_RsD(p, cell_line)
    u0[V.Raf] *= init_Raf(p, cell_line)
    u0[V.MEK] *= init_MEK(p, cell_line)
    u0[V.T] *= init_T(p, cell_line)
    u0[V.CREBn] *= init_CREBn(p, cell_line)
    u0[V.ERKc] *= init_ERKc(p, cell_line)
    u0[V.Elk1n] *= init_Elk1n(p, cell_line)
    u0[V.RSKc] *= init_RSKc(p, cell_line)

    return p, u0
end

function mul_ratio2mcf7!(p::Vector{Float64}, u0::Vector{Float64}, cell_line_2::String)
    p[C.PTEN] *= init_PTEN(p, cell_line_2) / init_PTEN(p, "MCF7")

    u0[V.E1] *= init_E1(p, cell_line_2) / init_E1(p, "MCF7")
    u0[V.E2] *= init_E2(p, cell_line_2) / init_E2(p, "MCF7")
    u0[V.E3] *= init_E3(p, cell_line_2) / init_E3(p, "MCF7")
    u0[V.E4] *= init_E4(p, cell_line_2) / init_E4(p, "MCF7")
    u0[V.G] *= init_G(p, cell_line_2) / init_G(p, "MCF7")
    u0[V.S] *= init_S(p, cell_line_2) / init_S(p, "MCF7")
    u0[V.I] *= init_I(p, cell_line_2) / init_I(p, "MCF7")
    u0[V.R] *= init_R(p, cell_line_2) / init_R(p, "MCF7")
    u0[V.O] *= init_O(p, cell_line_2) / init_O(p, "MCF7")
    u0[V.A] *= init_A(p, cell_line_2) / init_A(p, "MCF7")
    u0[V.Akt] *= init_Akt(p, cell_line_2) / init_Akt(p, "MCF7")
    u0[V.RsD] *= init_RsD(p, cell_line_2) / init_RsD(p, "MCF7")
    u0[V.Raf] *= init_Raf(p, cell_line_2) / init_Raf(p, "MCF7")
    u0[V.MEK] *= init_MEK(p, cell_line_2) / init_MEK(p, "MCF7")
    u0[V.T] *= init_T(p, cell_line_2) / init_T(p, "MCF7")
    u0[V.CREBn] *= init_CREBn(p, cell_line_2) / init_CREBn(p, "MCF7")
    u0[V.ERKc] *= init_ERKc(p, cell_line_2) / init_ERKc(p, "MCF7")
    u0[V.Elk1n] *= init_Elk1n(p, cell_line_2) / init_Elk1n(p, "MCF7")
    u0[V.RSKc] *= init_RSKc(p, cell_line_2) / init_RSKc(p, "MCF7")

    return p, u0
end


text:
export Dist2Set, LossTotal, IoU, g_CQ, MultipleSplitFeasCQ_fun_grad

function Dist2Set(input,P,TrOpts)
  #active_z_slice is only used for hyperspectral imaging where we map a 3D or 4D
  #data volume to a 2D map of the earth. this 2D map is located at slice number
  #active_z_slice in a x-y-z-nchan-n_ex tensor

  #initialize some stuff
  size_input = size(input)
  pen_grad   = zeros(Float32,size(input))
  pen_value  = 0f0

  #use softmax of final network state for constraining
  #the output for segmentation problems (don't use softmax for non-linear regression)
  #input      = softmax(Array(input)[:,:,TrainOpts.channels,:],dims=3)

  input      = softmax(input[:,:,TrOpts.active_channels,:],dims=3)

  for j in TrOpts.active_channels #loop over channels, each channel has a (different) corresponding projector
    input_slice        = input[:,:,j,1]
    pen_value_slice    = 0f0 #
    P_input            = P[j](input_slice) #project onto constraint set (can be an intersection)
    pen_value_slice    = 0.5f0*norm(vec(P_input)-vec(input_slice))^2 #squared point-to-set distance functino
    pen_grad[:,:,j,1] .= input_slice .- P_input # gradient of point-to-set distance function

    pen_value = pen_value + pen_value_slice #accumulate penalty value (loss) over channels
  end #end channels loop

  #at this point,the softmax of the network output has been projected onto the intersection. Now we need to 'complete' the gradient
  #via the last part of the chain-rule for distance-function-squared and softmax
  pen_grad[:,:,TrOpts.active_channels,1] .= ∇softmax(pen_grad[:,:,TrOpts.active_channels,1], input[:,:,TrOpts.active_channels,1]; dims=3)

  return pen_value, pen_grad
end

function g_CQ(input,TD_OP,P_sub,alpha_CQ,TrOpts)

  #initialize some stuff
  size_input = size(input)
  pen_grad   = zeros(Float32,size(input))
  pen_value  = 0f0

  #use softmax of final network state for constraining
  #the output for segmentation problems (don't use softmax for non-linear regression)
  input      = softmax(input[:,:,TrOpts.active_channels,:],dims=3)

  for j in TrOpts.active_channels #loop over channels, each channel has a (different) corresponding projector
    input_slice        = input[:,:,j,1]
    pen_value_slice    = 0f0 #
    #P_input            = P[j](input_slice) #project onto constraint set (can be an intersection)
    #pen_value_slice    = 0.5f0*norm(vec(P_input)-vec(input_slice))^2 #squared point-to-set distance functino
    f_temp,g_temp      = MultipleSplitFeasCQ_fun_grad(vec(deepcopy(input_slice)),TD_OP[j],P_sub[j],alpha_CQ[j])
    pen_grad[:,:,j,1] .= reshape(g_temp,size(input_slice)) # gradient of point-to-set distance function

    pen_value = pen_value + f_temp #accumulate penalty value (loss) over channels
  end #end channels loop

  #at this point,the softmax of the network output has been projected onto the intersection. Now we need to 'complete' the gradient
  #via the last part of the chain-rule for distance-function-squared and softmax
  pen_grad[:,:,TrOpts.active_channels,1] .= ∇softmax(pen_grad[:,:,TrOpts.active_channels,1], input[:,:,TrOpts.active_channels,1]; dims=3)

  return pen_value, pen_grad
end

function MultipleSplitFeasCQ_fun_grad(x::AbstractVector{T},A,P_sub,alpha::AbstractVector) where T

# returns the gradient of the sum of squared distance functions for the split-feasibility
#problem: A_i x = y_i , y_i \in C_i
# f = \sum_{i=1}^p \alpha_i/2 \| P_C_i (A_i x) - A_i x \|_2^2
# g = \sum_{i=1}^p \alpha_i/2 A_i^t (A_i x - P_C_i (A_i x) )

#serial implementation
#(each projection or mat-vec can still use multiple threads)

f    = zeros(T,length(alpha))#allocate function value per constraint/linear operator
ng   = zeros(T,length(alpha))#allocate norm of gradient per constraint/linear operator
grad = zeros(T,length(x))

for i = 1:length(A)
  y     = A[i]*x
  PcAx  = P_sub[i](deepcopy(y))
  f[i]  = norm(PcAx - y,2).^2
  grad .= grad .+ alpha[i].*A[i]'*(y.-PcAx)
end

return sum(f),grad
end

function LossTotal(HN,TrOpts,X0::AbstractArray{T, N},label,P,image_weights,active_z_slice::Int) where {T, N}
  #loss for hyperspectral imaging with 5d input
  alpha=TrOpts.alpha; use_gpu=TrOpts.use_gpu; lossf=TrOpts.lossf; lossg=TrOpts.lossg; flip_dims=TrOpts.flip_dims; permute_dims=TrOpts.permute_dims; active_channels=TrOpts.active_channels;

    if (isempty(flip_dims) && isempty(permute_dims)) == false
      X0, label, image_weights = AugmentDataLabel(X0, label, image_weights,TrOpts)#optional: augment data
    end

    Y_curr, Y_new, lgdet = HN.forward(X0,X0)
    if use_gpu == true
      Y_new = Y_new|>cpu
    end

    if isempty(label)==false
        lval         = lossf(Y_new[:,:,active_z_slice,active_channels,1],label[:,:,active_z_slice,active_channels,1],image_weights)
        (grad,dummy) = lossg(Y_new[:,:,active_z_slice,active_channels,1],label[:,:,active_z_slice,active_channels,1],image_weights)
      else
        lval = 0.0
    end

  if alpha>0
    dc2,dc2_grad = Dist2Set(Y_new[:,:,active_z_slice,:,:],P,TrOpts)

    #grad  = grad + alpha*dc2_grad[:,:,active_channels,1]
  else
    dc2 = 0f0
  end

  if isempty(label)==false && alpha>0
      grad  = grad + alpha*dc2_grad[:,:,active_z_slice,active_channels,1]
  elseif isempty(label)==true
      n = size(dc2_grad)
      #dc2_grad = reshape(dc2_grad,n[1],n[2],1,n[3],1)
      grad  = alpha*dc2_grad[:,:,active_channels,1]
  end

   ΔY_curr= zeros(Float32,size(Y_new))
   ΔY_curr[:,:,active_z_slice,active_channels,1] .= grad
   ΔY_new = zeros(Float32,size(Y_new))

   if use_gpu == true
     Y_new   = Y_new|>gpu
     ΔY_curr = ΔY_curr|>gpu
     ΔY_new  = ΔY_new|>gpu
   end
   ΔY_curr, ΔY_new, Y_curr, Y_new = HN.backward(ΔY_curr, ΔY_new, Y_curr, Y_new)

   return lval, dc2
end

function LossTotal(HN,TrOpts,X0::AbstractArray{T, N},label,P,image_weights,active_z_slice::Array{Any,1},P_mode::String,TD_OP,P_sub,alpha_CQ) where {T, N}
    alpha=TrOpts.alpha; use_gpu=TrOpts.use_gpu; lossf=TrOpts.lossf; lossg=TrOpts.lossg; active_channels=TrOpts.active_channels;

    Y_curr, Y_new, lgdet = HN.forward(X0,X0)
    if use_gpu == true
      Y_new = Y_new|>cpu
    end

    #initialize gradient
    if N==4
      grad = zeros(Float32,size(Y_new[:,:,active_channels,1]))
    elseif N==5
      grad = zeros(Float32,size(Y_new[:,:,:,active_channels,1]))
    end

    if isempty(label)==false
      if N==4
        lval         = TrOpts.lossf(Y_new[:,:,active_channels,1],label,image_weights)
        (grad,dummy) = TrOpts.lossg(Y_new[:,:,active_channels,1],label,image_weights)
      elseif N==5
        lval         = TrOpts.lossf(Y_new[:,:,:,active_channels,1],label,image_weights)
        (grad,dummy) = TrOpts.lossg(Y_new[:,:,:,active_channels,1],label,image_weights)
      end
    else
      lval = 0.0
    end

  if (alpha[1][1]>0 || alpha_CQ[1][1]>0)==true
    if P_mode == "Proj_intersection"
      dc2,dc2_grad = Dist2Set(Y_new,P,active_channels)
      dc2_grad .= alpha.*dc2_grad
    elseif P_mode == "g_CQ"
      dc2,dc2_grad = g_CQ(Y_new,TD_OP,P_sub,alpha_CQ,TrOpts)
    end

    if N==4
      grad  = grad + dc2_grad[:,:,active_channels,1]
    elseif N==5
      grad  = grad + dc2_grad[:,:,:,active_channels,1]
    end
  else
    dc2 = 0f0
  end

  #optional: randomly set parts of the gradient to zero
  if TrOpts.rand_grad_perc_zero!=0.0
    rand_mask = CartesianIndices(size(grad)[1:end-2])
    rand_mask = shuffle(vec(rand_mask))[1:round(Int,TrOpts.rand_grad_perc_zero*prod(size(rand_mask)))]
    if N==4;grad[rand_mask,:,:].=0f0;elseif N==5;grad[rand_mask,:,:,:].=0f0;end;
  end

   ΔY_curr= zeros(Float32,size(Y_new))
   if N==4
     ΔY_curr[:,:,active_channels,1] .= grad
   elseif N==5
     ΔY_curr[:,:,:,active_channels,1] .= grad
   end
   ΔY_new = zeros(Float32,size(Y_new))

   if use_gpu == true
     Y_new   = Y_new|>gpu
     ΔY_curr = ΔY_curr|>gpu
     ΔY_new  = ΔY_new|>gpu
   end
   ΔY_curr, ΔY_new, Y_curr, Y_new = HN.backward(ΔY_curr, ΔY_new, Y_curr, Y_new)

   return lval, dc2
end

function IoU(HN,data,labels)

if length(size(data[1]))==5
 return 1.0, 1.0
end

threshold = 0.65
IoU_pos = zeros(length(data))
IoU_neg = zeros(length(data))

  for i=1:length(data)
    ~, prediction, ~ = HN.forward(data[i],data[i])
    prediction = prediction |> cpu
    prediction[:,:,1:2,1].=softmax(prediction[:,:,1:2,1],dims=3);

    pred_thres = zeros(Int,size(prediction)[1:2])
    pos_inds   = findall(prediction[:,:,1] .> threshold)
    pred_thres[pos_inds] .= 1

    prediction_for_acc = pred_thres[5:end-4,5:end-4]
    mask_for_acc       = labels[i][5:end-4,5:end-4,1]

    pos_pred_inds = findall(prediction_for_acc.==1)
    neg_pred_inds = findall(prediction_for_acc.==0)

    true_pos_inds = findall(mask_for_acc .== 1)
    true_neg_inds = findall(mask_for_acc .== 0)

    #IoU
    IoU_pos[i] = length(intersect(pos_pred_inds,true_pos_inds))/length(union(pos_pred_inds,true_pos_inds))
    IoU_neg[i] = length(intersect(neg_pred_inds,true_neg_inds))/length(union(neg_pred_inds,true_neg_inds))
  end

return IoU_pos, IoU_neg
end

#Loss for applications other than Hyperspectral. Uses projection onto intersection.
function LossTotal(HN,TrOpts,X0::AbstractArray{T, N},label,P,image_weights,active_z_slice::Array{Any,1}) where {T, N}
    alpha=TrOpts.alpha; use_gpu=TrOpts.use_gpu; lossf=TrOpts.lossf; lossg=TrOpts.lossg; active_channels=TrOpts.active_channels;

    Y_curr, Y_new, lgdet = HN.forward(X0,X0)
    if use_gpu == true
      Y_new = Y_new|>cpu
    end

    #initialize gradient
    if N==4
      grad = zeros(Float32,size(Y_new[:,:,active_channels,1]))
    elseif N==5
      grad = zeros(Float32,size(Y_new[:,:,:,active_channels,1]))
    end

    if isempty(label)==false
      if N==4
        lval         = lossf(Y_new[:,:,active_channels,1],label,image_weights)
        (grad,dummy) = lossg(Y_new[:,:,active_channels,1],label,image_weights)
      elseif N==5
        lval         = lossf(Y_new[:,:,:,active_channels,1],label,image_weights)
        (grad,dummy) = lossg(Y_new[:,:,:,active_channels,1],label,image_weights)
      end
    else
      lval = 0.0
    end
    # # #take equal number of random pixels from each class for the gradient
    # pos_inds = findall(label[:,:,1,1].==1)
    # neg_inds = findall(label[:,:,1,1].==0)
    # npix = min(length(pos_inds),length(neg_inds))
    # npix = round(Int,npix/2)
    # pos_inds_select = shuffle(pos_inds)[1:npix]
    # neg_inds_select = shuffle(neg_inds)[1:npix]
    # random_mask = zeros(Float32,size(grad))
    # random_mask[pos_inds_select,1:2,1].=1
    # random_mask[neg_inds_select,1:2,1].=1

    #grad .= grad.*random_mask

  if alpha>0
    #if P_mode == "Proj_intersection"
      dc2,dc2_grad = Dist2Set(Y_new,P,TrOpts)
    #elseif P_mode == "g_CQ"
    #  dc2,dc2_grad = g_CQ(Y_new,TD_OP,P_sub,alpha_CQ,active_channels)
    #end
    # if (norm(alpha*dc2_grad[:,:,active_channels,1])/norm(grad)) > 10f0
    #   @warn "(norm(alpha*dc2_grad[:,:,active_channels,1])/norm(grad)) > 10f0"
    # elseif (norm(alpha*dc2_grad[:,:,active_channels,1])/norm(grad)) < 0.1f0
    #   @warn "(norm(alpha*dc2_grad[:,:,active_channels,1])/norm(grad)) < 0.1f0"
    # end
    if N==4
      grad  = grad + alpha*dc2_grad[:,:,active_channels,1]
    elseif N==5
      grad  = grad + alpha*dc2_grad[:,:,:,active_channels,1]
    end
  else
    dc2 = 0f0
  end

  #optional: randomly set parts of the gradient to zero
  if TrOpts.rand_grad_perc_zero!=0.0
    rand_mask = CartesianIndices(size(grad)[1:end-2])
    rand_mask = shuffle(vec(rand_mask))[1:round(Int,TrOpts.rand_grad_perc_zero*prod(size(rand_mask)))]
    if N==4;grad[rand_mask,:,:].=0f0;elseif N==5;grad[rand_mask,:,:,:].=0f0;end;
  end

   ΔY_curr= zeros(Float32,size(Y_new))
   if N==4
     ΔY_curr[:,:,active_channels,1] .= grad
   elseif N==5
     ΔY_curr[:,:,:,active_channels,1] .= grad
   end
   ΔY_new = zeros(Float32,size(Y_new))

   if use_gpu == true
     Y_new   = Y_new|>gpu
     ΔY_curr = ΔY_curr|>gpu
     ΔY_new  = ΔY_new|>gpu
   end
   ΔY_curr, ΔY_new, Y_curr, Y_new = HN.backward(ΔY_curr, ΔY_new, Y_curr, Y_new)

   return lval, dc2
end


text:
using Plots
plotlyjs()

include("../julia/QQplot.jl");
include("../julia/utils.jl");

x_true = readcsv("../data/spikes.csv")[:];
rs_estimated = readcsv("../data/firing_rate_estimated.csv")[:];

QQplot(rs_estimated, x_true)
title!("Renewal Gamma")
xlabel!("data")
ylabel!("Maximum likelihood")

max_stars_repo_name: ganow/julia-point-process

max_issues_repo_name: ganow/julia-point-process

max_forks_repo_name: ganow/julia-point-process


text:
############################################################
# joAbstractDAparallelLinearOperator subtypes ##############
############################################################

############################################################
# joDALinearOperator #######################################

export joDALinearOperator, joDALinearOperatorException
export joDAdistributedLinearOperator, joDAdistributedLinearOperatorException
export joDAdistributingLinearOperator, joDAdistributingLinearOperatorException
export joDAgatheringLinearOperator, joDAgatheringLinearOperatorException

# type definition
"""
joDALinearOperator is glueing type & constructor

!!! Do not use it to create the operators

"""
struct joDALinearOperator{DDT<:Number,RDT<:Number,N} <: joAbstractDAparallelLinearOperator{DDT,RDT,N}
    name::String
    m::Integer
    n::Integer
    nvc::Integer
    fop::Function              # forward
    fop_T::Nullable{Function}  # transpose
    fop_A::Nullable{Function}  # adjoint
    fop_C::Nullable{Function}  # conj
    iop::Nullable{Function}
    iop_T::Nullable{Function}
    iop_A::Nullable{Function}
    iop_C::Nullable{Function}
    PAs_in::joPAsetup    # input distributor
    PAs_out::joPAsetup   # output distributor
    fclean::Bool         # clean input vector post forward
    rclean::Bool         # clean input vector post reverse
end
"""
joDAdistributedLinearOperator is glueing type & constructor

!!! Do not use it to create the operators

"""
struct joDAdistributedLinearOperator{DDT<:Number,RDT<:Number,N} <: joAbstractDAparallelLinearOperator{DDT,RDT,N}
    name::String
    m::Integer
    n::Integer
    nvc::Integer
    fop::Function              # forward
    fop_T::Nullable{Function}  # transpose
    fop_A::Nullable{Function}  # adjoint
    fop_C::Nullable{Function}  # conj
    iop::Nullable{Function}
    iop_T::Nullable{Function}
    iop_A::Nullable{Function}
    iop_C::Nullable{Function}
    PAs_in::joPAsetup   # input distributor
    PAs_out::joPAsetup  # output distributor
    fclean::Bool        # clean input vector post forward
    rclean::Bool        # clean input vector post reverse
end
"""
joDAdistributingLinearOperator is glueing type & constructor

!!! Do not use it to create the operators

"""
struct joDAdistributingLinearOperator{DDT<:Number,RDT<:Number,N} <: joAbstractDAparallelLinearOperator{DDT,RDT,N}
    name::String
    m::Integer
    n::Integer
    nvc::Integer
    fop::Function              # forward
    fop_T::Nullable{Function}  # transpose
    fop_A::Nullable{Function}  # adjoint
    fop_C::Nullable{Function}  # conj
    iop::Nullable{Function}
    iop_T::Nullable{Function}
    iop_A::Nullable{Function}
    iop_C::Nullable{Function}
    PAs_out::joPAsetup  # output distributor
    gclean::Bool        # clean input vector post gathering
end
"""
joDAgatheringLinearOperator is glueing type & constructor

!!! Do not use it to create the operators

"""
struct joDAgatheringLinearOperator{DDT<:Number,RDT<:Number,N} <: joAbstractDAparallelLinearOperator{DDT,RDT,N}
    name::String
    m::Integer
    n::Integer
    nvc::Integer
    fop::Function              # forward
    fop_T::Nullable{Function}  # transpose
    fop_A::Nullable{Function}  # adjoint
    fop_C::Nullable{Function}  # conj
    iop::Nullable{Function}
    iop_T::Nullable{Function}
    iop_A::Nullable{Function}
    iop_C::Nullable{Function}
    PAs_in::joPAsetup  # input distributor
    gclean::Bool       # clean input vector post gathering
end

# type exception
struct joDALinearOperatorException <: Exception
    msg :: String
end
struct joDAdistributedLinearOperatorException <: Exception
    msg :: String
end
struct joDAdistributingLinearOperatorException <: Exception
    msg :: String
end
struct joDAgatheringLinearOperatorException <: Exception
    msg :: String
end

# unions
const joDALinOpsUnion{D,R,N}=Union{joDALinearOperator{D,R,N},joDAdistributedLinearOperator{D,R,N}} where {D,R,N}
const joDAdistributeLinOpsUnion{D,R,N}=Union{joDAdistribute{D,R,N},joDAdistributingLinearOperator{D,R,N}} where {D,R,N}
const joDAgatherLinOpsUnion{D,R,N}=Union{joDAgather{D,R,N},joDAgatheringLinearOperator{D,R,N}} where {D,R,N}


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

# Deterministic equivalent spconstructs #
# ===================================== #
function DEP(stochasticprogram::StochasticProgram{2}, structure::DeterministicEquivalent; optimizer = nothing)
    # Ensure stochastic program has been generated at this point
    if deferred(stochasticprogram)
        generate!(stochasticprogram)
    end
    if optimizer == nothing
        return structure.model
    end
    structure_model = copy(structure.model)
    set_optimizer(structure_model, optimizer)
    return structure_model
end


text:
fluxlimiter(::Type{FluxLimiter{minmod}}, rf)      = max(0.,min(1.,rf))
fluxlimiter(::Type{FluxLimiter{superbee}}, rf)    = max(0.,min(1.,2rf),min(rf,2.))
fluxlimiter(::Type{FluxLimiter{vanAlbada}}, rf)   = (rf^2+rf)/(rf^2+1)
fluxlimiter(::Type{FluxLimiter{vanAlbada2}}, rf)  = (2rf)/(rf^2+1)
fluxlimiter(::Type{FluxLimiter{vanLeer}}, rf)     = (rf+abs(rf))/(abs(rf)+1)
fluxlimiter(::Type{FluxLimiter{ospre}}, rf)       = 1.5(rf^2+rf)/(rf^2+rf+1)
fluxlimiter(::Type{FluxLimiter{hquick}}, rf)      = 2.0(rf+abs(rf))/(rf+2)
fluxlimiter(::Type{FluxLimiter{koren}}, rf)       = max(0,min(2rf,min((1+2rf)/3,2)))

@inline function updater!(RiverPhysics, RiverArrays, River)
    # unpack
    Kr, kappa_a, h, m = RiverPhysics.Kr, RiverPhysics.kappa_a,  RiverPhysics.h, RiverPhysics.m
    L, r = RiverArrays.L, RiverArrays.r
    dx, z, x = River.Δx^2, River.z, River.x

    # constant values
    cte1 = -Kr*(kappa_a^m)/dx
    cte2 = h*m
    # update river length
    river_length!(L, River)
    @tturbo for i ∈ 2:length(x)
        r[i] = cte1 * L[i]^(cte2)
    end
end

function river_length!(L, River)
    x, z, dx = River.x, River.z, River.Δx^2
    @tturbo for i in 2:length(x)
        L[i] = √(dx  + (z[i]-z[i-1])^2)
    end
    @inbounds @fastmath for i in 2:length(x)
        L[i] += L[i-1]
    end
end

@inline function rfactor!(rf, sol, ilast)
    """
    Wikipedia definition, paper has a slightly different one
    """
    @inbounds for i in 1:ilast
        s1 =  sol[i+1] - sol[i]
        if s1 == 0.0
            rf[i] = 1.0
        elseif s1 <= 0.0
            rf[i] = 0.0
        else
            rf[i] = (sol[i+2] - sol[i+1]) / s1
        end
	end
    @inbounds rf[ilast] = 1.0
    @inbounds rf[ilast-1] = 1.0
end

function TVD!(River, sol0, RiverPhysics::KnickPointsParameters, RiverArrays::KnickPointsArrays, dt, ilast)
    # Add reference to paper of Govers and Schwangart

    n, φ = RiverPhysics.n, RiverPhysics.φ
    r, rf = RiverArrays.r, RiverArrays.rf
    sol, dx = River.z, River.Δx

    rfactor!(rf, sol, ilast)
    copyto!(sol0, sol)
    # cache some constant values to avoid redundance
    invdx = 1/dx
    invdxdt = dt*invdx
    half_dt = dt*0.5

    @fastmath for i in 2:ilast
        # cache
        @inbounds soli, solplus, solminus = sol0[i], sol0[i+1], sol0[i-1]
        @inbounds ri = r[i]
        ri_half = ri*0.5
        dtdx = half_dt*invdx*ri*ri
        # solve
        α₀, α₁ = ri_half * (1+sign(ri)), ri_half*(1-sign(ri))
        Frlo = α₀*soli + α₁*solplus    # f right low
        Fllo = α₀*solminus + α₁*soli   # f left  low
        Frhi = ri_half*(soli + solplus) - dtdx*(solplus-soli)
        Flhi = ri_half*(soli + solminus) - dtdx*(soli-solminus)
        @inbounds Fᵣ = Frlo+fluxlimiter(φ, rf[i]) * (Frhi - Frlo )
        @inbounds Fₗ = Fllo+fluxlimiter(φ, rf[i-1]) * (Flhi - Fllo )
        # non linear term
        @inbounds multiplier  = ((sol0[i] - sol0[i+1])*invdx)^(n-1) # |-> pow32 speeds up the whole solver by 60%
        # update solution
        @inbounds sol[i] = max(sol[i] - invdxdt*(Fᵣ-Fₗ) * multiplier, sol[i+1])
        # @inbounds sol[i] -= invdxdt*(Fᵣ-Fₗ) * multiplier
    end
end


text:
module TestDecisionTreeWrapper

include(joinpath("..", "fixture_learners.jl"))
using .FixtureLearners
nfcp = NumericFeatureClassification()

using FactCheck

importall Orchestra.Transformers.DecisionTreeWrapper
using DecisionTree

facts("DecisionTree learners") do
  context("PrunedTree gives same results as its backend") do
    # Predict with Orchestra learner
    learner = PrunedTree()
    orchestra_predictions = fit_and_transform!(learner, nfcp)

    # Predict with original backend learner
    srand(1)
    model = build_tree(nfcp.train_labels, nfcp.train_instances)
    model = prune_tree(model, 1.0)
    original_predictions = apply_tree(model, nfcp.test_instances)

    # Verify same predictions
    @fact orchestra_predictions => original_predictions
  end

  context("RandomForest gives same results as its backend") do
    # Predict with Orchestra learner
    learner = RandomForest()
    orchestra_predictions = fit_and_transform!(learner, nfcp)

    # Predict with original backend learner
    srand(1)
    model = build_forest(
      nfcp.train_labels,
      nfcp.train_instances,
      size(nfcp.train_instances, 2),
      10,
      0.7
    )
    original_predictions = apply_forest(model, nfcp.test_instances)

    # Verify same predictions
    @fact orchestra_predictions => original_predictions
  end

  context("DecisionStumpAdaboost gives same results as its backend") do
    # Predict with Orchestra learner
    learner = DecisionStumpAdaboost()
    orchestra_predictions = fit_and_transform!(learner, nfcp)

    # Predict with original backend learner
    srand(1)
    model, coeffs = build_adaboost_stumps(
      nfcp.train_labels,
      nfcp.train_instances,
      7
    )
    original_predictions = apply_adaboost_stumps(
      model, coeffs, nfcp.test_instances
    )

    # Verify same predictions
    @fact orchestra_predictions => original_predictions
  end

  context("RandomForest handles training-dependent options") do
    # Predict with Orchestra learner
    learner = RandomForest({:impl_options => {:num_subfeatures => 2}})
    orchestra_predictions = fit_and_transform!(learner, nfcp)

    # Verify RandomForest didn't die
    @fact 1 => 1
  end
end

end # module


text:
# run with: julia -p auto compare_integrators.jl

@everywhere begin

    include("integrands.jl")
    using ucbature
    using PyPlot

    function randomize_integrand(dim, integrand, xmin, xmax, difficulty)
        c = rand(dim)
        w = rand(dim)
        cd = c .* (difficulty / sum(c))
        rand_integrand(x) = integrand(x, c, w)
    end

    function run_integrator(dim, integrand, xmin, xmax, difficulty, nb_evals, integrator, nb_runs, nb_evals_gt)
        mae_sum = 0.0
        for _ = 1:nb_runs
            # randomize integrand
            rand_integrand = randomize_integrand(dim, integrand, xmin, xmax, difficulty)
            # estimate ground truth
            ground_truth = ucbature.mc(rand_integrand, xmin, xmax, nb_evals_gt)
            # compute mae sum
            mae_sum += abs(ground_truth - integrator(rand_integrand, xmin, xmax, nb_evals))
        end
        mae_sum / nb_runs
    end

    function run_plot(dim, integrand, xmin, xmax, difficulty, list_nb_evals, integrator, nb_runs, nb_evals_gt, integrator_name)
        integrals = [run_integrator(dim, integrand, xmin, xmax, difficulty, nb_evals, integrator, nb_runs, nb_evals_gt)
                     for nb_evals = list_nb_evals]
        plot(list_nb_evals, integrals, label=integrator_name)
    end

    function run_full(integrand, xmin, xmax, difficulty, list_nb_evals, nb_runs, nb_evals_gt, integrand_name, outdir)
        fig = figure()
        dim = length(xmin)
        println("computing $(integrand_name) $(dim)D...")
        # run integrators and plot results
        run_plot(dim, integrand, xmin, xmax, difficulty, list_nb_evals, ucbature.mc, nb_runs, nb_evals_gt, "mc")
        run_plot(dim, integrand, xmin, xmax, difficulty, list_nb_evals, ucbature.str, nb_runs, nb_evals_gt, "str")
        run_plot(dim, integrand, xmin, xmax, difficulty, list_nb_evals, ucbature.ucb, nb_runs, nb_evals_gt, "ucb")
        # finish plot
        xscale("log")
        yscale("log")
        grid("on")
        legend(loc="upper right")
        xlabel("nb evaluations")
        ylabel("mean absolute error")
        title("integration $(integrand_name) $(dim)D")
        savefig("$(outdir)/out_integration_$(integrand_name)_$(dim)D.png")
        close(fig)
    end

    # global parameters
    nb_runs = 10
    list_nb_evals = [2^expo for expo = 10:16]
    nb_evals_gt = 2^20
    outdir = "output"
    if ~isdir(outdir)
        mkdir(outdir)
    end

    # auxiliary function
    function aux_run(data)
        f, x0, x1, difficulty, integrand_name = data
        run_full(f, x0, x1, difficulty, list_nb_evals, nb_runs, nb_evals_gt, integrand_name, outdir)
    end

end # @everywhere

# parallel runs
aux_data = [# (integrands.sinus, fill(-8.0, 2), fill(8.0, 2), 1.0, "sinus"),
            # (integrands.sinusc, fill(-8.0, 2), fill(8.0, 2), 1.0, "sinusc"),
            (integrands.sphere, fill(-2.0, 2), fill(2.0, 2), 1.0, "sphere"),
            # (integrands.oscillatory, fill(-5.0, 2), fill(5.0, 2), 6.0, "oscillatory"),
            # (integrands.productpeak, fill(-4.0, 2), fill(4.0, 2), 18.0, "productpeak"),
            # (integrands.cornerpeak, fill(0.0, 2), fill(1.0, 2), 2.2, "cornerpeak"),
            (integrands.gaussian, fill(-4.0, 2), fill(4.0, 2), 15.2, "gaussian"),
            # (integrands.continuous, fill(-4.0, 2), fill(4.0, 2), 16.1, "continuous"),
            # (integrands.discontinuous, fill(-4.0, 2), fill(4.0, 2), 16.4, "discontinuous"),
           ]

map(aux_run, aux_data)

max_stars_repo_name: florianLepretre/ucbature

max_issues_repo_name: florianLepretre/ucbature

max_forks_repo_name: florianLepretre/ucbature


text:
using Documenter, ECC

makedocs(sitename="ECC",
         doctest=true)


text:
import Optim

include("MV.jl")
include("wannierize_utils.jl")

filename = "free"
read_amn = true #read $file.amn as input
read_eig = true #read the eig file (can be set to false when not disentangling)
do_write_amn = true #write $file.optimize.amn at the end
nfrozen = 1 #will freeze if either n <= nfrozen or eigenvalue in frozen window
frozen_window_low = -Inf
frozen_window_high = -Inf
ftol = 1e-20 #tolerance on spread
gtol = 1e-6 #tolerance on gradient
maxiter = 3000 #maximum optimization iterations
m = 100 #history size of BFGS

# expert/experimental features
do_normalize_phase = false # perform a global rotation by a phase factor at the end
do_randomize_gauge = false #randomize initial gauge
cluster_size = 1e-6 #will also freeze additional eigenvalues if the freezing cuts a cluster. Set to 0 to disable
only_r2 = false #only minimize sum_n <r^2>_n, not sum_n <r^2>_n - <r>_n^2

srand(1)

p = read_system(filename,read_amn,read_eig)

# local number of frozen bands
function local_nfrozen(p, nfrozen, i, j, k, win_low, win_high)
    nf = 0
    @assert win_low == -Inf #for now
    K = p.ijk_to_K[i,j,k]
    nf = count(a -> a < win_high, p.eigs[K,:])
    ret = max(nf,nfrozen)
    @assert ret <= p.nwannier

    # if cluster of eigenvalues, take them all
    @assert issorted(p.eigs[K,:])
    while p.eigs[K,ret+1] < p.eigs[k,ret]+cluster_size
        ret = ret + 1
    end

    return ret
end

function local_frozen_sets(p, nfrozen, i, j, k, win_low, win_high)
    nf = 0
    K = p.ijk_to_K[i,j,k]

    l_frozen = (p.eigs[K,:] .>= win_low).&(p.eigs[K,:] .<= win_high)
    l_not_frozen = .!l_frozen

    # if cluster of eigenvalues and nfrozen > 0, take them all
    if nfrozen > 0
        @assert issorted(p.eigs[K,:])
        while nfrozen < p.nwannier && p.eigs[K,nfrozen+1] < p.eigs[k,nfrozen]+cluster_size
            nfrozen = nfrozen + 1
        end

        l_frozen[1:nfrozen] = true
        l_not_frozen = .!l_frozen
    end

    @assert count(l_frozen) <= p.nwannier
    return l_frozen, l_not_frozen
end
# There are three formats: A, (X,Y) and XY stored contiguously in memory
# A is the format used in the rest of the code, XY is the format used in the optimizer, (X,Y) is intermediate
# A is nb x nw
# X is nw x nw, Y is nb x nw and has first block equal to [I 0]

function XY_to_A(p,X,Y)
    A = zeros(Complex128,p.nband,p.nwannier,p.N1,p.N2,p.N3)
    for i=1:p.N1,j=1:p.N2,k=1:p.N3
        l_frozen, l_not_frozen = local_frozen_sets(p, nfrozen, i, j, k, frozen_window_low, frozen_window_high)
        lnf = count(l_frozen)
        @assert Y[:,:,i,j,k]'Y[:,:,i,j,k] ≈ eye(p.nwannier)
        @assert X[:,:,i,j,k]'X[:,:,i,j,k] ≈ eye(p.nwannier)
        @assert Y[l_frozen,1:lnf,i,j,k] ≈ eye(lnf)
        @assert vecnorm(Y[l_not_frozen,1:lnf,i,j,k]) ≈ 0
        @assert vecnorm(Y[l_frozen,lnf+1:end,i,j,k]) ≈ 0
        A[:,:,i,j,k] = Y[:,:,i,j,k]*X[:,:,i,j,k]
        # @assert normalize_and_freeze(A[:,:,i,j,k],lnf) ≈ A[:,:,i,j,k] rtol=1e-4
    end
    A
end

function A_to_XY(p,A)
    X = zeros(Complex128,p.nwannier,p.nwannier,p.N1,p.N2,p.N3)
    Y = zeros(Complex128,p.nband,p.nwannier,p.N1,p.N2,p.N3)
    for i=1:p.N1,j=1:p.N2,k=1:p.N3
        l_frozen, l_not_frozen = local_frozen_sets(p, nfrozen, i, j, k, frozen_window_low, frozen_window_high)
        lnf = count(l_frozen)
        Afrozen = normalize_and_freeze(A[:,:,i,j,k], l_frozen, l_not_frozen)
        Af = Afrozen[l_frozen,:]
        Ar = Afrozen[l_not_frozen,:]

        #determine Y
        if lnf != p.nwannier
            proj = Ar*Ar'
            proj = Hermitian((proj+proj')/2)
            D,V = eig(proj) #sorted by increasing eigenvalue
        end
        Y[l_frozen,1:lnf,i,j,k] = eye(lnf)
        if lnf != p.nwannier
            Y[l_not_frozen,lnf+1:end,i,j,k] = V[:,end-p.nwannier+lnf+1:end]
        end

        #determine X
        Xleft, S, Xright = svd(Y[:,:,i,j,k]'*Afrozen)
        X[:,:,i,j,k] = Xleft*Xright'

        @assert Y[:,:,i,j,k]'Y[:,:,i,j,k] ≈ eye(p.nwannier)
        @assert X[:,:,i,j,k]'X[:,:,i,j,k] ≈ eye(p.nwannier)
        @assert Y[l_frozen,1:lnf,i,j,k] ≈ eye(lnf)
        @assert vecnorm(Y[l_not_frozen,1:lnf,i,j,k]) ≈ 0
        @assert vecnorm(Y[l_frozen,lnf+1:end,i,j,k]) ≈ 0
        @assert Y[:,:,i,j,k]*X[:,:,i,j,k] ≈ Afrozen
    end
    X,Y
end

function XY_to_XY(p,XY) #XY to (X,Y)
    X = zeros(Complex128,p.nwannier,p.nwannier,p.N1,p.N2,p.N3)
    Y = zeros(Complex128,p.nband,p.nwannier,p.N1,p.N2,p.N3)
    for i=1:p.N1,j=1:p.N2,k=1:p.N3
        XYk = XY[:,i,j,k]
        X[:,:,i,j,k] = reshape(XYk[1:p.nwannier*p.nwannier], (p.nwannier, p.nwannier))
        Y[:,:,i,j,k] = reshape(XYk[p.nwannier*p.nwannier+1:end], (p.nband, p.nwannier))
    end
    X,Y
end
function obj(p,X,Y)
    A = XY_to_A(p,X,Y)

    res = omega(p,A,true,only_r2)
    func = res.Ωtot
    grad = res.gradient

    gradX = zeros(X)
    gradY = zeros(Y)
    for i=1:p.N1,j=1:p.N2,k=1:p.N3
        l_frozen, l_not_frozen = local_frozen_sets(p, nfrozen, i, j, k, frozen_window_low, frozen_window_high)
        lnf = count(l_frozen)
        gradX[:,:,i,j,k] = Y[:,:,i,j,k]'*grad[:,:,i,j,k]
        gradY[:,:,i,j,k] = grad[:,:,i,j,k]*X[:,:,i,j,k]'

        gradY[l_frozen,:,i,j,k] = 0.
        gradY[:,1:lnf,i,j,k] = 0.

        # to compute the projected gradient: redundant, taken care of by the optimizer
        # function proj_stiefel(G,X)
        #     G .- X*((X'G .+ G'X)./2)
        # end
        # gradX[:,:,i,j,k] = proj_stiefel(gradX[:,:,i,j,k],X[:,:,i,j,k])
        # gradY[:,:,i,j,k] = proj_stiefel(gradY[:,:,i,j,k],Y[:,:,i,j,k])
    end
    return func, gradX,gradY, res
end

function minimize(p,A)
    # initial X,Y
    X0,Y0 = A_to_XY(p,A)

    if do_randomize_gauge
        if read_amn
            error("don't set do_randomize_gauge and read_amn")
        end
        X0 = zeros(Complex128,p.nwannier,p.nwannier,p.N1,p.N2,p.N3)
        Y0 = zeros(Complex128,p.nband,p.nwannier,p.N1,p.N2,p.N3)
        for i=1:p.N1,j=1:p.N2,k=1:p.N3
            l_frozen, l_not_frozen = local_frozen_sets(p, nfrozen, i, j, k, frozen_window_low, frozen_window_high)
            lnf = count(l_frozen)
            X0[:,:,i,j,k] = normalize_matrix(randn(p.nwannier,p.nwannier) + im*randn(p.nwannier,p.nwannier))
            Y0[l_frozen,1:lnf,i,j,k] = eye(lnf)
            Y0[l_not_frozen, lnf+1:p.nwannier,i,j,k] = normalize_matrix(randn(p.nband-lnf,p.nwannier-lnf) + im*randn(p.nband-lnf,p.nwannier-lnf))
        end
    end

    M = p.nwannier*p.nwannier+p.nband*p.nwannier
    XY0 = zeros(Complex128, M, p.N1, p.N2, p.N3)
    for i=1:p.N1,j=1:p.N2,k=1:p.N3
        XY0[:,i,j,k] = vcat(vec(X0[:,:,i,j,k]),vec(Y0[:,:,i,j,k]))
    end

    # We have three formats:
    # (X,Y): Ntot x nw x nw, Ntot x nb x nw
    # A: ntot x nb x nw
    # XY: (nw*nw + nb*nw) x Ntot
    function fg!(G, XY)
        @assert size(G) == size(XY)
        X,Y = XY_to_XY(p,XY)

        f, gradX, gradY, res = obj(p,X,Y)

        for i=1:p.N1,j=1:p.N2,k=1:p.N3
            G[:,i,j,k] = vcat(vec(gradX[:,:,i,j,k]),vec(gradY[:,:,i,j,k]))
        end
        return f
    end

    function f(XY)
        return fg!(similar(XY),XY)
    end
    function g!(g, XY)
        fg!(g,XY)
        return g
    end

    # need QR orthogonalization rather than SVD to preserve the sparsity structure of Y
    XYkManif = Optim.ProductManifold(Optim.Stiefel_CholQR(), Optim.Stiefel_CholQR(), (p.nwannier, p.nwannier), (p.nband, p.nwannier))
    XYManif = Optim.PowerManifold(XYkManif, (M,), (p.N1,p.N2,p.N3))

    # stepsize_mult = 1
    # step = 0.5/(4*8*p.wb)*(p.N1*p.N2*p.N3)*stepsize_mult
    # ls = LineSearches.Static(step)
    ls = LineSearches.HagerZhang()
    # ls = LineSearches.BackTracking()
    # meth = Optim.GradientDescent
    # meth = Optim.ConjugateGradient
    meth = Optim.LBFGS
    res = Optim.optimize(f,g!,XY0, meth(manifold=XYManif, linesearch=ls,m=m), Optim.Options(show_trace=true,iterations=maxiter,f_tol=ftol, g_tol=gtol, allow_f_increases=true))
    display(res)
    XYmin = Optim.minimizer(res)

    Xmin,Ymin = XY_to_XY(p, XYmin)
    Amin = XY_to_A(p,Xmin,Ymin)
end

if read_amn
    A0 = copy(p.A)
else
    A0 = randn(size(p.A)) + im*randn(size(p.A))
end

for i=1:p.N1,j=1:p.N2,k=1:p.N3
    l_frozen, l_not_frozen = local_frozen_sets(p, nfrozen, i, j, k, frozen_window_low, frozen_window_high)
    A0[:,:,i,j,k] = normalize_and_freeze(A0[:,:,i,j,k],l_frozen,l_not_frozen)
end

A = minimize(p,A0)

# fix global phase
if do_normalize_phase
    for i=1:nwannier
        imax = indmax(abs.(A[:,i,1,1,1]))
        @assert abs(A[imax,i,1,1,1]) > 1e-2
        A[:,i,:,:,:] *= conj(A[imax,i,1,1,1]/abs(A[imax,i,1,1,1]))
    end
end

if do_write_amn
    write_amn(p,A,"$(p.filename).optimized")
end

max_stars_repo_name: VVitale/wannier

max_issues_repo_name: VVitale/wannier

max_forks_repo_name: VVitale/wannier


text:
Base.@kwdef mutable struct FileDialogDisplayProperties <: AbstractDisplayProperties
    caption::String
    action::String
    position::ImVec2 = ImVec2(0,0)
    width::Cfloat = Cfloat(320)
    height::Cfloat = Cfloat(200)
end

function get_caption(property::FileDialogDisplayProperties)
    property.caption
end

function get_action(property::FileDialogDisplayProperties)
    property.action
end

function get_position(property::FileDialogDisplayProperties)
    property.position
end

function get_width(property::FileDialogDisplayProperties)
    property.width
end

function get_height(property::FileDialogDisplayProperties)
    property.height
end

function set_caption!(property::FileDialogDisplayProperties, caption::String)
    property.caption = caption
end

function set_action!(property::FileDialogDisplayProperties, action::String)
    property.action = action
end

function set_position!(property::FileDialogDisplayProperties, position::ImVec2)
    property.position = position
end

function set_width!(property::FileDialogDisplayProperties, width::Cfloat)
    property.width = width
end

function set_height!(property::FileDialogDisplayProperties, height::Cfloat)
    property.height = height
end


text:
function pre_divide(points1, points2, K1, K2)
    p1 = Vector{SVector{2, Float64}}(undef, length(points1))
    p2 = Vector{SVector{2, Float64}}(undef, length(points1))

    @inbounds fx1, fy1, px1, py1 = K1[1, 1], K1[2, 2], K1[1, 3], K1[2, 3]
    @inbounds fx2, fy2, px2, py2 = K2[1, 1], K2[2, 2], K2[1, 3], K2[2, 3]
    @inbounds @simd for i in 1:length(points1)
        p1[i] = SVector{2, Float64}(
            (points1[i][1] - px1) / fx1,
            (points1[i][2] - py1) / fy1)
        p2[i] = SVector{2, Float64}(
            (points2[i][1] - px2) / fx2,
            (points2[i][2] - py2) / fy2)
    end
    p1, p2
end

"""
```julia
five_point(points1, points2, K1, K2; max_repr_error = 1.0)
```

Compute Essential matrix and recover pose from it for a given set of points.
This function accepts five points, if you have more of them,
consider using [`five_point_ransac`](@ref) version.

# Arguments:

- `points1`:
    Pixel coordinates of the matched points in `(x, y)` format
    in the first image.
- `points2`:
    Pixel coordinates of the matched points in `(x, y)` format
    in the second image.
- `K1`: Intrinsic matrix for the first set of points.
- `K2`: Intrinsic matrix for the second set of points.
- `max_repr_error`: Maximum allowed reprojection error for a point to be
    considered as inlier. Default is `1.0`.

# Returns:

`n_inliers, (E, P, inliers, repr_error)`:

- number of inliers
- tuple: essential matrix, pose matrix, boolean vector of inliers, error value.

!!! note

    Pose matrix transforms points from the first camera to the second camera.
"""
function five_point(points1, points2, K1, K2; max_repr_error = 1.0)
    p1, p2 = pre_divide(points1, points2, K1, K2)
    candidates = five_point_candidates(p1, p2)
    select_candidates(candidates, points1, points2, K1, K2; max_repr_error)
end

"""
```julia
five_point_ransac(
    points1, points2, K1, K2; max_repr_error = 1.0, ransac_kwargs...)
```

Compute Essential matrix and recover pose from it using RANSAC scheme.

# Arguments:

- `points1`:
    Pixel coordinates of the matched points in `(x, y)` format
    in the first image.
- `points2`:
    Pixel coordinates of the matched points in `(x, y)` format
    in the second image.
- `K1`: Intrinsic matrix for the first set of points.
- `K2`: Intrinsic matrix for the second set of points.
- `max_repr_error`: Maximum allowed reprojection error for a point to be
    considered as inlier. Default is `1.0`.
- `ransac_kwargs...`: Keyword arguments passed to [`ransac`](@ref).

# Returns:

`n_inliers, (E, P, inliers, repr_error)`:

- number of inliers
- tuple: essential matrix, pose matrix, boolean vector of inliers, error value.

!!! note

    Pose matrix transforms points from the first camera to the second camera.
"""
@inline five_point_ransac(pixels1, pixels2, K1, K2; max_repr_error = 1.0, ransac_kwargs...) =
    five_point_ransac(pixels1, pixels2, K1, K2, GEEV4x4Cache(); max_repr_error, ransac_kwargs...)

function five_point_ransac(
    pixels1, pixels2, K1, K2, cache::GEEV4x4Cache;
    max_repr_error = 1.0, ransac_kwargs...,
)
    pd1, pd2 = pre_divide(pixels1, pixels2, K1, K2)
    five_point_ransac(
        pixels1, pixels2, pd1, pd2, K1, K2, cache;
        max_repr_error, ransac_kwargs...)
end

function five_point_ransac(
    pixels1, pixels2, pd1, pd2, K1, K2, cache::GEEV4x4Cache;
    max_repr_error = 1.0, ransac_kwargs...,
)
    sample_selection(sample_ids) = (pd1[sample_ids], pd2[sample_ids])
    rank(models; sample_ids) = select_candidates(
        models, pixels1, pixels2, K1, K2, cache; max_repr_error, sample_ids)
    ransac(
        sample_selection, five_point_candidates, rank,
        length(pd1), 5; ransac_kwargs...)
end

"""
```julia
essential_ransac(pd1, pd2, focal_sum; threshold = 1.0, ransac_kwargs...)
```

Compute Essential matrix using the RANASC scheme.

# Arguments:

- `pd1`:
    Pixel coordinates **predivided** by `K^-1` of the matched points
    in `(x, y)` format in the first image.
- `pd2`:
    Pixel coordinates **predivided** by `K^-1` of the matched points
    in `(x, y)` format in the second image.
- `focal_sum`: `fx + fy`.
- `threshold`: Maximum error for the epipolar constraint. Default is `1.0`.
- `ransac_kwargs...`: Keyword arguments passed to [`ransac`](@ref).

# Returns:

`n_inliers, (E, inliers, repr_error)`:

- number of inliers
- tuple: essential matrix, boolean vector of inliers, error value.
"""
function essential_ransac(pd1, pd2, focal_sum; threshold = 1.0, ransac_kwargs...)
    threshold /= focal_sum * 0.5
    sample_selection(sample_ids) = (pd1[sample_ids], pd2[sample_ids])
    rank(models; sample_ids) = select_candidate(models, pd1, pd2; threshold)
    ransac(
        sample_selection, five_point_candidates, rank,
        length(pd1), 5; ransac_kwargs...)
end

struct FivePointCache
    F::Matrix{Float64}
    M::Matrix{Float64}
    Z::Matrix{Float64}
end

function FivePointCache(n_points)
    F = Matrix{Float64}(undef, n_points, 9)
    M = Matrix{Float64}(undef, 10, 20)
    Z = Matrix{Float64}(undef, 10, 10)
    FivePointCache(F, M, Z)
end

"""
Compute essential matrix using Five-Point algorithm.

# Arguments:

- `p1`: Vector of points, pre-divided by `K` matrix in `(x, y)` format.
- `p2`: Vector of points, pre-divided by `K` matrix in `(x, y)` format.

# Returns:

Vector of `3x3` essential matrix candidates.
"""
@inline five_point_candidates(p1, p2) =
    five_point_candidates!(FivePointCache(length(p1)), p1, p2)

function five_point_candidates!(cache::FivePointCache, p1, p2)
    @polyvar z

    V = null_space!(cache.F, p1, p2)
    E1 = SMatrix{3, 3, Float64, 9}(V[:, 6,]...)'
    E2 = SMatrix{3, 3, Float64, 9}(V[:, 7]...)'
    E3 = SMatrix{3, 3, Float64, 9}(V[:, 8]...)'
    E4 = SMatrix{3, 3, Float64, 9}(V[:, 9]...)'
    rref_M = compute_rref!(cache.M, E1, E2, E3, E4)

    B11, B12, B13 = to_polynoms(@view(rref_M[5, 11:20]), @view(rref_M[6, 11:20]), z)
    B21, B22, B23 = to_polynoms(@view(rref_M[7, 11:20]), @view(rref_M[8, 11:20]), z)
    B31, B32, B33 = to_polynoms(@view(rref_M[9, 11:20]), @view(rref_M[10, 11:20]), z)

    # Calculate determinant.
    P1 = B23 * B12 - B13 * B22
    P2 = B13 * B21 - B23 * B11
    P3 = B11 * B22 - B12 * B21
    det_B = P1 * B31 + P2 * B32 + P3 * B33

    # Normalize the coefficient of the highest order.
    coefficient(det_B, z^10) ≉ 0 && (det_B /= coefficient(det_B, z^10);)

    # Extract real roots of polynomial with companion matrix.
    fill!(cache.Z, 0.0)
    sub_Z = @view(cache.Z[1:9, 2:10])
    sub_Z[diagind(sub_Z)] .= 1.0
    @inbounds cache.Z[10, :] .= -coefficients(det_B)[end:-1:2]
    sol_z = Float64[real(s) for s in eigvals!(cache.Z) if isreal(s)]
    isempty(sol_z) && return SMatrix{3, 3, Float64, 9}[]

    # Compute x & y.
    z6 = Matrix{Float64}(undef, length(sol_z), 7)
    z7 = Matrix{Float64}(undef, length(sol_z), 8)
    @inbounds z6[:, 7] .= 1.0
    @inbounds z7[:, 8] .= 1.0
    @inbounds z7[:, 1] .= sol_z.^7
    @inbounds @simd for i in 1:6
        v = sol_z .^ i
        z6[:, 7 - i] .= v
        z7[:, 8 - i] .= v
    end

    P1z = z7 * coefficients(P1)
    P2z = z7 * coefficients(P2)
    P3z = z6 * coefficients(P3)

    sol_x = P1z ./ P3z
    sol_y = P2z ./ P3z
    candidates = Vector{SMatrix{3, 3, Float64, 9}}(undef, length(sol_x))
    @inbounds @simd for i in 1:length(sol_x)
        candidates[i] = @. sol_x[i] * E1 + sol_y[i] * E2 + sol_z[i] * E3 + E4
    end
    candidates
end

function select_candidate(candidates, pd1, pd2; threshold)
    best_n_inliers = 0
    best_error = maxintfloat()
    best_inliers = fill(true, length(pd1))

    E_res = SMatrix{3, 3, Float64, 9}(I)
    inliers = fill(false, length(pd1))
    for E in candidates
        fill!(inliers, false)
        n_inliers, avg_error = compute_essential_error!(
            inliers, pd1, pd2, E, threshold)

        n_inliers < 5 && continue
        n_inliers == best_n_inliers && best_error ≤ avg_error && continue

        best_error = avg_error
        best_n_inliers = n_inliers
        copy!(best_inliers, inliers)
        E_res = E
    end

    best_n_inliers, (E_res, best_inliers, best_error)
end

"""
Test candidates for the essential matrix and return the best candidate.
"""
function select_candidates(
    candidates, pixels1, pixels2, K1, K2, cache::GEEV4x4Cache;
    max_repr_error = 1.0, sample_ids,
)
    best_n_inliers = 0
    best_repr_error = maxintfloat()
    best_inliers = fill(false, length(pixels1))

    inliers = fill(true, length(pixels1))
    sample_inliers = fill(true, length(sample_ids))
    @inbounds sample_px1, sample_px2 = pixels1[sample_ids], pixels2[sample_ids]

    E_res = SMatrix{3, 3, Float64, 9}(I)
    P_res = SMatrix{3, 4, Float64, 12}(I)
    P_ref = SMatrix{3, 4, Float64, 12}(I)

    for E in candidates
        for P in compute_projections(E)
            sample_n_inliers, _ = chirality_test!(
                inliers, sample_px1, sample_px2,
                P_ref, P, K1, K2, cache; max_repr_error)
            sample_n_inliers < length(sample_ids) && continue

            fill!(inliers, true)
            n_inliers, repr_error = chirality_test!(
                inliers, pixels1, pixels2, P_ref, P,
                K1, K2, cache; max_repr_error)

            n_inliers < 5 && continue
            n_inliers == best_n_inliers &&
                best_repr_error ≤ repr_error && continue

            best_repr_error = repr_error
            best_n_inliers = n_inliers
            copy!(best_inliers, inliers)
            E_res = E
            P_res = P
        end
    end
    best_n_inliers, (E_res, P_res, best_inliers, best_repr_error)
end

@inline recover_pose(E, pixels1, pixels2, K1, K2; threshold = 1.0) =
    recover_pose(E, pixels1, pixels2, K1, K2, GEEV4x4Cache(); threshold)

function recover_pose(
    E, pixels1, pixels2, K1, K2, cache::GEEV4x4Cache; threshold = 1.0,
)
    best_n_inliers = 0
    best_error = maxintfloat()
    best_inliers = fill(true, length(pixels1))
    inliers = fill(true, length(pixels1))

    P_res = SMatrix{3, 4, Float64, 12}(I)
    P_ref = SMatrix{3, 4, Float64, 12}(I)

    for P in compute_projections(E)
        fill!(inliers, true)
        n_inliers, repr_error = chirality_test!(
            inliers, pixels1, pixels2, P_ref, P, K1, K2, cache;
            max_repr_error=threshold)

        n_inliers < 5 && continue
        n_inliers == best_n_inliers && threshold ≤ repr_error && continue

        best_error = repr_error
        best_n_inliers = n_inliers
        copy!(best_inliers, inliers)
        P_res = P
    end
    best_n_inliers, P_res, best_inliers, best_error
end


text:
module OnlineNMF

using DelimitedFiles:
    writedlm, readdlm
using Statistics:
    mean, var
using LinearAlgebra:
    Diagonal, lu!, qr!, svd, dot, norm, eigvecs
using Random:
    randperm
using ProgressMeter:
	Progress, next!
using ArgParse:
    ArgParseSettings, parse_args, @add_arg_table
using StatsBase:
	percentile
using DataFrames:
	DataFrame
using GLM:
	glm, coef, IdentityLink, @formula
using Distributions:
	Gamma, ccdf, Chisq
using CodecZstd:
	ZstdCompressorStream, ZstdDecompressorStream

export output, common_parse_commandline, csv2bin, mu, sinmf, ntd

include("Utils.jl")
include("csv2bin.jl")
include("mu.jl")
include("sinmf.jl")
include("ntd.jl")

end


text:
#!/usr/bin/env julia

mutable struct Instance
    num_depots::Int
    num_trips::Int
    depot_capacity::Vector{Int}
    adj_matrix::Matrix{Int}

    dh_forward::Vector{Vector{@NamedTuple{succ::Int, cost::Int}}}
    dh_backward::Vector{Vector{@NamedTuple{pred::Int, cost::Int}}}

    function Instance(num_depots, num_trips)
        L = num_depots + num_trips
        new(
            num_depots,
            num_trips,
            zeros(Int, num_depots),
            fill(-1, L, L),
            Vector{Vector{@NamedTuple{succ::Int, cost::Int}}}(undef, num_trips),
            Vector{Vector{@NamedTuple{pred::Int, cost::Int}}}(undef, num_trips)
        )
    end
end

# Helper functions
pullout_cost(inst::Instance, depot_id::Int, task_id::Int) = inst.adj_matrix[depot_id, inst.num_depots+task_id]

pullin_cost(inst::Instance, depot_id::Int, task_id::Int) = inst.adj_matrix[inst.num_depots+task_id, depot_id]

deadhead_cost(inst::Instance, pred_id::Int, succ_id::Int) = inst.adj_matrix[inst.num_depots+pred_id, inst.num_depots+succ_id]

function read_huisman_instance(fname::String)
    open(fname, "r") do fid
        header = split(readline(fid))

        num_depots = parse(Int, header[1])
        num_trips = parse(Int, header[2])

        inst = Instance(num_depots, num_trips)
        inst.depot_capacity = parse.(Int, header[3:end])

        row_idx = 1
        while !eof(fid)
            inst.adj_matrix[row_idx, :] = parse.(Int, split(readline(fid)))
            row_idx += 1
        end

        _create_deadhead_cache(inst)

        return inst
    end
end

function _create_deadhead_cache(inst::Instance)
    for i in 1:inst.num_trips
        fwd = Vector{@NamedTuple{succ::Int, cost::Int}}(undef, 0)
        bwd = Vector{@NamedTuple{pred::Int, cost::Int}}(undef, 0)

        for j in 1:inst.num_trips
            cost = deadhead_cost(inst, i, j)
            if cost != -1
                push!(fwd, (; :succ => j, :cost => cost))
            end

            cost = deadhead_cost(inst, j, i)
            if cost != -1
                push!(bwd, (; :pred => j, :cost => cost))
            end
        end

        inst.dh_forward[i] = fwd
        inst.dh_backward[i] = bwd
    end
end

max_stars_repo_name: afkummer/mdvsp-colgen-study

max_issues_repo_name: afkummer/mdvsp-colgen-study

max_forks_repo_name: afkummer/mdvsp-colgen-study


text:
# Julia wrapper for header: /usr/local/include/cpl_progress.h
# Automatically generated using Clang.jl wrap_c, version 0.0.0

"""
    GDALDummyProgress(double,
                      const char *,
                      void *) -> int

Stub progress function.
"""
function GDALDummyProgress(arg1::Cdouble, arg2, arg3)
    ccall((:GDALDummyProgress, libgdal), Cint, (Cdouble, Cstring, Ptr{Void}), arg1, arg2, arg3)
end

"""
    GDALTermProgress(double dfComplete,
                     const char * pszMessage,
                     void * pProgressArg) -> int

Simple progress report to terminal.

### Parameters
* **dfComplete**: completion ratio from 0.0 to 1.0.
* **pszMessage**: optional message.
* **pProgressArg**: ignored callback data argument.

### Returns
Always returns TRUE indicating the process should continue.
"""
function GDALTermProgress(arg1::Cdouble, arg2, arg3)
    ccall((:GDALTermProgress, libgdal), Cint, (Cdouble, Cstring, Ptr{Void}), arg1, arg2, arg3)
end

"""
    GDALScaledProgress(double dfComplete,
                       const char * pszMessage,
                       void * pData) -> int

Scaled progress transformer.
"""
function GDALScaledProgress(arg1::Cdouble, arg2, arg3)
    ccall((:GDALScaledProgress, libgdal), Cint, (Cdouble, Cstring, Ptr{Void}), arg1, arg2, arg3)
end

"""
    GDALCreateScaledProgress(double dfMin,
                             double dfMax,
                             GDALProgressFunc pfnProgress,
                             void * pData) -> void *

Create scaled progress transformer.

### Parameters
* **dfMin**: the value to which 0.0 in the sub operation is mapped.
* **dfMax**: the value to which 1.0 is the sub operation is mapped.
* **pfnProgress**: the overall progress function.
* **pData**: the overall progress function callback data.

### Returns
pointer to pass as pProgressArg to sub functions. Should be freed with GDALDestroyScaledProgress().
"""
function GDALCreateScaledProgress(arg1::Cdouble, arg2::Cdouble, arg3::GDALProgressFunc, arg4)
    ccall((:GDALCreateScaledProgress, libgdal), Ptr{Void}, (Cdouble, Cdouble, GDALProgressFunc, Ptr{Void}), arg1, arg2, arg3, arg4)
end

"""
    GDALDestroyScaledProgress(void * pData) -> void

Cleanup scaled progress handle.

### Parameters
* **pData**: scaled progress handle returned by GDALCreateScaledProgress().
"""
function GDALDestroyScaledProgress(arg1)
    ccall((:GDALDestroyScaledProgress, libgdal), Void, (Ptr{Void},), arg1)
end


text:
struct LinearFaltinsen78
    a::Float64
    h::Float64
    μ::Float64
    ω::Float64
    A::Float64
    g::Float64
    N::Int
    Aₙ::Vector{Float64}
    Bₙ::Vector{Float64}
    Cₙ::Vector{Float64}
    Dₙ::Vector{Float64}
    ωₙ²::Vector{Float64}
    Kₙ::Vector{Float64}
end

function LinearFaltinsen78(N, ω, a, h, μ; g=9.8, A=1.0)

    Kₙ = zeros(N+1)
    Aₙ = zeros(N+1)
    B\ = zeros(N+1)
    C = zeros(N+1)
    D = zeros(N+1)
    ωₙ = zeros(N+1)

    ω² = ω^2
    μ² = μ^2

    n = 0:N
    npia = @. (2n+1)*π / 2a

    ωₙ² = @.  g * npia * tanh( npia * h)
    Kₙ = @. ω*A/cosh(npia*h) * 2/a * (1/npia)^2 * (-1)^n

    Δω = @. (ωₙ² - ω²)
    den = @. Δω^2 + μ²*ω²

    Cₙ = @. Kₙ * (ω * Δω - μ²*ω) / den
    Dₙ = @. Kₙ * (Δω*μ + μ*ω²) / den

    Aₙ = @. -Cₙ - Kₙ*ω
    Bₙ = @. 1 / sqrt(ωₙ² - μ²/4) * (-μ/2*Aₙ - ω*Dₙ)

    return LinearFaltinsen78(a, h, μ, ω, A, g, N, Aₙ, Bₙ, Cₙ, Dₙ, ωₙ², Kₙ)
end

function Tₙ(f::LinearFaltinsen78, n, t)
    m = n-1
    μ = f.μ
    μ² = μ^2
    ωₙ² = f.ωₙ²[m]
    Aₙ = f.Aₙ[m]
    Bₙ = f.Bₙ[m]
    Cₙ = f.Cₙ[m]
    Dₙ = f.Dₙ[m]

    ω = f.ω

    ω₁ = sqrt(ωₙ²[n-1] - μ²/4)

    return exp(-μ/(2t)) * ( Aₙ * cos(ω₁*t) + Bₙ*sin(ω₁*t) ) +
        Cₙ * cos(ω*t) + Dₙ*cos(ω*t)
end

function potential₁(f::LinearFaltinse78, x, y, t)
    N = f.N
    ϕ₁ = 0.0
    ω = f.ω
    a = f.a
    h = f.h

    for n in 0:N
        npia = (2n+1)*π / 2a
        T = Tₙ(f, n, t)
        ϕ₁ = ϕ₁ + T * sin(npia*x) * cosh(npia*(y+h))
    end
    return ϕ₁

function potential(f::LinearFaltinse78, x, y, t)

    ϕ₁ = potential₁(f, x, y, t)
    ω = f.ω
    A = f.A
    ϕᶜ = A*x*cos(ω*t)

    return ϕ₁ + ϕᶜ

end


text:
const MOI = MathOptInterface

struct MOIGeneralProblem <: MOI.AbstractNLPEvaluator
    num_var::Int                 # number of decision variables
    num_con::Int                 # number of constraints
    primal_bounds
    constraint_bounds
end

struct MOIProblem <: MOI.AbstractNLPEvaluator
    num_var::Int                 # number of decision variables
    num_con::Int                 # number of constraints
    primal_bounds
    constraint_bounds
    prob
end

function moi_problem(prob)
    return MOIProblem(
        prob.num_var,
        prob.num_con,
        primal_bounds(prob),
        constraint_bounds(prob),
        prob)
end

pack(X, U, prob::MOIProblem) = pack(X, U, prob.prob)
unpack(Z, prob::MOIProblem) = unpack(Z, prob.prob)

primal_bounds(prob::MOI.AbstractNLPEvaluator) = prob.primal_bounds

constraint_bounds(prob::MOI.AbstractNLPEvaluator) = prob.constraint_bounds

function MOI.eval_objective(prob::MOIProblem, x)
    return eval_objective(prob.prob, x)
end

function MOI.eval_objective_gradient(prob::MOIProblem, grad_f, x)
    eval_objective_gradient!(grad_f, x, prob.prob)
end

function MOI.eval_constraint(prob::MOIProblem, g, x)
    eval_constraint!(g, x, prob.prob)
    return nothing
end

function MOI.eval_constraint_jacobian(prob::MOIProblem, jac, x)
    eval_constraint_jacobian!(jac, x, prob.prob)
    return nothing
end

function sparsity_jacobian(prob::MOIProblem)
    sparsity_jacobian(prob.prob)
end

MOI.features_available(prob::MOI.AbstractNLPEvaluator) = [:Grad, :Jac]
MOI.initialize(prob::MOI.AbstractNLPEvaluator, features) = nothing
MOI.jacobian_structure(prob::MOI.AbstractNLPEvaluator) = sparsity_jacobian(prob)
MOI.hessian_lagrangian_structure(prob::MOI.AbstractNLPEvaluator) = []
MOI.eval_hessian_lagrangian(prob::MOI.AbstractNLPEvaluator, H, x, σ, μ) = nothing

function solve(prob::MOI.AbstractNLPEvaluator, x0;
        tol = 1.0e-3,
        c_tol = 1.0e-2,
        max_iter = 1000,
        nlp = :ipopt,
        time_limit = 120,
        mipl = 0,
        mapl = 1)

    x_l, x_u = primal_bounds(prob)
    c_l, c_u = constraint_bounds(prob)

    nlp_bounds = MOI.NLPBoundsPair.(c_l, c_u)
    block_data = MOI.NLPBlockData(nlp_bounds, prob, true)

    if nlp == :ipopt
        solver = Ipopt.Optimizer()
        solver.options["max_iter"] = max_iter
        solver.options["tol"] = tol
        solver.options["constr_viol_tol"] = c_tol
        # solver.options["print_level"] = mapl
        # solver.options["linear_solver"] = "ma57"
    elseif nlp == :SNOPT7
        # solver = SNOPT7.Optimizer(
        #                           Major_feasibility_tolerance = c_tol,
        #                           Minor_feasibility_tolerance = tol,
        #                           Major_optimality_tolerance = tol,
        #                           Time_limit = time_limit,
        #                           Major_print_level = mapl,
        #                           Minor_print_level = mipl)

    solver = SNOPT7.Optimizer(MOI.MIN_SENSE,
        nothing, [], [], nothing,
        0, 0, 0,
        nothing,
        # Dict("Major_feasibility_tolerance" => c_tol,
        #      "Minor_feasibility_tolerance" => c_tol,
        #      "Minor_optimality_tolerance" => tol,
        #      "Major_optimality_tolerance" => tol,
        #      "Time_limit" => time_limit,
        #      "Major_print_level" => mapl,
        #      "Minor_print_level" => mipl))
        Dict("Major_feasibility_tolerance" => c_tol,
             "Major_optimality_tolerance" => tol,
             # "Minor_feasibility_tolerance" => c_tol,
             # "Minor_optimality_tolerance" => tol,
             "Time_limit" => time_limit,
             "Major_print_level" => mapl,
             "Minor_print_level" => mipl,
             "Major_iterations_limit" => 10000,
             "Minor_iterations_limit" => 200000,
             "Iterations_limit" => 1000000))
    else
        @error "nlp not setup"
    end

    x = MOI.add_variables(solver, prob.num_var)

    for i = 1:prob.num_var
        xi = MOI.SingleVariable(x[i])
        MOI.add_constraint(solver, xi, MOI.LessThan(x_u[i]))
        MOI.add_constraint(solver, xi, MOI.GreaterThan(x_l[i]))
        MOI.set(solver, MOI.VariablePrimalStart(), x[i], x0[i])
    end

    # Solve the problem
    MOI.set(solver, MOI.NLPBlock(), block_data)
    MOI.set(solver, MOI.ObjectiveSense(), MOI.MIN_SENSE)
    MOI.optimize!(solver)

    # Solver info
    info = [] #TODO add info from solve

    # Get the solution
    return MOI.get(solver, MOI.VariablePrimal(), x), info
end


text:
# Parts of this code were taken / derived from Graphs.jl. See LICENSE for
# licensing details.

# The Bellman Ford algorithm for single-source shortest path

struct NegativeCycleError <: Exception end

"""
    struct BellmanFord <: ShortestPathAlgorithm

The structure used to configure and specify that [`shortest_paths`](@ref)
should use the [Bellman-Ford algorithm](http://en.wikipedia.org/wiki/Bellman–Ford_algorithm).
No fields are specified or required.

### Implementation Notes
`BellmanFord` supports the following shortest-path functionality:
- negative distance matrices / weights
- (optional) multiple sources
- all destinations
"""
struct BellmanFord <: ShortestPathAlgorithm end
struct BellmanFordResult{T, U<:Integer} <: ShortestPathResult
    parents::Vector{U}
    dists::Vector{T}
end

function shortest_paths(
    graph::AbstractGraph{U},
    sources::AbstractVector{<:Integer},
    distmx::AbstractMatrix{T},
    ::BellmanFord
   ) where {T, U<:Integer}

    nvg = nv(graph)
    active = falses(nvg)
    active[sources] .= true
    dists = fill(typemax(T), nvg)
    parents = zeros(U, nvg)
    dists[sources] .= 0
    no_changes = false
    new_active = falses(nvg)

    @inbounds for i in vertices(graph)
        no_changes = true
        new_active .= false
        for u in vertices(graph)[active]
            for v in outneighbors(graph, u)
                relax_dist = distmx[u, v] + dists[u]
                if dists[v] > relax_dist
                    dists[v] = relax_dist
                    parents[v] = u
                    no_changes = false
                    new_active[v] = true
                end
            end
        end
        no_changes && break
        active, new_active = new_active, active
    end
    no_changes || throw(NegativeCycleError())
    return BellmanFordResult(parents, dists)
end

shortest_paths(g::AbstractGraph, v::Integer, distmx::AbstractMatrix, alg::BellmanFord) =
    shortest_paths(g, [v], distmx, alg)

has_negative_weight_cycle(g::AbstractGraph, ::BellmanFord) = false

function has_negative_weight_cycle(g::AbstractGraph, distmx::AbstractMatrix, alg::BellmanFord)
    try
        shortest_paths(g, vertices(g), distmx, alg)
    catch e
        isa(e, ShortestPaths.NegativeCycleError) && return true
    end
    return false
end


text:
# import ..Concepts:AbstractSamplingModels,
#     BernoulliModel,
#     VecOrMatOfNumbers,
#     provide

using ..Concepts

import StatsBase

struct Sampler{T<:AbstractSamplingModels}
    model::T
    draw

    function Sampler{BernoulliModel}()
        # abstract type constructor
        return new{BernoulliModel}()
    end

    function Sampler{UniformModel}()
        # abstract type constructor
        return new{UniformModel}()
    end

    function Sampler{NonUniformModel}()
        # abstract type constructor
        return new{NonUniformModel}()
    end

    function Sampler{BernoulliModel}(model::BernoulliModel)
        draw = begin
            function(x::VecOrMatOf{Number})
                if isa(x,Vector)
                    n = length(x)
                    mask = [rand(Distributions.Bernoulli(model.rate)) == 1 ? 1 : missing for i in 1:n]
                    return mask .* x
                end
                n,m = size(x);
                mask = [rand(Distributions.Bernoulli(model.rate)) == 1 ? 1 : missing for i in 1:n,j in 1:m]
                return mask .* x
            end
        end
        return new{BernoulliModel}(model,draw)
    end

    function Sampler{UniformModel}(model::UniformModel)
        draw = begin
            function(x::VecOrMatOf{T}) where T<:Any
                sampled_object = nothing
                if isa(x,Vector)
                    n = length(x)
                    mask  = [CartesianIndex(i) for i in StatsBase.sample(1:n,Int.(n * model.rate))]
                    # sampled_object = convert(MaybeMissing{T},Array{Missing}(undef,n))
                    sampled_object = type_conversion(MaybeMissing{T},Array{Missing}(undef,n))
                    for i in mask
                        sampled_object[i] = x[i]
                    end
                else
                    row,col = size(x)
                    mask = [CartesianIndex(StatsBase.sample(1:row),StatsBase.sample(1:col)) for i in 1:Int.(row * col * model.rate)]
#                    display(mask)
                    # sampled_object = convert(MaybeMissing{T},Array{Missing,2}(undef,row,col))
                    sampled_object = type_conversion(MaybeMissing{T},Array{Missing,2}(undef,row,col))
#                    display(sampled_object)
#                    display(x)
                    for i in mask
                        sampled_object[i] = x[i]
                    end
                end
                return sampled_object
            end
        end

        return new(model,draw)
    end
end

const Sampler(model::BernoulliModel) = Sampler{BernoulliModel}(model)
const Sampler(model::UniformModel)   = Sampler{UniformModel}(model)

@overload
function Concepts.provide(object::Sampler{Concepts.BernoulliModel};rate)
    return Sampler(BernoulliModel(rate))
end

@overload
function Concepts.provide(object::Sampler{Concepts.UniformModel};rate)
    return Sampler(UniformModel(rate))
end


text:
# This file is a part of SimilaritySearch.jl
using Test, SimilaritySearch, SparseArrays, LinearAlgebra

@testset "test database abstractions" begin
    X = rand(4, 100)
    A = MatrixDatabase(X)
    B = VectorDatabase(X)
    C = DynamicMatrixDatabase(X)
    D = B[1:100]
    @test D isa SubDatabase
    @test X === A.matrix
    @test X !== B.vecs
    @test X !== C.data
    @test length.([A, B, C, D]) == [100, 100, 100, 100]
    @test collect(A) == collect(B) == collect(C) == collect(D)
    for i in rand(1:100, 10)
        for v in getindex.([A, B, C, D], i)
            @test v == X[:, i]
        end
    end

    A[1] = 0
    @test sum(A[1]) == 0
    @test A[1] == X[:, 1]
    @test_throws DimensionMismatch A[2] = [1, 2]
    @test_throws DimensionMismatch C[2] = [1, 2]
    A[2] = [1, 2, 3, 4]
    @test A[2] == [1, 2, 3, 4]
    B[1] = [1, 2]
    @test typeof(B[1]) == Vector{Float64}
    @test B[1] == [1, 2]

    A = "ABCDEFGHIJKLMNÑOPQRSTUVWXYZαβγδι"
    F = VectorDatabase(type=String)
    for i in 1:100
        push!(F, String(rand(A, 5)))
    end

    @test length(F) == 100
    @test eltype(F) == String

    G = sparse(rand(1:10000, 30), rand(1:10000, 30), rand(30))
    H = MatrixDatabase(G)
    @test all([norm(H[i]) == norm(G[:, i]) for i in rand(1:size(G, 2), 100)])
end

max_stars_repo_name: sadit/SimilaritySearch

max_issues_repo_name: sadit/NNS

max_forks_repo_name: sadit/SimilaritySearch