text:
import Pkg
Pkg.activate("..")

using CrimsonSkyline
using Distributions: Beta, Geometric, Binomial, truncated
using StatsBase: mean, std

function bayes_net(t::Trace, data::Int64)
    prob = sample(t, "prob", Beta(2.0, 2.0))
    n = observe(t, "n", truncated(Geometric(prob), 1, Inf), data)
    sample(t, "count", Binomial(n, prob))
end

function main()
    @warn "~~~ generating kernels for closed-world inference (model is identifiable) ~~~"
    data = 7
    kernels = make_kernels(bayes_net; params = (data,))
    results = inference(
        bayes_net, Metropolis();
        params = (data,),
        inference_params = Dict(
            "kernels" => kernels,
            "num_iterations" => 55000,
            "burn" => 10000,
            "thin" => 500,
            "inverse_verbosity" => 5000
        )
    )
    prob_stats = map(f -> f(results, "prob"), [mean, std])
    count_stats = map(f -> f(results, "count"), [mean, std])
    @info "Posterior stats: prob (mean, std) = $prob_stats, count (mean, std) = $count_stats"
end

if abspath(PROGRAM_FILE) == @__FILE__
    main()
end


text:
using BasisFunctions: Dictionary, TensorProductDict, Dictionary1d, tensorproduct,
    extensionsize, gramdual, Weight, ProductWeight, DiscreteProductWeight, resize,
    isbasis, hastransform, dimensions, productmeasure
import Base: getindex
import BasisFunctions: components, factors, dictionary, measure, Measure

#################
# Platform
#################

export Platform, BasisPlatform, FramePlatform
"""
    abstract type Platform end

A platform represents a family of dictionaries.

A platform typically has a primal and a dual sequence. The platform maps an
index to a set of parameter values, which is then used to generate a primal
or dual dictionary (depending on the chosen measure).

See also [`BasisPlatform`](@ref), `FramePlatform`](@ref)
"""
abstract type Platform end

getindex(platform::Platform, param) = dictionary(platform, param)

correctparamformat(platform::Platform, _) = false

function dictionary(platform::Platform, param)
    if !correctparamformat(platform, param)
        throw(ArgumentError("Parameter $param not suited for platform $platform. "))
    end
    unsafe_dictionary(platform, param)
end

export dictionary, dualdictionary

"""
    dualdictionary(platform::Platform, param, measure::Measure; options...)

Return the dual dictionary of the platform.
"""
dualdictionary(platform::Platform, param, measure::Measure; dict=dictionary(platform, param), options...) =
    gramdual(dict, measure; options...)

include("platformadaptivity.jl")

"""
    abstract type BasisPlatform <: Platform end

A `BasisPlatform` represents a family of bases.
"""
abstract type BasisPlatform <: Platform end

"""
    abstract type FramePlatform <: Platform end

A `FramePlatform` represents a family of frames
(or ill-conditioned bases which are in the limit a frame).
"""
abstract type FramePlatform <: Platform end

include("styles.jl")

#################
# ModelPlatform
#################

"""
    struct ModelPlatform <: Platform

A `ModelPlatform` is a platform based on a model dictionary. The platform is
defined by resizing the dictionary, using its own implementation of `resize`.
All other operations are the defaults for the model dictionary.

This platform is convenient to compute adaptive approximations based on an
example of a dictionary from the desired family.
"""
struct ModelPlatform <: Platform
    model   ::  Dictionary
end

model(p::ModelPlatform) = p.model
unsafe_dictionary(p::ModelPlatform, param) = resize(model(p), param)

param_first(p::ModelPlatform) = dimensions(model(p))
DictionaryStyle(p::ModelPlatform) = isbasis(model(p)) ? BasisStyle() : UnknownDictionaryStyle()
SamplingStyle(p::ModelPlatform) = isbasis(model(p)) ? InterpolationStyle() : OversamplingStyle()
SolverStyle(p::ModelPlatform, samplingstyle::SamplingStyle) = SolverStyle(DictionaryStyle(p), p, samplingstyle)

export correctparamformat
correctparamformat(p::ModelPlatform, param)  =
    param isa typeof(dimensions(model(p)))

export measure
measure(platform::ModelPlatform) = measure(model(platform))
coefficienttype(platform::ModelPlatform) = coefficienttype(model(platform))

components(platform::ModelPlatform) = map(ModelPlatform, components(model(platform)))
component(platform::ModelPlatform, i) = ModelPlatform(component(model(platform), i))

#################
# ProductPlatform
#################

export ProductPlatform
"""
    struct ProductPlatform{N} <: Platform

A `ProductPlatform` corresponds to the product of two or more platforms. It results in a product
dictionary, product sampling operator, product solver, ... etcetera.
"""
struct ProductPlatform{N} <: Platform
    platforms   ::  NTuple{N,Platform}
end

ProductPlatform(platforms::Platform...) = ProductPlatform(platforms)

param_first(platform::ProductPlatform) = map(param_first, components(platform))

function DictionaryStyle(p::ProductPlatform)
    dict = TensorProductDict(map(x->dictionary(x,1), components(p))...)
    isbasis(dict) ? BasisStyle() : FrameStyle()
end
ProductPlatform(platform::Platform, n::Int) = ProductPlatform(ntuple(x->platform, n)...)
function dualdictionary(platform::ProductPlatform, param, measure::Measure; options...)
    @assert length(param)==length(components(platform))
    TensorProductDict(map((plati, parami, mi)->dualdictionary(plati, parami, mi; options...), components(platform), param, components(measure))...)
end
components(p::ProductPlatform) = p.platforms
factors(p::ProductPlatform) = components(p)

export productparameter
"""
    productparameter(p::ProductPlatform{N}, param)

Transform the parameter to a parameter accepted by a ProductPlatform
"""
productparameter(p::ProductPlatform{N}, n::Int) where {N} = ntuple(x->n, Val(N))
productparameter(p::ProductPlatform{N}, n::NTuple{N,Any}) where {N} = n

SamplingStyle(platform::ProductPlatform) = ProductSamplingStyle(map(SamplingStyle, components(platform)))
function SolverStyle(p::ProductPlatform, samplingstyle::ProductSamplingStyle)
    @assert nfactors(p) == ncomponents(samplingstyle)
    ProductSolverStyle(map(SolverStyle, factors(p), factors(samplingstyle)))
end

dictionary(p::ProductPlatform, n::Int) = dictionary(p, productparameter(p, n))

unsafe_dictionary(p::ProductPlatform, n) = tensorproduct(map(dictionary, components(p), n)...)

dualdictionary(platform::ProductPlatform, param, measure::Union{ProductWeight,DiscreteProductWeight}; options...) =
    tensorproduct(map((platformi, parami, mi)->dualdictionary(platformi, parami, mi; options...),
        components(platform), param, components(measure))...)

measure(platform::ProductPlatform) = productmeasure(map(measure, components(platform))...)

correctparamformat(p::ProductPlatform{N}, plt_par::NTuple{N,Int}) where N =
    all(map(correctparamformat, p.platforms, plt_par))

# we allow an integer platform parameter for convenience, but in that case
# the `productparameter` function gives the correct tuple
correctparamformat(p::ProductPlatform, plt_par::Int) = true

correctparamformat(::ProductPlatform, plt_par) = false

## Default platforms

export platformparameter, platform
"""
    platformparameter(dict::Dictionary) = dimensions(dict)

Return the parameter that is given to a model-based platform to obtain this dictionary.
"""
platformparameter(dict::Dictionary) = dimensions(dict)

"""
    platform(dict::Dictionary)

Return a platform that generates dictionaries of the type of dict.
"""
platform(dict::Dictionary) = ModelPlatform(dict)
platform(dict::TensorProductDict) = ProductPlatform(map(platform, factors(dict))...)

ext: jl


text:
using Zarr: AbstractStore, Zarr, zopen, zgroup, zcreate
export DiskArrayShuffleStore
struct DiskArrayShuffleStore{S<:DiskArrayShuffler} <: AbstractStore
    s::S
    name::String
    metadata::Dict{String,Vector{UInt8}}
end
Base.show(io::IO,s::DiskArrayShuffleStore) = print(io,"Shuffled Store")
function DiskArrayShuffleStore(s::DiskArrayShuffler; name = "batches", attrs = Dict(),nbatches=100)
    onear = s.currentcoll[].chunkviews |> first |> first
    soutput = size(onear)
    cs = (soutput..., s.sampler.batchsize)
    soutput = (cs...,nbatches)
    #Create a mock Zarr array to contain the metadata
    g = zgroup(Zarr.DictStore())
    zcreate(eltype(onear),g,name,soutput..., compressor=Zarr.NoCompressor(), attrs=attrs)
    Zarr.consolidate_metadata(g)
    DiskArrayShuffleStore(s,name,g.storage.a)
end

function Base.getindex(s::DiskArrayShuffleStore,i::AbstractString)
    if i in (s.name * "/.zarray", s.name*"/.zattrs",".zgroup",".zmetadata")
        return get(s.metadata,i,nothing)
    else
        schunk = split(i,'/')[end]
        chunk = parse.(Int,split(schunk,'.'))
        all(iszero,chunk[2:end]) || return nothing
        r = readbatch(s.s)
        nd = ndims(first(r))
        rone = reduce((i,j)->cat(i,j,dims=nd+1),r)
        return collect(vec(reinterpret(UInt8,rone)))
    end
end


text:
"""
    struct User

Represents a user on Reddit.
"""
struct User
    name::String
end

"""
    about()

Get information about default user(s).
"""
function about()
    about(default())
end

"""
    about(a::Authorized)

Get information about the current user.
"""
function about(a::Authorized)
    JSON.parse(get("/user/$(a.username)/about", a))["data"]
end

"""
    about(user::User)

Get information about a specified user using default credentials.
"""
function about(user::User)
    about(user, default())
end

"""
    about(user::User, a::Authorized)

Get information about a specified user.
"""
function about(user::User, a::Authorized)
    JSON.parse(get("/user/$(user.name)/about", a))["data"]
end

"""
    blocked()

Get users blocked by default user(s).
"""
function blocked()
    blocked(default())
end

"""
    blocked(a::Authorized)

Get an Array of users blocked by current user.
"""
function blocked(a::Authorized)
    JSON.parse(get("/prefs/blocked", a))["data"]["children"]
end

"""
    comments(; count=nothing, sort="new")

Get all comments by default user(s).
"""
function comments(; count=1000, sort="new")
    a = default()
    comments(User(a.username), a, count=count, sort=sort)
end

"""
    comments(a::Authorized; count=nothing, sort="new")

Get all comments by current user.
"""
function comments(a::Authorized; count=1000, sort="new")
    comments(User(a.username), count=count, sort=sort)
end

"""
    comments(user::User; count=nothing, sort="new)

Get all comments by default user(s).
"""
function comments(user::User; count=1000, sort="new")
    a = default()
    comments(User(user.name), default(), count=count, sort=sort)
end

"""
    comments(user::User, a::Authorized; count=nothing, sort="new")

Get all comments by a user.
"""
function comments(user::User, a::Authorized;
                  count=1000, sort="new")
    fetchedcomments = Array{Dict{String, Any}, 1}()
    fetchedcount = 0
    limit = 100
    after = ""
    if count < 100
        limit = count
    end
    api = "/user/$(user.name)/comments.json?sort=$sort&limit=$limit"
    firstfetched = JSON.parse(get(api, a))
    for fetched in firstfetched["data"]["children"]
        push!(fetchedcomments, fetched["data"])
    end
    after = firstfetched["data"]["after"]
    fetchedcount += firstfetched["data"]["dist"]
    done = false
    while !done && fetchedcount < count
        afterapi = api*"&after=$after&count=$fetchedcount"
        nextfetched = JSON.parse(get(afterapi, a))
        for fetched in nextfetched["data"]["children"]
            push!(fetchedcomments, fetched["data"])
        end
        fetchedcount += nextfetched["data"]["dist"]
        after = nextfetched["data"]["after"]
        if after == nothing
            done = true
        end
    end
    return fetchedcomments
end

"""
    deletecomment(id::AbstractString)

Delete a comment made by the default user(s) by id.
"""
function deletecomment(id::AbstractString)
    comment(id, default())
end

"""
    delete(id::AbstractString, a::Authorized)

Delete a comment by id.
"""
function deletecomment(id::AbstractString, a::Authorized)
    body = "id=$id"
    JSON.parse(post("/api/del", body, a))
end

"""
    downvoted()

Get all threads downvoted by default user(s).
"""
function downvoted()
    downvoted(default())
end

"""
    downvoted(a::Authorized)

Get all threads a user has downvoted.
"""
function downvoted(a::Authorized)
    JSON.parse(get("/user/$(a.username)/downvoted", a))
end

"""
    friends()

Get friends of default user(s).
"""
function friends()
    friends(default())
end

"""
    friends(a::Authorized)

Get friends of current user.
"""
function friends(a::Authorized)
    JSON.parse(get("/api/v1/me/friends", a))["data"]["children"]
end

"""
    gilded()

Get all threads gilded by default user(s).
"""
function gilded()
    gilded(default())
end

"""
    gilded(a::Authorized)

Get all threads current user has gilded.
"""
function gilded(a::Authorized)
    gilded(a.username, a)
end

"""
    gilded(user::User)

Get all threads a user has gilded, using default credentials.
"""
function gilded(user::User)
    gilded(user.name, default())
end

"""
    gilded(user::AbstractString)

Get all threads a user has gilded, using default credentials.
"""
function gilded(user::AbstractString)
    gilded(user, default())
end

"""
    gilded(user::User, a::Authorized)

Get all threads a user has gilded, using default credentials.
"""
function gilded(user::User, a::Authorized)
    gilded(user.name, a)
end

"""
    gilded(user::AbstractString, a::Authorized)

Get all threads a user has gilded.
"""
function gilded(user::AbstractString, a::Authorized)
    JSON.parse(get("/user/$user/gilded", a))["data"]["children"]
end

"""
    hidden()

Get all threads hidden by default user(s).
"""
function hidden()
    hidden(default())
end

