function getModuloList(inputfile)
    lines = readlines(inputfile)
    idsStirng = lines[2]
    ids::Array{Int} = []
    for idSubString in split(idsStirng, ',')
        idString = String(idSubString)
        if idString != "x"
            push!(ids, parse(Int, idString))
        else
            push!(ids, 0)
        end
    end
    # println(ids)
    return ids
end

function printRow(t, moduloList)
    row = []
    for i in 1:length(moduloList)
        value = (t % moduloList[i] > 0) ? 0 : 1
        push!(row, value)
    end
    # println(row)
end

function found(timestamp, moduloList)
    t = timestamp
    result = true
    for i in 1:length(moduloList)
        # print(t, " ")
        if (moduloList[i] > 0)
            if t % moduloList[i] == 0
                result = result && true
                # print(" D")
            else
                result = result && false
                return false
            end
        else
            # print(" never mind")
        end
        # println()
        t += 1
    end
    # println(result)
    return result
end

function findFirst(inputFile)
    moduloList = getModuloList(inputFile)
    println(moduloList)
    firstMod = moduloList[1]
    t = firstMod
    isFound = false
    while !isFound
        if found(t, moduloList)
            break
        else
            t += firstMod
        end
    end
    println(t)
end

findFirst("data_test.txt")


using Test
using MARY_fRG.Basics

@testset "点相关的测试" begin
    pt1 = Point2D(1., 1.)
    println(pt1)
    pt2 = Point2D(2., 3.)
    pt3 = pt1 + pt2
    @test (pt3.y == 4) && (pt3.x == 3)
    pt4 = middle_point(pt1, pt3)
    @test (pt4.x == 2) && (pt4.y == 2.5)
    ang1 = absolute_angle(pt1)
    @test ang1 == pi / 4
    ang1 = absolute_angle(Point2D(-1., 1.))
    @test ang1 == 3pi / 4
    ang1 = absolute_angle(Point2D(1., -1.))
    @test ang1 == 7pi / 4
    pt5 = -pt1
    println(pt5)
    pt6 = pt5 - pt1
    println(pt6)
    pt7 = pt5 - 2pt1
    println(pt7)
end


module SplittingOrderConditions

import Base: start, next, done

export Lyndon, MultiFor, multinomial_coeff, generate_equations, call_form
export generate_equations_ABC, get_ABC_coeffs
export compute_residua, compute_residua_ABC, get_lem, get_p_lem

using Combinatorics

immutable Lyndon
    s::Int
    n::Int
end

Base.start(::Lyndon) = Int[-1]

function Base.next(L::Lyndon, w::Vector{Int})
    if w == [-1]
        w[end] += 1
        return (copy(w), w)
    end
    m = length(w)
    while length(w) < L.n
        push!(w, w[end-m+1])
    end
    while length(w) > 0 && w[end] == L.s - 1
        pop!(w)
    end
    w[end] += 1
    (copy(w), w)
end

Base.done(L::Lyndon, w::Vector{Int}) = w == [L.s-1]

immutable MultiFor
    k::Array{Int,1}
end

Base.start(MF::MultiFor) = Int[]

Base.done(MF::MultiFor, k::Array{Int,1}) = MF.k==k

function Base.next(MF::MultiFor, k::Array{Int,1})
    if k==Int[]
        k = zeros(Int, length(MF.k))
        return(copy(k), k)
    end
    for i=1:length(k)
        if k[i]<MF.k[i]
            k[i] += 1
            for j = 1:i-1
                k[j] = 0
            end
            return (copy(k), k)
        end
    end
end

multinomial_coeff(q::Int, k::Array{Int,1}) = div(factorial(q), prod([factorial(i) for i in k]))

function generate_equations(q::Int, s::Int; symmetric::Bool=false, palindromic::Bool=false,
                            last_b_is_zero::Bool=false, with_brackets::Bool=true)
    eqs = Dict{Array{Int,1}, AbstractString}()
    if symmetric && iseven(q)
        return eqs
    end
    if symmetric
        palindromic = false
        last_b_is_zero = true
        s2 = ceil(Int,s/2)
        s3 = isodd(s)?s2-1:s2
        if with_brackets
            a = vcat([string("a[", j, "]") for j=1:s2], [string("a[", j, "]") for j=s3:-1:1])
        else
            a = vcat([string("a", j) for j=1:s2], [string("a", j) for j=s3:-1:1])
        end
        s2 = ceil(Int,(s-1)/2)
        s3 = iseven(s)?s2-1:s2
        if with_brackets
            b = vcat([string("b[", j, "]") for j=1:s2], [string("b[", j, "]") for j=s3:-1:1],["0"])
        else
            b = vcat([string("b", j) for j=1:s2], [string("b", j) for j=s3:-1:1],["0"])
        end
    elseif palindromic
        last_b_is_zero = false
        s2 = ceil(Int,s/2)
        s3 = isodd(s)?s2-1:s2
        if with_brackets
            a = vcat([string("a[", j, "]") for j=1:s2], [string("b[", j, "]") for j=s3:-1:1])
        else
            a = vcat([string("a", j) for j=1:s2], [string("b", j) for j=s3:-1:1])
        end
        b = reverse(a)
    else
        if with_brackets
            a = [string("a[", j, "]") for j=1:s]
            b = [string("b[", j, "]") for j=1:s]
        else
            a = [string("a", j) for j=1:s]
            b = [string("b", j) for j=1:s]
        end
    end
    if palindromic
        for j in Lyndon(2, q)
            if length(j)==q && !haskey(eqs, [x==0?1:0 for x in reverse(j)])
                eqs[j] = ""
            end
        end
    else
        for j in Lyndon(2, q)
            if length(j)==q
                eqs[j] = ""
            end
        end
    end

    for k0 in Combinatorics.WithReplacementCombinations(1:s,q)
        k = zeros(Int,s)
        for j in k0
            k[j]+=1
        end
        for l = MultiFor(k)
            W=Int[]
            for j=s:-1:1
                append!(W, ones(Int, l[j]))
                append!(W, zeros(Int, k[j]-l[j]))
            end
            if W in keys(eqs) && !(last_b_is_zero && l[s]>0)
                coeff =  multinomial_coeff(q, k)
                for j=1:s
                    coeff = coeff * binomial(k[j], l[j])
                end
                h = string(coeff)
                for j=1:s
                    if k[j]-l[j]>0
                        h = string(h,"*",a[j])
                    end
                    if k[j]-l[j]>1
                        h = string(h,"^",k[j]-l[j])
                    end
                    if l[j]>0
                        h = string(h,"*",b[j])
                    end
                    if l[j]>1
                        h = string(h,"^",l[j])
                    end
                end
                eqs[W] = string(eqs[W],"+",h)
            end
        end
    end
    for (W, h) in eqs
        eqs[W] = string(h,"-1")
    end
    eqs
end

function generate_equations_ABC(q::Int, s::Int; symmetric::Bool=false, palindromic::Bool=false,
                            last_c_is_zero::Bool=false, last_b_is_zero::Bool=false, with_brackets::Bool=true)
    eqs = Dict{Array{Int,1}, AbstractString}()
    if symmetric && iseven(q)
        return eqs
    end
    if last_b_is_zero
       last_c_is_zero = true
    end
    if symmetric
        error("symmertric not implemented.")
        #palindromic = false
        #last_c_is_zero = true
        #last_b_is_zero = true
        #s2 = ceil(Int,s/2)
        #s3 = isodd(s)?s2-1:s2
        #if with_brackets
        #    a = vcat([string("a[", j, "]") for j=1:s2], [string("a[", j, "]") for j=s3:-1:1])
        #else
        #    a = vcat([string("a", j) for j=1:s2], [string("a", j) for j=s3:-1:1])
        #end
        #s2 = ceil(Int,(s-1)/2)
        #s3 = iseven(s)?s2-1:s2
        #if with_brackets
        #    b = vcat([string("b[", j, "]") for j=1:s2], [string("b[", j, "]") for j=s3:-1:1],["0"])
        #else
        #    b = vcat([string("b", j) for j=1:s2], [string("b", j) for j=s3:-1:1],["0"])
        #end
    elseif palindromic
        last_c_is_zero = false
        last_b_is_zero = false
        s2 = ceil(Int,s/2)
        s3 = isodd(s)?s2-1:s2
        if with_brackets
            a = vcat([string("a[", j, "]") for j=1:s2], [string("c[", j, "]") for j=s3:-1:1])
            b = vcat([string("b[", j, "]") for j=1:s2], [string("b[", j, "]") for j=s3:-1:1])
        else
            a = vcat([string("a", j) for j=1:s2], [string("c", j) for j=s3:-1:1])
            b = vcat([string("b", j) for j=1:s2], [string("b", j) for j=s3:-1:1])
        end
        c = reverse(a)
    else
        if with_brackets
            a = [string("a[", j, "]") for j=1:s]
            b = [string("b[", j, "]") for j=1:s]
            c = [string("c[", j, "]") for j=1:s]
        else
            a = [string("a", j) for j=1:s]
            b = [string("b", j) for j=1:s]
            c = [string("c", j) for j=1:s]
        end
    end
    if palindromic
        rr = (2, 1, 0)
        for j in Lyndon(3, q)
            if length(j)==q && !haskey(eqs, [rr[x+1] for x in reverse(j)])
                eqs[j] = ""
            end
        end
    else
        for j in Lyndon(3, q)
            if length(j)==q
                eqs[j] = ""
            end
        end
    end

    for k0 in Combinatorics.WithReplacementCombinations(1:s,q)
        k = zeros(Int,s)
        for j in k0
            k[j]+=1
        end
        for lA = MultiFor(k)
        for lB = MultiFor(k-lA)
        for lC = MultiFor(k-lA-lB)
            W=Int[]
            for j=s:-1:1
                append!(W, 2*ones(Int, lC[j]))
                append!(W, ones(Int, lB[j]))
                append!(W, zeros(Int, lA[j]))
            end
            if W in keys(eqs) && !(last_b_is_zero && lB[s]>0)  && !(last_c_is_zero && lC[s]>0)
                coeff =  multinomial_coeff(q, k)
                for j=1:s
                    coeff = coeff * multinomial_coeff(k[j], [lA[j], lB[j], lC[j]])
                end
                h = string(coeff)
                for j=1:s
                    if lA[j]>0
                        h = string(h,"*",a[j])
                    end
                    if lA[j]>1
                        h = string(h,"^",lA[j])
                    end
                    if lB[j]>0
                        h = string(h,"*",b[j])
                    end
                    if lB[j]>1
                        h = string(h,"^",lB[j])
                    end
                    if lC[j]>0
                        h = string(h,"*",c[j])
                    end
                    if lC[j]>1
                        h = string(h,"^",lC[j])
                    end
                end
                eqs[W] = string(eqs[W],"+",h)
            end
        end
        end
        end
    end
    for (W, h) in eqs
        eqs[W] = string(h,"-1")
    end
    eqs
end

function compute_residua(q::Int, a::Vector, b::Vector)
    s = length(a)
    T = typeof(a[1])
    eqs = Dict{Array{Int,1}, T}()

    aa = zeros(T, s, q+1)
    bb = zeros(T, s, q+1)
    aa[:,1] = ones(T, s)
    bb[:,1] = ones(T, s)
    for k=1:q
        aa[:,k+1] = aa[:,k].*a
        bb[:,k+1] = bb[:,k].*b
    end

    for j in Lyndon(2, q)
        if length(j)==q
            eqs[j] = -one(T)
        end
    end

    for k0 in Combinatorics.WithReplacementCombinations(1:s,q)
        k = zeros(Int,s)
        for j in k0
            k[j]+=1
        end
        for l = MultiFor(k)
            W=Int[]
            for j=s:-1:1
                append!(W, ones(Int, l[j]))
                append!(W, zeros(Int, k[j]-l[j]))
            end
            if W in keys(eqs)
                h =  multinomial_coeff(q, k)
                for j=1:s
                    h *= binomial(k[j], l[j])
                end
                for j=1:s
                    h *= aa[j, k[j]-l[j]+1] * bb[j, l[j]+1]
                end
                eqs[W] += h
            end
        end
    end
    eqs
end

get_lem(q, a, b) = norm(collect(values(compute_residua(q, a, b))))

function get_p_lem(a,b)
    q = 1
    while true
        lem = get_lem(q,a,b)
        if lem>1e-11
            return q-1, lem
        end
        q += 1
    end
end

function get_ABC_coeffs(a,b)
    s=length(a)
    T = typeof(a[1])
    S = 2*s-1
    A = zeros(T, S)
    B = zeros(T, S)
    C = zeros(T, S)
    h = a[1]
    for k=1:s
        K = 2*k-1
        A[K] = a[k]
        B[K] = h
        C[K] = b[k]
        h = b[k]-h
        if k<s
            A[K+1] = 0.0
            B[K+1] = h
            C[K+1] = 0.0
            h = a[k+1]-h
        end
    end
    A, B, C
end

function compute_residua_ABC(q::Int, a::Vector, b::Vector, c::Vector)
    s = length(a)
    T = typeof(a[1])
    eqs = Dict{Array{Int,1}, T}()

    aa = zeros(T, s, q+1)
    bb = zeros(T, s, q+1)
    cc = zeros(T, s, q+1)
    aa[:,1] = ones(T, s)
    bb[:,1] = ones(T, s)
    cc[:,1] = ones(T, s)
    for k=1:q
        aa[:,k+1] = aa[:,k].*a
        bb[:,k+1] = bb[:,k].*b
        cc[:,k+1] = cc[:,k].*c
    end

    for j in Lyndon(3, q)
        if length(j)==q
            eqs[j] = -one(T)
        end
    end

    for k0 in Combinatorics.WithReplacementCombinations(1:s,q)
        k = zeros(Int,s)
        for j in k0
            k[j]+=1
        end
        for lA = MultiFor(k)
        for lB = MultiFor(k-lA)
        for lC = MultiFor(k-lA-lB)
            W=Int[]
            for j=s:-1:1
                append!(W, 2*ones(Int, lC[j]))
                append!(W, ones(Int, lB[j]))
                append!(W, zeros(Int, lA[j]))
            end
            if W in keys(eqs)
                h =  multinomial_coeff(q, k)
                for j=1:s
                    h *= multinomial_coeff(k[j], [lA[j], lB[j], lC[j]])
                end
                for j=1:s
                    h *= aa[j, lA[j]+1] * bb[j, lB[j]+1] * cc[j, lC[j]+1]
                end
                eqs[W] += h
            end
        end
        end
        end
    end
    eqs
end

get_lem(q, a, b, c) = norm(collect(values(compute_residua_ABC(q, a, b, c))))

function get_p_lem(a,b,c)
    q = 1
    while true
        lem = get_lem(q,a,b,c)
        if lem>1e-11
            return q-1, lem
        end
        q += 1
    end
end