"""
    hidden(a::Authorized)

Get all threads current user has hidden.
"""
function hidden(a::Authorized)
    JSON.parse(get("/user/$(a.username)/hidden", a))["data"]["children"]
end

"""
    karma()

Get karma breakdown for default user(s).
"""
function karma()
    karma(default())
end

"""
    karma(a::Authorized)

Get karma breakdown for current user.
"""
function karma(a::Authorized)
    JSON.parse(get("/api/v1/me/karma", a))["data"]
end

"""
    me()

Get identity information for default user(s).
"""
function me()
    me(default())
end

"""
    me(a::Authorized)

Get identity information for current user.
"""
function me(a::Authorized)
    JSON.parse(get("/api/v1/me", a))
end

"""
    overview()

Get an overview for default user(s).
"""
function overview()
    overview(default())
end

"""
    overview(a::Authorized)

Get an overview for current user.
"""
function overview(a::Authorized)
    overview(a.username, a)
end

"""
    overview(user::User)

Get an overview for a user, using default credentials.
"""
function overview(user::User)
    overview(user.name, default())
end

"""
    overview(user::AbstractString)

Get an overview for a user, using default credentials.
"""
function overview(user::AbstractString)
    overview(user, default())
end

"""
    overview(user::User, a::Authorized)

Get an overview for a user.
"""
function overview(user::User, a::Authorized)
    overview(user.name, a)
end

"""
    overview(user::AbstractString, a::Authorized)

Get an overview for a user.
"""
function overview(user::AbstractString, a::Authorized)
    JSON.parse(get("/user/$user/overview", a))["data"]["children"]
end

"""
    preferences()

Get all preferences for default user(s).
"""
function preferences()
    preferences(default())
end

"""
    preferences(a::Authorized)

Get all preferences for current user.
"""
function preferences(a::Authorized)
    JSON.parse(get("/api/v1/me/prefs", a))
end

"""
    reply(parent::AbstractString, text::AbstractString)

Reply to a comment by name, as default user(s).
"""
function reply(parent::AbstractString, text::AbstractString)
    reply(parent, text, default())
end

"""
    reply(parent::AbstractString, text::AbstractString, a::Authorized)

Reply to a comment by name.
"""
function reply(parent::AbstractString, text::AbstractString, a::Authorized)
    body = "api_type=json&text=$text&thing_id=$parent"
    JSON.parse(post("/api/comment", body, a))
end

"""
    saved()

Get all posts saved by default user(s).
"""
function saved()
    saved(default())
end

"""
    saved(a::Authorized)

Get all posts a user has saved.
"""
function saved(a::Authorized)
    JSON.parse(get("/user/$(a.username)/saved", a))["data"]["children"]
end

"""
    submit(sub::Subreddit,
           title::AbstractString,
           text::AbstractString;
           ad=false,
           flair_id="",
           flair_text="",
           kind="self",
           nsfw=false,
           resubmit=true,
           sendreplies=true,
           spoiler=false)

Submit a post to a Subreddit using the defalut credentials.
"""
function submit(sub::Subreddit,
                title::AbstractString,
                body::AbstractString;
                ad=false,
                flair_id="",
                flair_text="",
                kind="self",
                nsfw=false,
                resubmit=true,
                sendreplies=true,
                spoiler=false)
    submit(sub, title, body, default(), ad=ad, flair_id=flair_id,
           flair_text=flair_text, kind=kind, nsfw=nsfw, resubmit=resubmit,
           sendreplies=sendreplies, spoiler=spoiler)
end

"""
    submit(sub::Subreddit,
           title::AbstractString,
           text::AbstractString,
           a::Authorized;
           ad=false,
           flair_id="",
           flair_text="",
           kind="self",
           nsfw=false,
           resubmit=true,
           sendreplies=true,
           spoiler=false)

Submit a post to a Subreddit.
"""
function submit(sub::Subreddit,
                title::AbstractString,
                body::AbstractString,
                a::Authorized;
                ad=false,
                flair_id="",
                flair_text="",
                kind="self",
                nsfw=false,
                resubmit=true,
                sendreplies=true,
                spoiler=false)
    params = "api_type=json&sr=$(sub.name)&title=$title&kind=$kind"
    if kind == "self"
        params *= "&text=$body"
    elseif kind == "link"
        params *= "&url=$body"
    else
        println("Error: Invalid kind")
        return nothing
    end
    if flair_id != ""
        params *= "&flair_id=$flair_id"
    end
    if flair_text != ""
        params *= "&flair_text=$flair_text"
    end
    params *= "&ad=$ad&nsfw=$nsfw&sendreplies=$sendreplies&spoiler=$spoiler"
    JSON.parse(post("/api/submit", params, a))
end

"""
    submitted()

Get all posts current submitted by default user(s).
"""
function submitted()
    submitted(default())
end

"""
    submitted(a::Authorized)

Get all posts current user has submitted.
"""
function submitted(a::Authorized)
    submitted(User(a.username), a)
end

"""
    submitted(user::AbstractString)

Get all posts a user has submitted, using default credentials.
"""
function submitted(user::AbstractString)
    submitted(User(user), default())
end

"""
    submitted(user::User)

Get all posts a user has submitted, using default credentials.
"""
function submitted(user::User)
    submitted(user, default())
end

"""
    submitted(user::User, a::Authorized)

Get all posts a user has submitted.
"""
function submitted(user::User, a::Authorized)
    JSON.parse(get("/user/$(user.name)/submitted", a))["data"]["children"]
end

"""
    trophies()

Get all trophies for default user(s).
"""
function trophies()
    trophies(default())
end

"""
    trophies(a::Authorized)

Get all trophies for current user.
"""
function trophies(a::Authorized)
    JSON.parse(get("/api/v1/me/trophies", a))["data"]["trophies"]
end

"""
    upvoted()

Get all posts upvoted by default user(s).
"""
function upvoted()
    upvoted(default())
end

"""
    upvoted(a::Authorized)

Get all posts a user has upvoted.
"""
function upvoted(a::Authorized)
    JSON.parse(get("/user/$(a.username)/upvoted", a))["data"]["children"]
end


text:
module TestTemperatureProfile
using Test
using SARA
using SARA: TemperatureProfile, inverse_profile

@testset "temperature profile" begin
    P = TemperatureProfile()
    # 1. correctness of inverse profile
    T_max = 800
    log10_τ = 3
    f = P(T_max, log10_τ)
    g = inverse_profile(P, T_max, log10_τ)
    n = 16
    for T in T_max*rand(n)
        @test f(g(T)) ≈ T
    end
    for x in 1e-1randn(n)
        @test g(f(x)) ≈ abs(x)
    end

    T_max = 1025.0
    log10_τ = 4.0
    P = TemperatureProfile()
    f = P(T_max, log10_τ) # maps position of measured optical data to temperature
    g = inverse_profile(P, T_max, log10_τ)
    T = T_max / 2
    @test T ≈ f(g(T))
    x = 1e-1randn()
    @test abs(x) ≈ g(f(x))

    # test differentiation of g
    using ForwardDiff: derivative
    d = derivative(g, T)
    @test d isa Real
end

# d = (t->ForwardDiff.derivative(mean(C), t)).(T) # WARNING: investigate numerical issue (NaN)

end

# using JSON
# 2. testing against Bi2O3 data
# include("../NatCom2020/inner_load_data.jl")
#
# path = "/Users/sebastianament/Documents/SEA/XRD Analysis/SARA/Bi2O3_19F44_01/"
# # path = ""
# file = "Bi2O3_19F44_01_inner_loop_data.json"
# f = JSON.parsefile(path * file)
# position, optical, rescaling_parameters = load_data(f)
#
# tol = 1e-6
# P = TemperatureProfile()
# for i in 1:length(f)
#     x = 1000 * (f[i]["pos_wafer"] .- f[i]["center_real"])
#     T = f[i]["temp_profile"]
#     # T = f[i]["temp_profile_si"]
#     T_max, τ = f[i]["meta"]["temp"], f[i]["meta"]["dwell"]
#     log10_τ = log10(τ)
#     if maximum(abs, P(T_max, log10_τ).(x) .- T) > tol
#         println("Hey!")
#     end
# end

#### old plot of fwhm
# mn = [2.39794001e+00, 4.00000000e+02, 2.50000000e+01, 1.77686356e-02, 0.00000000e+00]
# mx = [   4.,         1300.,          350.,            9.92052851,    5.,        ]
#fig = plt.figure()
#ax = fig.add_subplot(111, projection='3d')
#plt.xlabel('log tau')
#plt.ylabel('temp')
#X,Y = meshgrid(range(mn[0], mx[0], 20), range(mn[1], mx[1], 20))

# X = collect(range(mn[1], stop=mx[1], length=20))
# Y = collect(range(mn[2], stop=mx[2], length=20))

# using Plots
# pyplot()
# Z = ffwhm(Y, X)
# surface(X[:], Y[:], ffwhm.(Y, X'), rstride=1, cstride=1, alpha=0.2)


text:
import Base: parse

#this is all pretty hacky with regards to the new structure and atom api. can for sure be a lot better!
"Quantum espresso card option parser"
function cardoption(line)
	sline = split(line)
	if length(sline) < 2 && lowercase(sline[1])=="k_points"
		return :tpiba
	else
	  	return Symbol(match(r"((?:[a-z][a-z0-9_]*))", sline[2]).match)
  	end
end

"""
    qe_read_output(filename::String, T=Float64)

Reads a generic quantum espresso input, returns a dictionary with all found data in the file.
Possible keys:
 - `:fermi`
 - `:polarization`
 - `:pol_mod`
 - `:k_cryst`
 - `:k_cart`
 - `:alat`
 - `:cell_parameters`
 - `:pos_option`
 - `:atomic_positions`
 - `:total_force`
 - `:colin_mag_moments`
 - `:bands`
 - `:accuracy`
 - `:converged`
"""
function qe_read_output(filename::String, T=Float64)
    out = Dict{Symbol,Any}()
    open(filename, "r") do f
        prefac_k     = nothing
        k_eigvals    = Array{Array{T,1},1}()
        lowest_force = T(1000000)
        atsyms = Symbol[]
        while !eof(f)
            line = readline(f)

            #polarization
            if occursin("C/m^2", line)
                s_line = split(line)
                P      = parse(T, s_line[3])
                mod    = parse(T, s_line[5][1:end-1])
                readline(f)
                s_line = parse.(T, split(readline(f))[6:2:10])
                out[:polarization] = Point3{T}(P * s_line[1], P * s_line[2], P * s_line[3])
                out[:pol_mod]      = mod

            #input structure parameters
            elseif occursin("lattice parameter", line)
                out[:in_alat] = uconvert(Ang, parse(Float64, split(line)[5])*1a₀)

            elseif occursin("crystal axes", line)
                cell_1 = parse.(Float64, split(readline(f))[4:6]) .* out[:in_alat]
                cell_2 = parse.(Float64, split(readline(f))[4:6]) .* out[:in_alat]
                cell_3 = parse.(Float64, split(readline(f))[4:6]) .* out[:in_alat]
                out[:in_cell] = Mat3([cell_1 cell_2 cell_3])

            elseif occursin("reciprocal axes", line)
                cell_1 = parse.(Float64, split(readline(f))[4:6]) .* 2π/out[:in_alat]
                cell_2 = parse.(Float64, split(readline(f))[4:6]) .* 2π/out[:in_alat]
                cell_3 = parse.(Float64, split(readline(f))[4:6]) .* 2π/out[:in_alat]
                out[:in_recip_cell] = Mat3([cell_1 cell_2 cell_3])

            elseif occursin("atomic species", line)
                line = readline(f)
                while !isempty(line)
                    push!(atsyms, Symbol(strip_split(line)[1]))
                    line = readline(f)
                end

                #PseudoPot
            elseif occursin("PseudoPot", line)
                !haskey(out, :pseudos) && (out[:pseudos] = Dict{Symbol, Pseudo}())
                pseudopath = readline(f) |> strip |> splitdir
                out[:pseudos][Symbol(split(line)[5])] = Pseudo(pseudopath...)
                #fermi energy
            elseif occursin("Fermi", line)
                out[:fermi]        = parse(T, split(line)[5])
            elseif occursin("lowest unoccupied", line) && occursin("highest occupied", line)
                out[:fermi]        = parse(T, split(line)[7])

            elseif occursin("lowest unoccupied", line) || occursin("highest occupied", line)
                out[:fermi]        = parse(T, split(line)[5])

            elseif occursin("cryst.", line) && length(split(line)) == 2
                out[:k_cryst] = Vector{Vec3{T}}()
                line = readline(f)
                while line != "" && !occursin("--------", line)
                    push!(out[:k_cryst], parse_k_line(line, T))
                    line = readline(f)
                end

                #k_cart
            elseif occursin("cart.", line) && length(split(line)) == 5
                line = readline(f)
                alat = out[:in_alat]
                out[:k_cart] = Vector{Vec3{typeof(2π/alat)}}()
                while line != "" && !occursin("--------", line)
                    push!(out[:k_cart], parse_k_line(line, T) .* 2π/alat)
                    line = readline(f)
                end

                #bands
            elseif occursin("k", line) && occursin("PWs)", line)
                tmp = T[]
                readline(f)
                line = readline(f)
                while line != "" && !occursin("--------", line)
                    append!(tmp, parse_line(T, line))
                    line = readline(f)
                end
                push!(k_eigvals, tmp)

                #errors
            elseif occursin("mpirun noticed", line)
                @warn "File ended unexpectedly, returning what info has been gathered so far."
                return out
                break
                #vcrel outputs
            elseif occursin("Begin final coordinates", line)
                line = readline(f)
                while !occursin("End final coordinates", line)
                    if occursin("CELL_PARAMETERS", line)
                        out[:alat]            = occursin("angstrom", line) ? :angstrom : parse(T, split(line)[end][1:end-1])
                        out[:cell_parameters] = reshape(T[parse.(T, split(readline(f))); parse.(T, split(readline(f))); parse.(T, split(readline(f)))], (3,3))'
                    elseif occursin("ATOMIC_POSITIONS", line)
                        out[:pos_option]      = cardoption(line)
                        line  = readline(f)
                        atoms = []
                        while !occursin("End", line)
                            s_line = split(line)
                            key    = Symbol(s_line[1])
                            push!(atoms, key=>Point3{T}(parse.(T, s_line[2:end])...))
                            line = readline(f)
                        end
                        posdict = Dict{Symbol, Vector{Point3{T}}}()
                        for (atsym, pos) in atoms
                            if haskey(posdict, atsym)
                                push!(posdict[atsym], pos)
                            else
                                posdict[atsym] = [pos]
                            end
                        end
                        out[:atomic_positions] = posdict
                        break
                    end
                    line = readline(f)
                end
                if haskey(out, :alat) && haskey(out, :cell_parameters)
                    pseudo_data = InputData(:atomic_species, :none, out[:pseudos])
                    tmp_flags = Dict{Symbol, Any}(:ibrav => 0)
                    if haskey(out, :alat)
                        tmp_flags[:A] = out[:alat] == :angstrom ? 1.0 : conversions[:bohr2ang] * out[:alat]
                    else
                        tmp_flags[:A] = 1.0
                    end
                    cell_data = InputData(:cell_parameters, :alat, Mat3(out[:cell_parameters]))
                    atoms_data = InputData(:atomic_positions, out[:pos_option], out[:atomic_positions])
                    out[:final_structure] = extract_structure!("newstruct", tmp_flags, cell_data, atsyms, atoms_data, pseudo_data)
                end

            elseif occursin("Total force", line)
                force = parse(T, split(line)[4])
                if force <= lowest_force
                    lowest_force      = force
                    out[:total_force] = force
                end
            elseif occursin("Magnetic moment per site", line)
                key = :colin_mag_moments
                out[key] = T[]
                line = readline(f)
                while !isempty(line)
                    push!(out[key], parse.(Float64, split(line)[6]))
                    line = readline(f)
                end
            elseif occursin("estimated scf accuracy", line)
                key = :accuracy
                acc = parse(T, split(line)[end-1])
                if haskey(out, key)
                    push!(out[key], acc)
                else
                    out[key] = [acc]
                end
            elseif occursin("convergence NOT achieved", line)
                out[:converged] = false
            elseif occursin("convergence has been achieved", line)
                out[:converged] = true
            end
        end

        #process bands
        if !isempty(k_eigvals)
            out[:bands] = Vector{DFBand{T}}()
            for i1=1:length(k_eigvals[1])
                eig_band = T[]
                for i = 1:length(out[:k_cart])
                    push!(eig_band, k_eigvals[i][i1])
                end
                if !haskey(out, :k_cryst)
                    out[:k_cryst] = (out[:in_recip_cell]^-1,) .* out[:k_cart]
                end

                push!(out[:bands], DFBand(out[:k_cart], out[:k_cryst], eig_band))
            end
        end

        return out
    end