function call_form(input::ASCIIString; threads=1)
    attempts = 0
    FORM_PID=0
    path = joinpath(dirname(@__FILE__), "../deps/bin")
    out = "Error initializing preset external channels\n"
    while out == "Error initializing preset external channels\n"
        (so,si,pr) = readandwrite(`$path/tform -w$(threads) -t /tmp -M  -q -pipe 0,1 $path/pipe.frm`)
        FORM_PID = readuntil(so,'\n')[1:end-1];
        print(si, FORM_PID,',',getpid(),"\n\n", input,"\n\n")
        close(si)
        out = readall(so)
        close(si)
        close(so)
        close(pr)
        attempts +=1
    end
    println(STDERR, "# of attemps to call form = ",attempts)
    try
        rm("/tmp/xform$FORM_PID.str") # delete generated temporary file
    end
    out
end

function compile_fg(fun::AbstractString, n::Integer; threads=1)
    fun = replace(replace(fun,"[", "("), "]", ")")
    input = string("""
Off Statistics;
V x;
S u;
""","Local f0 = ",
    fun,
    ";\n",
"""
#write <> "begin"
#write <> "  if length(g)==0"
Print;
Format O3;
Format maple;
.sort;
#write <> "    return f0"
#write <> "  else"
I i;
S m, xx, u;
Hide f0;
#do i=1,$(n)
    Local f'i' = f0;
    id x('i') = xx;
    id xx^m? = m*xx^(m-1);
    id xx = x('i');
    .sort
    Hide f'i';
#enddo
.sort;
""","Local H = ",
    join(["u^$(j+1)*f$(j)" for j=0:n],:+),
    ";\n",
"""
B u;
.sort
#optimize H
B u;
.sort
""",
    join(["Local F$(j) = H[u^$(j+1)];\n" for j=0:n], ""),
"""
.sort
#write <> "%4O",
""", "#write <> \"\n   f=%e   ",
join(["g($(j))=%e" for j=1:n],"    "),
"\", ",
join(["F$(j)" for j=0:n],","),"\n",
"""
#write <> "    return f"
#write <> "  end"
#write <> "end"
.end
""")
    out = call_form(input, threads=threads)
    out = replace(replace(out,"(", "["), ")", "]")
    out = replace(out,"length[g]", "length(g)")
    eval(parse(string("(x,g)->",out)))
end

end


include("autodiff.jl")
include("utils.jl")
include("numbers.jl")
include("groups/group_utils.jl")
include("sized_abstract_array.jl")
include("orthonormal_bases.jl")

# starting with tests of simple manifolds
include("circle.jl")
include("cholesky_space.jl")
include("euclidean.jl")
include("fixed_rank.jl")
include("grassmann.jl")
include("hyperbolic.jl")
include("sphere.jl")
include("stiefel.jl")
include("symmetric.jl")
include("rotations.jl")
include("symmetric_positive_definite.jl")
include("torus.jl")

#meta manifolds
include("product_manifold.jl")
include("power_manifold.jl")
include("vector_bundle.jl")
include("graph.jl")

include("metric.jl")
include("statistics.jl")

# Lie groups and actions
include("groups/groups_general.jl")
include("groups/array_manifold.jl")
include("groups/circle_group.jl")
include("groups/translation_group.jl")
include("groups/special_orthogonal.jl")
include("groups/product_group.jl")
include("groups/semidirect_product_group.jl")
include("groups/special_euclidean.jl")
include("groups/group_operation_action.jl")
include("groups/rotation_action.jl")
include("groups/translation_action.jl")


# Code for computing with domains.

############################
# The union of two domains
############################

# DD can be any collection type: a Tuple, or a Vector
# Eventually it will likely be forced to be an AbstractSet
#
"""
A `UnionDomain` represents the union of a set of domains.
"""
struct UnionDomain{DD,T} <: Domain{T}
    domains  ::  DD
end

"""
The `UnionDomain` constructor can be invoked in one of three ways:
- with a list of arguments: UnionDomain(d1, d2)
- with a single domain: UnionDomain(d::Domain)
- or with any iterable list of domains: UnionDomain(domains)
"""
UnionDomain(domains...) = UnionDomain(domains)
UnionDomain(d::Domain) = UnionDomain((d,))

UnionDomain(domains) = _UnionDomain(expand_union(promote_domains(domains)...))
_UnionDomain(domains) = UnionDomain{typeof(domains),eltype(first(domains))}(domains)

"Return all arguments in a tuple of domains, with each `UnionDomain` replaced by its elements."
expand_union() = ()
expand_union(domain, domains...) = (domain, expand_union(domains...)...)
expand_union(domain::UnionDomain, domains...) = (elements(domain)..., expand_union(domains...)...)

convert(::Type{Domain}, v::AbstractVector{<:Domain}) = UnionDomain(v)
convert(::Type{Domain}, v::AbstractSet{<:Domain}) = UnionDomain(v)

convert(::Type{Domain{T}}, d::UnionDomain{S}) where {S,T} =
    UnionDomain(convert_domain.(T, elements(d)))

Domain(v::AbstractVector) = convert(Domain, v)
Domain(v::AbstractSet) = convert(Domain, v)

elements(d::UnionDomain) = d.domains

hash(d::UnionDomain, h::UInt) = hash(Set(d.domains), h)

union(domains::Domain...) = UnionDomain(domains...)

# Catch one specific case
union(d1::Domain, d2::Domain)  = d1 == d2 ? d1 : UnionDomain(d1, d2)

# The union of domains corresponds to a logical OR of their characteristic functions
indomain(x, d::UnionDomain) = mapreduce(d->in(x, d), |, elements(d))
approx_indomain(x, d::UnionDomain, tol = default_tolerance(d)) = mapreduce(d->approx_in(x, d, tol), |, elements(d))

point_in_domain(d::UnionDomain) = point_in_domain(element(d,1))

==(a::UnionDomain, b::UnionDomain) = Set(elements(a)) == Set(elements(b))

isempty(d::UnionDomain) = all(isempty, d.domains)

function show(io::IO, d::UnionDomain)
    print(io, "a union of $(numelements(d)) domains:\n")
    for (i,e) in enumerate(elements(d))
        print(io, "\t$i.\t: ", e, "\n")
    end
end

## ualgebra

for op in (:+, :-, :*, :/)
    @eval begin
        $op(domain::UnionDomain, x::Number) = UnionDomain(broadcast($op, elements(domain), x))
        $op(x::Number, domain::UnionDomain) = UnionDomain(broadcast($op, x, elements(domain)))
    end
end

\(x::Number, domain::UnionDomain) = UnionDomain(broadcast(\, x, elements(domain)))

for (op, mop) in ((:minimum, :min), (:maximum, :max), (:infimum, :min), (:supremum, :max))
    @eval $op(d::UnionDomain) = mapreduce($op, $mop, elements(d))
end

setdiff(d1::UnionDomain, d2::UnionDomain) = UnionDomain(setdiff.(elements(d1), Ref(d2)))

function setdiff(d1::UnionDomain, d2::Domain)
    s = Set(elements(d1))
    # check if any element is in d1 and just remove
    s2 = Set(setdiff(s, tuple(d2)))
    s2 ≠ s && return UnionDomain(s2)

    UnionDomain(setdiff.(elements(d1), Ref(d2)))
end

function setdiff(d1::Domain, d2::UnionDomain)
    ret = d1
    for d in elements(d2)
        ret = setdiff(ret, d)
    end
    ret
end

# use \ as a synomym for setdiff, in the context of domains (though, generically,
# \ means left division in Julia)
\(d1::Domain, d2::Domain) = setdiff(d1, d2)

##############################
# The intersection of domains
##############################

"""
An `IntersectionDomain` represents the intersection of a set of domains.
"""
struct IntersectionDomain{DD,T} <: Domain{T}
    domains ::  DD
end

"""
The `IntersectionDomain` constructor can be invoked in one of three ways:
- with a list of arguments: IntersectionDomain(d1, d2)
- with a single domain: IntersectionDomain(d::Domain)
- or with any iterable list of domains: IntersectionDomain(domains)
"""
IntersectionDomain(domains...) = IntersectionDomain(domains)
IntersectionDomain(d::Domain) = IntersectionDomain((d,))

IntersectionDomain(domains) = _IntersectionDomain(expand_intersection(promote_domains(domains)...))
_IntersectionDomain(domains) = IntersectionDomain{typeof(domains),eltype(first(domains))}(domains)

"Return all arguments in a tuple of domains, with each `UnionDomain` replaced by its elements."
expand_intersection() = ()
expand_intersection(domain, domains...) = (domain, expand_intersection(domains...)...)
expand_intersection(domain::IntersectionDomain, domains...) = (elements(domain)..., expand_intersection(domains...)...)

elements(d::IntersectionDomain) = d.domains

intersect(d1::Domain, d2::Domain) = d1 == d2 ? d1 : IntersectionDomain(d1, d2)
function intersect(d1::UnionDomain, d2::UnionDomain)
    d1 == d2 && return d1
    union(intersect.(Ref(d1), elements(d2))...)
end
intersect(d1::UnionDomain, d2::Domain) = union(intersect.(d1.domains, Ref(d2))...)
intersect(d1::Domain, d2::UnionDomain) = union(intersect.(Ref(d1), d2.domains)...)

# The intersection of domains corresponds to a logical AND of their characteristic functions
indomain(x, d::IntersectionDomain) = mapreduce(d->in(x, d), &, elements(d))

(&)(d1::Domain, d2::Domain) = intersect(d1,d2)

function intersect(d1::ProductDomain, d2::ProductDomain)
    if numelements(d1) == numelements(d2)
        ProductDomain([intersect(element(d1,i), element(d2,i)) for i in 1:numelements(d1)]...)
    else
        IntersectionDomain(d1, d2)
    end
end

convert(::Type{Domain{T}}, d::IntersectionDomain{S}) where {S,T} =
    IntersectionDomain(convert_domain.(T, elements(d)))

function show(io::IO, d::IntersectionDomain)
    print(io, "the intersection of $(numelements(d)) domains:\n")
    print(io, "\t1.\t: ", element(d,1), "\n")
    print(io, "\t2.\t: ", element(d,2), "\n")
end

#########################################
### The difference between two domains
#########################################

struct DifferenceDomain{D1,D2,T} <: Domain{T}
    d1    ::  D1
    d2    ::  D2
end

DifferenceDomain(d1, d2) = _DifferenceDomain(promote_domain(d1, d2)...)
_DifferenceDomain(d1, d2) = DifferenceDomain{typeof(d1),typeof(d2),eltype(d1)}(d1, d2)

convert(::Type{Domain{T}}, d::DifferenceDomain{S}) where {S,T} =
    DifferenceDomain(convert_domain(T, d.d1), convert_domain(T, d.d2))

function setdiff(d1::Domain{T}, d2::Domain) where T
    d1 == d2 && return EmptySpace{T}()
    DifferenceDomain(d1, d2)
end

setdiff(d1::Domain, d2) = DifferenceDomain(d1, d2)
setdiff(d1, d2::Domain) = DifferenceDomain(d1, d2)

# The difference between two domains corresponds to a logical AND NOT of their characteristic functions
indomain(x, d::DifferenceDomain) = in(x, d.d1) && (~in(x, d.d2))

\(d1::Domain, d2) = setdiff(d1, d2)

function show(io::IO, d::DifferenceDomain)
    print(io, "the difference of 2 domains:\n")
    print(io, "\t1.\t: ", d.d1, "\n")
    print(io, "\t2.\t: ", d.d2, "\n")
end


# verify transformation semigroup generated by wrappers
#
REAL_ID = [4 5 6; 3 4 5; 2 3 4]
COMPLEX_ID = REAL_ID * (1 + im)

#identity
#transpose
#adjoint
symmup(A) = Symmetric(A, :U)
symmlo(A) = Symmetric(A, :L)
hermup(A) = Hermitian(A, :U)
hermlo(A) = Hermitian(A, :L)
triaup(A) = UpperTriangular(A)
trialo(A) = LowerTriangular(A)
triaupu(A) = UnitUpperTriangular(A)
trialou(A) = UnitLowerTriangular(A)
diagonal(A) = Diagonal(A)
bidiup(A) = Bidiagonal(A, :U)
bidilo(A) = Bidiagonal(A, :L)
tridi(A) = Tridiagonal(A)
symmtridiup(A) = SymTridiagonal(Symmetric(A, :U))
symmtridilo(A) = SymTridiagonal(Symmetric(A, :L))
hermtridiup(A) = Hermitian(Tridiagonal(A), :U)
hermtridilo(A) = Hermitian(Tridiagonal(A), :L)

REAL_GEN = [identity, transpose, symmup, symmlo, triaup, trialo, triaupu, trialou,
            diagonal, bidiup, bidilo, tridi, symmtridiup, symmtridilo]

COMPLEX_GEN = [REAL_GEN..., adjoint, hermup, hermlo, hermtridiup, hermtridilo]

function closegen(ID::T, GEN::Vector) where T<:AbstractMatrix
    dic = Dict{T,Any}(ID => Set{Any}())
    modified = true
    len(set::Set) = isempty(set) ? 0 : length(first(set))
    setkg(k, gen) = isempty(dic[k]) ? (gen,) : (gen, first(dic[k])...)
    while modified
        modified = false
        keysprev = collect(keys(dic))
        for gen in GEN
            for k in keysprev
                gk = Matrix(gen(k))
                if !haskey(dic, gk)
                    dic[gk] = Set{Any}(Ref(setkg(k, gen)))
                    modified = true
                elseif len(dic[gk]) == len(dic[k]) + 1
                    push!(dic[gk], setkg(k, gen))
                end
            end
        end
    end
    dic
end

closereal() = closegen(REAL_ID, REAL_GEN)
closecomplex() = closegen(COMPLEX_ID, COMPLEX_GEN)


@symbol_func function K_TB(cur_reactor::AbstractReactor)
  cur_K = 4.027e-2

  cur_K *= K_n(cur_reactor)

  cur_K *= cur_reactor.epsilon

  cur_K /= cur_reactor.max_beta_N

  cur_K *= 1 + cur_reactor.f_D

  cur_K *= 1 + cur_reactor.nu_n

  cur_K *= 1 + cur_reactor.nu_T

  cur_K /= 1 + cur_reactor.nu_n + cur_reactor.nu_T

  cur_K
end


# The purpose of this script is to visually compare the emirical performance
# of various algorithms for line spetral estimation for separating
# two closeby sinusoids.
#
# Parameters:
# N: number of measurements per snapshot
# K: number of sinusoids
# L: number of snapshots
# m: number of sparse measurements which can be used in some of the algorithms;
#    set m = N for full measurements
# SNR: signal-to-noise-ratio [dB] of the AWGN model
# f: location of the frequencies of the sinusoids

include("MUSIC.jl")
include("rootMUSIC.jl")
include("ESPRIT.jl")
include("TLSesprit.jl")
include("MPM.jl")
include("ANM.jl")
include("EMAC.jl")
include("periodogram.jl")
include("structHMirls.jl")
using Plots, Random

# Parameters
N = 10
K = 2
L = 1
m = 8
SNR = 30

#f = rand(K)
f = [0.3; 0.40]