end

function qe_read_output(input::DFInput{QE}, args...; kwargs...)
    out = Dict{Symbol, Any}()
    if isprojwfccalc(input)
        pdos_files = searchdir(dir(input), ".pdos_")
        if flag(input, :kresolveddos) == true
            out[:heatmaps]   = Vector{Matrix{Float64}}()
            out[:ytickvals] = Vector{Vector{Float64}}()
            out[:yticks]    = Vector{Vector{Float64}}()
            for f in pdos_files
                th, vals, ticks = qe_read_kpdos(joinpath(dir(input), f), args...; kwargs...)
                push!(out[:heatmaps], th)
                push!(out[:ytickvals], vals)
                push!(out[:yticks], ticks)
            end
        else
            out[:pdos] = NamedTuple{(:energies, :values), Tuple{Vector{Float64}, Matrix{Float64}}}[]
            for f in pdos_files
                energs, vals = qe_read_pdos(joinpath(dir(input), f), args...; kwargs...)
                push!(out[:pdos], (energies=energs, values=vals))
            end
        end
        out[:states], out[:bands] = qe_read_projwfc(outpath(input))
        return out
    else
        return qe_read_output(outpath(input))
    end
end
"""
    qe_read_bands(filename::String, T=Float64)

Reads the output file of a 'bands' calculation in Quantum Espresso.
Returns an array of DFBands each with the same k_points and their respective energies.
"""
qe_read_bands_file(filename::String, T=Float64) = qe_read_output(filename, T)[:bands]

"""
    qe_read_ks_from_output(filename::String, T=Float64)

Read k-points from a Quantum Espresso bands output file in cartesian (2pi/alat in Angstrom^-1!) and crystalline coordinates.
Returns (k_points_cart,k_points_cryst).
"""
function qe_read_ks_from_output(filename::String, T=Float64)
    t = qe_read_output(filename, T)
    return t[:k_cart], t[:k_cryst]
end

"""
    qe_read_fermi_from_output(filename::String,T=Float64)

Reads the Fermi level from a Quantum Espresso scf calculation output file
(if there is one).
"""
qe_read_fermi_from_output(filename::String, T=Float64) = qe_read_output(filename, T)[:fermi]

"""
    qe_read_kpdos(filename::String,column=1;fermi=0)

Reads the k_resolved partial density of states from a Quantum Espresso projwfc output file.
Only use this if the flag kresolveddos=true in the projwfc input file!! The returned matrix can be readily plotted using heatmap() from Plots.jl!
Optional input: column = 1 (column of the output, 1 = first column after ik and E)
fermi  = 0 (possible fermi offset of the read energy values)
Return:         Array{Float64,2}(length(k_points),length(energies)) ,
(ytickvals,yticks)
"""
function qe_read_kpdos(filename::String, column=1; fermi=0)
    read_tmp = readdlm(filename, Float64, comments=true)
    zmat     = zeros(typeof(read_tmp[1]), Int64(read_tmp[end, 1]), div(size(read_tmp)[1], Int64(read_tmp[end, 1])))
    for i1 = 1:size(zmat)[1]
        for i2 = 1:size(zmat)[2]
            zmat[i1, i2] = read_tmp[size(zmat)[2] * (i1 - 1) + i2, 2 + column]
        end
    end

    yticks    = collect(Int(div(read_tmp[1, 2] - fermi, 1)):1:Int(div(read_tmp[end, 2] - fermi, 1)))
    ytickvals = [findfirst(x -> norm(yticks[1] + fermi - x) <= 0.1, read_tmp[:, 2])]
    for (i, tick) in enumerate(yticks[2:end])
        push!(ytickvals, findnext(x -> norm(tick + fermi - x) <= 0.1, read_tmp[:, 2], ytickvals[i]))
    end

    return  zmat', ytickvals, yticks
end

"""
    qe_read_pdos(filename::String)

Reads partial dos file.
"""
function qe_read_pdos(filename::String)
    read_tmp = readdlm(filename, skipstart=1)
    energies = read_tmp[:, 1]
    values   = read_tmp[:, 2:end]

    return energies, values
end

"""
    qe_read_projwfc(filename::String)

Reads the output file of a projwfc.x calculation.
Each kpoint will have as many energy dos values as there are bands in the scf/nscf calculation that
generated the density upon which the projwfc.x was called.
Returns:
    states: [(:atom_id, :wfc_id, :j, :l, :m),...] where each j==0 for a non spin polarized calculation.
    kpdos : kpoint => [(:e, :ψ, :ψ²), ...] where ψ is the coefficient vector in terms of the states.
"""
function qe_read_projwfc(filename::String)
    lines  = readlines(filename) .|> strip

    state_tuple = NamedTuple{(:atom_id, :wfc_id, :j, :l, :m), Tuple{Int, Int, Float64, Float64, Float64}}
    states = state_tuple[]
    istart = findfirst(x -> x == "Atomic states used for projection", lines) + 3
    istop  = findnext(isempty, lines, istart) - 1
    for i = istart:istop
        l = replace_multiple(lines[i], "(" => " ", ")" => " ", "," => "", "=" => " ", ":" => "", "#" => " ") |> split
        if length(l) == 11 #spinpolarized
            push!(states, state_tuple((parse.(Int,(l[4], l[7]))..., 0.0, parse.(Float64,( l[9], l[11]))...)))
        else #not spin polarized
            push!(states, state_tuple((parse.(Int,(l[4], l[7]))..., parse.(Float64, (l[9], l[11], l[13]))...)))
        end
    end

    ilowd = findnext(x -> x == "Lowdin Charges:", lines, istop)
    ETuple = NamedTuple{(:e, :ψ, :ψ²), Tuple{Float64, Vector{Float64}, Float64}}
    kdos  = Pair{Vec3{Float64}, Vector{ETuple}}[]
    while istop < ilowd - 2
        istart = findnext(x -> occursin("k = ", x), lines, istop)
        istop  = findnext(isempty, lines, istart) - 1
        k = Vec3(parse.(Float64, split(lines[istart])[3:end]))
        etuples = ETuple[]
        istop_ψ  = istart - 1
        istart_ψ = istart
        while istop_ψ < istop - 1
            eline = replace_multiple(lines[istop_ψ + 2], "="=>"", "(" => " ", ")"=>" ")
            e = parse(Float64, split(eline)[end-1])
            coeffs = zeros(length(states))
            istart_ψ = findnext(x -> x[1:3] == "psi", lines, istop_ψ + 1)
            istop_ψ  = findnext(x -> x[2:4] == "psi", lines, istart_ψ) - 1

            for i=istart_ψ:istop_ψ
                l = replace_multiple(lines[i], "psi =" => " ", "*[#" => " ", "]+" => " ") |> strip |> split
                for k=1:2:length(l)
                    coeffs[parse(Int, l[k+1])] = parse(Float64, l[k])
                end
            end
            ψ² = parse(Float64, split(lines[istop_ψ + 1])[end])
            push!(etuples, (e = e, ψ = coeffs, ψ² = ψ²))
        end
        push!(kdos, k => etuples)
    end
    nkstot = length(kdos)
    nbnd   = length(last(kdos[1]))
    bands = [DFBand(nkstot) for i = 1:nbnd]
    for b in bands
        b.extra[:ψ]  = Vector{Vector{Float64}}(undef, nkstot)
        b.extra[:ψ²] = Vector{Float64}(undef, nkstot)
    end

    for (i, (k, energies)) in enumerate(kdos)
        for (ie, etuple) in enumerate(energies)
            bands[ie].k_points_cryst[i] = zero(Vec{3, Float64})
            bands[ie].k_points_cart[i]  = k
            bands[ie].eigvals[i]    = etuple.e
            bands[ie].extra[:ψ][i]  = etuple.ψ
            bands[ie].extra[:ψ²][i] = etuple.ψ²
        end
    end
    return states, bands
end

"""
    qe_read_polarization(filename::String, T=Float64)

Returns the polarization and modulus.
"""
function qe_read_polarization(filename::String, T=Float64)
    t = qe_read_output(filename, T)
    return t[:polarization], t[:pol_mod]
end

qe_read_vcrel(filename::String, T=Float64) =
	qe_read_output(filename, T) do x
		return x[:cell_parameters], x[:alat], x[:atomic_positions], x[:pos_option]
    end

function alat(flags, pop=false)
    if haskey(flags, :A)
        alat = pop ? pop!(flags, :A) : flags[:A]
		alat *= 1Ang
    elseif haskey(flags, :celldm_1)
        alat = pop ? pop!(flags, :celldm_1) : flags[:celldm_1]
        alat *= 1a₀
    elseif haskey(flags, :celldm)
        alat = pop ? pop!(flags, :celldm)[1] : flags[:celldm][1]
        alat *= 1a₀
    else
        error("Cell option 'alat' was found, but no matching flag was set. \n
               The 'alat' has to  be specified through 'A' or 'celldm(1)'.")
    end
    return alat
end

#TODO handle more fancy cells
function extract_cell!(flags, cell_block)
    if cell_block != nothing
        _alat = 1.0Ang
        if cell_block.option == :alat
            @assert pop!(flags, :ibrav) == 0 "Only ibrav = 0 allowed for now."
            _alat = alat(flags)

        elseif cell_block.option == :bohr
            _alat = 1u"a₀"
        end

        return (_alat .* cell_block.data)'
    end
end

function qe_DFTU(speciesid::Int, parsed_flags::SymAnyDict)
	U  = 0.0
	J0 = 0.0
	J  = [0.0]
	α  = 0.0
	β  = 0.0
	if haskey(parsed_flags, :Hubbard_U) && length(parsed_flags[:Hubbard_U]) >= speciesid
		U = parsed_flags[:Hubbard_U][speciesid]
	end
	if haskey(parsed_flags, :Hubbard_J0) && length(parsed_flags[:Hubbard_J0]) >= speciesid
		J0 = parsed_flags[:Hubbard_J0][speciesid]
	end
	if haskey(parsed_flags, :Hubbard_J) && length(parsed_flags[:Hubbard_J]) >= speciesid
		J = Float64.(parsed_flags[:Hubbard_J][:, speciesid])
	end
	if haskey(parsed_flags, :Hubbard_alpha) && length(parsed_flags[:Hubbard_alpha]) >= speciesid
		α = parsed_flags[:Hubbard_alpha][speciesid]
	end
	if haskey(parsed_flags, :Hubbard_beta) && length(parsed_flags[:Hubbard_beta]) >= speciesid
		β = parsed_flags[:Hubbard_beta][speciesid]
	end
	return DFTU{Float64}(U=U, J0=J0, α=α, β=β, J=J)
end

degree2π(ang) = ang / 180 * π

function qe_magnetization(atid::Int, parsed_flags::SymAnyDict)
	θ = haskey(parsed_flags, :angle1) && length(parsed_flags[:angle1]) >= atid ? parsed_flags[:angle1][atid] : 0.0
	θ = degree2π(θ)
	ϕ = haskey(parsed_flags, :angle2) && length(parsed_flags[:angle2]) >= atid ? parsed_flags[:angle2][atid] : 0.0
	ϕ = degree2π(ϕ)

	start = haskey(parsed_flags, :starting_magnetization) && length(parsed_flags[:starting_magnetization]) >= atid ?
		parsed_flags[:starting_magnetization][atid] : 0.0
	if start isa AbstractVector
		return Vec3{Float64}(start...)
	else
		return start * Vec3{Float64}(sin(θ) * cos(ϕ), sin(θ) * sin(ϕ), cos(θ))
	end
end

function extract_atoms!(parsed_flags, atsyms, atom_block, pseudo_block, cell::Mat3{LT}) where {LT <: Length}
    atoms = Atom{Float64, LT}[]

    option = atom_block.option
    if option == :crystal || option == :crystal_sg
        primv = cell
        cell  = Mat3(Matrix(1.0I, 3, 3))
    elseif option == :alat
        primv = alat(parsed_flags, true) * Mat3(Matrix(1.0I, 3, 3))
    elseif option == :bohr
        primv = 1a₀ .* Mat3(Matrix(1.0I, 3, 3))
    else
        primv = 1Ang .* Mat3(Matrix(1.0I, 3, 3))
    end
    for (at_sym, positions) in atom_block.data
        pseudo = haskey(pseudo_block.data, at_sym) ? pseudo_block.data[at_sym] : (@warn "No Pseudo found for atom '$at_sym'.\nUsing Pseudo()."; Pseudo())
        speciesid = findfirst(isequal(at_sym), atsyms)
        for pos in positions
            push!(atoms, Atom(name=at_sym, element=element(at_sym), position_cart=primv * pos, position_cryst=ustrip.(inv(cell) * pos), pseudo=pseudo, magnetization=qe_magnetization(speciesid, parsed_flags), dftu=qe_DFTU(speciesid, parsed_flags)))
        end
    end

    return atoms
end

function extract_structure!(name, parsed_flags, cell_block, atsyms, atom_block, pseudo_block)
    if atom_block == nothing
        return nothing
    end
    cell = extract_cell!(parsed_flags, cell_block)
    atoms = extract_atoms!(parsed_flags, atsyms, atom_block, pseudo_block, cell)
    return Structure(name, cell, atoms)
end

function separate(f, A::AbstractVector{T}) where T
    true_part = T[]
    false_part = T[]
    while length(A) > 0
        t = pop!(A)
        if f(t)
            push!(true_part, t)
        else
            push!(false_part, t)
        end
    end
    return reverse(true_part), reverse(false_part)
end

"""
    qe_read_input(filename, T=Float64; execs=[Exec("pw.x")], run=true, structure_name="noname")

Reads a Quantum Espresso input file. The `QE_EXEC` inside execs gets used to find which flags are allowed in this input file, and convert the read values to the correct Types.
Returns a `DFInput{QE}` and the `Structure` that is found in the input.
"""
function qe_read_input(filename; execs=[Exec("pw.x")], run=true, structure_name="noname")
    @assert ispath(filename) "$filename is not a valid path."
    t_lines = read(filename)          |>
        String                        |>
        x -> split(x, "\n")          .|>
        x -> cut_after(x, '!')       .|>
        x -> cut_after(x, '#')
    lines = join(t_lines, "\n")       |>
        x -> replace(x, ", " => "\n") |>
        x -> replace(x, "," => " ")   |>
        x -> replace(x, "'" => " ")   |>
        x -> split(x, "\n")          .|>
        strip                         |>
        x -> filter(y -> !occursin("&", y), x) |>
        x -> filter(y -> !(occursin("/", y) && length(y) == 1), x) |>
        x -> filter(!isempty, x)

    exec = getfirst(x->x.exec ∈ QE_EXECS, execs)

    flaglines, lines = separate(x -> occursin("=", x), lines)
    flaglines = strip_split.(flaglines, "=")
    easy_flaglines, difficult_flaglines = separate(x-> !occursin("(", x[1]), flaglines)
    parsed_flags = SymAnyDict()
    #easy flags
    for (f, v) in easy_flaglines
        sym = Symbol(f)
        typ = flagtype(QE, exec, sym)
        if eltype(typ) <: Bool
            v = strip(lowercase(v), '.')
        elseif eltype(typ) <: Number
            v = replace(v, "d" => "e")
        end
        if typ <: AbstractArray
	        typ = eltype(typ)
        end
        tval = typ != String ? parse.((typ,), split(v)) : v
        parsed_flags[sym] = length(tval) == 1 ? tval[1] : tval
    end

    used_lineids = Int[]
    findcard(s) = findfirst(l -> occursin(s, lowercase(l)), lines)
    i = findcard("atomic_species")
    if i!=nothing
        push!(used_lineids, i)
        nat  = parsed_flags[:nat]
        ntyp = parsed_flags[:ntyp]

        pseudos = InputData(:atomic_species, :none, Dict{Symbol, Pseudo}())
        pseudo_dir = pop!(parsed_flags, :pseudo_dir, "./")
        atsyms = Symbol[]
        for k=1:ntyp
            push!(used_lineids, i + k)
            sline = strip_split(lines[i+k])
            atsym = Symbol(sline[1])
            pseudos.data[atsym] = Pseudo(sline[end], pseudo_dir)
            push!(atsyms, atsym)
        end

        i = findcard("cell_parameters")
        append!(used_lineids, [i, i+1, i+2, i+3])
        cell_block = InputData(:cell_parameters,
                               cardoption(lines[i]),
                               Mat3([parse(Float64, split(lines[i+k])[j]) for k=1:3, j=1:3]))

        i = findcard("atomic_positions")
        push!(used_lineids, i)
        atom_block = InputData(:atomic_positions,
                               cardoption(lines[i]),
                               Dict{Symbol, Vector{Point3{Float64}}}() )

        for k=1:nat
            push!(used_lineids, i+k)
            sline = split(lines[i+k])
            atsym = Symbol(sline[1])
            point = Point3(parse.(Float64, sline[2:4]))
            if !haskey(atom_block.data, atsym)
                atom_block.data[atsym] = [point]
            else
                push!(atom_block.data[atsym], point)
            end
        end

        #the difficult flags, can only be present if atomic stuff is found
        for (f, v) in difficult_flaglines
            _s = split(replace(replace(replace(f, "(" => " "), ")" => " "), "," => " "))

            sym = Symbol(_s[1])
            ids = parse.(Int, _s[2:end])
            typ = flagtype(QE, exec, sym)
            v = replace(v, "d" => "e")
            parsedval = parse.((eltype(typ),), split(v))
            if !haskey(parsed_flags, sym)
                if typ <: AbstractMatrix
                    parsed_flags[sym] = length(parsedval) == 1 ? zeros(eltype(typ), ntyp, 10) : fill(zeros(eltype(typ), length(parsedval)), ntyp, 10) #arbitrary limit
                else
                    parsed_flags[sym] = length(parsedval) == 1 ? zeros(eltype(typ), ntyp) : fill(zeros(eltype(typ), length(parsedval)), ntyp)
                end
            end
            if length(ids) == 1
                parsed_flags[sym][ids[1]] = length(parsedval) == 1 ? parsedval[1] : parsedval
            else
                parsed_flags[sym][ids[1], ids[2]] = length(parsedval) == 1 ? parsedval[1] : parsedval
            end
        end
        structure = extract_structure!(structure_name, parsed_flags, cell_block, atsyms, atom_block, pseudos)
        delete!.((parsed_flags,), [:ibrav, :nat, :ntyp, :A, :celldm_1, :celldm])
        delete!.((parsed_flags,), [:Hubbard_U, :Hubbard_J0, :Hubbard_alpha, :Hubbard_beta, :Hubbard_J])
        delete!.((parsed_flags,), [:starting_magnetization, :angle1, :angle2]) #hubbard and magnetization flags
    else
        structure = nothing
    end

    datablocks = InputData[]
    i = findcard("k_points")
    if i!=nothing
        append!(used_lineids, [i, i+1])
        k_option = cardoption(lines[i])
        if k_option == :automatic
            s_line = split(lines[i+1])
            k_data = parse.(Int, s_line)
        else
            nks    = parse(Int, lines[i+1])
            k_data = Vector{Vector{Float64}}(undef, nks)
            for k = 1:nks
                push!(used_lineids, i+1+k)
                k_data[k] = parse.(Float64, split(lines[i+1+k]))
            end
        end
        push!(datablocks, InputData(:k_points, k_option, k_data))
    end
    #Now we need to deal with lines that were not flaglines and not inside the usual cards.
    #Ultimately the goal would be to actually take the info read from the documentation to
    #understand the syntax of the cards, and read them accordingly. In the meantime we do
    #it like this.
    remlines = String[]
    for i=1:length(lines)
        if i ∈ used_lineids
            continue
        end
        push!(remlines, lines[i])
    end
    !isempty(remlines) && push!(datablocks, InputData(:noname, :nooption, remlines))

    dir, file = splitdir(filename)
    return DFInput{QE}(splitext(file)[1], dir, parsed_flags, datablocks, execs, run), structure
end

function qe_writeflag(f, flag, value)
    if isa(value, Matrix)
        for i=1:size(value)[1], j=1:size(value)[2]
            if !iszero(value[i,j])
                write(f, "  $flag($i,$j) = $(value[i, j])\n")
            end
        end
    elseif isa(value, Vector)
        for i=1:length(value)
            if !iszero(value[i])
                if length(value[i]) == 1
                    write(f, "  $flag($i) = $(value[i])\n")
                else
                    write(f, "  $flag($i) =")
                    for v in value[i]
                        write(f, " $v")
                    end
                    write(f, "\n")
                end
            end
        end
    elseif isa(value, String)
        write(f, "  $flag = '$value'\n")
    else
        write(f, "  $flag = $value\n")
    end
end

"""
    save(input::DFInput{QE}, structure, filename::String=inpath(input))

Writes a Quantum Espresso input file.
"""
function save(input::DFInput{QE}, structure, filename::String=inpath(input); relative_positions=true)
    if haskey(flags(input), :calculation)
        setflags!(input, :calculation => replace(input[:calculation], "_" => "-"), print=false)
    end
    open(filename, "w") do f
        writeflag(flag_data) = qe_writeflag(f, flag_data[1], flag_data[2])
        write_dat(data)       = write_data(f, data)

        controls = Dict{Symbol, Dict{Symbol, Any}}()

        for (flag, val) in input.flags
            block, variable = qe_block_variable(input, flag)
            if !haskey(controls, block.name)
                controls[block.name] = Dict{Symbol, Any}()
            end
            controls[block.name][flag] = val
        end

        #Here we try to figure out the correct order of the control blocks
        # first we find the order of the pw.x inputs, the rest should follow.
        blocks2file = []
        for name in [:control, :system, :electrons, :ions, :cell]
            push!(blocks2file, name => pop!(controls, name, nothing))
        end
        for name in keys(controls)
            push!(blocks2file, name => pop!(controls, name, nothing))
        end
        filter!(x->x[2]!=nothing, blocks2file)
        for (name, flags) in blocks2file
            write(f, "&$name\n")
            if name == :system
                nat   = length(atoms(structure))
                ntyp  = length(unique(atoms(structure)))
                # A     = 1.0
                ibrav = 0
                write(f,"  ibrav = $ibrav\n")
                # write(f,"  A = $A\n")
                write(f,"  nat = $nat\n")
                write(f,"  ntyp = $ntyp\n")
            end
            map(writeflag, [(flag, data) for (flag, data) in flags])
            write(f, "/\n\n")
        end
        if exec(input, "pw.x") != nothing
            write_structure(f, input, structure, relative_positions=relative_positions)
        end
        for dat in input.data
            if dat.name != :noname
                if dat.option != :none
                    write(f, "$(uppercase(String(dat.name))) ($(dat.option))\n")
                else
                    write(f, "$(uppercase(String(dat.name)))\n")
                end
            end
            if dat.name == :k_points && dat.option != :automatic
                write(f, "$(length(dat.data))\n")
                write_dat(dat.data)
            else
                write_dat(dat.data)
            end
            write(f, "\n")
        end
    end
    #TODO handle writing hubbard and magnetization better
    delete!.((input.flags,), (:Hubbard_U, :Hubbard_J0, :Hubbard_J, :Hubbard_alpha, :Hubbard_beta,
    						  :starting_magnetization, :angle1, :angle2))
end

function write_structure(f, input::DFInput{QE}, structure; relative_positions=true)
    unique_at = unique(atoms(structure))
    pseudo_lines = String[]
    atom_lines   = String[]
    for at in unique_at
        push!(pseudo_lines, "$(name(at)) $(element(at).atomic_weight)   $(pseudo(at).name)\n")
    end

    for at in atoms(structure)
        push!(atom_lines, position_string(QE, at, relative=relative_positions))
    end

    write(f, "ATOMIC_SPECIES\n")
    write.((f, ), pseudo_lines)

    write(f, "\n")
    write(f, "CELL_PARAMETERS (angstrom)\n")
    write_cell(f, (ustrip.(uconvert.(Ang, cell(structure))))')
    write(f, "\n")

    if relative_positions
	    write(f, "ATOMIC_POSITIONS (crystal) \n")
    else
	    write(f, "ATOMIC_POSITIONS (angstrom) \n")
    end
    write.((f, ), atom_lines)
    write(f, "\n")
end

function qe_generate_pw2waninput(input::DFInput{Wannier90}, qeprefix, runexecs)
    flags = Dict()
    flags[:prefix] = qeprefix
    flags[:seedname] = "$(name(input))"
    flags[:outdir] = "$(joinpath(dir(input), "outputs"))"
    flags[:wan_mode] = "standalone"
    flags[:write_mmn] = true
    flags[:write_amn] = true
    if flag(input, :spin) != nothing
        flags[:spin_component] = flag(input, :spin)
    end
    if flag(input, :wannier_plot) != nothing
        flags[:write_unk] = flag(input, :wannier_plot)
    end
    if any(flag(input, :berry_task) .== ("morb"))
        flags[:write_uHu] = true
    end
    pw2wanexec = Exec("pw2wannier90.x", runexecs[2].dir)
    run = get(input.flags, :preprocess, false) && input.run
    return DFInput{QE}("pw2wan_$(flags[:seedname])", dir(input), flags, InputData[], [runexecs[1], pw2wanexec], run)
end


text:
using JSON
include("./lib/SignalGenerator.jl")
using SignalGenerator

# function genSig(n, mu, sigma, changeStep, sigStd)
function gen_sig1()
  jcont = JSON.parsefile("./config.json")
  saveto = jcont["out_art_sig1"]
  # n, mu, sigma, changeStep, sigStd)
  sig = genSig(n=10, mu=100, sigma=10, changeStep=1.0, sigStd = 0.3)
  f = open(saveto, "w")
  JSON.print(f, sig)
  close(f)
  println("Saving into $saveto")
end
gen_sig1()

max_stars_repo_name: av-maslov/Recurrency

max_issues_repo_name: av-maslov/Recurrency

max_forks_repo_name: av-maslov/Recurrency


text:
# ------------------------------------------------------------------
# Copyright (c) 2018, Júlio Hoffimann Mendes <juliohm@stanford.edu>
# Licensed under the ISC License. See LICENCE in the project root.
# ------------------------------------------------------------------

"""
    returnlevels(xs)

Return periods and levels of data `xs`.
"""
function returnlevels(xs::AbstractVector)
  ms = sort(xs)
  n = length(ms)
  p = (1:n) ./ (n + 1)
  δt = 1 ./ (1 .- p)

  δt, ms
end

"""
    returnlevels(gev, mmin, mmax; nlevels=50)

Return `nlevels` periods and levels of generalized extreme
value distribution `gev` with maxima in the interval `[mmin,mmax]`.
"""
function returnlevels(gev::GeneralizedExtremeValue,
                      mmin::Real, mmax::Real;
                      nlevels::Int=50)
  ms = linspace(mmin, mmax, nlevels)
  δt = 1 ./ (1 .- cdf.(gev, ms))

  δt, ms
end

"""
    meanexcess(xs, k)

Return mean excess of the data `xs` using previous `k` values.
"""
meanexcess(xs::AbstractVector, k::Int) = meanexcess(xs, [k])[1]

function meanexcess(xs::AbstractVector, ks::AbstractVector{Int})
  ys = sort(xs, rev=true)
  [mean(log.(ys[1:k-1])) - log(ys[k]) for k in ks]
end


text:
# Abstract types

## Spatial sampling types

abstract type AbstractSpatialSampling{T} end
abstract type AbstractCartesianSpatialSampling{T}<:AbstractSpatialSampling{T} end

## k-space sampling types

abstract type AbstractKSpaceSampling{T} end
abstract type AbstractKSpaceFixedSizeSampling{T}<:AbstractKSpaceSampling{T} end
abstract type AbstractKSpaceCartesianSampling{T}<:AbstractKSpaceFixedSizeSampling{T} end

max_stars_repo_name: grizzuti/UtilitiesForMRI

max_issues_repo_name: grizzuti/UtilitiesForMRI

max_forks_repo_name: grizzuti/UtilitiesForMRI


text:
function ispangram(input)
  seen = Dict()
  r = r"[a-z]"

  for i in input
    seen[lowercase(i)] = 1
  end

  if length(collect(eachmatch(r,join(keys(seen))))) == 26
    return true
  else
    return false
  end
end


text: # TODO: Need a simple environment for test here!


text:
#
# Replace default random number generator by other that aleviates the
# mutiple-threading problems
#

import Random
using StableRNGs
using Future: randjump

function init_random(options)
    if options.seed > 0
        if options.StableRNG == true
            RNG = StableRNGs.StableRNG(options.seed)
        else
            RNG = [randjump(Random.MersenneTwister(options.seed), big(10)^20)]
            foreach(_ -> push!(RNG, randjump(last(RNG), big(10)^20)), 2:Threads.nthreads())
        end
    else
        if options.StableRNG == true
            seed = abs(rand(Int))
            RNG = StableRNGs.StableRNG(seed)
        else
            RNG = [randjump(Random.MersenneTwister(), big(10)^20)]
            foreach(_ -> push!(RNG, randjump(last(RNG), big(10)^20)), 2:Threads.nthreads())
        end
    end
    return RNG
end

random(RNG::StableRNGs.LehmerRNG) = rand(RNG)
random(RNG::StableRNGs.LehmerRNG, arg) = rand(RNG, arg)

random(RNG::Array{T}) where {T} = rand(RNG[Threads.threadid()])
random(RNG::Array{T}, arg) where {T} = rand(RNG[Threads.threadid()], arg)


text:
function rescale!(
    dist::Distributions.DiscreteNonParametric,
    examinees::Vector{<:AbstractExaminee};
    metric::Vector{Float64} = [0.0, 1.0]
)
#check observed metric
    N = size(examinees, 1)
    Wk = mapreduce(e -> e.latent.posterior.p, +, examinees) ./ N
    Xk = dist.support
    observed =
        [LinearAlgebra.dot(Wk, Xk), sqrt(LinearAlgebra.dot(Wk, Xk .^ 2))]
    observed = [
        observed[1] - metric[1],
        observed[2] / metric[2],
    ]
    Xk2, Wk2 = my_rescale(Xk, Wk, observed)
    Wk = cubic_spline_int(Xk, Xk2, Wk2)
    dist = Distributions.DiscreteNonParametric(Xk2[2:(end-1)], Wk; check_args=false)
    return nothing
end

function _rescale(
    dist::Distributions.DiscreteNonParametric,
    latents::Vector{<:AbstractLatent};
    metric::Vector{Float64} = [0.0, 1.0]
)
#check observed metric
    N = size(latents, 1)
    Wk = mapreduce(l -> l.posterior.p, +, latents) ./ N
    Xk = dist.support
    observed =
        [LinearAlgebra.dot(Wk, Xk), sqrt(LinearAlgebra.dot(Wk, Xk .^ 2))]
    observed = [
        observed[1] - metric[1],
        observed[2] / metric[2],
    ]
    Xk2, Wk2 = my_rescale(Xk, Wk, observed)
    Wk = cubic_spline_int(Xk, Xk2, Wk2)
    return Distributions.DiscreteNonParametric(Xk2[2:(end-1)], Wk; check_args=false)::Distributions.DiscreteUnivariateDistribution
end


text:
arrowlength = maxd/30
arrowheight = (hlim[2]-hlim[1])/20

struct ACASRender
    m::UnresponsiveACASMDP
    s::ACASState
end

POMDPModelTools.render(m::UnresponsiveACASMDP, step::NamedTuple) = render(m, step.s)
POMDPModelTools.render(m::UnresponsiveACASMDP, s::ACASState) = ACASRender(m, s)
POMDPModelTools.render(m::UnresponsiveACASMDP, s::AbstractVector) = render(m, convert(ACASState, s))
POMDPModelTools.render(m::UnresponsiveACASMDP) = ACASRender(m, rand(initialstate(m)))

ownarrow(s::ACASState) = [-arrowlength/2, arrowlength/2, -arrowlength/2], [s.h_o+arrowheight/2, s.h_o, s.h_o-arrowheight/2]
intarrow(s::ACASState) = [s.d+arrowlength/2, s.d-arrowlength/2, s.d+arrowlength/2], [s.h_i+arrowheight/2, s.h_i, s.h_i-arrowheight/2]
title(s::ACASState) = "[hₒ=$(s.h_o) ft, ḣₒ=$(s.hdot_o) ft/min, hᵢ=$(s.h_i) ft, d=$(s.d) ft]"

function Base.show(io::IO, m::MIME"text/plain", r::ACASRender)
    s = r.s
    ox, oy = ownarrow(s)
    plt = UnicodePlots.lineplot(ox, oy,
                                name="Ownship >",
                                title=title(s),
                                width=70,
                                grid=false,
                                xlabel="d (ft)",
                                ylabel="h (ft)",
                                ylim=hlim,
                                xlim=(-0.1*maxd,1.1*maxd))
    ix, iy = intarrow(s)
    UnicodePlots.lineplot!(plt, ix, iy, name="Intruder <")
    show(io, m, plt)
end

function transform(xy)
    x, y = xy
    r = x/maxd
    d = (maximum(hlim) - y)/(hlim[2]-hlim[1])
    return (r, d)
end

function Base.show(io::IO, m::MIME"text/html", r::ACASRender)
    s = r.s
    g = compose(context(),
            text(0.5, 0.05, title(s), hcenter, vtop),
            (context(), polygon(transform.(collect(zip(ownarrow(s)...)))),
                        text(transform((0, s.h_o - arrowheight))..., "Ownship", hleft, vtop),
                        fill("blue")),
            (context(), polygon(transform.(collect(zip(intarrow(s)...)))),
                        text(transform((s.d, s.h_i - arrowheight))..., "Intruder", hright, vtop),
                        fill("green")),
           )
    show(io, m, g)
end


text:
module CassetteSpeen

using ..Ahorn, Maple

@mapdef Entity "brokemiahelper/cassetteSpinner" CassetteSpinner(x::Integer, y::Integer, index::Integer=0, tempo::Number=1, attachToSolid::Bool=false)

colorNames = Dict{String, Int}(
    "Blue" => 0,
    "Rose" => 1,
    "Bright Sun" => 2,
    "Malachite" => 3
)

const placements = Ahorn.PlacementDict(
    "Cassette Spinner ($index - $color)" => Ahorn.EntityPlacement(
        CassetteSpinner,
        "point",
        Dict{String, Any}(
            "index" => index,
        )
    ) for (color, index) in colorNames
)

Ahorn.editingOptions(entity::CassetteSpinner) = Dict{String, Any}(
    "index" => colorNames
)

#function Ahorn.render(ctx::Ahorn.Cairo.CairoContext, entity::CassetteCar, room::Maple.Room)
#     # x, y = Ahorn.position(entity)

#     # width = 64
#     # height = 18

#     # # index = Int(get(entity.data, "index", 0))
#     # # color = get(colors, index, defaultColor)

#     # Ahorn.drawRectangle(ctx, 0, 0, width, height, defaultColor, defaultColor)
#     x = Int(get(entity.data, "x", 0))
#     y = Int(get(entity.data, "y", 0))

#     width = Int(get(entity.data, "width", 32))
#     height = Int(get(entity.data, "height", 32))

     #Ahorn.drawRectangle(ctx, 0, 0, 64, 18, (0.0, 0.0, 1.0, 0.4), (0.0, 0.0, 1.0, 1.0))
#end

Ahorn.resizable(entity::CassetteSpinner) = false, false

function carSelection(entity::CassetteSpinner)
	rectangle = Ahorn.getEntityRectangle(entity)
	rectangle = Ahorn.Rectangle(rectangle.x-8, rectangle.y-8, 16, 16)
	return rectangle
end

Ahorn.selection(entity::CassetteSpinner) = carSelection(entity)

colors = Dict{Integer, Ahorn.colorTupleType}(
    1 => (240, 73, 190, 255) ./ 255,
	2 => (252, 220, 58, 255) ./ 255,
	3 => (56, 224, 78, 255) ./ 255,
)

defaultColor = (73, 170, 240, 255) ./ 255
borderMultiplier = (0.9, 0.9, 0.9, 1)

end

max_stars_repo_name: Brokemia/BrokemiaHelper

max_issues_repo_name: Brokemia/BrokemiaHelper

max_forks_repo_name: Brokemia/BrokemiaHelper


text:
include("initialize.jl")
include("shuffle_config.jl")

"""
Calculates the pair- and triplet-SRO sum separately.
Inputs are a 'Backpack' and the sequence to be calculated.
"""
function sro(X::Backpack, sequence)
    A = copy(X.alpha)
    B = copy(X.beta)
    alpha_sum = 0
    beta_sum = 0
# Pair part: Going through all combinations (without double counts)
    for i = 1:length(sequence), j = i+1:length(sequence)
        if sequence[i][2] != 0 && sequence[j][2] != 0
            A[sequence[i][2], sequence[j][2], X.shell_matrix[i, j]] -=
            X.pair_prefactors[sequence[i][2], sequence[j][2], X.shell_matrix[i, j]]
        end
    end
# Summation (and substracting all pairs between to equal elements)
    for i = 1:size(A)[3]
        alpha_sum += 1/i*(sum(broadcast(abs, A[:, :, i])) - sum([abs(A[n, n, i]) for n in 1:size(A)[1]]))
    end

# Triplet part: This time, the tuple values of each element not only indicate
# the index, but also act as a key for the right mole fraction:
    for i = 1:length(sequence), j = 1:length(sequence), k = 1:length(sequence)
        if sequence[i][2] != 0 && sequence[j][2] != 0 && sequence[k][2] != 0
            si, sj, sk = X.triplet_shell_matrix[i, j, k]
            (si, sj, sk) == (0,0,0) && continue
            spa = sequence[i][2]
            spb = sequence[j][2]
            spc = sequence[k][2]
            xa = X.mole_fractions[sequence[i]]
            xb = X.mole_fractions[sequence[j]]
            xc = X.mole_fractions[sequence[k]]
            prefactor = X.trip_prefactors[(si, sj, sk)]/(xa*xb*xc)
            B[spa, spb, spc, si, sj, sk] -= prefactor
        end
    end
# Summation (and substracting all triplets between to equal elements)
    for i = 1:size(B)[4], j = 1:size(B)[5], k = 1:size(B)[6]
        beta_sum += 1/(3*i*j*k)*(sum(broadcast(abs, B[:, :, :, i, j, k])) - sum([abs(B[n, n, n, i, j, k]) for n in 1:size(B)[1]]))
    end
    return alpha_sum, beta_sum
end

# Helping structure for the results
struct SRO_results
    pairs::Vector
    triplets::Vector
    sro_sum::Vector
    pair_extrema::Vector{Tuple{Float64,Int64}}
    triplet_extrema::Vector{Tuple{Float64,Int64}}
    overall_extrema::Vector{Tuple{Float64,Int64}}
end

"""
Helps to evaluate the desired calculations and outputs the minima and maxima.
"""
function evaluate(X::Backpack, permutations, weight_factor::Float64)
    perms = permutations
    my_pairs = []
    my_triplets = []
    my_sros = []
    for i = 1:length(perms)
        current_sro = sro(X, perms[i])
        push!(my_pairs, current_sro[1])
        push!(my_triplets, current_sro[2])
        push!(my_sros, current_sro[1]+weight_factor*current_sro[2])
    end
    pair_extrema = [findmin(my_pairs), findmax(my_pairs)]
    triplet_extrema = [findmin(my_triplets), findmax(my_triplets)]
    overall_extrema = [findmin(my_sros), findmax(my_sros)]
    return SRO_results(my_pairs, my_triplets, my_sros, pair_extrema, triplet_extrema, overall_extrema)
end

max_stars_repo_name: ulipoto/WCSRO

max_issues_repo_name: ulipoto/WCSRO

max_forks_repo_name: ulipoto/WCSRO


text:
using Base.Test

# meta.jl
import WaveletScattering: get_γ, get_θ, get_χ, get_j,
    get_bandwidth, get_centerfrequency, get_qualityfactor, get_scale

# spec.jl
import WaveletScattering: default_ɛ, uncertainty

# spec1d.jl
import WaveletScattering: Spec1D

# bandwidths, centerfrequencies, qualityfactors, scales, uncertainty
numerictypes = [Float16, Float32, Float64]
nfos = [1, 2, 4, 8, 12, 24, 32]
for T in numerictypes, nfo in nfos, max_q in nfos[nfos.<=nfo],
    log2_s in (7+ceil(Int, log2.(nfo)):18), max_s in [max_q*exp2.(5:14); Inf]
    machine_precision = max(1e-10, default_ɛ(T))
    spec = Spec1D(
            signaltype=T,
            n_filters_per_octave=nfo,
            max_qualityfactor=max_q,
            log2_size=log2_s,
            max_scale=max_s)
    γs = map(get_γ, spec.ψmetas)
    θs = map(get_θ, spec.ψmetas)
    χs = map(get_χ, spec.ψmetas)
    js = map(get_j , spec.ψmetas)
    bws = map(get_bandwidth, spec.ψmetas)
    ξs = map(get_centerfrequency, spec.ψmetas)
    qs = map(get_qualityfactor, spec.ψmetas)
    scs = map(get_scale, spec.ψmetas)
    # gammas
    @test eltype(γs) <: Int16
    # chis
    @test eltype(χs) <: Int8
    # js
    @test eltype(js) <: Int8
    @test map(Int16, χs) + nfo * map(Int16, js) == γs
    # thetas
    @test all(θs .== 0)
    @test eltype(θs) <: Int8
    # bandwidths
    @test_approx_eq bws ξs./qs
    # centerfrequencies
    @test_approx_eq ξs[1] spec.motherfrequency
    difflogξs = diff(log2.(ξs[:]))
    @test difflogξs ≈ (-ones(difflogξs)/spec.n_filters_per_octave)
    @test all(ξs.>0.0)
    @test ξs ≈ bws.*qs
    # qualityfactors
    @test all(qs.>=0.0)
    @test all(qs.<=max_q)
    @test qs ≈ ξs./bws
    # scales
    @test all(scs.>0.0)
    @test all(scs[qs.>1.0] .< (max_s+machine_precision))
    @test all(scs .< (exp2.(spec.log2_size)+machine_precision))
    # uncertainty
    empirical_uncertainty = bws .* scs
    @test all(abs.(diff(empirical_uncertainty[:])) .< machine_precision)
    @test all(abs.(uncertainty(spec)-empirical_uncertainty).< machine_precision)
end


text:
using Test
using Scruff
using Scruff.Utils
using Scruff.RTUtils
using AbstractTrees
using DataStructures

@testset "Util" begin
    @testset "Cartesian product" begin
        @testset "With no arrays" begin
            @test cartesian_product(Array{Int,1}[]) == [[]]
        end
        @testset "With single array" begin
            @test cartesian_product([[1,2]]) == [[1], [2]]
        end
        @testset "With two arrays" begin
            @test cartesian_product([[1,2], [3,4,5]]) ==
                    [[1,3], [1,4], [1,5], [2,3], [2,4], [2,5]]
        end
        @testset "With empty array" begin
            @test cartesian_product([[1,2], [], [3,4,5]]) == []
        end
    end

    @testset "Topsort" begin
        @testset "Acyclic" begin
            graph = Dict(2 => [1], 5 => [2, 1], 3 => [2], 4 => [1])
            ord = topsort(graph)
            @test length(ord) == 5
            @test 1 in ord
            @test 2 in ord
            @test 3 in ord
            @test 4 in ord
            @test 5 in ord
            i1 = indexin(1, ord)[1]
            i2 = indexin(2, ord)[1]
            i3 = indexin(3, ord)[1]
            i4 = indexin(4, ord)[1]
            i5 = indexin(5, ord)[1]
            @test i1 < i2
            @test i2 < i5
            @test i2 < i3
            @test i1 < i4
        end

        @testset "Cyclic" begin
            graph = Dict(2 => [1, 3], 5 => [2, 1], 3 => [2], 4 => [1])
            ord = topsort(graph)
            @test length(ord) == 5
            @test 1 in ord
            @test 2 in ord
            @test 3 in ord
            @test 4 in ord
            @test 5 in ord
            i1 = indexin(1, ord)[1]
            i2 = indexin(2, ord)[1]
            i3 = indexin(3, ord)[1]
            i4 = indexin(4, ord)[1]
            i5 = indexin(5, ord)[1]
            # i2 and i3 could be in any order, but must be greater than i1 and less than i5
            @test i1 < i2
            @test i1 < i3
            @test i2 < i5
            @test i3 < i5
            @test i1 < i4
        end
    end

    @testset "Continuous utilities" begin
        @testset "Intervals" begin
            @test make_intervals([0, 1, 3]) == [(-Inf,0.5), (0.5,2), (2,Inf)]
            @test make_intervals(([1])) == [(-Inf, Inf)]
        end
        @testset "Linear value" begin
            @test linear_value([2,-3], -1, [4,-2]) == 13
            @test linear_value([], -1, []) == -1
        end
        @testset "Bounded linear value" begin
            is1 = [(3,5), (-3,-1)]
            cs1 = [[3,-1], [3,-3], [5,-1], [5,-3]]
            vs1 = map(x -> linear_value([2,-3], -1, x), cs1)
            @test bounded_linear_value([2,-3], -1, is1) ==
                (minimum(vs1), maximum(vs1))
        end
        @testset "Normal density" begin
            @test isapprox(normal_density(0, 0, 1), 0.3989, atol = 0.0001)
            @test isapprox(normal_density(1, 1, 1), 0.3989, atol = 0.0001)
            @test isapprox(normal_density(-1, 0, 1), 0.2420, atol = 0.0001)
            @test isapprox(normal_density(-2, 0, 4), 0.2420 / 2, atol = 0.0001)
        end
    end

    @testset "Memo" begin
        count = 0
        function f(x)
            count += 1
            return x
        end
        g = memo(f)
        @test g(1) == 1
        @test count == 1
        @test g(1) == 1
        @test count == 1
        @test g(2) == 2
        @test count == 2
    end

    @testset "factor" begin
        fact1 = Factor{1}((2,), (1,), [0.1, 0.9])
        fact2 = Factor{1}((3,), (2,), [0.2, 0.3, 0.5])
        prod12 = Factor{2}((3, 2), (2, 1), [0.02, 0.18, 0.03, 0.27, 0.05, 0.45])
        prod21 = Factor{2}((2, 3), (1, 2), [0.02, 0.03, 0.05, 0.18, 0.27, 0.45])
        fact23 = Factor{2}((3, 2), (2, 3), [0.2, 0.8, 0.3, 0.7, 0.4, 0.6])
        prod1223 = Factor{3}((3, 2, 2), (2, 3, 1),
                                [0.004, 0.036, 0.016, 0.144,
                                0.009, 0.081, 0.021, 0.189,
                                0.02, 0.18, 0.03, 0.27])
        prod2123 = prod1223 # Even though the order of variables in one of the
                            # arguments is different, the result is the same
        prod122321 = Factor{3}((2, 3, 2), (1, 2, 3),
                                [0.00008, 0.00032, 0.00027, 0.00063,
                                0.001, 0.0015, 0.00648, 0.02592,
                                0.02187,   0.05103,   0.081,  0.1215])
        sum1 = Factor{2}((3, 2), (2, 3),
                            [0.00008 + 0.00648, 0.00032 + 0.02592, 0.00027 + 0.02187,
                            0.00063 + 0.05103, 0.001 + 0.081, 0.0015 + 0.1215])
        sum2 = Factor{2}((2, 2), (1, 3),
                            [0.00008 + 0.00027 + 0.001, 0.00032 + 0.00063 + 0.0015,
                            0.00648 + 0.02187 + 0.081, 0.02592 + 0.05103 + 0.1215])
        sum3 = Factor{2}((2, 3), (1, 2),
                            [0.00008 + 0.00032, 0.00027 + 0.00063, 0.001 + 0.0015,
                            0.00648 + 0.02592, 0.02187 + 0.05103, 0.081 + 0.1215])
        @testset "Product" begin
            @testset "Multiplying two independent factors" begin
                @test isapprox(product(fact1, fact2), prod12)
                @test isapprox(product(fact2, fact1), prod21)
            end
            @testset "Multiplying two factors with a shared variable" begin
                @test isapprox(product(prod12, fact23), prod1223)
                @test isapprox(product(prod21, fact23), prod2123)
            end
            @testset "Multiplying factors with two shared variables" begin
                @test isapprox(product(prod1223, prod21), prod122321)
            end
        end

        @testset "Summing out a variable" begin
            @testset "With variables" begin
                @test isapprox(sum_over(prod122321, 1), sum1)
                @test isapprox(sum_over(prod122321, 2), sum2)
                @test isapprox(sum_over(prod122321, 3), sum3)
            end
            @testset "When last variable" begin
                f = sum_over(fact1, 1)
                @test isnothing(f.dims)
                @test isnothing(f.keys)
                @test f.entries == [1.0]
            end
        end
    end
end


text:
"""
*ProspectTheory*

`ProspectTheory` constructs a model object for cummulative prospect theory.
By default, parameters for utility curvature and probability weigting are equal gains and losses.

- `α`: utility curvature for gains
- `β`: utility curvature for losses
- `γg`: probability weighting parameter for gains
- `γl`: probability weighting parameter for losses
- `λ`: loss aversion parameter

Constructor
````julia
ProspectTheory(;α=.80, β=α, γg=.70, γl=γg, λ=2.25)
````
*References*

Fennema, H., & Wakker, P. (1997). Original and cumulative prospect theory: A discussion of empirical differences. Journal of Behavioral Decision Making, 10(1), 53-64.

Tversky, A., & Kahneman, D. (1992). Advances in prospect theory: Cumulative representation of uncertainty. Journal of Risk and uncertainty, 5(4), 297-323.
"""
mutable struct ProspectTheory{T1,T2,T3,T4,T5} <:UtilityModel
    α::T1
    β::T2
    γg::T3
    γl::T4
    λ::T5
end

function ProspectTheory(;α=.80, β=α, γg=.70, γl=γg, λ=2.25)
    return ProspectTheory(α, β, γg, γl, λ)
end

"""
*compute_utility*

`compute_utility` computes utility of gamble outcomes according to prospect theory

- `model`: a model object for prospect theory
- `gamble`: a gamble object

Function Signature
````julia
compute_utility(model::ProspectTheory, gamble::Gamble)
````
"""
function compute_utility(model::ProspectTheory, gamble)
    @unpack α,β,λ = model
    vl,vg = split_values(gamble)
    utilg = vg.^α
    utill = @. -λ*abs(vl)^β
    return [utill; utilg]
end

"""
*compute_weights*

`compute_weights` computes decision weights based on cummulative outcomes

- `model`: a model object for prospect theory
- `gamble`: a gamble object

Function Signature
````julia
compute_weights(model::ProspectTheory, gamble::Gamble)
````
"""
function compute_weights(model::ProspectTheory, gamble::Gamble)
    pl,pg = split_probs(gamble)
    @unpack γg,γl = model
    ω = [_compute_weights(pl, γl); _compute_weights(pg, γg)]
    return ω
end

"""
*_compute_weights*

`_compute_weights` computes decision weights based on cummulative outcomes

- `p`: a probability vector
- `γ`: parameter that controls weighting of low and high probabilities

Function Signature
````julia
_compute_weights(p, γ)
````
"""
function _compute_weights(p, γ)
    n = length(p)
    f(i) = weight(sum(p[i:n]), γ) - weight(sum(p[(i+1):n]), γ)
    ω = [f(i) for i in 1:n-1]
    isempty(p) ? nothing : push!(ω, weight(p[n], γ))
    return ω
end

weight(p, γ) = (p^γ)/(p^γ + (1-p)^γ)^(1/γ)

function sort!(model::ProspectTheory, gamble)
    @unpack p,v = gamble
    i = sortperm(v)
    p .= p[i]; v .= v[i]
    gains = v .>= 0
    pl = @view p[.!gains]
    vl = @view v[.!gains]
    reverse!(vl); reverse!(pl)
    return nothing
end

function split_values(gamble)
    @unpack v = gamble
    gains = v .>= 0
    vg = @view v[gains]
    vl = @view v[.!gains]
    return vl,vg
end

function split_probs(gamble)
    @unpack v,p = gamble
    gains = v .>= 0
    pg = @view p[gains]
    pl = @view p[.!gains]
    return pl,pg
end


text:
module DiscreteEntropyEstimators

import Base: convert
import StatsBase

export
    # generic functions
    entropy,

    # helper functions
    convert,

    # generic types
    Frequencies,
    Samples,
    Observations,

    # estimators
    MaximumLikelihoodEntropyEstimator,
    MLEntropyEstimator,
    MLEE,
    MillerMaddowBiasCorrectedEntropyEstimator,
    MillerMaddowBiasCorrected,
    MMBCEE,

    # convenient functions
    entropy_ml

include("generic.jl")
include("estimators.jl")

end # module


text:
struct Support; node::Vector{Int}; type::Array{Bool,2}; end
struct NodalLoad; node::Vector{Int}; val::Array{Float64,2}; end

function StructDomain(Nx::Int, Ny::Int, Lx::Float64, Ly::Float64, Supp, Load)

	# generate NODE using meshgrid
	X, Y = LinRange(0., Lx, Nx+1), LinRange(0., Ly, Ny+1)
	X, Y = MeshGrid2D(X, Y); NODE = [X[:] Y[:]]; Nn = size(NODE, 1)

	# generate 4-node structured ELEM
	k = 0; ELEM = zeros(Int, Nx*Ny, 4)
	for j=1:Ny, i=1:Nx
		n1, n2 = (i-1)*(Ny+1)+j, i*(Ny+1)+j
		k+=1; ELEM[k, :] = [n1 n2 n2+1 n1+1]
	end

	# tolerance for coords search
	tol = 0.01*min(Lx/Nx, Ly/Ny)

	# generate SUPP
	SNode=zeros(Int,Nn); SType=zeros(Bool,Nn,2); NSupp=0
	for s in Supp
		node = FindNodeSet(NODE, s[1:2], tol)
		for j in node; NSupp+=1; SNode[NSupp]=j; SType[NSupp,:]=s[3:4]; end
	end
	SUPP = Support(SNode[1:NSupp], SType[1:NSupp,:])

	# generate LOAD
	LNode=zeros(Int,Nn); LVal=zeros(Nn,2); NLoad=0
	for L in Load
		node = FindNodeSet(NODE, L[1:2], tol)
		for j in node; NLoad+=1; LNode[NLoad]=j; LVal[NLoad,:]=L[3:4]; end
	end
	LOAD = NodalLoad(LNode[1:NLoad], LVal[1:NLoad,:])

	return NODE, ELEM, SUPP, LOAD
end

#------------------------------------------------------------------------------
# Implementation of 2D mesh grid
function MeshGrid2D(x, y)
	nx, ny = length(x), length(y)
	return x' .* ones(ny), ones(nx)' .* y
end

#------------------------------------------------------------------------------
# Find node-set function
function FindNodeSet(Node, pos, tol::Float64)

	x, y = pos
	X, Y = Node[:,1], Node[:,2]

	xMin, xMax, yMin, yMax = -Inf, +Inf, -Inf, +Inf
	if x>=0; xMin=x-tol; xMax=x+tol; end
	if y>=0; yMin=y-tol; yMax=y+tol; end

	return findall((X.>xMin).&(X.<xMax).&(Y.>yMin).&(Y.<yMax))
end

max_stars_repo_name: ricardoaf/jutop

max_issues_repo_name: ricardoaf/jutop

max_forks_repo_name: ricardoaf/jutop


text:
include("line_and_plane.jl")
include("circle_and_sphere.jl")


text:
import Base: eachindex, AbstractVector, AbstractMatrix, map, ==
export FiniteMean, FiniteKernel, LhsFiniteCrossKernel, RhsFiniteCrossKernel,
    FiniteCrossKernel

const IntVec = AbstractVector{<:Integer}

"""
    FiniteMean <: Function

A mean function defined on a finite index set. Has a method of `mean` which requires no
additional data.
"""
struct FiniteMean <: MeanFunction
    μ::MeanFunction
    X::AbstractVector
end
(μ::FiniteMean)(n) = μ.μ(getindex(μ.X, n))
==(μ::FiniteMean, μ′::FiniteMean) = μ.μ == μ′.μ && μ.X == μ′.X
eachindex(μ::FiniteMean) = eachindex(μ.X)
length(μ::FiniteMean) = length(μ.X)
_map(μ::FiniteMean, q::IntVec) = map(μ.μ, μ.X[q])
function print(io::IO, μ::FiniteMean)
    println(io, "FiniteMean ($(length(μ))):")
    print(io, Shunted(Shunt(4, ' '), μ.μ))
end

# Sugar
map(μ::FiniteMean, ::Colon) = map(μ, eachindex(μ))

"""
    FiniteKernel <: Kernel

A kernel valued on a finite index set. Has a method of `cov` which requires no additional
data.
"""
struct FiniteKernel <: Kernel
    k::Kernel
    X::AbstractVector
end
(k::FiniteKernel)(n, n′) = k.k(getindex(k.X, n), getindex(k.X, n′))
(k::FiniteKernel)(n) = k.k(k.X[n], k.X[n])
==(k::FiniteKernel, k′::FiniteKernel) = k.k == k′.k && k.X == k′.X
eachindex(k::FiniteKernel) = eachindex(k.X)
length(k::FiniteKernel) = length(k.X)
function print(io::IO, k::FiniteKernel)
    println(io, "FiniteKernel $(size(k)):")
    print(io, Shunted(Shunt(4, ' '), k.k))
end
_map(k::FiniteKernel, q::IntVec) = map(k.k, k.X[q])
_map(k::FiniteKernel, q::IntVec, q′::IntVec) = map(k.k, k.X[q], k.X[q′])
function _pairwise(k::FiniteKernel, q::IntVec)
    return pairwise(k.k, k.X[q])
end
_pairwise(k::FiniteKernel, q::IntVec, q′::IntVec) = pairwise(k.k, k.X[q], k.X[q′])

# Sugar
map(k::FiniteKernel, ::Colon) = map(k, eachindex(k))
map(k::FiniteKernel, ::Colon, ::Colon) = map(k, eachindex(k), eachindex(k))
pairwise(k::FiniteKernel, ::Colon) = pairwise(k, eachindex(k))
pairwise(k::FiniteKernel, ::Colon, ::Colon) = pairwise(k, eachindex(k), eachindex(k))
pairwise(k::FiniteKernel, ::Colon, q′::IntVec) = pairwise(k, eachindex(k), q′)
pairwise(k::FiniteKernel, q::IntVec, ::Colon) = eachindex(k, q, eachindex(k))

"""
    LhsFiniteCrossKernel <: CrossKernel

A cross kernel whose first argument is defined on a finite index set. Useful for defining
cross-covariance between a Finite kernel and other non-Finite kernels.
"""
struct LhsFiniteCrossKernel <: CrossKernel
    k::CrossKernel
    X::AbstractVector
end
(k::LhsFiniteCrossKernel)(n, x) = k.k(getindex(k.X, n), x)
==(k::LhsFiniteCrossKernel, k′::LhsFiniteCrossKernel) = k.k == k′.k && k.X == k′.X
size(k::LhsFiniteCrossKernel, N::Int) = N == 1 ? length(k.X) : size(k.k, N)
eachindex(k::LhsFiniteCrossKernel, N::Int) = N == 1 ? eachindex(k.X) : eachindex(k.k, 2)
_map(k::LhsFiniteCrossKernel, q::IntVec, X′::AV) = map(k.k, k.X[q], X′)
_pairwise(k::LhsFiniteCrossKernel, q::IntVec, X′::AV) = pairwise(k.k, k.X[q], X′)
function print(io::IO, k::LhsFiniteCrossKernel)
    println(io, "LhsFiniteCrossKernel $(size(k)):")
    print(io, Shunted(Shunt(4, ' '), k.k))
end

# Sugar
map(k::LhsFiniteCrossKernel, ::Colon, X′::AV) = map(k, eachindex(k, 1), X′)
pairwise(k::LhsFiniteCrossKernel, ::Colon, X′::AV) = pairwise(k, eachindex(k, 1), X′)

"""
    RhsFiniteCrossKernel <: CrossKernel

A cross kernel whose second argument is defined on a finite index set. You can't really do
anything with this object other than use it to construct other objects.
"""
struct RhsFiniteCrossKernel <: CrossKernel
    k::CrossKernel
    X′::AbstractVector
end
(k::RhsFiniteCrossKernel)(x, n′) = k.k(x, getindex(k.X′, n′))
==(k::RhsFiniteCrossKernel, k′::RhsFiniteCrossKernel) = k.k == k′.k && k.X′ == k′.X′
size(k::RhsFiniteCrossKernel, N::Int) = N == 2 ? length(k.X′) : size(k.k, N)
eachindex(k::RhsFiniteCrossKernel, N::Int) = N == 1 ? eachindex(k.k, 1) : eachindex(k.X′)
_map(k::RhsFiniteCrossKernel, X::AV, q′::IntVec) = map(k.k, X, k.X′[q′])
_pairwise(k::RhsFiniteCrossKernel, X::AV, q′::IntVec) = pairwise(k.k, X, k.X′[q′])
function print(io::IO, k::RhsFiniteCrossKernel)
    println(io, "RhsFiniteCrossKernel $(size(k)):")
    print(io, Shunted(Shunt(4, ' '), k.k))
end

# Sugar
map(k::RhsFiniteCrossKernel, X::AV, ::Colon) = map(k, X, eachindex(k, 2))
pairwise(k::RhsFiniteCrossKernel, X::AV, ::Colon) = pairwise(k, X, eachindex(k, 2))

"""
    FiniteCrossKernel <: CrossKernel

A cross kernel valued on a finite index set. Has a method of `xcov` which requires no
additional data.
"""
struct FiniteCrossKernel <: CrossKernel
    k::CrossKernel
    X::AbstractVector
    X′::AbstractVector
end
(k::FiniteCrossKernel)(n, n′) = k.k(getindex(k.X, n), getindex(k.X′, n′))
==(k::FiniteCrossKernel, k′::FiniteCrossKernel) = k.k == k′.k && k.X == k′.X && k.X′ == k′.X′
size(k::FiniteCrossKernel, N::Int) = N == 1 ? length(k.X) : (N == 2 ? length(k.X′) : 1)
eachindex(k::FiniteCrossKernel, N::Int) = N == 1 ? eachindex(k.X) : eachindex(k.X′)
_map(k::FiniteCrossKernel, q::IntVec, q′::IntVec) = map(k.k, k.X[q], k.X′[q′])
_pairwise(k::FiniteCrossKernel, q::IntVec, q′::IntVec) = pairwise(k.k, k.X[q], k.X′[q′])
function print(io::IO, k::FiniteCrossKernel)
    println(io, "FiniteCrossKernel $(size(k)):")
    print(io, Shunted(Shunt(4, ' '), k.k))
end

function zero(k::CrossKernel)
    if size(k, 1) < Inf && size(k, 2) < Inf
        return FiniteZeroCrossKernel(eachindex(k, 1), eachindex(k, 2))
    elseif size(k, 1) < Inf
        return LhsFiniteZeroCrossKernel(eachindex(k, 1))
    elseif size(k, 2) < Inf
        return RhsFiniteZeroCrossKernel(eachindex(k, 2))
    else
        return ZeroKernel{Float64}()
    end
end

# Sugar
map(k::FiniteCrossKernel, ::Colon, ::Colon) = map(k, eachindex(k, 1), eachindex(k, 2))
function pairwise(k::FiniteCrossKernel, ::Colon, ::Colon)
    return pairwise(k, eachindex(k, 1), eachindex(k, 2))
end
pairwise(k::FiniteCrossKernel, ::Colon, X′::IntVec) = pairwise(k, eachindex(k, 1), X′)
pairwise(k::FiniteCrossKernel, X::IntVec, ::Colon) = pairwise(k, X, eachindex(k, 2))

################################## Optimisations for zeros #################################

struct FiniteZeroMean{TX} <: MeanFunction
    X::TX
end
length(μ::FiniteZeroMean) = length(μ.X)
==(μ::FiniteZeroMean, μ′::FiniteZeroMean) = length(μ) == length(μ′)
eachindex(μ::FiniteZeroMean) = eachindex(μ.X)
print(io::IO, μ::FiniteZeroMean) = print(io, "FiniteZeroMean ($(length(μ)))")
_map(μ::FiniteZeroMean, q::AV) = Zeros(length(q))

struct FiniteZeroKernel{TX} <: Kernel
    X::TX
end
length(k::FiniteZeroKernel) = length(k.X)
==(k::FiniteZeroKernel, k′::FiniteZeroKernel) = length(k) == length(k′)
eachindex(k::FiniteZeroKernel) = eachindex(k.X)
print(io::IO, k::FiniteZeroKernel) = print(io, "FiniteZeroKernel $(size(k))")
_map(k::FiniteZeroKernel, q::AV) = Zeros(length(q))
_map(k::FiniteZeroKernel, q::AV, q′::AV) = Zeros(length(q))
_pairwise(k::FiniteZeroKernel, q::AV) = Zeros(length(q), length(q))
_pairwise(k::FiniteZeroKernel, q::AV, q′::AV) = Zeros(length(q), length(q′))

+(x::FiniteZeroMean, x′::FiniteZeroMean) = zero(x)

struct FiniteZeroCrossKernel{TX, TX′} <: CrossKernel
    X::TX
    X′::TX′
end
size(k::FiniteZeroCrossKernel, dim::Int) = dim == 1 ? length(k.X) : length(k.X′)
==(k::FiniteZeroCrossKernel, k′::FiniteZeroCrossKernel) = size(k) == size(k′)
eachindex(k::FiniteZeroCrossKernel, dim::Int) = dim == 1 ? eachindex(k.X) : eachindex(k.X′)
print(io::IO, k::FiniteZeroCrossKernel) = print(io, "FiniteZeroCrossKernel $(size(k))")
_map(k::FiniteZeroCrossKernel, q::AV) = Zeros(length(q))
_map(k::FiniteZeroCrossKernel, q::AV, q′::AV) = Zeros(length(q))
_pairwise(k::FiniteZeroCrossKernel, q::AV) = Zeros(length(q), length(q))
_pairwise(k::FiniteZeroCrossKernel, q::AV, q′::AV) = Zeros(length(q), length(q′))

struct LhsFiniteZeroCrossKernel{TX} <: CrossKernel
    X::TX
end
size(k::LhsFiniteZeroCrossKernel, dim::Int) = dim == 1 ? length(k.X) : Inf
==(k::LhsFiniteZeroCrossKernel, k′::LhsFiniteZeroCrossKernel) = size(k) == size(k′)
function eachindex(k::LhsFiniteZeroCrossKernel, dim::Int)
    return dim == 1 ? eachindex(k.X) : eachindex(ZeroKernel{Float64}(), 2)
end
function print(io::IO, k::LhsFiniteZeroCrossKernel)
    print(io, "LhsFiniteZeroCrossKernel $(size(k))")
end
_map(k::LhsFiniteZeroCrossKernel, q::AV, X′::AV) = Zeros(length(q))
_pairwise(k::LhsFiniteZeroCrossKernel, q::AV, X′::AV) = Zeros(length(q), length(X′))

struct RhsFiniteZeroCrossKernel{TX′} <: CrossKernel
    X′::TX′
end
size(k::RhsFiniteZeroCrossKernel, dim::Int) = dim == 1 ? Inf : length(k.X′)
==(k::RhsFiniteZeroCrossKernel, k′::RhsFiniteZeroCrossKernel) = size(k) == size(k′)
function eachindex(k::RhsFiniteZeroCrossKernel, dim::Int)
    return dim == 1 ? eachindex(ZeroKernel{Float64}(), 1) : eachindex(k.X′)
end
function print(io::IO, k::RhsFiniteZeroCrossKernel)
    print(io, "RhsFiniteZeroCrossKernel $(size(k))")
end
_map(k::RhsFiniteZeroCrossKernel, X::AV, q′::AV) = Zeros(length(X))
_pairwise(k::RhsFiniteZeroCrossKernel, X::AV, q′::AV) = Zeros(length(X), length(q′))

# More sugar.
finite(μ::MeanFunction, X::AbstractVector) = FiniteMean(μ, X)
finite(μ::ZeroMean, X::AbstractVector) = FiniteZeroMean(X)
finite(μ::FiniteZeroMean, q::AbstractVector) = FiniteZeroMean(q)

finite(k::Kernel, X::AbstractVector) = FiniteKernel(k, X)
finite(k::ZeroKernel, X::AbstractVector) = FiniteZeroKernel(X)
finite(k::FiniteZeroKernel, q::AbstractVector) = FiniteZeroKernel(q)

finite(k::CrossKernel, X::AV, X′::AV) = FiniteCrossKernel(k, X, X′)
function finite(k::ZeroKernel, X::AV, X′::AV)
    return length(X) == length(X′) ? FiniteZeroKernel(X) : FiniteZeroCrossKernel(X, X′)
end
finite(k::FiniteZeroCrossKernel, q::AV, q′::AV) = FiniteZeroCrossKernel(q, q′)

const LhsFinite = Union{LhsFiniteCrossKernel, LhsFiniteZeroCrossKernel}
const RhsFinite = Union{RhsFiniteCrossKernel, RhsFiniteZeroCrossKernel}

lhsfinite(k::CrossKernel, X::AbstractVector) = LhsFiniteCrossKernel(k, X)
lhsfinite(k::ZeroKernel, X::AbstractVector) = LhsFiniteZeroCrossKernel(X)
lhsfinite(k::RhsFiniteCrossKernel, X::AbstractVector) = finite(k.k, X, k.X′)
lhsfinite(k::RhsFiniteZeroCrossKernel, X::AV) = FiniteZeroCrossKernel(X, k.X′)

rhsfinite(k::CrossKernel, X′::AbstractVector) = RhsFiniteCrossKernel(k, X′)
rhsfinite(k::ZeroKernel, X′::AbstractVector) = RhsFiniteZeroCrossKernel(X′)
rhsfinite(k::LhsFiniteCrossKernel, X′::AbstractVector) = finite(k.k, k.X, X′)
rhsfinite(k::LhsFiniteZeroCrossKernel, X′::AV) = FiniteZeroCrossKernel(k.X, X′)


text:
# MadNLP.jl
# Created by Sungho Shin (sungho.shin@wisc.edu)

abstract type AbstractInteriorPointSolver end

include("restoration.jl")

mutable struct InteriorPointSolver{KKTSystem} <: AbstractInteriorPointSolver
    nlp::AbstractNLPModel
    kkt::KKTSystem

    opt::Options
    cnt::Counters
    logger::Logger

    n::Int # number of variables (after reformulation)
    m::Int # number of cons
    nlb::Int
    nub::Int

    x::Vector{Float64} # primal (after reformulation)
    l::Vector{Float64} # dual
    zl::Vector{Float64} # dual (after reformulation)
    zu::Vector{Float64} # dual (after reformulation)
    xl::Vector{Float64} # primal lower bound (after reformulation)
    xu::Vector{Float64} # primal upper bound (after reformulation)

    obj_val::Float64
    f::Vector{Float64}
    c::Vector{Float64}

    jacl::Vector{Float64}

    d::AbstractKKTVector{Float64, Vector{Float64}}
    p::AbstractKKTVector{Float64, Vector{Float64}}

    _w1::AbstractKKTVector{Float64, Vector{Float64}}
    _w2::AbstractKKTVector{Float64, Vector{Float64}}

    _w3::AbstractKKTVector{Float64, Vector{Float64}}
    _w4::AbstractKKTVector{Float64, Vector{Float64}}

    x_trial::Vector{Float64}
    c_trial::Vector{Float64}
    obj_val_trial::Float64

    x_slk::StrideOneVector{Float64}
    c_slk::SubVector{Float64}
    rhs::Vector{Float64}

    ind_ineq::Vector{Int}
    ind_fixed::Vector{Int}
    ind_llb::Vector{Int}
    ind_uub::Vector{Int}

    x_lr::SubVector{Float64}
    x_ur::SubVector{Float64}
    xl_r::SubVector{Float64}
    xu_r::SubVector{Float64}
    zl_r::SubVector{Float64}
    zu_r::SubVector{Float64}

    dx_lr::SubVector{Float64}
    dx_ur::SubVector{Float64}
    x_trial_lr::SubVector{Float64}
    x_trial_ur::SubVector{Float64}

    linear_solver::AbstractLinearSolver
    iterator::AbstractIterator

    obj_scale::Vector{Float64}
    con_scale::Vector{Float64}
    con_jac_scale::Vector{Float64}
    inf_pr::Float64
    inf_du::Float64
    inf_compl::Float64

    theta_min::Float64
    theta_max::Float64
    mu::Float64
    tau::Float64

    alpha::Float64
    alpha_z::Float64
    ftype::String

    del_w::Float64
    del_c::Float64
    del_w_last::Float64

    filter::Vector{Tuple{Float64,Float64}}

    RR::Union{Nothing,RobustRestorer}
    status::Status
    output::Dict
end

function InteriorPointSolver(nlp::AbstractNLPModel;
    option_dict::Dict{Symbol,Any}=Dict{Symbol,Any}(), kwargs...
)
    opt = Options(linear_solver=default_linear_solver())
    set_options!(opt,option_dict,kwargs)
    check_option_sanity(opt)

    KKTSystem = if opt.kkt_system == SPARSE_KKT_SYSTEM
        MT = (opt.linear_solver.INPUT_MATRIX_TYPE == :csc) ? SparseMatrixCSC{Float64, Int32} : Matrix{Float64}
        SparseKKTSystem{Float64, MT}
    elseif opt.kkt_system == SPARSE_UNREDUCED_KKT_SYSTEM
        MT = (opt.linear_solver.INPUT_MATRIX_TYPE == :csc) ? SparseMatrixCSC{Float64, Int32} : Matrix{Float64}
        SparseUnreducedKKTSystem{Float64, MT}
    elseif opt.kkt_system == DENSE_KKT_SYSTEM
        MT = Matrix{Float64}
        VT = Vector{Float64}
        DenseKKTSystem{Float64, VT, MT}
    elseif opt.kkt_system == DENSE_CONDENSED_KKT_SYSTEM
        MT = Matrix{Float64}
        VT = Vector{Float64}
        DenseCondensedKKTSystem{Float64, VT, MT}
    end
    return InteriorPointSolver{KKTSystem}(nlp, opt; option_linear_solver=option_dict)
end

# Inner constructor
function InteriorPointSolver{KKTSystem}(nlp::AbstractNLPModel, opt::Options;
    option_linear_solver::Dict{Symbol,Any}=Dict{Symbol,Any}(),
) where {KKTSystem<:AbstractKKTSystem}
    cnt = Counters(start_time=time())

    logger = Logger(print_level=opt.print_level,file_print_level=opt.file_print_level,
                    file = opt.output_file == "" ? nothing : open(opt.output_file,"w+"))
    @trace(logger,"Logger is initialized.")

    # generic options
    opt.disable_garbage_collector &&
        (GC.enable(false); @warn(logger,"Julia garbage collector is temporarily disabled"))
    set_blas_num_threads(opt.blas_num_threads; permanent=true)

    @trace(logger,"Initializing variables.")
    ind_cons = get_index_constraints(nlp; fixed_variable_treatment=opt.fixed_variable_treatment)
    ns = length(ind_cons.ind_ineq)
    n = get_nvar(nlp)+ns
    m = get_ncon(nlp)

    # Initialize KKT
    kkt = KKTSystem(nlp, ind_cons)

    xl = [get_lvar(nlp);view(get_lcon(nlp),ind_cons.ind_ineq)]
    xu = [get_uvar(nlp);view(get_ucon(nlp),ind_cons.ind_ineq)]
    x = [get_x0(nlp);zeros(ns)]
    l = copy(get_y0(nlp))
    zl= zeros(get_nvar(nlp)+ns)
    zu= zeros(get_nvar(nlp)+ns)

    f = zeros(n) # not sure why, but seems necessary to initialize to 0 when used with Plasmo interface
    c = zeros(m)

    n_jac = nnz_jacobian(kkt)

    nlb = length(ind_cons.ind_lb)
    nub = length(ind_cons.ind_ub)

    x_trial=Vector{Float64}(undef,n)
    c_trial=Vector{Float64}(undef,m)

    x_slk= view(x,get_nvar(nlp)+1:n)
    c_slk= view(c,ind_cons.ind_ineq)
    rhs = (get_lcon(nlp).==get_ucon(nlp)).*get_lcon(nlp)

    x_lr = view(x, ind_cons.ind_lb)
    x_ur = view(x, ind_cons.ind_ub)
    xl_r = view(xl, ind_cons.ind_lb)
    xu_r = view(xu, ind_cons.ind_ub)
    zl_r = view(zl, ind_cons.ind_lb)
    zu_r = view(zu, ind_cons.ind_ub)
    x_trial_lr = view(x_trial, ind_cons.ind_lb)
    x_trial_ur = view(x_trial, ind_cons.ind_ub)

    aug_vec_length = is_reduced(kkt) ? n+m : n+m+nlb+nub

    _w1 = is_reduced(kkt) ? ReducedKKTVector(n, m) : UnreducedKKTVector(n, m, nlb, nub)
    _w2 = is_reduced(kkt) ? ReducedKKTVector(n, m) : UnreducedKKTVector(n, m, nlb, nub)
    _w3 = is_reduced(kkt) ? ReducedKKTVector(n, m) : UnreducedKKTVector(n, m, nlb, nub)
    _w4 = is_reduced(kkt) ? ReducedKKTVector(n, m) : UnreducedKKTVector(n, m, nlb, nub)

    jacl = zeros(n) # spblas may throw an error if not initialized to zero

    d = UnreducedKKTVector(n, m, nlb, nub)
    dx_lr = view(d.xp, ind_cons.ind_lb) # TODO
    dx_ur = view(d.xp, ind_cons.ind_ub) # TODO

    p = UnreducedKKTVector(n, m, nlb, nub)

    obj_scale = [1.0]
    con_scale = ones(m)
    con_jac_scale = ones(n_jac)

    @trace(logger,"Initializing linear solver.")
    cnt.linear_solver_time =
        @elapsed linear_solver = opt.linear_solver.Solver(get_kkt(kkt) ; option_dict=option_linear_solver,logger=logger)

    n_kkt = size(kkt, 1)
    buffer_vec = similar(full(d), n_kkt)
    @trace(logger,"Initializing iterative solver.")
    iterator = opt.iterator(linear_solver, kkt, buffer_vec)

    @trace(logger,"Initializing fixed variable treatment scheme.")

    if opt.inertia_correction_method == INERTIA_AUTO
        opt.inertia_correction_method = is_inertia(linear_solver) ? INERTIA_BASED : INERTIA_FREE
    end

    !isempty(option_linear_solver) && print_ignored_options(logger, option_linear_solver)

    return InteriorPointSolver{KKTSystem}(nlp,kkt,opt,cnt,logger,
        n,m,nlb,nub,x,l,zl,zu,xl,xu,0.,f,c,
        jacl,
        d, p,
        _w1, _w2, _w3, _w4,
        x_trial,c_trial,0.,x_slk,c_slk,rhs,
        ind_cons.ind_ineq,ind_cons.ind_fixed,ind_cons.ind_llb,ind_cons.ind_uub,
        x_lr,x_ur,xl_r,xu_r,zl_r,zu_r,dx_lr,dx_ur,x_trial_lr,x_trial_ur,
        linear_solver,iterator,
        obj_scale,con_scale,con_jac_scale,
        0.,0.,0.,0.,0.,0.,0.,0.,0.," ",0.,0.,0.,
        Vector{Float64}[],nothing,INITIAL,Dict(),
    )
end

include("utils.jl")
include("kernels.jl")
include("callbacks.jl")
include("factorization.jl")
include("solver.jl")


text:
using Base.Test

include("optunity.jl")

testit_dict(d::Dict) = -(d[:x]-1)^2 - (d[:y]-2)^2 - (d[:z]+3)^4

vars, details = maximize(testit_dict, num_evals=10000, solver_name="grid search", x=[-5,5], y=[-5,5], z=[-5,5])

@test_approx_eq_eps vars["x"]  1.0 .2
@test_approx_eq_eps vars["y"]  2.0 .2
@test_approx_eq_eps vars["z"] -3.0 .2

vars, details = maximize(testit_dict, num_evals=10000, z=[-5,5], y=[-5,5], x=[-5,5])

@test_approx_eq_eps vars["x"]  1.0 .2
@test_approx_eq_eps vars["y"]  2.0 .2
@test_approx_eq_eps vars["z"] -3.0 .2

vars, details = maximize(testit_dict, num_evals=10000, solver_name="nelder-mead", x=[-5,5], y=[-5,5], z=[-5,5])

@test_approx_eq_eps vars["x"]  0.8 .1
@test_approx_eq_eps vars["y"]  1.9 .1
@test_approx_eq_eps vars["z"] -3.4 .1

max_stars_repo_name: xrounder/optunity

max_issues_repo_name: xrounder/optunity

max_forks_repo_name: xrounder/optunity