# Create sparse samples at m uniformly random locations \in [N]
omega = sort(randperm(N)[1:m])
a(f) = exp.(2*pi*1im*f*(0:N-1))
A(f) = hcat(a.(f')...)

# Initialize data matrix
Y = Array{Complex{Float64},2}(undef,N,L)

# Draw a random noise sample from a multivariate Gaussian distribution
noise = (randn(N,L)+1im*randn(N,L))/ (sqrt(2)*10^(SNR/10))

# Draw signal coefficients from the uniform distribution
# on the complex unit sphere
S = exp.(2*pi*rand(K,L)*1im)
#S = [1, 1]

Y = A(f)*S + noise

f_range = range(0, 1, length=1000)
spec_periodogram = periodogram(Y,f_range)
(f_MUSIC, specMUSIC) = MUSIC(Y,K,m)

f_ESPRIT = ESPRIT(Y,K)
f_rootMUSIC = rootMUSIC(Y,K)
f_TLSesprit = TLSesprit(Y,K)
#println(f_TLSesprit)

f_MPM = MPM(Y,K)
#f_ANM = ANM(Y[omega,:],N,K,omega,1/SNR)
#f_EMAC = EMAC(Y[omega,:],N,K,omega,1/SNR)
f_structHMirls = structHMirls(Y[omega,1],N,K,omega,1/SNR)

plot(f,seriestype="vline",label="ground truth",
xlims = (0,1),
ylims = (0,1.1),
xticks = 0:0.1:1,
xlabel = "f",
ylabel = "normalized spectra"
)
#plot!(f_range,spec_periodogram,label="Periodogram")
#plot!(f_range,specMUSIC,label="MUSIC")
plot!(f_rootMUSIC,seriestype="vline",label="root MUSIC")
plot!(f_structHMirls,seriestype="vline",label="structHMirls")
#plot!(f_MUSIC,seriestype="vline",label="MUSIC")
#plot!(f_ESPRIT,seriestype="vline",label="ESPRIT")
plot!(f_TLSesprit,seriestype="vline",label="TLSesprit")
#plot!(f_MPM,seriestype="vline",label="MPM")
#plot!(f_ANM,seriestype="vline",label="ANM")
#plot!(f_EMAC,seriestype="vline",label="EMAC")
#savefig("/home/thomas/atom-amd64/usr/share/atom/example/periodogram_vs_MUSIC")


# exp(-x^2/2) H_n(x) / sqrt(π*prod(1:n))

struct ForwardWeightedHermitePlan{T}
    Vtw::Matrix{T} # vandermonde
end

struct BackwardWeightedHermitePlan{T}
    V::Matrix{T} # vandermonde
end

function _weightedhermite_vandermonde(n)
    V = Array{Float64}(undef, n, n)
    x,w = unweightedgausshermite(n)
    for k=1:n
        V[k,:] = FastGaussQuadrature.hermpoly_rec(0:n-1, sqrt(2)*x[k])
    end
    V,w
end

function ForwardWeightedHermitePlan(n::Integer)
    V,w = _weightedhermite_vandermonde(n)
    ForwardWeightedHermitePlan(V' * Diagonal(w / sqrt(π)))
end

BackwardWeightedHermitePlan(n::Integer) = BackwardWeightedHermitePlan(_weightedhermite_vandermonde(n)[1])

*(P::ForwardWeightedHermitePlan, v::AbstractVector) = P.Vtw*v
*(P::BackwardWeightedHermitePlan, v::AbstractVector) = P.V*v

weightedhermitetransform(v) = ForwardWeightedHermitePlan(length(v))*v
iweightedhermitetransform(v) = BackwardWeightedHermitePlan(length(v))*v


### Generate linear data through Born scattering modeling for each num_vintage
## author: Ziyi Yin (ziyi.yin@gatech.edu)

using DrWatson
@quickactivate "CompassTimeLapseCCS"

using PyPlot, JUDI
using JLD2, LinearAlgebra
using JOLI, FFTW
using Printf, Statistics

JLD2.@load "../model/CompassTimeLapse.jld2"
JLD2.@load "../model/MarineSourceLoc.jld2" xrec_stack zrec_stack
JLD2.@load "../model/OBN.jld2" OBN_loc

model0_stack = [Model(n,d,o,m0_stack[i]; nb=100) for i = 1:num_vintage]

### reciprocity

nsrc = length(OBN_loc)
nrec = length(xrec_stack[1])

dtS = dtR = 1f0
timeS = timeR = 2000f0

xsrc = convertToCell(OBN_loc)
ysrc = convertToCell(range(0f0,stop=0f0,length=nsrc))
zsrc = convertToCell(range(12*d[2],stop=12*d[2],length=nsrc))

srcGeometry = Geometry(xsrc, ysrc, zsrc; dt=dtS, t=timeS)
wavelet = ricker_wavelet(timeS,dtS,0.025f0)

q = judiVector(srcGeometry, wavelet)

yrec = 0f0

recGeometry_stack = [Geometry(xrec_stack[i],yrec,zrec_stack[i]; dt=dtR, t=timeR, nsrc=nsrc) for i = 1:num_vintage]

ntComp = get_computational_nt(srcGeometry, recGeometry_stack[1], model0_stack[1])
info = Info(prod(n), nsrc, ntComp)

opt = Options(isic=true)

F0_stack = [judiProjection(info, recGeometry_stack[i])*judiModeling(info, model0_stack[i]; options=opt)*adjoint(judiProjection(info, srcGeometry)) for i = 1:num_vintage]
J_stack = [judiJacobian(F0_stack[i],q) for i = 1:num_vintage]
d_lin = [J_stack[i]*dm_stack[i] for i = 1:num_vintage]    # generate linear data through Born scattering

JLD2.@save "../data/d_lin.jld2" d_lin q


using OrdinaryDiffEq, RecursiveArrayTools, Unitful
using LinearAlgebra, Test

@testset "Algorithms" begin
algs = [Euler(),Midpoint(),Heun(),Ralston(),RK4(),SSPRK104(),SSPRK22(),SSPRK33(),
        SSPRK43(),SSPRK432(),BS3(),BS5(),DP5(),DP8(),Feagin10(),Feagin12(),
        Feagin14(),TanYam7(),Tsit5(),TsitPap8(),Vern6(),Vern7(),Vern8(),Vern9()]

@testset "Scalar units" begin
    f(y,p,t) = 0.5*y / 3.0u"s"
    u0 = 1.0u"N"
    prob = ODEProblem(f,u0,(0.0u"s",1.0u"s"))

    for alg in algs
        @show alg
        sol = solve(prob,alg,dt=1u"s"/10)
    end

    sol = solve(prob,ExplicitRK())
end

@testset "2D units" begin
    f(dy,y,p,t) = (dy .= 0.5.*y ./ 3.0u"s")
    u0 = [1.0u"N" 2.0u"N"
          3.0u"N" 1.0u"N"]
    prob = ODEProblem(f,u0,(0.0u"s",1.0u"s"))

    for alg in algs
        @show alg
        sol = solve(prob,alg,dt=1u"s"/10)
    end

    sol = solve(prob,ExplicitRK())
end
end

@testset "Mixed units" begin
@testset "With ArrayPartition" begin
    r0 = [1131.340, -2282.343, 6672.423]u"km"
    v0 = [-5.64305, 4.30333, 2.42879]u"km/s"
    Δt = 86400.0*365u"s"
    μ = 398600.4418u"km^3/s^2"
    rv0 = ArrayPartition(r0,v0)

    function f(dy, y, μ, t)
        r = norm(y.x[1])
        dy.x[1] .= y.x[2]
        dy.x[2] .= -μ .* y.x[1] / r^3
    end

    prob = ODEProblem(f,rv0,(0.0u"s",1.0u"s"),μ)
    for alg in [Tsit5()]
        @show alg
        sol = solve(prob,alg)
    end

    for alg in [    AutoVern6(Rodas5(autodiff=false)),
                    AutoVern7(Rodas5(autodiff=false)),
                    AutoVern8(Rodas5(autodiff=false)),
                    AutoVern9(Rodas5(autodiff=false))]
        @show alg
        @test_broken sol = solve(prob,alg)
    end
end

@testset "Without ArrayPartition" begin
    # coordinate: u = [position, momentum]
    # parameters: p = [mass, force constanst]
    function f_harmonic!(du,u,p,t)
        du[1] = u[2]/p[1]
        du[2] = -p[2]*u[1]
    end

    mass = 1.0u"kg"
    k = 1.0u"N/m"
    p = [mass, k]

    u0 = [1.0u"m", 0.0u"kg*m/s"] # initial values (position, momentum)
    tspan = (0.0u"s", 10.0u"s")
    prob = ODEProblem(f_harmonic!, u0, tspan, p)
    @test_broken solve(prob, Tsit5()).retcode == :Success
end

@testset "Unitful time with unitless state" begin
  using OrdinaryDiffEq, Unitful
  u0 = 30.0
  tspan = (0.0u"s",10.0u"s")
  prob1 = ODEProblem((du,u,t,p)->(du[1] = -0.2u"1/s"*u[1]),[u0],tspan)
  prob2 = ODEProblem((u,t,p)->(-0.2u"1/s"*u[1]),u0,tspan)
  prob3 = ODEProblem((u,t,p)->[-0.2u"1/s"*u[1]],[u0],tspan)
  for prob in [prob1, prob2, prob3]
    solve(prob, Tsit5()).retcode === :Success
  end
end

end


@testset "DouglasRachford" begin
    # Though this seems a strange way, it is a way to compute the mid point
    M = Sphere(2)
    p = [1., 0., 0.]
    q = [0., 1., 0.]
    r = [0., 0., 1.]
    start = [0.,0.,1.]
    result = geodesic(M, p, q, 0.5)
    F(x) = distance(M, x, p)^2 + distance(M, x, q)^2
    prox1 = (η,x) -> prox_distance(M, η, p, x)
    prox2 = (η,x) -> prox_distance(M, η, q, x)
    @test_throws ErrorException DouglasRachford(M,F,Array{Function,1}([prox1,]),start) # we need more than one prox
    xHat = DouglasRachford(M,F,[prox1,prox2],start)
    @test F(start) > F(xHat)
    @test_broken distance(M, xHat, result) ≈ 0
    # but we can also compute the riemannian center of mass (locally) on Sn
    # though also this is not that useful, but easy to test that DR works
    F2(x) = distance(M,x,p)^2 + distance(M,x,q)^2 + distance(M,x,r)^2
    prox3 = (η,x) -> prox_distance(M,η,r,x)
    o = DouglasRachford(M,F2,[prox1,prox2,prox3],start;
    debug = [DebugCost(), DebugIterate(), DebugProximalParameter(),100],
    record = [RecordCost(), RecordProximalParameter()],
    return_options=true
    )
    xHat2 = get_solver_result(o)
    drec2 = get_record(o)
    xCmp = 1/sqrt(3)*ones(3)
    # since the default does not run that long -> rough estimate
    @test_broken distance(M,xHat2,xCmp) ≈ 0 atol = 10^(-5)
end


import NLPModels: grad, cons, jac

"""
unconstrained_check: return the infinite norm of the gradient of the objective function

`unconstrained_check( :: AbstractNLPModel, :: NLPAtX; pnorm :: Real = Inf, kwargs...)`

Require `state.gx` (filled if not provided)

See also `unconstrained2nd_check`, `optim_check_bounded`, `KKT`
"""
function unconstrained_check(pb::AbstractNLPModel, state::NLPAtX; pnorm::Real = Inf, kwargs...)
  if state.gx == _init_field(typeof(state.gx)) # should be filled if empty
    update!(state, gx = grad(pb, state.x))
  end

  return norm(state.gx, pnorm)
end

"""
unconstrained2nd_check: check the norm of the gradient and the smallest
                   eigenvalue of the hessian.

`unconstrained2nd_check( :: AbstractNLPModel, :: NLPAtX; pnorm :: Real = Inf, kwargs...)`

Require are `state.gx`, `state.Hx` (filled if not provided).

See also `unconstrained_check`, `optim_check_bounded`, `KKT`
"""
function unconstrained2nd_check(pb::AbstractNLPModel, state::NLPAtX; pnorm::Real = Inf, kwargs...)
  if state.gx == _init_field(typeof(state.gx)) # should be filled if empty
    update!(state, gx = grad(pb, state.x))
  end
  if state.Hx == _init_field(typeof(state.Hx))
    update!(state, Hx = hess(pb, state.x).data)
  end

  res = max(norm(state.gx, pnorm), max(-eigmin(Symmetric(state.Hx, :L)), 0.0))

  return res
end

"""
optim\\_check\\_bounded: gradient of the objective function projected

`optim_check_bounded( :: AbstractNLPModel, :: NLPAtX; pnorm :: Real = Inf, kwargs...)`

Require `state.gx` (filled if not provided).

See also `unconstrained_check`, `unconstrained2nd_check`, `KKT`
"""
function optim_check_bounded(pb::AbstractNLPModel, state::NLPAtX; pnorm::Real = Inf, kwargs...)
  if state.gx == _init_field(typeof(state.gx)) # should be filled if void
    update!(state, gx = grad(pb, state.x))
  end

  proj = max.(min.(state.x - state.gx, pb.meta.uvar), pb.meta.lvar)
  gradproj = state.x - proj

  return norm(gradproj, pnorm)
end

"""
constrained: return the violation of the KKT conditions
length(lambda) > 0
"""
function _grad_lagrangian(pb::AbstractNLPModel, state::NLPAtX)
  if (pb.meta.ncon == 0) & !has_bounds(pb)
    return state.gx
  elseif pb.meta.ncon == 0
    return state.gx + state.mu
  else
    return state.gx + state.mu + state.Jx' * state.lambda
  end
end

function _sign_multipliers_bounds(pb::AbstractNLPModel, state::NLPAtX)
  if has_bounds(pb)
    return vcat(
      min.(max.(state.mu, 0.0), -state.x + pb.meta.uvar),
      min.(max.(-state.mu, 0.0), state.x - pb.meta.lvar),
    )
  else
    return zeros(0)
  end
end

function _sign_multipliers_nonlin(pb::AbstractNLPModel, state::NLPAtX)
  if pb.meta.ncon == 0
    return zeros(0)
  else
    return vcat(
      min.(max.(state.lambda, 0.0), -state.cx + pb.meta.ucon),
      min.(max.(-state.lambda, 0.0), state.cx - pb.meta.lcon),
    )
  end
end

function _feasibility(pb::AbstractNLPModel, state::NLPAtX)
  if pb.meta.ncon == 0
    return vcat(max.(state.x - pb.meta.uvar, 0.0), max.(-state.x + pb.meta.lvar, 0.0))
  else
    return vcat(
      max.(state.cx - pb.meta.ucon, 0.0),
      max.(-state.cx + pb.meta.lcon, 0.0),
      max.(state.x - pb.meta.uvar, 0.0),
      max.(-state.x + pb.meta.lvar, 0.0),
    )
  end
end

"""
KKT: verifies the KKT conditions

`KKT( :: AbstractNLPModel, :: NLPAtX; pnorm :: Real = Inf, kwargs...)`

Note: `state.gx` is mandatory + if bounds `state.mu` + if constraints `state.cx`, `state.Jx`, `state.lambda`.

See also `unconstrained_check`, `unconstrained2nd_check`, `optim_check_bounded`
"""
function KKT(pb::AbstractNLPModel, state::NLPAtX; pnorm::Real = Inf, kwargs...)

  #Check the gradient of the Lagrangian
  gLagx = _grad_lagrangian(pb, state)
  #Check the complementarity condition for the bounds
  dual_res_bounds = _sign_multipliers_bounds(pb, state)
  #Check the complementarity condition for the constraints
  res_nonlin = _sign_multipliers_nonlin(pb, state)
  #Check the feasibility
  feas = _feasibility(pb, state)

  res = vcat(gLagx, feas, dual_res_bounds, res_nonlin)

  return norm(res, pnorm)
end


"""
    A mixed continuous-discrete distribution
"""
struct MixedDistribution{DT<:Dst.ContinuousUnivariateDistribution} <: Dst.ContinuousUnivariateDistribution
    mass_probs::Vector{Float64}
    mass_points::Vector{Float64}
    cont_weights::Vector{Float64}
    cont_dists::Vector{DT}
    function MixedDistribution(mass_probs::Vector{Float64}, mass_points::Vector{Float64}, cont_weights::Vector{Float64}, cont_dists::Vector{DT}) where {DT <: Dst.ContinuousUnivariateDistribution}
        length(mass_probs) == length(mass_points)  || ArgumentError("Mass length inconsistency")
        length(cont_weights) == length(cont_dists) || ArgumentError("Continuous length inconsistency")
        # sort mass points
        pairs = sort(zip(mass_points, mass_probs) |> collect, by = t -> t[1])
        mass_points = [t[1] for t in pairs]
        mass_probs =  [t[2] for t in pairs]
        s = sum(mass_probs) + sum(cont_weights)
        s ≈ one(eltype(mass_probs)) && return new{DT}(mass_probs, mass_points, cont_weights, cont_dists)
        @warn "Prior probabilities not normalized"
        new{DT}(mass_probs./s, mass_points, cont_weights./s, cont_dists)
    end
end

function Base.maximum(md::MixedDistribution)
    max_mass = isempty(md.mass_points) ? -Inf64 : maximum(md.mass_points)
    max_cont = isempty(md.cont_dists)  ? -Inf64 : maximum.(md.cont_dists) |> maximum
    return max(max_mass, max_cont)
end

function Base.minimum(md::MixedDistribution)
    min_mass = isempty(md.mass_points) ? +Inf64 : minimum(md.mass_points)
    min_cont = isempty(md.cont_dists)  ? +Inf64 : minimum.(md.cont_dists) |> minimum
    return min(min_mass, min_cont)
end

function Dst.ncategories(md::MixedDistribution)
    return length(md.mass_probs)
end

function Statistics.mean(md::MixedDistribution)
	μ = sum(md.mass_points.*md.mass_probs)
	for idx in eachindex(md.cont_dists)
		μ += mean(md.cont_dists[idx]) * md.cont_weights[idx]
	end
	return μ
end

"""
	Variance of mixed continuous discrete computed as a mixture:
	Var[X] = E[Var[X|θ]] + Var[E[X|θ]]
"""
function Statistics.var(md::MixedDistribution)
	disc_weight = sum(md.mass_probs)
    μ_discrete = sum(md.mass_points.*md.mass_probs)/disc_weight
    σ2d = sum((md.mass_points.-μ_discrete).^2 .* md.mass_probs)
    μ = mean(md)
    tot_var = isempty(md.mass_probs) ? 0.0 : disc_weight * (σ2d + μ_discrete^2)
    for idx in eachindex(md.cont_dists)
		tot_var += md.cont_weights[idx] * (var(md.cont_dists[idx]) + mean(md.cont_dists[idx])^2)
    end
    return tot_var - μ^2
end

function Dst.insupport(md::MixedDistribution, x::Real)
    return minimum(md) <= x && x <= maximum(md)
end

function Dst.pdf(md::MixedDistribution, x::Real)
    any(md.mass_points .≈ x) && DomainError("PDF undefined at a mass point")
end

function Dst.cdf(md::MixedDistribution, x::Real)
    cdf_point = 0.0
    for idx in eachindex(md.mass_probs)
        if md.mass_points[idx] <= x
            cdf_point += md.mass_probs[idx]
        end
    end
    for idx in eachindex(md.cont_dists)
        cdf_point += md.cont_weights[idx] * Dst.cdf(md.cont_dists[idx], x)
    end
    return cdf_point
end

"""
	Number of components of the mixed distribution:
	number of discrete points + continuous distributions
"""
Dst.ncomponents(md::MixedDistribution) = Dst.ncategories(md) + length(md.cont_dists)

"""
	Returns a vector of prior probabilities [discrete priors, continuous priors]
"""
Dst.probs(md::MixedDistribution) = vcat(md.mass_probs, md.cont_weights)

Base.rand(rng::AbstractRNG, md::MixedDistribution) = ordered_rand(md, rand(rng))

Base.rand(md::MixedDistribution) = rand(Random.GLOBAL_RNG, md)

"""
Integration of function including Dirac peaks
∫x⋅f(x)dx
Uses the trapezoid method with given step dx
"""
function partial_expectancy(md::MixedDistribution; xlow = minimum(md), xhigh = maximum(md), dx = min((xhigh - xlow) * 0.05, 0.01))
    xhigh > xlow || throw(ArgumentError("xlow = $xlow greater than xhigh = $xhigh"))
    disc_idx = 1
    while disc_idx <= length(md.mass_points) && md.mass_points[disc_idx] < xlow
        disc_idx += 1
    end
    still_mass = disc_idx <= length(md.mass_points) && md.mass_points[disc_idx] <= xhigh
    x = xlow
    integrator = 0.0
    while x <= xhigh
        δ = min(dx, xhigh - x)
        # check if mass point in chunk
        if still_mass && md.mass_points[disc_idx] <= x + δ
            integrator += md.mass_points[disc_idx] * md.mass_probs[disc_idx]
            disc_idx += 1
            if disc_idx > length(md.mass_points) || md.mass_probs[disc_idx] > xhigh
                still_mass = false
            end
        end
        # add continuous chunk
        for (d,p) in zip(md.cont_dists, md.cont_weights)
            if Dst.insupport(d, x) && Dst.insupport(d, x + δ)
                integrator += p * 0.5 * δ * (x + δ*0.5) * (Dst.pdf(d,x) + Dst.pdf(d,x+δ))
            end
        end
        x += dx
    end
    return integrator
end

"""
Computing integral of function including Dirac peaks in intervals between eval_points
"""
function partial_expectancy(md::MixedDistribution, eval_points::V; dx = min((xhigh - xlow) * 0.05, 0.01)) where {V<:AbstractVector{<:Real}}
    intervals = begin
        sorted_eval = sort(eval_points)
        collect(zip(sorted_eval[1:end-1], sorted_eval[2:end]))
    end
    ys = [partial_expectancy(md, xlow = xl, xhigh = xh, dx = dx) for (xl,xh) in intervals]
    (intervals, ys)
end

"""
Exception used when Cumulated Distribution Function does not integrate to 1 in a computation
"""
struct CDFException{D<:Dst.Distribution} <: Exception
	d::D
	CDFException(d::D) where {D<:Dst.Distribution} = new{D}(d)
end

"""
	Ordered_rand returns the quantity corresponding to
	an already-generated random number
"""
function ordered_rand(md::MixedDistribution, r)
	cs_discrete = cumsum(md.mass_probs)
	for idx in eachindex(cs_discrete)
		if r <= cs_discrete[idx]
			return md.mass_points[idx]
		end
	end
	prob_offset = isempty(cs_discrete) ? 0.0 : last(cs_discrete)
	cs_cont = vcat(cs_discrete, cumsum(md.cont_weights) .+ prob_offset)
	for idx in length(cs_discrete)+1:length(cs_cont)
		if r <= cs_cont[idx]
			cont_idx = idx-length(cs_discrete)
			last_prob = idx == 1 ? 0.0 : cs_cont[idx-1]
			scaled_point = (r - last_prob) / (cs_cont[idx] - last_prob)
			return Dst.quantile(md.cont_dists[cont_idx], scaled_point)
		end
	end
	# all points reached beforehand
	throw(CDFException(md))
end


include("../demos/lasso.jl")
include("../demos/rpca.jl")


@testset "Basic 3D Orientation Tests" begin
	w = (0f0,0f0); x = (1,1,2,2); y = (2,2,1,1); z = (1,2,1,2);
	p = (2f0,2f0,2f0); q = (3,4,3); r = (3f0,1f0,1f0,1f0); s = (2,3,4); t = (-4f0,-1f0,-1f0,-7f0);

	@test plane(x,y,z) == (-3, 3, 3, -3)
	@test dot(q,s) == 30
	@test toPoint3H(q) == (1,q[1],q[2],q[3])
	@test abovePl(t,plane(x,y,z))
	@test !belowPl(toPoint3H(p),plane(x,y,z))
	@test plane(p, q, s) == plane(toPoint3H(p), toPoint3H(q), toPoint3H(s))

#	pts = [Point3H(1f0, a, b, c) for a in [2f0 4f0], b in [2f0 4f0], c in [2f0 4f0]]

	a = (1f0,3f0,3f0,3f0); b = (1f0,2f0,2f0,2f0); c = (1f0,2f0,2f0,4f0); d = (1f0,4f0,2f0,2f0); e = (1f0,2f0,4f0,2f0); f = (1f0,2f0,4f0,4f0); g = (1f0,4f0,4f0,2f0);
	h = (1f0); i = (1f0);

#	@test CH([a;b;c;d;e;f;g])

end;


module SlimeMoldOptim

using  InvertedIndices
using  Random
using  Statistics
export slimeMoldAlgo, SlimeMolds

  struct SlimeMolds
    S::Matrix{Float32};
    fitnessS::Vector{Float32};
    x_b_global::Vector{Float32};
    minFGlobal::Float32;
    x_w_global::Vector{Float32};
    maxFGlobal::Float32;
    iter::Int;
  end

  function drawMoldsPositionsBoundary!( ;rng::RandomDevice,
                                        X::Matrix{Float32},
                                        rand_mat::Matrix{Float32},
                                        upper::Vector{Float32},
                                        lower::Vector{Float32} )
    @assert length( upper ) == length( lower )
    @assert all( size( X )  == size( rand_mat ) )

    nFireworks = size( X, 2 );
    d          = length( lower );
    rand_mat  .= Float32.( rand( rng, d, nFireworks ) );
    for iter ∈ 1:1:nFireworks
      @inbounds x_sp     = view( X,        :, iter );
      @inbounds rand_vec = view( rand_mat, :, iter );
      for k ∈ 1:1:d
        @inbounds lu      = lower[k]::Float32;
        @inbounds x_sp[k] = ( rand_vec[k] * ( upper[k] - lu ) + lu )::Float32;
      end
    end
  end

  function smellIndex!( fitnessS::Vector{Float32},
                        W::Vector{Float32},
                        rng::RandomDevice )

    nMolds = length( fitnessS )::Int;
    τ      = median( fitnessS )::Float32;
    bF     = minimum( fitnessS )::Float32;
    wF     = maximum( fitnessS )::Float32;
    denom  = ( bF - wF )::Float32;
    if denom == 0.0f0
      denom = 1f-20;
    end
    rr     = rand( rng, Float32, nMolds );
    @simd for ii ∈ 1:1:nMolds
      if fitnessS[ii] <= τ
        @inbounds W[ii] = 1.0f0 + rr[ii] * log( ( ( bF - fitnessS[ii] ) / denom ) + 1.0f0 );
      else
        @inbounds W[ii] = 1.0f0 - rr[ii] * log( ( ( bF - fitnessS[ii] ) / denom ) + 1.0f0 );
      end
    end
  end

  function updateMoldPositions!( S::Matrix{Float32},
                                 vb::Float32,
                                 vc::Float32,
                                 W::Vector{Float32},
                                 z::Float32,
                                 rng::RandomDevice,
                                 lower::Vector{Float32},
                                 upper::Vector{Float32},
                                 x_b_global::Vector{Float32},
                                 minFGlobal::Float32,
                                 fitnessS::Vector{Float32} )
    nMolds = size( S, 2 );
    d      = size( S, 1 );
    rr     = rand( rng, Float32, nMolds );
    Threads.@threads for ii ∈ 1:1:nMolds
      r = rr[ii];
      if r < z
        ## reinit slime mold
        @inbounds S[:,ii] .= ( rand( rng, Float32, d ) .* ( upper .- lower ) ) .+ lower;
      else
        p = tanh( abs( fitnessS[ii] - minFGlobal  ) );
        if r < p
          ## execute exploitation step
          idx                = collect( 1:1:nMolds )[Not( ii )][rand( 1:d-1, 1 )][1];
          @inbounds S[:,ii] .= x_b_global .+ vb * ( W[idx] * ( S[:,idx] .- S[:,ii] ) )
        else
          ## execute exploration step
          @inbounds S[:,ii] .= vc * S[:,ii];
        end
        ## check if we are out of bounds, if true, reinit
        if any( S[:,ii] .< lower ) || any( S[:,ii] .> upper )
          @inbounds S[:,ii] .= ( rand( rng, Float32, d ) .* ( upper .- lower ) ) .+ lower;
        end
      end
    end
  end

  function evaluateCandidatesParallel( ;kwargs::Tuple,
                                       f::Function,
                                       X::Matrix{Float32},
                                       fitness::Vector{Float32} )
    @assert size( X, 2 ) == size( fitness, 1 )
    n = size( X, 2 )::Int;
    @assert n > 0
    Threads.@threads for ii ∈ 1:1:n
      x            = view( X, :, ii );
      fitness[ii]  = f( x; kwargs=kwargs );
    end
  end

  """

           slimeMoldAlgo( nMolds::Int,
                          n_vc::Int,
                          kwargs...;
                          lower::Vector{Float32},
                          upper::Vector{Float32},
                          objFunction::Function,
                          maxiter::Int,
                          ϵ_conv::Float32,
                          z::Float32=0.03f0 )

        minimize objective function objFunction, the solution space is limited by lower and upper bound.
        - The optimization algorithm utilized is the slime mold algorithm.
        - The nMolds parameter governs the number of molds being modeled.
        - n_vc is the number of steps learning parameter vb needs to reach 0.0.
        - maxiter denotes the overall number of algo iteraions, it has to be >= n_vc.

        the number of sparks per firework, in remains constant foreach firework. ϵ_A is the smoothing parameter
        controlling the variance of amplitudes computed foreach fw. C_a ist the upscaling parameter for explosion
        amplitudes. C_r is the downscaling parameter for explosion amplitudes. XPrimary is the feature set of the
        primary algorithm to be tuned. yPrimary is the target set of the primary algorithm. maxiter is the maximum
        number iteraions.  ϵ_conv denotes the convergence parameter.
  """
  function slimeMoldAlgo( nMolds::Int,
                          n_vc::Int,
                          kwargs...;
                          lower::Vector{Float32},
                          upper::Vector{Float32},
                          objFunction::Function,
                          maxiter::Int,
                          ϵ_conv::Float32,
                          z::Float32=0.03f0 )
    @assert nMolds >  0
    @assert n_vc   >  0
    @assert n_vc   >= maxiter
    @assert length( lower ) == length( upper )
    @assert all( lower .< upper )
    @assert maxiter >  0
    @assert ϵ_conv  >  0.0f0
    @assert z       >= 0.0f0
    @assert z       <  1.0f0

    d = length( lower );
    @assert d > 0
    rng = RandomDevice();

    δ        = 1.0f0 / Float32( n_vc - 1 );
    vc_full  = collect( 1.0f0:-δ:0.0f0 );
    S        = Matrix{Float32}( undef, d, nMolds );
    rand_mat = Matrix{Float32}( undef, d, nMolds );
    fitnessS = Vector{Float32}( undef, nMolds );
    W        = Vector{Float32}( undef, nMolds );
    x_b      = Vector{Float32}( undef, d );
    p        = -1.0f0;
    drawMoldsPositionsBoundary!( ;rng=rng, X=S, rand_mat=rand_mat, upper=upper, lower=lower );
    maxFGlobal = Float32( objFunction( S[:,1]; kwargs=() ) );
    minFGlobal = Float32( objFunction( S[:,1]; kwargs=() ) );
    x_b_global = zeros( Float32, d );
    x_w_global = zeros( Float32, d );
    iter       = 1;
    while iter <= maxiter
      evaluateCandidatesParallel(;kwargs   = kwargs,
                                 f         = objFunction,
                                 X         = S,
                                 fitness   = fitnessS );
      smellIndex!( fitnessS, W, rng );
      if maximum( fitnessS ) > maxFGlobal
        @inbounds maxFGlobal  = fitnessS[argmax( fitnessS )]::Float32;
        @inbounds x_w_global .= S[:,argmax( fitnessS )];
      end
      if minimum( fitnessS ) < minFGlobal
        @inbounds minFGlobal  = fitnessS[argmin( fitnessS )]::Float32;
        @inbounds x_b_global .= S[:,argmin( fitnessS )];
      end
      a             = Float32( atanh( -( iter / maxiter ) + 1.0f0 ) );
      vb            = ( ( rand( rng, Float32, 1 )[1] * 2.0f0 * a ) - a )::Float32;
      @inbounds vc  = vc_full[iter]::Float32;
      updateMoldPositions!( S, vb, vc, W, z, rng, lower, upper, x_b_global, minFGlobal, fitnessS );
      iter += 1;
    end
    return SlimeMolds( S, fitnessS, x_b_global, minFGlobal, x_w_global, maxFGlobal, iter - 1 )
  end

end


# This implements an efficient Gibbs sampler for the multivariate truncated
# normal distribution subject to linear constraints from:

# Y. Li and S. Ghosh (2015)
# "Efficient Sampling Methods for Truncated Multivariate Normal and Student-t
#  Distributions Subject to Linear Inequality Constraints"
# Journal of Statistical Theory & Practice
# doi: 10.1080/15598608.2014.996690
# http://dx.doi.org/10.1080/15598608.2014.996690

function GibbsTMVN(μ, Σ, c, d, R̃, NumSamples = 10000,
                   x_init = zeros(length(μ)))
  p = length(μ)
  m = size(R̃, 1)
  samples = zeros(p, NumSamples)
  Σhalf = chol(Σ)'
  Σhalfinv = Σhalf \ eye(m)
  R = R̃ * Σhalf
  a = c - R̃ * μ
  b = d - R̃ * μ
  x = Σhalfinv * (x_init - μ)
  Rx = R * x

  # Make sure starting value is feasible:
  if !(all(Rx .>= a) && all(Rx .<= b))
    x = getFeasiblePoint(a, b, R)
  end

  bounds = [-Inf, Inf]
  tmp_l = zeros(m)
  tmp_u = zeros(m)

  for iter = 1:NumSamples
    for i = 1:p
      bounds[1] = -Inf
      bounds[2] = Inf
      GetBounds!(bounds, tmp_l, tmp_u, a, b, R, x, i)
      Spl = TruncatedNormalSampler(bounds[1], bounds[2])
      x[i] = rand(Spl)
    end
    samples[:, iter] = x
  end

  for iter = 1:NumSamples
     samples[:, iter] = Σhalf * samples[:, iter] + μ
  end
  return samples
end

function getFeasiblePoint(a::Vector{Float64}, b::Vector{Float64}, R::Matrix{Float64})
  set_default_solver(SCSSolver(verbose=false))
  m = size(R, 2)
  x = Variable(m)
  # Convex solver doesn't work with Infinity constraints
  finiteidx_a = (1:m)[a .> -Inf]
  finiteidx_b = (1:m)[b .< Inf]
  problem = minimize(0)
  for i in finiteidx_a
    problem.constraints += R[i, :] * x >= a
  end
  for i in finiteidx_b
    problem.constraints += R[i, :] * x <= b
  end
  solve!(problem)
  return vec(x.value)
end

function GetBounds!(bounds::Vector{Float64}, tmp_l::Vector{Float64},
                    tmp_u::Vector{Float64}, a::Array{Float64},
                    b::Vector{Float64}, R::Matrix{Float64},
                    x::Vector{Float64}, i::Int)
  m = size(R, 1)
  p = size(R, 2)
  tmp_l[:] = 0.0
  tmp_u[:] = 0.0

  for j = 1:m
    for k = 1:m
      if k != i
        tmp_l[j] -= R[j, k] * x[k]
      end
    end
    tmp_u[j] = copy(tmp_l[j])
    tmp_l[j] += a[j]
    tmp_l[j] /= R[j, i]
    tmp_u[j] += b[j]
    tmp_u[j] /= R[j, i]
  end

  for j = 1:m
    if R[j, i] > 0
      bounds[1] = max(tmp_l[j], bounds[1])
      bounds[2] = min(tmp_u[j], bounds[2])
    elseif R[j, i] < 0
      bounds[1] = max(tmp_u[j], bounds[1])
      bounds[2] = min(tmp_l[j], bounds[2])
    end
  end
  nothing
end


include("utils.jl")

@testset "fixed Rank" begin
    M = FixedRankMatrices(3, 2, 2)
    M2 = FixedRankMatrices(3, 2, 1)
    Mc = FixedRankMatrices(3, 2, 2, ℂ)
    x = SVDMPoint([1.0 0.0; 0.0 1.0; 0.0 0.0])
    x2 = SVDMPoint([1.0 0.0; 0.0 1.0; 0.0 0.0], 1)
    v = UMVTVector([0.0 0.0; 0.0 0.0; 1.0 1.0], [1.0 0.0; 0.0 1.0], zeros(2, 2))
    @test repr(M) == "FixedRankMatrices(3, 2, 2, ℝ)"
    @test repr(Mc) == "FixedRankMatrices(3, 2, 2, ℂ)"
    @test sprint(show, "text/plain", x) == """
    $(sprint(show, SVDMPoint{Matrix{Float64}, Vector{Float64}, Matrix{Float64}}))
    U factor:
     3×2 $(sprint(show, Matrix{Float64})):
      1.0  0.0
      0.0  1.0
      0.0  0.0
    singular values:
     2-element $(sprint(show, Vector{Float64})):
      1.0
      1.0
    Vt factor:
     2×2 $(sprint(show, Matrix{Float64})):
      1.0  0.0
      0.0  1.0"""
    @test sprint(show, "text/plain", v) == """
    $(sprint(show, UMVTVector{Matrix{Float64}, Matrix{Float64}, Matrix{Float64}}))
    U factor:
     3×2 $(sprint(show, Matrix{Float64})):
      0.0  0.0
      0.0  0.0
      1.0  1.0
    M factor:
     2×2 $(sprint(show, Matrix{Float64})):
      1.0  0.0
      0.0  1.0
    Vt factor:
     2×2 $(sprint(show, Matrix{Float64})):
      0.0  0.0
      0.0  0.0"""

    @test inner(M, x, v, v) == norm(M, x, v)^2
    @test x == SVDMPoint(x.U, x.S, x.Vt)
    @test v == UMVTVector(v.U, v.M, v.Vt)
    @testset "Fixed Rank Matrices – Basics" begin
        @test representation_size(M) == (3, 2)
        @test representation_size(Mc) == (3, 2)
        @test manifold_dimension(M) == 6
        @test manifold_dimension(Mc) == 12
        @test !is_manifold_point(M, SVDMPoint([1.0 0.0; 0.0 0.0], 2))
        @test_throws DomainError is_manifold_point(
            M,
            SVDMPoint([1.0 0.0; 0.0 0.0], 2),
            true,
        )
        @test is_manifold_point(M2, x2)

        @test !is_tangent_vector(
            M,
            SVDMPoint([1.0 0.0; 0.0 1.0; 0.0 0.0]),
            UMVTVector(zeros(2, 1), zeros(1, 2), zeros(2, 2)),
        )
        @test !is_tangent_vector(M, SVDMPoint([1.0 0.0; 0.0 0.0], 2), v)
        @test_throws DomainError is_tangent_vector(
            M,
            SVDMPoint([1.0 0.0; 0.0 0.0], 2),
            v,
            true,
        )
        @test !is_tangent_vector(M, x, UMVTVector(x.U, v.M, x.Vt, 2))
        @test_throws DomainError is_tangent_vector(
            M,
            x,
            UMVTVector(x.U, v.M, x.Vt, 2),
            true,
        )
        @test !is_tangent_vector(M, x, UMVTVector(v.U, v.M, x.Vt, 2))
        @test_throws DomainError is_tangent_vector(
            M,
            x,
            UMVTVector(v.U, v.M, x.Vt, 2),
            true,
        )

        @test is_manifold_point(M, x)
        @test is_tangent_vector(M, x, v)
    end
    types = [[Matrix{Float64}, Vector{Float64}, Matrix{Float64}]]
    TEST_FLOAT32 && push!(types, [Matrix{Float32}, Vector{Float32}, Matrix{Float32}])

    for T in types
        @testset "Type $T" begin
            y = retract(M, x, v, PolarRetraction())
            z = SVDMPoint([1/sqrt(2) 1/sqrt(2); 1/sqrt(2) -1/sqrt(2); 0.0 0.0])
            pts = []
            for p in [x, y, z]
                push!(pts, SVDMPoint(convert.(T, [p.U, p.S, p.Vt])...))
            end
            for p in pts
                @test is_manifold_point(M, p)
            end
            @testset "SVD MPoint Basics" begin
                s = svd(x.U * Diagonal(x.S) * x.Vt)
                x2 = SVDMPoint(s)
                x3 = SVDMPoint(s.U, s.S, s.Vt)
                @test SVDMPoint(x.U, x.S, x.Vt) == x
                @test x.S == x2.S
                @test x.U == x2.U
                @test x.Vt == x2.Vt
                @test x == x2
                @test x2.U == x3.U
                @test x2.S == x3.S
                @test x2.Vt == x3.Vt
                @test x2 == x3
                y = SVDMPoint([1.0 0.0; 0.0 0.0; 0.0 0.0], 1)
                s2 = svd([1.0 0.0; 0.0 0.0; 0.0 0.0])
                y2 = SVDMPoint(s2, 1)
                y3 = SVDMPoint(s2.U, s2.S, s2.Vt, 1)
                @test y.S == y2.S
                @test y.U == y2.U
                @test y.Vt == y2.Vt
                @test y == y2
                @test y2.U == y3.U
                @test y2.S == y3.S
                @test y2.Vt == y3.Vt
                @test y2 == y3

                @test is_manifold_point(M, x)
                xM = x.U * Diagonal(x.S) * x.Vt
                @test is_manifold_point(M, xM)
                @test !is_manifold_point(M, xM[1:2, :])
                @test_throws DomainError is_manifold_point(M, xM[1:2, :], true)
                @test_throws DomainError is_manifold_point(
                    FixedRankMatrices(3, 2, 1),
                    x,
                    true,
                )
                @test_throws DomainError is_manifold_point(
                    FixedRankMatrices(3, 2, 1),
                    xM,
                    true,
                )
                xF1 = SVDMPoint(2 * x.U, x.S, x.Vt)
                @test !is_manifold_point(M, xF1)
                @test_throws DomainError is_manifold_point(M, xF1, true)
                xF2 = SVDMPoint(x.U, x.S, 2 * x.Vt)
                @test !is_manifold_point(M, xF2)
                @test_throws DomainError is_manifold_point(M, xF2, true)
            end
            @testset "UMV TVector Basics" begin
                w = UMVTVector(v.U, 2 * v.M, v.Vt)
                @test v + w == UMVTVector(2 * v.U, 3 * v.M, 2 * v.Vt)
                @test v - w == UMVTVector(0 * v.U, -v.M, 0 * v.Vt)
                @test 2 * v == UMVTVector(2 * v.U, 2 * v.M, 2 * v.Vt)
                @test v * 2 == UMVTVector(v.U * 2, v.M * 2, v.Vt * 2)
                @test 2 \ v == UMVTVector(2 \ v.U, 2 \ v.M, 2 \ v.Vt)
                @test v / 2 == UMVTVector(v.U / 2, v.M / 2, v.Vt / 2)
                @test +v == v
                @test -v == UMVTVector(-v.U, -v.M, -v.Vt)
                w = UMVTVector(v.U, v.M, v.Vt)
                @test v == w
                w = allocate(v, number_eltype(v))
                zero_tangent_vector!(M, w, x)
                oneP = SVDMPoint(one(zeros(3, 3)), ones(2), one(zeros(2, 2)), 2)
                @test oneP == one(x)
                oneV = UMVTVector(one(zeros(3, 3)), one(zeros(2, 2)), one(zeros(2, 2)), 2)
                @test oneV == one(v)
            end
            test_manifold(
                M,
                pts,
                test_exp_log=false,
                default_inverse_retraction_method=nothing,
                test_injectivity_radius=false,
                default_retraction_method=PolarRetraction(),
                test_is_tangent=false,
                test_project_tangent=true,
                test_default_vector_transport=false,
                test_forward_diff=false,
                test_reverse_diff=false,
                test_vector_spaces=false,
                test_vee_hat=false,
                test_tangent_vector_broadcasting=false, #broadcast not so easy for 3 matrix type
                projection_atol_multiplier=15,
                retraction_methods=[PolarRetraction()],
                mid_point12=nothing,
            )
        end
    end
end


# Model.
@sugar function hmm(N::Int)
    z₀ ~ Normal(0.0, 1.0)
    z = z₀
    for i in 1 : N
        rand(:z => i, Normal(z, 1.0))
        rand(:x => i, Normal(z, 1.0))
    end
end

# Data.
data_len = 100
x = map(x) do k
    3.0 + randn()
end
obs = target(map(1 : data_len) do i
                 (:x => i, ) => x[i]
             end)

@testset "Inference library smoke test 1 - Bayesian linear regression" begin

    # Importance sampling.
    is_test = () -> begin
        n_samples = 50000
        ps, lnw = importance_sampling(obs, n_samples, bayesian_linear_regression, (x, ))
        zipped = zip(ps.calls, lnw)

        est_σ = sum(map(zipped) do (cl, w)
                        (cl[:σ]) * exp(w)
                    end)

        est_β = sum(map(zipped) do (cl, w)
                        (cl[:β]) * exp(w)
                    end)
        @test isapprox(0.0, est_β - 3.0, atol=1e-1)
    end

    # MH random walk kernel.
    mh_test = () -> begin
        n_iters = 50000
        ret, cl = generate(obs, bayesian_linear_regression, x)
        calls = []
        for i in 1 : n_iters
            cl, _ = mh(target([(:σ, ), (:β, )]), cl)
            i % 30 == 0 && begin
                push!(calls, cl)
            end
        end

        est_σ = sum(map(calls) do cl
                        (cl[:σ])
                    end) / length(calls)

        est_β = sum(map(calls) do cl
                        (cl[:β])
                    end) / length(calls)
        @test isapprox(0.0, est_β - 3.0, atol=1e-1)
    end

    # MH random walk kernel with proposal.
    proposal = (cl, y) -> begin
        σ = rand(:σ, Normal(6.0, 1.0))
        β = rand(:β, Normal(mean(y) / length(y), 3.0))
    end

    mh_test_with_proposal = () -> begin
        n_iters = 10000
        ret, cl = generate(obs, bayesian_linear_regression, x)
        calls = []
        for i in 1 : n_iters
            cl, _ = mh(cl, proposal, (y, ))
            i % 30 == 0 && begin
                push!(calls, cl)
            end
        end

        est_σ = sum(map(calls) do cl
                        (cl[:σ])
                    end) / length(calls)

        est_β = sum(map(calls) do cl
                        (cl[:β])
                    end) / length(calls)
        @test isapprox(0.0, est_β - 3.0, atol=1e-1)
    end

    # ADVI.
    surrogate = () -> begin
        μ₁ = learnable(:μ₁)
        μ₂ = learnable(:μ₂)
        σ₂ = learnable(:σ₂)
        σ = rand(:σ, Normal(1.0 + μ₁, 0.3))
        β = rand(:β, Normal(μ₂, σ₂ ^ 2))
    end

    advi_test = () -> begin
        ps = learnables([(:μ₁, ) => 5.0,
                         (:μ₂, ) => 1.0,
                         (:σ₂, ) => 1.0])
        @time ps, elbows, _ = advi(obs, ps,
                                   surrogate, (),
                                   bayesian_linear_regression, (x, );
                                   opt = ADAM(0.05, (0.9, 0.8)),
                                   iters = 1000)
    end

    # ADGV.
    surrogate = () -> begin
        μ₁ = learnable(:μ₁)
        μ₂ = learnable(:μ₂)
        σ₂ = learnable(:σ₂)
        σ = rand(:σ, Normal(1.0 + μ₁, 0.3))
        β = rand(:β, Normal(μ₂, σ₂ ^ 2))
    end

    adgv_test = () -> begin
        ps = learnables([(:μ₁, ) => 5.0,
                         (:μ₂, ) => 1.0,
                         (:σ₂, ) => 1.0])
        ps, elbows, _ = adgv(obs, ps, 10,
                             surrogate, (),
                             bayesian_linear_regression, (x, );
                             opt = ADAM(0.05, (0.9, 0.8)),
                             iters = 1000)
    end

    # HMC kernel.
    hmc_test = () -> begin
        n_iters = 1000
        ret, cl = generate(obs, bayesian_linear_regression, x)
        calls = []
        for i in 1 : n_iters
            cl, _ = hmc(target([(:σ, ), (:β, )]), cl)
            i % 10 == 0 && begin
                push!(calls, cl)
            end
        end

        est_σ = sum(map(calls) do cl
                        (cl[:σ])
                    end) / length(calls)

        est_β = sum(map(calls) do cl
                        (cl[:β])
                    end) / length(calls)
        @test isapprox(0.0, est_β - 3.0, atol=1e-1)
    end

    # Combination kernel.
    combo_kernel_test = () -> begin
        n_iters = 1000
        ret, cl = generate(obs, bayesian_linear_regression, x)
        calls = []
        for i in 1 : n_iters
            cl, _ = mh(cl, proposal, (y, ))
            cl, _ = hmc(target([(:σ, ), (:β, )]), cl)
            i % 10 == 0 && begin
                push!(calls, cl)
            end
        end

        est_σ = sum(map(calls) do cl
                        (cl[:σ])
                    end) / length(calls)

        est_β = sum(map(calls) do cl
                        (cl[:β])
                    end) / length(calls)
        @test isapprox(0.0, est_β - 3.0, atol=1e-1)
    end

    # Boomerang kernel.
    boomerang_test = () -> begin
        n_iters = 500
        ret, cl = generate(obs, bayesian_linear_regression, x)
        sel = get_target(cl)
        sel_values, _ = get_choice_gradients(sel, cl, 1.0)
        d = length(array(sel_values, Float64))
        flow = Boomerang(sparse(I, d, d), zeros(d), 1.0)
        θ = rand(MvNormal(d, 1.0))
        calls = []
        for i in 1 : n_iters
            cl, _ = pdmk(target([(:σ, ), (:β, )]), cl, flow, θ)
            i % 10 == 0 && begin
                push!(calls, cl)
            end
        end

        est_σ = sum(map(calls) do cl
                        (cl[:σ])
                    end)

        est_β = sum(map(calls) do cl
                        (cl[:β])
                    end)
    end

    @time is_test()
    @time mh_test()
    @time mh_test_with_proposal()
    @time hmc_test()
    @time combo_kernel_test()
    #@time boomerang_test()
    #@time advi_test()
    #@time adgv_test()

end


using VegaLite
using DataFrames

# a simple histogram of random standard normal draws

data_values(DataFrame(x=randn(200))) +
    mark_bar() +
    encoding_x_quant(:x; bin=Dict(:maxbins=>20), axis=Dict(:title=>"values")) +
    encoding_y_quant(:*, aggregate="count", axis=Dict(:title=>"number of draws"))


# Tests boteSalvatICX, boteSalvatEdgeEnergy and boteSalvatAvailable
using Test
using BoteSalvatICX

@testset "ICX" begin

    # Comparing with values generated by xion.f
    @test isapprox(ionizationcrosssection(12,4,7.71792E+01),3.85315E-18,atol=3e-20)
    @test isapprox(ionizationcrosssection(12,4,1.29569E+04),1.24124E-18,atol=1e-20)
    @test isapprox(ionizationcrosssection(12,4,5.78762E+06),1.30393E-19,atol=1e-21)

    @test isapprox(ionizationcrosssection(23,2,2.90068E+08),5.64142E-21,atol=1e-23)
    @test isapprox(ionizationcrosssection(23,2,1.53993E+04),2.63509E-20,atol=1e-22)

    @test isapprox(ionizationcrosssection(45,7,6.68344E+08),1.62343E-20,atol=1e-22)
    @test isapprox(ionizationcrosssection(45,7,1.22321E+03),2.11571E-19,atol=2e-21)

    @test isapprox(ionizationcrosssection(67,1,1.63117E+05),1.36404E-23,atol=1e-25)
    @test isapprox(ionizationcrosssection(78,9,2.98538E+03),2.02808E-20,atol=1e-22)

    @test isapprox(ionizationcrosssection(99,3,1.00000E+09),1.39024E-22,atol=1e-24)
    @test isapprox(ionizationcrosssection(99,3,2.66073E+04),3.85171E-24,atol=1e-25)

    @test isapprox(ionizationcrosssection(1,1,6.87860E+02),1.57729E-17,atol=1e-19)

end

@testset "EdgeEnergy" begin
    # Not a very good set of values...
    @test isapprox(edgeenergy(67,9),1.36386E+3,atol=1.0)
    @test isapprox(edgeenergy(57,4),5.48366E+3,atol=1.0)
    @test isapprox(edgeenergy(42,1),1.99691E+4,atol=1.0)
    @test isapprox(edgeenergy(42,1),1.99691E+4,atol=1.0)
end

@testset "Available" begin
    @test hasedge(1,1)==true
    @test hasedge(1,2)==false
    @test hasedge(11,2)==true
    @test hasedge(12,4)==true
    @test hasedge(12,5)==false
    @test hasedge(26,6)==true
    @test all(i->hasedge(33,i),1:9)
    @test hasedge(33,10)==false
    @test all(z->hasedge(z,9),33:99)
end


using Primes

@testset "Rational reconstruction" begin
    ms = [107, 199, 509, 2^31 - 1]
    as = [QQ(0), QQ(1), QQ(1, 2), QQ(3, 2), QQ(4, 5), QQ(1, 7), QQ(6)]

    for m in ms
        for a in as
            gf = GF(m)
            ac = lift( numerator(a) // gf(denominator(a)) )
            ar1 = Groebner.rational_reconstruction(ac, BigInt(m))
            ar2 = Groebner.rational_reconstruction(Int(ac), m)
            @test ar1 == ar2 == a
        end
    end

    # random tests
    ms = Primes.primes(2^20, 2^21)[1:3]
    as = [QQ(rand(1:100), rand(1:100)) for _ in 1:10]
    for m in ms
        for a in as
            gf = GF(m)
            ac = lift( numerator(a) // gf(denominator(a)) )
            ar1 = Groebner.rational_reconstruction(ac, BigInt(m))
            ar2 = Groebner.rational_reconstruction(Int(ac), m)
            @test ar1 == ar2 == a
        end
    end

end


@testset "vector1" begin
    env = SymbolicEnv()
    cache = SymbolicCache()
    x = [@expr(x), 1, 1.0]
    println(x)
end

@testset "vector2" begin
    env = SymbolicEnv()
    cache = SymbolicCache()
    x = [@expr(x), 1, 1.0]
    println(convert(AbstractSymbolic{Float64}, x))
end

@testset "SymbolicVector1" begin
    @bind x = 10
    v = [x + i for i = 1:10]
    @test seval(v) == [seval(x)+i for i = 1:10]
end

@testset "SymbolicVector1" begin
    @bind x = 10
    v = [x^i + i for i = 1:10]
    @test seval(v, :x) == [i * seval(x)^(i-1) for i = 1:10]
end

@testset "SymbolicVec3" begin
    m = @expr [1, x, 3]
    @test m[1].val == 1
    @test m[2].var == :x
    @test m[3].val == 3
end

@testset "SymbolicVec4" begin
    test = SymbolicEnv()
    m = @expr [1, x, y]
    @bind test x = 10
    @bind test y = 20
    result = seval(m[1:2], test)
    @test result == [1, 10]
end

@testset "SymbolicVec5" begin
    test = SymbolicEnv()
    m = @expr [1, x, y]
    expr = sum(m)
    println(expr)
    @bind test x = 10
    @bind test y = 20
    result = seval(expr, test)
    @test result == 31.0
    result = seval(expr, :x, test)
    @test result == 1.0
end

@testset "SymbolicDot1" begin
    test = SymbolicEnv()
    e1 = @expr [x^2, y, 10]
    e2 = @expr [x^2, y, 10]
    expr = dot(e1, e2)
    println(expr)
    x = 0.5
    y = 0.8
    @bind test x = x
    @bind test y = y
    @test seval(expr, test, SymbolicCache()) == dot([seval(x,test)^2, seval(y,test), 10], [seval(x,test)^2, seval(y,test), 10])
end

@testset "SymbolicDot2" begin
    test = SymbolicEnv()
    e1 = @expr [x^2, y, 10]
    e2 = @expr [x^2, y, 10]
    expr = dot(e1, e2)
    x0 = 0.5
    y0 = 0.8
    @bind test x = x0
    @bind test y = y0
    h = 0.0001
    test1 = SymbolicEnv()
    @bind test1 begin
        x = x0 + h
        y = y0
    end
    test2 = SymbolicEnv()
    @bind test2 begin
        x = x0 - h
        y = y0
    end
    ex = (seval(expr, test1, SymbolicCache()) - seval(expr, test2, SymbolicCache())) / (2*h)
    @test isapprox(seval(expr, :x, test, SymbolicCache()), ex, atol=1.0e-5)
end

@testset "SymbolicDot3" begin
    test = SymbolicEnv()
    e1 = @expr [x^2, y, 10]
    e2 = @expr [x^2, y, 10]
    expr = dot(e1, e2)
    x0 = 10.0
    y0 = 0.8
    @bind test begin
        x = x0
        y = y0
    end
    # h = 0.00001
    # @env test1 begin
    #     x = x+h
    #     y = y
    # end
    # @env test2 begin
    #     x = x-h
    #     y = y
    # end
    # ex = (seval(expr, test1, SymbolicCache()) - 2*seval(expr, test, SymbolicCache()) + seval(expr, test2, SymbolicCache())) / (h^2)
    # @test isapprox(ex, 4*3*x^2)
    @test isapprox(seval(expr, (:x,:x), test), 4*3*x0^2)
end

@testset "SymbolicPlus1" begin
    test = SymbolicEnv()
    e1 = convert(AbstractVectorSymbolic{Float64}, @expr [x^2, y, 10])
    e2 = convert(AbstractVectorSymbolic{Float64}, @expr [x, y^2, 10])
    expr = e1 + e2
    println(expr)
    x0 = 10.0
    y0 = 0.8
    @bind test begin
        x = x0
        y = y0
    end
    println(seval(expr, test))
    println(seval(expr[1], test))
    println(seval(expr[1:2], test))
    println(seval(expr, :x, test))
    println(seval(expr[1], :x, test))
    println(seval(expr[1:2], :x, test))
    println(seval(expr, (:x,:y), test))
    println(seval(expr[1], (:x,:y), test))
    println(seval(expr[1:2], (:x,:y), test))
    # @test isapprox(seval(expr, (:x,:x), test), 4*3*x0^2)
end

@testset "SymbolicDot5" begin
    @bind begin
        x = 4
        y = 3
    end
    v1 = [x, y]
    v2 = [x, 8]
    println(seval(dot(v1, v2)))
    println(seval(dot(v1, v2), :x))
    println(seval(dot(v1, v2), :y))
    println(seval(dot(v1, v2), (:x,:y)))
end

@testset "SymbolicPlus2" begin
    test = SymbolicEnv()
    e1 = convert(AbstractVectorSymbolic{Float64}, @expr [x^2, y, 10])
    e2 = convert(AbstractVectorSymbolic{Float64}, @expr [x, y^2, 10])
    expr = e1 + e2
    println(expr)
    x0 = 10.0
    y0 = 0.8
    @bind test begin
        x = x0
        y = y0
    end
    println(seval(expr, test))
    println(seval(expr[1], test))
    println(seval(expr[[true, false, false]], test))
end


using WaterLily
using LinearAlgebra: norm2
include("ThreeD_Plots.jl")

function TGV(p=6,Re=1e5)
    # Define vortex size, velocity, viscosity
    L = 2^p; U = 1; ν = U*L/Re

    # Taylor-Green-Vortex initial velocity field
    function uλ(i,vx)
        x,y,z = @. (vx-1.5)*π/L                # scaled coordinates
        i==1 && return -U*sin(x)*cos(y)*cos(z) # u_x
        i==2 && return  U*cos(x)*sin(y)*cos(z) # u_y
        return 0.                              # u_z
    end

    # Initialize simulation
    return Simulation((L+2,L+2,L+2),zeros(3),L;U,uλ,ν)
end

function ω_mag_data(sim)
    # plot the vorticity modulus
    @inside sim.flow.σ[I] = WaterLily.ω_mag(I,sim.flow.u)*sim.L/sim.U
    return @view sim.flow.σ[2:end-1,2:end-1,2:end-1]
end

sim,fig = volume_video!(TGV(),ω_mag_data,name="TGV.mp4",duration=10)


import Pkg; Pkg.activate(".")

using Distributions
using AdvancedMH
import Random
import LinearAlgebra

eye(n::Int) = Matrix(LinearAlgebra.I(n) * 1.0)

function update_mean!(current_mean, x, iter)
  current_mean .+= (x - current_mean) / iter
end

function update_cov!(current_cov, current_mean, x, iter)
  d = x - current_mean
  current_cov .= current_cov * (iter - 1)/iter + (d*d') * (iter - 1)/iter^2
  current_cov .= Matrix(LinearAlgebra.Symmetric(current_cov))
end

# FIXME: How to write this properly?
mutable struct RandomWalkAdaptiveProposal{P} <: AdvancedMH.Proposal{P}
  proposal::P
  iter::Int
end

d(p::RandomWalkAdaptiveProposal) = length(p.proposal.components[1].μ)

function RWAP(d::Int; beta=0.05, iter=1)
  p = MixtureModel([MvNormal(zeros(d), eye(d) * .01/d),
                    MvNormal(zeros(d), eye(d) * .01/d)],
                   [1-beta, beta])
  return RandomWalkAdaptiveProposal(p, iter)
end
function RWAP(d::Distribution; beta=0.05, iter=1)
  K = length(d)
  p = MixtureModel([d, MvNormal(zeros(K), eye(K) * .01/K)],
                   [1-beta, beta])
  return RandomWalkAdaptiveProposal(p, iter)
end

function update!(p::RandomWalkAdaptiveProposal, x)
  p.iter += 1
  _mean = p.proposal.components[1].μ
  _cov = Matrix(p.proposal.components[1].Σ)

  update_mean!(_mean, x, p.iter)
  update_cov!(_cov, _mean, x, p.iter)

  # Is there a better way?
  p.proposal.components[1] = MvNormal(_mean, _cov)

  return
end

function AdvancedMH.propose(rng::Random.AbstractRNG, p::RandomWalkAdaptiveProposal, m::DensityModel)
  return AdvancedMH.propose(rng, StaticProposal(p.proposal), m)
end
function AdvancedMH.propose(
  rng::Random.AbstractRNG,
  proposal::RandomWalkAdaptiveProposal{<:Union{Distribution,AbstractArray}},
  model::DensityModel,
  t
)
  update!(proposal, t)
  if proposal.iter < 2 * d(proposal)
    out = rand(rng, MvNormal(t, proposal.proposal.components[2].Σ))
  else
    _cov = proposal.proposal.components[1].Σ * 2.38^2 / d(proposal)
    out = rand(rng, MvNormal(t, _cov))
  end
  return out
end
function AdvancedMH.q(
  proposal::RandomWalkAdaptiveProposal{<:Union{Distribution,AbstractArray}},
  t,
  t_cond
)
  # return logpdf(proposal, t - t_cond)
  return 0  # because the proposal is symmetric.
end

#= TEST
p = RWAP(3)
update!(p, randn(3))
=#

RWAM(d::Int; beta=0.05, iter=1) = AdvancedMH.MetropolisHastings(RWAP(d, beta=beta, iter=iter))
function RWAM(d::Distribution; beta=0.05, iter=1)
  return AdvancedMH.MetropolisHastings(RWAP(d, beta=beta, iter=iter))
end


############################################################
# Square Loss
#
# L(\hat{y},y) = 1/2N \sum_{i=1}^N (\hat{y}_i - y_i)^2
############################################################
@defstruct SquareLossLayer Layer (
  name :: AbstractString = "square-loss",
  (weight :: AbstractFloat = 1.0, weight >= 0),
  (bottoms :: Vector{Symbol} = Symbol[], length(bottoms) == 2),
)
@characterize_layer(SquareLossLayer,
  has_loss  => true,
  can_do_bp => true,
  is_sink   => true,
  has_stats => true,
)

type SquareLossLayerState{T} <: LayerState
  layer      :: SquareLossLayer
  loss       :: T

  loss_accum :: T
  n_accum    :: Int

  # a helper blob used to compute the loss without destroying
  # the pred results passed up
  pred_copy :: Blob
end

function setup(backend::Backend, layer::SquareLossLayer, inputs::Vector{Blob}, diffs::Vector{Blob})
  data_type = eltype(inputs[1])
  pred_copy = make_blob(backend, data_type, size(inputs[1])...)

  state = SquareLossLayerState(layer, zero(data_type), zero(data_type), 0, pred_copy)
  return state
end
function shutdown(backend::Backend, state::SquareLossLayerState)
  destroy(state.pred_copy)
end

function reset_statistics(state::SquareLossLayerState)
  state.n_accum = 0
  state.loss_accum = zero(typeof(state.loss_accum))
end
function dump_statistics(storage, state::SquareLossLayerState, show::Bool)
  update_statistics(storage, "$(state.layer.name)-square-loss", state.loss_accum)

  if show
    loss = @sprintf("%.4f", state.loss_accum)
    m_info("  Square-loss (avg over $(state.n_accum)) = $loss")
  end
end

function forward(backend::CPUBackend, state::SquareLossLayerState, inputs::Vector{Blob})
  pred  = inputs[1]
  label = inputs[2]

  data_type = eltype(pred)
  n = length(pred)

  copy!(state.pred_copy, pred)
  BLAS.axpy!(n, convert(data_type, -1), label.data, 1, state.pred_copy.data, 1)
  state.loss = state.layer.weight * 0.5/get_num(pred)*BLAS.dot(state.pred_copy.data, state.pred_copy.data)

  # accumulate statistics
  state.loss_accum = (state.loss_accum*state.n_accum + state.loss*get_num(pred)) / (state.n_accum+get_num(pred))
  state.n_accum += get_num(pred)
end

function backward(backend::CPUBackend, state::SquareLossLayerState, inputs::Vector{Blob}, diffs::Vector{Blob})
  diff = diffs[1]
  if isa(diff, CPUBlob)
    pred  = inputs[1]
    label = inputs[2]

    data_type = eltype(pred)
    n = length(pred)
    num = get_num(pred)

    erase!(diff)
    BLAS.axpy!(n, convert(data_type, state.layer.weight/num), pred.data, 1, diff.data, 1)
    BLAS.axpy!(n, convert(data_type, -state.layer.weight/num), label.data, 1, diff.data, 1)
  end

  # the "label" also needs gradient
  if isa(diffs[2], CPUBlob)
    copy!(diffs[2], diff)
    BLAS.scal!(n, -one(data_type), diffs[2].data, 1)
  end
end


function initialize!(integrator, cache::DImplicitEulerCache) end
function initialize!(integrator, cache::DImplicitEulerConstantCache) end

@muladd function perform_step!(integrator, cache::DImplicitEulerConstantCache, repeat_step=false)
  @unpack t,dt,uprev,u,f,p = integrator
  alg = unwrap_alg(integrator, true)

  # initial guess
  guess = zero(u)

  # Wont work for non-mutable types
  zero_func! = (out, x) -> begin
                out = f((x .- uprev) ./ dt, x, p, t)
              end
  u = NLsolve.nlsolve(zero_func!, guess).zero

  if integrator.opts.adaptive && integrator.success_iter > 0
    # local truncation error (LTE) bound by dt^2/2*max|y''(t)|
    # use 2nd divided differences (DD) a la SPICE and Shampine

    # TODO: check numerical stability
    uprev2 = integrator.uprev2
    tprev = integrator.tprev

    dt1 = dt*(t+dt-tprev)
    dt2 = (t-tprev)*(t+dt-tprev)
    c = 7/12 # default correction factor in SPICE (LTE overestimated by DD)
    r = c*dt^2 # by mean value theorem 2nd DD equals y''(s)/2 for some s

    tmp = r*integrator.opts.internalnorm.((u - uprev)/dt1 - (uprev - uprev2)/dt2,t)
    atmp = calculate_residuals(tmp, uprev, u, integrator.opts.abstol, integrator.opts.reltol,integrator.opts.internalnorm,t)
    integrator.EEst = integrator.opts.internalnorm(atmp,t)
  else
    integrator.EEst = 1
  end

  integrator.u = u
end

@muladd function perform_step!(integrator, cache::DImplicitEulerCache, repeat_step=false)
  @unpack t,dt,uprev,u,f,p = integrator
  @unpack atmp,tmp = cache
  alg = unwrap_alg(integrator, true)

  # initial guess
  guess = tmp
  guess .= zero(eltype(u))

  zero_func! = (out, x) -> begin
                f(out, (x .- uprev) ./ dt, x, p, t)
              end
  u .= NLsolve.nlsolve(zero_func!, guess).zero

  if integrator.opts.adaptive && integrator.success_iter > 0
    # local truncation error (LTE) bound by dt^2/2*max|y''(t)|
    # use 2nd divided differences (DD) a la SPICE and Shampine

    # TODO: check numerical stability
    uprev2 = integrator.uprev2
    tprev = integrator.tprev

    dt1 = dt*(t+dt-tprev)
    dt2 = (t-tprev)*(t+dt-tprev)
    c = 7/12 # default correction factor in SPICE (LTE overestimated by DD)
    r = c*dt^2 # by mean value theorem 2nd DD equals y''(s)/2 for some s

    @.. tmp = r*integrator.opts.internalnorm((u - uprev)/dt1 - (uprev - uprev2)/dt2,t)
    calculate_residuals!(atmp, tmp, uprev, u, integrator.opts.abstol, integrator.opts.reltol,integrator.opts.internalnorm,t)
    integrator.EEst = integrator.opts.internalnorm(atmp,t)
  else
    integrator.EEst = 1
  end
end


using ADCME
using PyCall
using LinearAlgebra
using SparseArrays
using PyPlot
using Random
using Test
Random.seed!(233)

sparse_mat_mul = load_op("build/libSparseMatMul", "sparse_sparse_mat_mul")
diag_sparse_mat_mul = load_op("build/libSparseMatMul", "diag_sparse_mat_mul")
sparse_diag_mat_mul = load_op("build/libSparseMatMul", "sparse_diag_mat_mul")

function Base.:*(s1::SparseTensor, s2::SparseTensor)
    ii1, jj1, vv1 = find(s1)
    ii2, jj2, vv2 = find(s2)
    m, n = size(s1)
    n_, k = size(s2)
    if n!=n_
        error("IGACS: matrix size mismatch: ($m, $n) vs ($n_, $k)")
    end
    mat_mul_fn = sparse_mat_mul
    if s1._diag
        mat_mul_fn = diag_sparse_mat_mul
    elseif s2._diag
        mat_mul_fn = sparse_diag_mat_mul
    end
    ii3, jj3, vv3 = mat_mul_fn(ii1-1,jj1-1,vv1,ii2-1,jj2-1,vv2,m,n,k)
    SparseTensor(ii3, jj3, vv3, m, k)
end
# sparse_mat_mul = py"sparse_mat_mul"
################## End Load Operator ##################
A = sprand(10,5,0.3)
B = sprand(5,20,0.3)
C = A*B
CC = SparseTensor(A)*SparseTensor(B)
# TODO: specify your input parameters
sess = Session()
init(sess)
C_ = run(sess, CC)
@test C_≈C

A = spdiagm(0=>[1.;2.;3;4;5])
B = sprand(5,20,0.3)
C = A*B
CC = SparseTensor(A)*SparseTensor(B)
# TODO: specify your input parameters
sess = Session()
init(sess)
C_ = run(sess, CC)
@test C_≈C

A = sprand(10,5,0.5)
B = spdiagm(0=>[1.;2.;3;4;5])
C = A*B
CC = SparseTensor(A)*SparseTensor(B)
# TODO: specify your input parameters
sess = Session()
init(sess)
C_ = run(sess, CC)
@test C_≈C


@inline function Vreman_eddy_viscosity(S::SymTen,ω::Vec,c::Real,Δ2::Real)
    α = S + AntiSymTen(-0.5*ω)
    β = Δ2 * symmetric(α' ⋅ α)
    B = β.xx*(β.yy + β.zz) + β.yy*β.zz - (β.xy^2 + β.xz^2 + β.yz^2)
    B = ifelse(B<0.0,0.0,B)
    return c*fsqrt(B/(α:α))
end


# =========================================================================================
"""
    function φΩ{T}(
        ξ    ::Vector{T},
        model::NonlocalXieModel1{T}
    )

Computes the interior nonlocal electrostatic potential ``φ_Ω`` for the given observation
point ``ξ``.

# Unit
``V = \\frac{C}{F}``

# Return type
`T`

!!! warning
    This function does not verify whether ξ is located inside of the sphere!
"""
function φΩ(ξ::Vector{T}, model::NonlocalXieModel1{T}) where T
    a  = model.radius
    λ  = model.params.λ
    εΩ = model.params.εΩ
    εΣ = model.params.εΣ
    ε∞ = model.params.ε∞
    A₃ = model.A₃
    r  = norm(ξ)

    φ = zero(T)
    for (qi, q) in enumerate(model.charges)   # Eq. (17a)

        # if ξ is close to the origin, all terms for n > 0 become negligible
        if r < 1e-10
            φ += A₃[1, qi]
            continue
        end

        # if q is close to the origin, compute nonlocal Born potential
        # (see born/potentials.jl)
        if norm(q.pos) < 1e-10
            ν = √(εΣ/ε∞) * a/λ
            φ += q.val * (1/r + 1/a/εΣ * (1 - εΣ + (εΣ - ε∞)/ε∞ * sinh(ν)/ν * exp(-ν)))
            continue
        end

        # otherwise, use Eq. (18)
        P = legendre(model.len, cos(ξ, q.pos, r))
        φj = zero(T)
        for n in 0:model.len-1
            φj += A₃[n+1, qi] * r^n * P(n)
        end
        φ += φj * q.val
    end

    (φ + φmol(ξ, model.charges) / 4π / εΩ) * T(ec/ε0)
end

# =========================================================================================
"""
    function φΣ{T}(
        ξ    ::Vector{T},
        model::NonlocalXieModel1{T}
    )

Computes the exterior nonlocal electrostatic potential ``φ_Σ`` for the given observation
point ``ξ``.

# Unit
``V = \\frac{C}{F}``

# Return type
`T`

!!! warning
    This function does not verify whether ξ is located outside of the sphere!
"""
function φΣ(ξ::Vector{T}, model::NonlocalXieModel1{T}) where T
    a  = model.radius
    λ  = model.params.λ
    εΣ = model.params.εΣ
    ε∞ = model.params.ε∞
    A₁ = model.A₁
    A₂ = model.A₂
    κ  = λ \ √(εΣ/ε∞)
    r  = norm(ξ)

    kᵣ = spherical_besselk(model.len, κ * r)

    φ = zero(T)
    for (qi, q) in enumerate(model.charges)   # Eq. (17a)

        # if q is close to origin, compute nonlocal Born potential
        # (see born/potentials.jl)
        if norm(q.pos) < 1e-10
            ν = √(εΣ/ε∞) * a/λ
            φ += q.val / εΣ / r * (1 + (εΣ - ε∞)/ε∞ * sinh(ν)/ν * exp(-ν * r/a))
            continue
        end

        # otherwise, use Eq. (18)
        P = legendre(model.len, cos(ξ, q.pos, r))
        φj = zero(T)
        for n in 0:model.len-1
            φj += (ε∞ - εΣ)/ε∞ * A₂[n+1, qi] * kᵣ(n) * P(n) +
                  A₁[n+1, qi] / r^(n+1) * P(n)
        end
        φ += φj * q.val
    end

    φ * T(ec/ε0)
end


@testset "Linear differential equations" begin

    T, t = Nemo.PowerSeriesRing(Nemo.GF(2^31 - 1), 300, "t"; model=:capped_absolute)

    for d = 1:5
        for c = 1:5
            S = Nemo.MatrixSpace(T, d, d)
            A = random_ps_matrix(T, S)
            B = random_ps_matrix(T, S)
            Sconst = Nemo.MatrixSpace(Nemo.GF(2^31 - 1), d, d)
            Y0 = Sconst([rand(Int) % 100 for i = 1:d, j = 1:d])
            @time sol = ps_matrix_linear_de(A, B, Y0)
            to_be_zero = map(ps_diff, sol) - A * sol - B
            @test truncate_matrix(to_be_zero, 299) == truncate_matrix(zero(S), 299)
        end
    end

end


#=
Miscellaneous physics
=#

# provides c0, G, m_H, k_B, sigma_sb, a_rad, Msun, Rsun, Lsun
include("constants.jl")

function mu_from_composition(X, Y)
    #=
    Calculates the dimensionless mean molecular weight for a fully ionized gas.

    Parameters
    ----------
    X : hydrogen mass fraction
    Y : helium mass fraction

    Returns
    -------
    mu : mean molecular weight
    =#
    return 2 / (1 + 3 * X + 0.5 * Y)
end # mu_from_composition

function beta(Pgas, T)
    #=
    Calculates the ratio between gas pressure and total pressure.

    Parameters
    ----------
    Pgas : gas pressure in dyne cm^-2
    T : temperature in K

    Returns
    -------
    beta : Pgas / P
    =#
    Prad = a_rad * T .^ 4 / 3
    return Pgas / (Pgas + Prad)
end # beta

function rho_ideal(Pgas, T, mu)
    #=
    Calculates the density of an ideal fully ionized gas.

    Parameters
    ----------
    Pgas : gas pressure in dyne cm^-2
    T : temperature in K
    mu : dimensionless mean molecular weight

    Returns
    -------
    rho : density in g cm^-3
    =#
    Prad = a_rad * T .^ 4 / 3
    P = Pgas .+ Prad
    rho = P .* m_H .* mu ./ (k_B * T)
    return rho
end # rho_ideal

function acceleration(m, r)
    #=
    Calculates the local acceleration due to gravity.
    =#
    return G * m ./ (r .^ 2)

end # acceleration

function scale_height(m, r, P, T, mu)
    #=
    Calculates the pressure scale height.
    =#
    g = acceleration(m, r)
    rho = rho_ideal(P, T, mu)
    return P ./ (rho .* g)
end # scale_height

function T_surface(R, L_surface)
    #=
    Calculates the effective temperature, given the surface luminosity and the
    stellar radius.
    =#
    return (L_surface ./ (4 * pi * R .^ 2 * sigma_sb)) .^ (1 / 4)
end # T_surface

function P_surface(M, R, kappa)
    #=
    Calculates the surface pressure, given the total mass, radius, and surface
    opacity.
    =#
    g = acceleration(M, R)
    return 2 * g ./ (3 * kappa)
end # P_surface

function del_ad()
    return 0.4
end # nabla_ad

function del_rad(m, l, P, T, kappa)
    #=
    Gradient (d logP / d logT) in fully radiative case.
    =#
    return 3 * kappa .* l .* P ./ (16 * pi * a_rad * c0 * G * m .* T .^ 4)
end # nabla_rad

function epsilon_pp(rho, T, X)
    #=
    Proton-proton chain luminosity density.
    Kippenhahn chapter 18 page 194, from Agnulo+1999
    =#
    # approximate the pp2, pp3 correction as a logistic curve from 1 to 1.5
    # compare with Fig 18.7 in Kippenhahn for small 4He fractions
    f(x) = 1 + 0.5 ./ (1 + e .^ (-3 .* (x - 2)))
    phi = f(T / 1e7)     # pp2, pp3 correction
    f11 = 1       # weak screening approximation
    T9 = T / 1e9
    g11 = 1 + 3.82 * T9 + 1.51 * T9 .^ 2 + 0.144 * T9 .^ 3 - 0.0114 * T9 .^ 4
    factor = 2.57e4 .* phi .* f11 .* g11 .* rho .* X ^ 2
    return factor .* T9 .^ (-2 / 3) .* e .^ (-3.381 ./ (T9 .^ (1/3)))
end # epsilon_pp

function epsilon_cno(rho, T, X, Y)
    #=
    CNO luminosity density
    Kippanhahn chapter 18 page 196
    =#
    T9 = T / 1e9
    X_CNO = (1 - X - Y) * 0.7
    g141 = 1 - 2 .* T9 + 3.41 .* T9 .^ 2 - 2.43 .* T9 .^ 3
    factor = 8.24e25 .* g141 .* X_CNO .*  X .* rho
    return factor .* T9 .^ (-2 / 3) .* e .^ (-15.231 .* T9 .^ (-1 / 3) - (T9 / 0.8) .^ 2)
end # epsilon_cno


export CategoryWithProducts, ob, terminal, delete, product, proj1, proj2, pair,
  CategoryWithCoproducts, initial, create, coproduct, coproj1, coproj2, copair,
  CompleteCategory, equalizer, incl, factorize,
  CocompleteCategory, coequalizer, proj

# Limits
########

""" Theory of a *category with (finite) products*

Finite products are presented in biased style, via the nullary case (terminal
objects) and the binary case (binary products). The equational axioms are
standard, especially in type theory (Lambek & Scott, 1986, Section 0.5 or
Section I.3).

For a monoidal category axiomatization, see [`CartesianCategory`](@ref).
"""
@theory Category(Ob,Hom) => CategoryWithProducts(Ob,Hom,Terminal,Product) begin
  Terminal()::TYPE
  Product(foot1::Ob, foot2::Ob)::TYPE

  # Terminal object.
  terminal()::Terminal()
  ob(⊤::Terminal())::Ob
  delete(⊤::Terminal(), C::Ob)::(C → ob(⊤))

  # Binary products.
  product(A::Ob, B::Ob)::Product(A,B)
  ob(Π::Product(A,B))::Ob ⊣ (A::Ob, B::Ob)
  proj1(Π::Product(A,B))::(ob(Π) → A) ⊣ (A::Ob, B::Ob)
  proj2(Π::Product(A,B))::(ob(Π) → B) ⊣ (A::Ob, B::Ob)
  (pair(Π::Product(A,B), f::(C → A), g::(C → B))::(C → ob(Π))
    ⊣ (A::Ob, B::Ob, C::Ob))

  # Projection axioms.
  (pair(Π,f,g) ⋅ proj1(Π) == f
    ⊣ (A::Ob, B::Ob, C::Ob, Π::Product(A,B), f::(C → A), g::(C → B)))
  (pair(Π,f,g) ⋅ proj2(Π) == g
    ⊣ (A::Ob, B::Ob, C::Ob, Π::Product(A,B), f::(C → A), g::(C → B)))

  # Uniqueness axioms.
  f == g ⊣ (C::Ob, ⊤::Terminal(), f::(C → ob(⊤)), g::(C → ob(⊤)))
  (pair(h⋅proj1(Π), h⋅proj2(Π)) == h
    ⊣ (A::Ob, B::Ob, C::Ob, Π::Product(A,B), h::(C → ob(Π))))
end

""" Theory of a *(finitely) complete category*

Finite limits are presented in biased style, via finite products and equalizers.
The equational axioms for equalizers are obscure, but can found in (Lambek &
Scott, 1986, Section 0.5), apparently following "Burroni's pioneering ideas".
"""
@theory CategoryWithProducts(Ob,Hom,Terminal,Product) =>
    CompleteCategory(Ob,Hom,Terminal,Product,Equalizer) begin
  Equalizer(f::(A → B), g::(A → B))::TYPE ⊣ (A::Ob, B::Ob)

  # Equalizers.
  equalizer(f::(A → B), g::(A → B))::Equalizer(f,g) ⊣ (A::Ob, B::Ob)
  ob(eq::Equalizer(f,g))::Ob ⊣ (A::Ob, B::Ob, f::(A → B), g::(A → B))
  (incl(eq::Equalizer(f,g))::(ob(eq) → A)
    ⊣ (A::Ob, B::Ob, f::(A → B), g::(A → B)))
  (factorize(eq::Equalizer(f,g), h::(C → A),
             eq_h::Equalizer(h⋅f,h⋅g))::(ob(eq_h) → ob(eq))
    ⊣ (A::Ob, B::Ob, C::Ob, f::(A → B), g::(A → B)))

  # Equalizer axioms.
  (incl(eq)⋅f == incl(eq)⋅g
    ⊣ (A::Ob, B::Ob, f::(A → B), g::(A → B), eq::Equalizer(f,g)))
  (incl(eq) == id(A)
    ⊣ (A::Ob, B::Ob, f::(A → B), eq::Equalizer(f,f)))
  (factorize(eq,h,eq_h) ⋅ incl(eq) == incl(eq_h) ⋅ h
    ⊣ (A::Ob, B::Ob, C::Ob, f::(A → B), g::(A → B), h::(C → A),
       eq::Equalizer(f,g), eq_h::Equalizer(h⋅f, h⋅g)))
  (factorize(eq, k⋅incl(eq), eq_k) == k
    ⊣ (A::Ob, B::Ob, D::Ob, f::(A → B), g::(A → B), eq::Equalizer(f,g),
       k::(D → ob(eq)), eq_k::Equalizer(k⋅incl(eq)⋅f, k⋅incl(eq)⋅g)))
end

# Colimits
##########

""" Theory of a *category with (finite) coproducts*

Finite coproducts are presented in biased style, via the nullary case (initial
objects) and the binary case (binary coproducts). The axioms are dual to those
of [`CategoryWithProducts`](@ref).

For a monoidal category axiomatization, see [`CocartesianCategory`](@ref).
"""
@theory Category(Ob,Hom) => CategoryWithCoproducts(Ob,Hom,Initial,Coproduct) begin
  Initial()::TYPE
  Coproduct(foot1::Ob, foot2::Ob)::TYPE

  # Initial object.
  initial()::Initial()
  ob(⊥::Initial())::Ob
  create(⊥::Initial(), C::Ob)::(ob(⊥) → C)

  # Binary coproducts.
  coproduct(A::Ob, B::Ob)::Coproduct(A,B)
  ob(⨆::Coproduct(A,B))::Ob ⊣ (A::Ob, B::Ob)
  coproj1(⨆::Coproduct(A,B))::(A → ob(⨆)) ⊣ (A::Ob, B::Ob)
  coproj2(⨆::Coproduct(A,B))::(B → ob(⨆)) ⊣ (A::Ob, B::Ob)
  (copair(⨆::Coproduct(A,B), f::(A → C), g::(B → C))::(ob(⨆) → C)
    ⊣ (A::Ob, B::Ob, C::Ob))

  # Coprojection axioms.
  (coproj1(⨆) ⋅ copair(⨆,f,g) == f
    ⊣ (A::Ob, B::Ob, C::Ob, ⨆::Coproduct(A,B), f::(A → C), g::(B → C)))
  (coproj2(⨆) ⋅ copair(⨆,f,g) == g
    ⊣ (A::Ob, B::Ob, C::Ob, ⨆::Coproduct(A,B), f::(A → C), g::(B → C)))

  # Uniqueness axioms.
  f == g ⊣ (C::Ob, ⊥::Initial(), f::(ob(⊥) → C), g::(ob(⊥) → C))
  (copair(coproj1(⨆)⋅h, coproj2(⨆)⋅h) == h
    ⊣ (A::Ob, B::Ob, C::Ob, ⨆::Coproduct(A,B), h::(ob(⨆) → C)))
end

""" Theory of a *(finitely) cocomplete category*

Finite colimits are presented in biased style, via finite coproducts and
coequalizers. The axioms are dual to those of [`CompleteCategory`](@ref).
"""
@theory CategoryWithCoproducts(Ob,Hom,Initial,Coproduct) =>
    CocompleteCategory(Ob,Hom,Initial,Coproduct,Coequalizer) begin
  Coequalizer(f::(A → B), g::(A → B))::TYPE ⊣ (A::Ob, B::Ob)

  # Coequalizers.
  coequalizer(f::(A → B), g::(A → B))::Coequalizer(f,g) ⊣ (A::Ob, B::Ob)
  ob(eq::Coequalizer(f,g))::Ob ⊣ (A::Ob, B::Ob, f::(A → B), g::(A → B))
  (proj(eq::Coequalizer(f,g))::(B → ob(eq))
    ⊣ (A::Ob, B::Ob, f::(A → B), g::(A → B)))
  (factorize(coeq::Coequalizer(f,g), h::(B → C),
             coeq_h::Coequalizer(f⋅h,g⋅h))::(ob(coeq) → ob(coeq_h))
    ⊣ (A::Ob, B::Ob, C::Ob, f::(A → B), g::(A → B)))

  # Coequalizer axioms.
  (f⋅proj(coeq) == g⋅proj(coeq)
    ⊣ (A::Ob, B::Ob, f::(A → B), g::(A → B), coeq::Coequalizer(f,g)))
  (proj(coeq) == id(B)
    ⊣ (A::Ob, B::Ob, f::(A → B), coeq::Coequalizer(f,f)))
  (proj(coeq) ⋅ factorize(coeq,h,coeq_h) == h ⋅ proj(coeq_h)
    ⊣ (A::Ob, B::Ob, C::Ob, f::(A → B), g::(A → B), h::(B → C),
       coeq::Coequalizer(f,g), coeq_h::Coequalizer(f⋅h, g⋅h)))
  (factorize(coeq, proj(coeq)⋅k, coeq_k) == k
    ⊣ (A::Ob, B::Ob, D::Ob, f::(A → B), g::(A → B),
       coeq::Coequalizer(f,g), k::(ob(coeq) → D),
       coeq_k::Coequalizer(f⋅proj(coeq)⋅k, g⋅proj(coeq)⋅k)))
end


module HyperCubicRoots

# # Obtain the Gauss hypergeometric function ₂F₁(a, b; c; z):
import HypergeometricFunctions
HgF = HypergeometricFunctions._₂F₁

export HgF,
       solve_real_cubic_roots, solve_all_cubic_roots,
       solve_quadratic_roots,
       solve_all_quartic_roots

# Tolerance to maintain accuracy on the evaluation of the root (p(root) ≈ 0)
const CUBIC_ATOL = 10.0*sqrt(eps(Float64))
const QUARTIC_ATOL = 10.0*sqrt(eps(Float64))

poor_conditioning(v, compare_tol) = any(abs.(v) .> compare_tol) ? true : false

include("quadratic.jl")
include("cubic.jl")
include("quartic.jl")

end