"""
    flatten 2D matrix into (m*n, 1) matrix

        mainly used for images to flatten images

    inputs:
        x := 3D (rgp, m, n) matrix

    outputs:
        y := 2D (rgp*m*n, 1) matrix
"""
function flatten(x)
    rgp, n, m = size(x)
    return reshape(x, (rgp*m*n, 1))
end

"""
    perform the forward propagation using

    input:
        x := (n0, m) matrix
        y := (c,  m) matrix where c is the number of classes

    return cache of A, Z, Yhat, Cost
"""
function forwardProp(X::Matrix{T},
                     Y::Matrix{T},
                     model::Model) where {T}

    W::AbstractArray{Matrix{T},1},
    B::AbstractArray{Matrix{T},1},
    layers::AbstractArray{Layer,1},
    lossFun = model.W, model.B, model.layers, model.lossFun
    regulization = model.regulization
    λ = model.λ
    m = size(X)[2]
    c = size(Y)[1]
    L = length(layers)
    A = Vector{Matrix{eltype(X)}}()
    Z = Vector{Matrix{eltype(X)}}()
    push!(Z, W[1]*X .+ B[1])
    actFun = layers[1].actFun
    push!(A, eval(:($actFun.($Z[1]))))
    for l=2:L
        push!(Z, W[l]*A[l-1] .+ B[l])
        actFun = layers[l].actFun

        if isequal(actFun, :softmax)
            a = Matrix{eltype(X)}(undef, c, 0)
            for i=1:m
                zCol = Z[l][:,i]
                a = hcat(a, eval(:($(actFun)($zCol))))
            end
            push!(A, a)
        else
            push!(A, eval(:($(actFun).($Z[$l]))))
        end
    end

    if isequal(lossFun, :categoricalCrossentropy)
        cost = sum(eval(:($lossFun.($A[$L], $Y))))/ m
    else
        cost = sum(eval(:($lossFun.($A[$L], $Y)))) / (m*c)
    end

    if regulization > 0
        cost += (λ/2m) * sum([norm(w, regulization) for w in W])
    end

    return Dict("A"=>A,
                "Z"=>Z,
                "Yhat"=>A[L],
                "Cost"=>cost)
end #forwardProp

# export forwardProp

"""
    predict Y using the model and the input X and the labels Y
    inputs:
        model := the trained model
        X := the input matrix
        Y := the input labels to compare with

    output:
        a Dict of
        "Yhat" := the predicted values
        "Yhat_bools" := the predicted labels
        "accuracy" := the accuracy of the predicted labels
"""
function predict(model::Model, X, Y)
    Ŷ = forwardProp(X, Y, model)["Yhat"]
    T = eltype(Ŷ)
    layers = model.layers
    c, m = size(Y)
    # if isbool(Y)
    acc = 0
    if isequal(layers[end].actFun, :softmax)
        Ŷ_bool = BitArray(undef, c, 0)
        p = Progress(size(Ŷ)[2], 0.01)
        for v in eachcol(Ŷ)
            Ŷ_bool = hcat(Ŷ_bool, v .== maximum(v))
            next!(p)
        end

        acc = sum([Ŷ_bool[:,i] == Y[:,i] for i=1:size(Y)[2]])/m
        println("Accuracy is = $acc")
    end

    if isequal(layers[end].actFun, :σ)
        Ŷ_bool = BitArray(undef, c, 0)
        p = Progress(size(Ŷ)[2], 0.01)
        for v in eachcol(Ŷ)
            Ŷ_bool = hcat(Ŷ_bool, v .> T(0.5))
            next!(p)
        end

        acc = sum(Ŷ_bool .== Y)/(c*m)
        println("Accuracy is = $acc")
    end
    return Dict("Yhat"=>Ŷ,
                "Yhat_bool"=>Ŷ_bool,
                "accuracy"=>acc)
end #predict

function backProp(X,Y,
                  model::Model,
                  cache::Dict{})

    layers::AbstractArray{Layer, 1} = model.layers
    lossFun = model.lossFun
    m = size(X)[2]
    L = length(layers)
    A, Z = cache["A"], cache["Z"]
    W, B, regulization, λ = model.W, model.B, model.regulization, model.λ

    D = [rand(size(A[i])...) .< layers[i].keepProb for i=1:L]

    if layers[L].keepProb < 1.0 #to save time of multiplication in case keepProb was one
        A = [A[i] .* D[i] for i=1:L]
        A = [A[i] ./ layers[i].keepProb for i=1:L]
    end
    # init all Arrays
    dA = Vector{Matrix{eltype(A[1])}}([similar(mat) for mat in A])
    dZ = Vector{Matrix{eltype(Z[1])}}([similar(mat) for mat in Z])
    dW = Vector{Matrix{eltype(W[1])}}([similar(mat) for mat in W])
    dB = Vector{Matrix{eltype(B[1])}}([similar(mat) for mat in B])

    dlossFun = Symbol("d",lossFun)
    actFun = layers[L].actFun
    if ! isequal(actFun, :softmax) && !(actFun==:σ &&
                                        lossFun==:binaryCrossentropy)
        dA[L] = eval(:($dlossFun.($A[$L], $Y))) #.* eval(:($dActFun.(Z[L])))
    end
    if layers[L].keepProb < 1.0 #to save time of multiplication in case keepProb was one
        dA[L] = dA[L] .* D[L]
        dA[L] = dA[L] ./ layers[L].keepProb
    end
    for l=L:-1:2
        actFun = layers[l].actFun
        dActFun = Symbol("d",actFun)
        if l==L && (isequal(actFun, :softmax) ||
                    (actFun==:σ && lossFun==:binaryCrossentropy))
            dZ[l] = A[l] .- Y
        else
            dZ[l] = dA[l] .* eval(:($dActFun.($Z[$l])))
        end

        dW[l] = dZ[l]*A[l-1]' ./m

        if regulization > 0
            if regulization==1
                dW[l] .+= (λ/2m)
            else
                dW[l] .+= (λ/m) .* W[l]
            end
        end
        dB[l] = 1/m .* sum(dZ[l], dims=2)
        dA[l-1] = W[l]'dZ[l]
        if layers[l-1].keepProb < 1.0 #to save time of multiplication in case keepProb was one
            dA[l-1] = dA[l-1] .* D[l-1]
            dA[l-1] = dA[l-1] ./ layers[l-1].keepProb
        end
    end

    l=1 #shortcut cause I just copied from the for loop
    actFun = layers[l].actFun
    dActFun = Symbol("d",actFun)
    dZ[l] = dA[l] .* eval(:($dActFun.($Z[$l])))

    dW[l] = 1/m .* dZ[l]*X' #where X = A[0]

    if regulization > 0
        if regulization==1
            dW[l] .+= (λ/2m)
        else
            dW[l] .+= (λ/m) .* W[l]
        end
    end
    dB[l] = 1/m .* sum(dZ[l], dims=2)

    grads = Dict("dW"=>dW,
                 "dB"=>dB)
    return grads
end #backProp

function updateParams(model::Model, grads::Dict, tMiniBatch::Integer)
    W, B, V, S, layers = model.W, model.B, model.V, model.S, model.layers
    optimizer = model.optimizer
    dW, dB = grads["dW"], grads["dB"]
    L = length(layers)
    α = model.α
    β1, β2, ϵAdam = model.β1, model.β2, model.ϵAdam

    #initialize the needed variables to hold the corrected values
    #it is being init here cause these are not needed elsewhere
    VCorrected, SCorrected = deepInitVS(W,B,optimizer)
    if optimizer==:adam || optimizer==:momentum
        for i=1:L
            V[:dw][i] .= β1 .* V[:dw][i] .+ (1-β1) .* dW[i]
            V[:db][i] .= β1 .* V[:db][i] .+ (1-β1) .* dB[i]

            ##correcting
            VCorrected[:dw][i] .= V[:dw][i] ./ (1-β1^tMiniBatch)
            VCorrected[:db][i] .= V[:db][i] ./ (1-β1^tMiniBatch)

            if optimizer==:adam
                S[:dw][i] .= β2 .* S[:dw][i] .+ (1-β2) .* (dW[i].^2)
                S[:db][i] .= β2 .* S[:db][i] .+ (1-β2) .* (dB[i].^2)

                ##correcting
                SCorrected[:dw][i] .= S[:dw][i] ./ (1-β2^tMiniBatch)
                SCorrected[:db][i] .= S[:db][i] ./ (1-β2^tMiniBatch)

                ##update parameters with adam
                W[i] .-= (α .* (VCorrected[:dw][i] ./ (sqrt.(SCorrected[:dw][i]) .+ ϵAdam)))
                B[i] .-= (α .* (VCorrected[:db][i] ./ (sqrt.(SCorrected[:db][i]) .+ ϵAdam)))
            else#if optimizer==:momentum
                W[i] .-= (α .* VCorrected[:dw][i])
                B[i] .-= (α .* VCorrected[:db][i])
            end #if optimizer==:adam
        end #for i=1:L
    else
        W .-= (α .* dW)
        B .-= (α .* dB)
    end #if optimizer==:adam || optimizer==:momentum
    return W, B
end #updateParams

function updateParams!(model::Model, grads::Dict, tMiniBatch::Integer)
    W, B, V, S, layers = model.W, model.B, model.V, model.S, model.layers
    optimizer = model.optimizer
    dW, dB = grads["dW"], grads["dB"]
    L = length(layers)
    α = model.α
    β1, β2, ϵAdam = model.β1, model.β2, model.ϵAdam

    #initialize the needed variables to hold the corrected values
    #it is being init here cause these are not needed elsewhere
    VCorrected, SCorrected = deepInitVS(W,B,optimizer)
    if optimizer==:adam || optimizer==:momentum
        for i=1:L
            V[:dw][i] .= β1 .* V[:dw][i] .+ (1-β1) .* dW[i]
            V[:db][i] .= β1 .* V[:db][i] .+ (1-β1) .* dB[i]

            ##correcting
            VCorrected[:dw][i] .= V[:dw][i] ./ (1-β1^tMiniBatch)
            VCorrected[:db][i] .= V[:db][i] ./ (1-β1^tMiniBatch)

            if optimizer==:adam
                S[:dw][i] .= β2 .* S[:dw][i] .+ (1-β2) .* (dW[i].^2)
                S[:db][i] .= β2 .* S[:db][i] .+ (1-β2) .* (dB[i].^2)

                ##correcting
                SCorrected[:dw][i] .= S[:dw][i] ./ (1-β2^tMiniBatch)
                SCorrected[:db][i] .= S[:db][i] ./ (1-β2^tMiniBatch)

                ##update parameters with adam
                W[i] .-= (α .* (VCorrected[:dw][i] ./ (sqrt.(SCorrected[:dw][i]) .+ ϵAdam)))
                B[i] .-= (α .* (VCorrected[:db][i] ./ (sqrt.(SCorrected[:db][i]) .+ ϵAdam)))
            else#if optimizer==:momentum
                W[i] .-= (α .* VCorrected[:dw][i])
                B[i] .-= (α .* VCorrected[:db][i])
            end #if optimizer==:adam
        end #for i=1:L
    else
        W .-= (α .* dW)
        B .-= (α .* dB)
    end #if optimizer==:adam || optimizer==:momentum
    model.W, model.B = W, B
    return
end #updateParams!

function updateParams!(model::Model,
                       cLayer::Layer,
                       cnt::Integer = -1;
                       tMiniBatch::Integer = 1)

    optimizer = model.optimizer
    α = model.α
    β1, β2, ϵAdam = model.β1, model.β2, model.ϵAdam

    if cLayer.updateCount >= cnt
        return nothing
    end #if cLayer.updateCount >= cnt

    cLayer.updateCount += 1
    #initialize the needed variables to hold the corrected values
    #it is being init here cause these are not needed elsewhere
    VCorrected = Dict(:dw=>similar(cLayer.dW), :db=>similar(cLayer.dB))
    SCorrected = Dict(:dw=>similar(cLayer.dW), :db=>similar(cLayer.dB))
    if optimizer==:adam || optimizer==:momentum

        cLayer.V[:dw] .= β1 .* cLayer.V[:dw] .+ (1-β1) .* cLayer.dW
        cLayer.V[:db] .= β1 .* cLayer.V[:db] .+ (1-β1) .* cLayer.dB

        ##correcting
        VCorrected[:dw] .= cLayer.V[:dw] ./ (1-β1^tMiniBatch)
        VCorrected[:db] .= cLayer.V[:db] ./ (1-β1^tMiniBatch)

        if optimizer==:adam
            cLayer.S[:dw] .= β2 .* cLayer.S[:dw] .+ (1-β2) .* (cLayer.dW.^2)
            cLayer.S[:db] .= β2 .* cLayer.S[:db] .+ (1-β2) .* (cLayer.dB.^2)

            ##correcting
            SCorrected[:dw] .= cLayer.S[:dw] ./ (1-β2^tMiniBatch)
            SCorrected[:db] .= cLayer.S[:db] ./ (1-β2^tMiniBatch)

            ##update parameters with adam
            cLayer.W .-= (α .* (VCorrected[:dw] ./ (sqrt.(SCorrected[:dw]) .+ ϵAdam)))
            cLayer.B .-= (α .* (VCorrected[:db] ./ (sqrt.(SCorrected[:db]) .+ ϵAdam)))

        else#if optimizer==:momentum

            cLayer.W .-= (α .* VCorrected[:dw])
            cLayer.B .-= (α .* VCorrected[:db])

        end #if optimizer==:adam
    else
        cLayer.W .-= (α .* cLayer.dW)
        cLayer.B .-= (α .* cLayer.dB)
    end #if optimizer==:adam || optimizer==:momentum

    return nothing
end #updateParams!


abstract type KayRuleModel <: MixingRule end

struct KayRuleParam <: EoSParam
end

@newmodelsimple KayRule KayRuleModel KayRuleParam

export KayRule

function KayRule(components::Vector{String}; activity=nothing, userlocations::Vector{String}=String[], activity_userlocations::Vector{String}=String[], verbose::Bool=false)
    #params = getparams(components, ["properties/critical.csv"]; userlocations=userlocations, verbose=verbose)
    #acentricfactor = SingleParam(params["w"],"acentric factor")
    packagedparams = KayRuleParam()
    model = KayRule(packagedparams, verbose=verbose)
    return model
end

function mixing_rule(model::ABCubicModel,V,T,z,mixing_model::KayRuleModel,α,a,b,c)
    n = sum(z)
    invn = (one(n)/n)
    invn2 = invn^2
    #b̄ = (dot(z,Symmetric(b.^(1/3)),z) * invn2).^3
    #ā = √(dot(z,Symmetric(a .* sqrt.(α*α') ./ b).^2,z)* invn2)  * b̄
    b̄ = zero(eltype(z))
    ab2 = zero(T+first(z))
    for i in 1:length(z)
        zi = z[i]
        zi2 = zi*zi
        αi = α[i]
        b̄ += cbrt(b[i,i])*zi2
        abij2 = (a[i,i]*αi/b[i,i])^2
        ab2 += abij2*zi2
        for j in 1:(i-1)
            zij =zi*z[j]
            abij2 = (a[i,j]*sqrt(αi*α[j])/b[i,j])^2
            ab2 += 2*abij2*zij
            b̄ += 2*cbrt(b[i,j])*zij
        end
    end
    b̄ = (b̄*invn2)^3
    ā = sqrt(ab2*invn2)*b̄
    c̄ = dot(z,c)*invn
    return ā,b̄,c̄
end

is_splittable(::KayRule) = false


"""
PropDist(args...)--> UnivariateDistribution

Fit the given data to a specific statistical distribution
Distributions.jl package is used to fit

The next table show the list of variables allowed:

|PropertyName|Args|Default|Input|Description
|---|---|---|---|---|
|Prop|Mandatory|--|Array{Number,1}| Data to fit|
|Distr|Mandatory|--|Distr=Distribution|Distribution to fit the data given|
"""
@userplot prophist

@recipe function f(h::prophist)

    x, d=h.args

    MinProp=minimum(x)
    MaxProp=maximum(x)
    PropRange=range(MinProp,stop=MaxProp,length=100)
    DisProp=Distributions.fit(d, x[:]*1)
    DisParam=params(DisProp)
    FitProp=pdf.(DisProp, PropRange);
    Quant=quantile.(DisProp,[0.1,0.5,0.9])

@series begin
        seriestype := :histogram
        normalize := true
        bins := bins=round(Int,sqrt(size(x,1)))
        seriescolor --> :red
        label --> "Histogram"
        x
    end

@series begin
        seriestype := :path
        linestyle := :dash
        linewidth := 3
        linecolor --> :black
        label --> "Probabilitic Distribution Fitted"
        PropRange, FitProp

 end
println("Percentiles 10% - %50 - 90% \n         $Quant ")
end


#== # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #
#
# Description
#
#   Convert a satellite state vector from an Earth-Centered Inertial (ECI)
#   reference frame to an Earth-Centered, Earth-Fixed (ECEF) reference frame.
#
# # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #
#
# References
#
#   [1] Vallado, D. A (2013). Fundamentals of Astrodynamics and Applications.
#       Microcosm Press, Hawthorn, CA, USA.
#
# # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # ==#

export svECItoECEF

"""
    function svECItoECEF(sv::SatelliteStateVector, ECI, ECEF, JD_UTC [, eop_data])

Convert the satellite state vector `sv` from the Earth-Centered Inertial (ECI)
reference frame `ECI` to the Earth-Centered, Earth-Fixed (ECEF) reference frame
at the Julian day `JD_UTC` [UTC]. The `eop_data` may be required depending on
the selection of the input and output reference system. For more information,
see the documentation of the function `rECItoECEF`.

!!! info

    It is assumed that the input velocity and acceleration in `sv` are obtained
    by an observer on the ECI frame. Thus, the output will contain the velocity
    and acceleration as measured by an observer on the ECEF frame.

"""
function svECItoECEF(sv::SatelliteStateVector, T_ECI::T_ECIs,
                     T_ECEF::Type{Val{:ITRF}}, JD_UTC::Number,
                     eop_data::EOPData_IAU1980)

    # First, convert from ECI to PEF.
    sv_pef = svECItoECEF(sv, T_ECI, Val{:PEF}, JD_UTC, eop_data)

    # Now, convert from PEF to ITRF.
    return svECEFtoECEF(sv_pef, Val{:PEF}, Val{:ITRF}, JD_UTC, eop_data)
end

function svECItoECEF(sv::SatelliteStateVector, T_ECI::T_ECIs_IAU_2006,
                     T_ECEF::Type{Val{:ITRF}}, JD_UTC::Number,
                     eop_data::EOPData_IAU2000A)

    # First, convert from ECI to TIRS.
    sv_tirs = svECItoECEF(sv, T_ECI, Val{:TIRS}, JD_UTC, eop_data)

    # Now, convert from TIRS to ITRF.
    return svECEFtoECEF(sv_tirs, Val{:TIRS}, Val{:ITRF}, JD_UTC, eop_data)
end

function svECItoECEF(sv::SatelliteStateVector,
                     T_ECI::Union{T_ECIs, T_ECIs_IAU_2006},
                     T_ECEF::Union{Type{Val{:PEF}},Type{Val{:TIRS}}},
                     JD_UTC::Number,
                     eop_data::Union{Nothing,EOPData_IAU1980, EOPData_IAU2000A} = nothing)

    # Get the matrix that converts the ECI to the ECEF.
    if eop_data == nothing
        D = rECItoECEF(DCM, T_ECI, T_ECEF, JD_UTC)
    else
        D = rECItoECEF(DCM, T_ECI, T_ECEF, JD_UTC, eop_data)
    end

    # Since the ECI and ECEF frames have a relative velocity between them, then
    # we must account from it when converting the velocity and acceleration. The
    # angular velocity between those frames is computed using `we` and corrected
    # by the length of day (LOD) parameter of the EOP data, if available.
    ω  = we * ( 1 - (eop_data != nothing ? eop_data.LOD(JD_UTC)/86400 : 0 ) )
    vω = SVector{3}(0,0,ω)

    # Compute the position in the ECEF frame.
    r_ecef = D*sv.r

    # Compute the velocity in the ECEF frame.
    vω_x_r = vω × r_ecef
    v_ecef = D*sv.v - vω_x_r

    # Compute the acceleration in the ECI frame.
    a_ecef = D*sv.a - vω × vω_x_r - 2vω × v_ecef

    return satsv(sv.t, r_ecef, v_ecef, a_ecef)
end


#=
Created on Tuesday 04 February 2020
Last update: Wedneday 05 February 2020

@author: Michiel Stock
michielfmstock@gmail.com

Dijkstra, A* and Bellman FOrd
=#

module ShortestPath

export reconstruct_path, dijkstra, a_star

using DataStructures, STMO

"""
    reconstruct_path(previous::Dict{T,T}, source::T, sink::T) where {T}

Reconstruct the path from the output of the Dijkstra algorithm.

Inputs:
        - previous : a Dict with the previous node in the path
        - source : the source node
        - sink : the sink node
Ouput:
        - the shortest path from source to sink
"""
function reconstruct_path(previous::Dict{T,T}, source::T, sink::T) where {T}
    path = T[sink]
    current = sink
    while current != source
        current = previous[current]
        push!(path, current)
    end
    return reverse!(path)
end

"""
    dijkstra(graph::AdjList{R,T}, source::T) where {R<:Real,T}

Dijkstra without a specified sink. Give the `graph` and a `source` and this
function uses Dijkstra's algorithm to compute all distances between the nodes
and the source. Returns
    - `distances`: a dictionary with the distances
    - `previous`: a dictionary representing the tree of the shortest paths.
"""
function dijkstra(graph::AdjList{R,T}, source::T) where {R<:Real,T}
    # initialize the tentative distances
    distances = Dict(v => Inf for v in keys(graph))
    distances[source] = 0.0
    previous = Dict{T,T}()
    vertices_to_check = [(0.0, source)]
    while length(vertices_to_check) > 0
        dist, u = heappop!(vertices_to_check)
        for (dist_u_v, v) in graph[u]
            new_dist = dist + dist_u_v
            if distances[v] > new_dist
                distances[v] = new_dist
                previous[v] = u
                heappush!(vertices_to_check, (new_dist, v))
            end
        end
    end
    return distances, previous
end

"""
    dijkstra(graph::AdjList{R,T}, source::T, sink::T) where {R<:Real,T}

Dijkstra's shortest path algorithm.

Inputs:
    - `graph` : adjacency list representing a weighted directed graph
    - `source`
    - `sink`

Outputs:
    - the shortest path
    - the cost of this shortest path
"""
function dijkstra(graph::AdjList{R,T}, source::T, sink::T) where {R<:Real,T}
    # initialize the tentative distances
    distances = Dict(v => Inf for v in keys(graph))
    distances[source] = 0.0
    previous = Dict{T,T}()
    vertices_to_check = [(0.0, source)]
    while length(vertices_to_check) > 0
        dist, u = heappop!(vertices_to_check)
        if u == sink
            return reconstruct_path(previous, source, sink), dist
        end
        for (dist_u_v, v) in graph[u]
            new_dist = dist + dist_u_v
            if distances[v] > new_dist
                distances[v] = new_dist
                previous[v] = u
                heappush!(vertices_to_check, (new_dist, v))
            end
        end
    end
end

"""
    a_star(graph::AdjList{R,T}, source::T, sink::T, heuristic) where {R<:Real,T}

A* shortest path algorithm.

Inputs:
    - `graph` : adjacency list representing a weighted directed graph
    - `source`
    - `sink`
    - `heuristic` : a function that inputs a node and returns an lower bound
            for the distance to the source. Note that a distance can be turned into
            a heuristic using `n -> d(n, sink)`

Outputs:
    - the shortest path
    - the cost of this shortest path
"""
function a_star(graph::AdjList{R,T}, source::T, sink::T, heuristic) where {R<:Real,T}
    # initialize the tentative distances
    distances = Dict(v => Inf for v in keys(graph))
    distances[source] = 0.0
    previous = Dict{T,T}()
    vertices_to_check = [(heuristic(source), source)]
    while length(vertices_to_check) > 0
        dist, u = heappop!(vertices_to_check)
        if u == sink
            return reconstruct_path(previous, source, sink), distances[sink]
        end
        for (dist_u_v, v) in graph[u]
            new_dist = dist + dist_u_v
            if distances[v] > new_dist
                distances[v] = new_dist
                min_dist_to_sink = new_dist + heuristic(v)
                previous[v] = u
                heappush!(vertices_to_check, (min_dist_to_sink, v))
            end
        end
    end
end

end  # module ShortestPaths


"""
    AeroState(alpha, beta, alpha_dot, tas, eas, cas, ias, qinf, mach)

Aerodynamic state struct storing relevant aerodynamic data.

# Constructors

Different constructors are provided:

    AeroState(alpha, beta, alpha_dot, tas, eas, cas, ias, qinf, mach)
    AeroState(state::State, env::Environment)
    AeroState(tas::Number, alpha::Number, beta::Number, height::Number)

"""
struct AeroState
    alpha :: Number  # AOA [rad]
    beta :: Number  # AOS [rad]

    alpha_dot::Number  # rad/s

    tas :: Number  # True Airspeed [m/s]
    eas :: Number  # Equivalent Airspeed [m/s]
    cas :: Number  # Calibrated Airspeed [m/s]
    ias :: Number  # Indicated Airspeed [m/s]

    qinf :: Number  # Dynamic Pressure [Pa]
    mach :: Number  # Mach number
end

function AeroState(state::State, env::Environment)

    p = get_pressure(env)
    ρ = get_density(env)
    a = get_sound_velocity(env)

    # coming from
    wind_hor = get_wind_NED(env)
    wind_body = hor2body(wind_hor..., get_euler_angles(state)...)

    aero_speed = get_body_velocity(state) + wind_body

    tas, alpha, beta = tas_alpha_beta_from_uvw(aero_speed...)

    cas = tas2cas(tas, ρ, p)
    ias = tas
    eas = tas2eas(tas, ρ)

    qinf = incompressible_qinf(tas, ρ)
    mach = tas / a

    AeroState(alpha, beta, 0, tas, eas, cas, ias, qinf, mach)
end

function AeroState(tas::Number, alpha::Number, beta::Number, height::Number)

    T, p, ρ, a = atmosphere_isa(height)

    cas = tas2cas(tas, ρ, p)
    ias = tas
    eas = tas2eas(tas, ρ)

    qinf = incompressible_qinf(tas, ρ)
    mach = tas / a

    AeroState(alpha, beta, 0, tas, eas, cas, ias, qinf, mach)
end

"""
    generate_state_aerostate(pos::Position, att::Attitude, tas::Number, alpha::Number,
                                beta::Number, env::Environment,
                                ang_vel=[0.0, 0.0, 0.0], accel=[0.0, 0.0, 0.0],
                                ang_accel=[0.0, 0.0, 0.0])

returns both state and aerostate provided the values of aircrft states, environment and
aerodynamic angles.
"""
function generate_state_aerostate(pos::Position, att::Attitude, tas::Number, alpha::Number,
                         beta::Number, env::Environment,
                         ang_vel=[0.0, 0.0, 0.0], accel=[0.0, 0.0, 0.0],
                         ang_accel=[0.0, 0.0, 0.0])

    aerostate = AeroState(tas, alpha, beta, get_height(pos))

    aero_vel_body = wind2body([tas, 0, 0]..., alpha, beta)
    wind_vel_body = hor2body(get_wind_NED(env)..., get_euler_angles(att)...)
    vel = aero_vel_body - wind_vel_body

    state = State(pos, att, vel, ang_vel, accel, ang_accel)
    return state, aerostate
end

"""
    check_state_aerostate_env_coherence(state::State, aerost::AeroState,
                                        env::Environment)

Check if aerodynamic velocity in aerostate is coherent with inertial velocity in
state and wind velocity in environment.
"""
function check_state_aerostate_env_coherence(state::State, aerost::AeroState,
                                             env::Environment)

    att = get_attitude(state)
    # wind vel from environment
    wind = get_wind_NED(env)
    # ground speed from state
    vel = get_horizon_velocity(state)
    # Composition (wind posititve blowing towards)
    aero_vel_estimation = vel + wind

    # Transform aerostate velocity to horizon
    tas, α, β = get_tas(aerost), get_alpha(aerost), get_beta(aerost)
    aero_vel = body2hor(wind2body(tas, 0, 0, α, β)..., get_euler_angles(state)...)

    isapprox(aero_vel_estimation, aero_vel)
end

get_alpha(aerost::AeroState) = aerost.alpha
get_beta(aerost::AeroState) = aerost.beta
get_aero_angles(aerost::AeroState) = [aerost.alpha, aerost.beta]
get_alpha_dot(aerost::AeroState) = aerost.alpha_dot

get_aero_speeds(aerost::AeroState) = [aerost.tas, aerost.eas, aerost.cas, aerost.ias]
get_tas(aerost::AeroState) = aerost.tas
get_eas(aerost::AeroState) = aerost.eas
get_cas(aerost::AeroState) = aerost.cas
get_ias(aerost::AeroState) = aerost.ias

get_qinf(aerost::AeroState) = aerost.qinf
get_mach(aerost::AeroState) = aerost.mach


module MoreMatrices

const ɛ = ([0 0 0; 0 0 1; 0 -1 0], [0 0 -1; 0 0 0; 1 0 0], [0 1 0; -1 0 0; 0 0 0])

export
    ɛ,

    # operations.jl

    kron_shuffle!,
    kron_shuffle,
    kron_unshuffle!,
    kron_unshuffle,
    kron_factor,
    kron_sum,
    join_blocks,
    split_blocks,
    rotation!,
    rotation,
    euler_zyz!,
    euler_zyz,
    nearest_orthogonal!,
    nearest_orthogonal,
    nearest_special_orthogonal!,
    nearest_special_orthogonal,
    cayley,
    rotation2skew,
    vec2skew!,
    vec2skew,
    skew2vec!,
    skew2vec,

    # random_matrices.jl

    random_orthogonal!,
    random_orthogonal,
    random_rotation,
    random_permutation

include("operations.jl")
include("random_matrices.jl")

end


using Test
using Mods

@testset "Constructors" begin
    @test one(Mod{17}) == Mod{17}(1)
    @test oneunit(Mod{17}) == Mod{17}(1)
    @test zero(Mod{17}) == 0
    @test iszero(zero(Mod{17}))
    @test Mod{17}(1) == Mod{17}(1, 0)
    @test Mod{17}(1, 2) == 1 + 2im
    a = Mod{17}(3)
    @test typeof(a) == Mod{17,Int}
    a = Mod{17}(3, 2)
    @test typeof(a) == GaussMod{17,Int}
    a = zero(Mod{17})
    @test typeof(a) == Mod{17,Int}
    a = Mod{17}(1 // 2 + (3 // 4)im)
    @test typeof(a) == GaussMod{17,Int}
end

@testset "Mod arithmetic" begin
    p = 23
    a = Mod{p}(2)
    b = Mod{p}(25)
    @test a == b
    @test a == 2
    @test a == -21

    b = Mod{p}(20)
    @test a + b == 22
    @test a - b == -18
    @test a + a == 2a
    @test 0 - a == -a

    @test a * b == Mod{p}(17)
    @test (a / b) * b == a
    @test (b // a) * (2 // 1) == b
    @test a * (2 // 3) == (2a) * inv(Mod{p}(3))

    @test is_invertible(a)
    @test !is_invertible(Mod{10}(4))

    @test a^(p - 1) == 1
    @test a^(-1) == inv(a)

    @test Mod{p}(3 // 7) == 3 // 7
    @test isequal(3 // 7, Mod{p}(3 // 7))

    @test -Mod{13,Int}(typemin(Int)) == -Mod{13}(mod(typemin(Int), 13))
end

@testset "Mod arithmetic with UInt" begin
    p = UInt(23)
    a = Mod{p,UInt}(2)
    b = Mod{p,UInt}(25)
    @test a == b
    @test a == 2
    @test a == 25

    b = Mod{p,UInt}(20)
    @test a + b == 22
    @test a - b == 28
    @test a + a == 2a
    @test 0 - a == -a

    @test a * b == Mod{p,UInt}(17)
    @test (a / b) * b == a
    @test (b // a) * (2 // 1) == b
    @test a * (2 // 3) == (2a) * inv(Mod{p,UInt}(3))

    @test is_invertible(a)
    @test !is_invertible(Mod{10}(4))

    @test a^(p - 1) == 1
    @test a^(-1) == inv(a)
end

@testset "GaussMod arithmetic" begin
    @test one(GaussMod{6}) == Mod{6}(1 + 0im)
    @test zero(GaussMod{6}) == Mod{6}(0 + 0im)
    @test GaussMod{6}(2 + 3im) == Mod{6}(2 + 3im)
    @test GaussMod{6,Int}(2 + 3im) == Mod{6}(2 + 3im)
    @test rand(GaussMod{6}) isa GaussMod
    @test eltype(rand(GaussMod{6}, 4, 4)) <: GaussMod
    p = 23
    a = GaussMod{p}(3 - im)
    b = Mod{p}(5 + 5im)

    @test a + b == 8 + 4im
    @test a + Mod{p}(11) == Mod{p}(14, 22)
    @test -a == 20 + im
    @test a - b == Mod{p}(3 - im - 5 - 5im)

    @test a * b == Mod{p}((3 - im) * (5 + 5im))
    @test a / b == Mod{p}((3 - im) // (5 + 5im))

    @test a^(p * p - 1) == 1
    @test is_invertible(a)
    @test a * inv(a) == 1

    @test a / (1 + im) == a / Mod{p}(1 + im)
    @test imag(a * a') == 0
end

@testset "Large Modulus" begin

    p = 9223372036854775783   # This is a large prime
    x = Mod{p}(-2)
    @test x * x == 4
    @test x + x == -4
    @test x / x == 1
    @test x / 3 == x / Mod{p}(3)
    @test (x / 3) * (3 // x) == 1
    @test x // x == value(x) / x
    @test x^4 == 16
    @test x^(p - 1) == 1   # Fermat Little Theorem test
    @test 2x == x + x
    @test x - x == 0
    y = inv(x)
    @test x * y == 1
    @test x + p == x
    @test x * p == 0
    @test p - x == x - 2x

    p = 9223372036854775783   # This is a large prime
    x = Mod{p}(-2 + 0im)
    @test x * x == 4
    @test x + x == -4
    @test x / x == 1
    @test x / 3 == x / Mod{p}(3)
    @test (x / 3) * (3 // x) == 1
    @test x // x == value(x) / x
    @test x^4 == 16
    @test x^(p - 1) == 1   # Fermat Little Theorem test
    @test 2x == x + x
    @test x - x == 0
    y = inv(x)
    @test x * y == 1
    @test x + p == x
    @test x * p == 0
    @test p - x == x - 2x

    @test 0 <= value(rand(Mod{p})) < p

end

@testset "CRT" begin
    p = 23
    q = 91
    a = Mod{p}(17)
    b = Mod{q}(32)

    @test CRT(Int32, [], []) === Int32(0)
    x = CRT(a, b)
    @test typeof(x) <: BigInt
    @test a == mod(x, p)
    @test b == mod(x, q)

    c = Mod{101}(86)
    x = CRT(Int, a, b, c)
    @test typeof(x) <: Int

    @test a == mod(x, p)
    @test b == mod(x, q)
    @test c == mod(x, 101)

    @test CRT(
        BigInt,
        Mod{9223372036854775783}(9223372036854775782),
        Mod{9223372036854775643}(9223372036854775642),
    ) == 85070591730234614113402964855534653468
end

@testset "Matrices" begin
    M = zeros(Mod{11}, 3, 3)
    @test sum(M) == 0

    M = ones(Mod{11}, 5, 5)
    @test sum(M) == 3

    M = rand(Mod{11}, 5, 6)
    @test size(M) == (5, 6)

    A = rand(Mod{17}, 5, 5)
    X = values.(A)
    @test sum(X) == sum(Mod{17}.(A))
end

@testset "Hashing/Iterating" begin
    x = Mod{17}(11)
    y = x + 0im
    @test x == y
    @test hash(x) == hash(y)
    @test typeof(x) !== typeof(y)

    A = Set([x, y])
    @test length(A) == 1

    v = [Mod{10}(t) for t = 1:15]
    w = [Mod{10}(t + 0im) for t = 1:15]
    S = Set(v)
    T = Set(w)
    @test length(S) == 10
    @test S == T
    @test union(S, T) == intersect(S, T)
end

@testset "Iteration" begin
    @test sum(k for k in Mod{7}) == zero(Mod{7})
    @test collect(Mod{7}) == Mod{7}.(0:6)
    @test prod(k for k in Mod{7} if k != 0) == -1  # Wilson's theorem

    @test sum(GaussMod{7}) == 0
    @test length(GaussMod{5}) == 25
end

@testset "Linear system solution" begin
    A = Mod{13}.([1 2; 3 -4])
    b = Mod{13}.([3, 4])
    x = A \ b
    @test A * x == b
    @test inv(A) * b == x
end


# Copyright © 2021, UChicago Argonne, LLC
# All Rights Reserved.

# This source code is licensed under the BSD-style license found in the
# LICENSE file in the root directory of this source tree.

function get_vbounds(obs, operators)
    # Initialization
    operators = Set(operators)
    (v_lb, v_ub) = (-1e+09, 1e+09)
    factor1 = 10
    factor2 = 2
    eps = 1e-09

    # if !isempty(intersect(operators, Set("EL")))
    #     factor1 = 5
    #     factor2 = 2
    # end

    ## Compute min/max of indep/dep variables
    # (z_lb, z_ub) = (minimum(obs[:,end]), maximum(obs[:,end]))
    # (x_lb, x_ub) = (minimum(obs[:,1:end-1]), maximum(obs[:,1:end-1]))
    (obs_lb, obs_ub) = (minimum(obs), maximum(obs))

    ## (1) Shrink bounds
    (v_lb, v_ub) = (maximum([v_lb, factor1^(-sign(obs_lb)) * obs_lb]), minimum([v_ub, factor1 * obs_ub]))

    ## (2) Expand bounds if it is too tight
    (v_lb, v_ub) = (minimum([v_lb, factor2^(-sign(obs_lb)) * obs_lb]), maximum([v_ub, factor2 * obs_ub]))

    ## (3) make bounds symmetric and +- epsilon
    magnitude = maximum([abs(v_lb), abs(v_ub)]) + eps
    (v_lb, v_ub) = (-magnitude, magnitude)

    v_lb, v_ub
end

function get_bounds(obs, operators, params)
    bounds = Dict()
    v_lb, v_ub = get_vbounds(obs, operators)
    merge!(bounds, Dict("v_lb"=>v_lb, "v_ub"=>v_ub))
    merge!(bounds, Dict("c_lb"=>get(params, "DBL_CONSTANT_LB", -2.0), "c_ub"=>get(params, "DBL_CONSTANT_UB", 2.0)))

    bounds["cst_abslb"] = 1/(4*pi)
    bounds["mult_child_abslb"] = 0.01
    bounds["div_child_abslb"] = 0.01
    bounds["sqrt_rhs_lb"] = 0.01
    bounds["exp_rhs_lb"] = v_lb > 0 ? max(v_lb, log(v_lb)) : v_lb
    bounds["exp_rhs_ub"] = min(v_ub, log(v_ub))
    bounds["log_rhs_lb"] = max(v_lb, exp(v_lb))
    bounds["log_rhs_ub"] = min(v_ub, exp(v_ub))

    bounds
end

function get_df_vars(nodes, y_indexes, num_obs)
    df_vars = DataFrame(vartype=String[], nodeid=Int[], dataid=Int[], operator=Any[], sol=Float64[])
    for (n,o) in y_indexes
        append!(df_vars, Dict("vartype"=>"y", "nodeid"=>n, "dataid"=>0, "operator"=>o, "sol"=>0.0))
    end

    for n in nodes, i=1:num_obs
        append!(df_vars, Dict("vartype"=>"v", "nodeid"=>n, "dataid"=>i, "operator"=>0, "sol"=>0.0))
    end

    for n in nodes
        append!(df_vars, Dict("vartype"=>"c", "nodeid"=>n, "dataid"=>0, "operator"=>0, "sol"=>0.0))
    end

    for i=1:num_obs
        append!(df_vars, Dict("vartype"=>"e", "nodeid"=>0, "dataid"=>i, "operator"=>0, "sol"=>0.0))
    end

    df_vars.index = 1:nrow(df_vars)
    select!(df_vars, ["index", "vartype", "nodeid", "dataid", "operator", "sol"])
    df_vars
end

function create_decision_variables(model, nodes, operA, y_indexes, num_obs, bounds, params)
    @variable(model, y[n in nodes, o in operA; (n,o) in y_indexes], Bin)
    @variable(model, bounds["v_lb"] <= v[1:num_obs, nodes] <= bounds["v_ub"])
    if get(params, "INT_CONSTANT_ISINTEGER", 0) == 1
        @variable(model, bounds["c_lb"] <= c[nodes] <= bounds["c_ub"], Int)
    else
        @variable(model, bounds["c_lb"] <= c[nodes] <= bounds["c_ub"])
    end
    @variable(model, -1e+09 <= e[1:num_obs] <= 1e+09)
    df_vars = get_df_vars(nodes, y_indexes, num_obs)
    model, y, c, v, e, df_vars
end

function set_objective_function(model, e, num_obs)
    @objective(model, Min, 1 / num_obs * sum(e.^2))
    model
end

function limit_number_of_constant(model, y, y_indexes, numcst)
    @constraint(model, sum(y for (n,o) in y_indexes if o == 'c') <= numcst)
    model
end

function add_e_deining_constraints(model, v, e, obs, num_obs)
    @constraint(model, edef[i in 1:num_obs], e[i] == v[i,1] - obs[i,end])
    model
end

function add_tree_defining_constraints(model, formulation, y, y_indexes, nodes, nleaves, operBU, operB, operU, operL, indvars)
    if occursin("Coz", formulation)
        @constraints(model, begin
            # [Cozad]-(1d) at most one operator is active
            active_at_most[n1 in nodes], sum(y[n,o] for (n,o) in y_indexes if n == n1) <= 1
            # [Cozad]-(1i) at least one independent variable is active
            active_indvar, sum(y[n,o] for (n,o) in y_indexes if o in indvars) >= 1
            # [Cozad]-(1e) if binary/unary is active, then rch is active
            active_rch[n1 in nleaves],
                sum(y[n1,o] for o in operBU) <=
                sum(y[n,o] for (n,o) in y_indexes if n == 2*n1+1)
            # [Cozad]-(1f) if binary is active, then lch is active
            active_lch[n1 in nleaves],
                sum(y[n1,o] for o in operB) <=
                sum(y[n,o] for (n,o) in y_indexes if n == 2*n1)
            # [Cozad]-(1h)
            inactive_rch[n1 in nleaves],
                sum(y[n1,o] for o in operL) <=
                1 - sum(y[n,o] for (n,o) in y_indexes if n == 2*n1+1)
            # [Cozad]-(1g)
            inactive_lch[n1 in nleaves],
                sum(y[n1,o] for o in union(operU, operL)) <=
                1 - sum(y[n,o] for (n,o) in y_indexes if n == 2*n1)
        end)
    elseif occursin("Imp", formulation)
        @constraints(model, begin
            # [Cozad]-(1d) at most one operator is active
            active_at_most[n1 in nodes], sum(y[n,o] for (n,o) in y_indexes if n == n1) <= 1
            # [Cozad]-(1i) at least one independent variable is active
            active_indvar, sum(y[n,o] for (n,o) in y_indexes if o in indvars) >= 1
            # improvement of (1e-1h). Also, improvement of [Newman]'s (18-23)
            active_inactive_rch[n1 in nleaves],
                sum(y[n1,o] for o in operBU) ==
                sum(y[n,o] for (n,o) in y_indexes if n == 2*n1+1)
            active_inactive_lch[n1 in nleaves],
                sum(y[n1,o] for o in operB) ==
                sum(y[n,o] for (n,o) in y_indexes if n == 2*n1)
        end)
    else
        nothing
    end

    model
end

function add_value_defining_constraints(model, formulation, y, c, v, y_indexes, bounds, nodes, nleaves, obs, num_obs, operA, operB, indvars)
    ### Constrs for defining domain bounds
    @info "Constrs for defining domain bounds"

    v_lb = bounds["v_lb"]
    v_ub = bounds["v_ub"]
    c_lb = bounds["c_lb"]
    c_ub = bounds["c_ub"]

    if 'C' in operA && bounds["cst_abslb"] > 0
        @constraint(model, domain_cst[n in nodes], c[n]^2 >= bounds["cst_abslb"]^2 * y[n,'C'])
    end

    if ('*' in operA && bounds["mult_child_abslb"] > 0) || ('D' in operA && bounds["div_child_abslb"] > 0)
        @constraint(model, domain_lch_abslb[i in 1:num_obs, n in nleaves],
        v[i,2*n]^2 >=
            ( '*' in operA ? bounds["mult_child_abslb"]^2 * y[n,'*'] : 0 ) +
            ( 'D' in operA ? bounds["div_child_abslb"]^2 * y[n,'D'] : 0 ) )

        @constraint(model, domain_rch_abslb[i in 1:num_obs, n in nleaves],
            v[i,2*n+1]^2 >=
                ( '*' in operA ? bounds["mult_child_abslb"]^2 * y[n,'*'] : 0 ) +
                ( 'D' in operA ? bounds["div_child_abslb"]^2 * y[n,'D'] : 0 ) )
    end

    if ('R' in operA && bounds["sqrt_rhs_lb"] > 0) || ('E' in operA && bounds["exp_rhs_lb"] > v_lb) ||
        ('L' in operA && bounds["log_rhs_lb"] > v_lb)
        @constraint(model, domain_rhs_lb[i in 1:num_obs, n in nleaves],
            v[i,2*n+1] >= v_lb +
                ( ('R' in operA && bounds["sqrt_rhs_lb"] > 0) ? (bounds["sqrt_rhs_lb"] - v_lb) * y[n,'D'] : 0 ) +
                ( ('E' in operA && bounds["exp_rhs_lb"] > v_lb) ? (bounds["exp_rhs_lb"] - v_lb) * y[n,'E'] : 0 ) +
                ( ('L' in operA && bounds["log_rhs_lb"] > v_lb) ? (bounds["log_rhs_lb"] - v_lb) * y[n,'L'] : 0 ) )
    end

    if ('E' in operA && bounds["exp_rhs_ub"] < v_ub) || ('L' in operA && bounds["log_rhs_ub"] < v_ub)
        @constraint(model, domain_rhs_ub[i in 1:num_obs, n in nleaves],
            v[i,2*n+1] <= v_ub +
                ( ('E' in operA && bounds["exp_rhs_ub"] < v_ub) ? (bounds["exp_rhs_ub"] - v_ub) * y[n,'E'] : 0 ) +
                ( ('L' in operA && bounds["log_rhs_ub"] < v_ub) ? (bounds["log_rhs_ub"] - v_ub) * y[n,'L'] : 0 ) )
    end

    if 'S' in operA
        @constraints(model, begin
            domain_asincos_node_lb[i in 1:num_obs, n in nleaves], v[i,n] >= v_lb*(1 - y[n,'S'])
            domain_asincos_node_ub[i in 1:num_obs, n in nleaves], v[i,n] <= y[n,'S'] + v_ub*(1 - y[n,'S'])
            domain_asincos_rch_lb[i in 1:num_obs, n in nleaves], v[i,2*n+1] >= -y[n,'S'] + v_lb*(1 - y[n,'S'])
            domain_asincos_rch_ub[i in 1:num_obs, n in nleaves], v[i,2*n+1] <= y[n,'S'] + v_ub*(1 - y[n,'S'])
        end)
    end

    ### Constrs for defining v
    if occursin("Coz", formulation)
        @info "Constrs for defining v (Cozad) $(formulation)"

        @constraints(model, begin
            v_inactive_ub[i in 1:num_obs, n in nodes], v[i,n] <= v_ub * sum(y[n1,o] for (n1,o) in y_indexes if n1 == n)
            v_inactive_lb[i in 1:num_obs, n in nodes], v[i,n] >= v_lb * sum(y[n1,o] for (n1,o) in y_indexes if n1 == n)
        end)

        @constraints(model, begin
            v_indvars_ub[i in 1:num_obs, n in nodes, o in indvars], v[i,n] - obs[i,o] <= (v_ub - obs[i,o]) * (1 - y[n,o])
            v_indvars_lb[i in 1:num_obs, n in nodes, o in indvars], v[i,n] - obs[i,o] >= (v_lb - obs[i,o]) * (1 - y[n,o])
        end)

        if 'P' in operA
            @constraints(model, begin
                v_pi_ub[i in 1:num_obs, n in nodes], v[i,n] - pi <= (v_ub - v_lb) * (1 - y[n,'P'])
                v_pi_lb[i in 1:num_obs, n in nodes], v[i,n] - pi >= (v_lb - v_ub) * (1 - y[n,'P'])
            end)
        end

        if 'C' in operA
            @constraints(model, begin
                v_cst_ub[i in 1:num_obs, n in nodes], v[i,n] - c[n] <= (v_ub - c_lb) * (1 - y[n,'C'])
                v_cst_lb[i in 1:num_obs, n in nodes], v[i,n] - c[n] >= (v_lb - c_ub) * (1 - y[n,'C'])
            end)
        end
    elseif occursin("Imp", formulation)
        @info "Constrs for defining v (New) $(formulation)"

        @constraints(model, begin
            v_indvars_ub[i in 1:num_obs, n1 in nodes],
                v[i,n1] <=
                    v_ub * sum(y[n,o] for (n,o) in y_indexes if n==n1 && o in operB) +
                    min(v_ub, sqrt(v_ub)) * ((n1,'R') in y_indexes ? y[n1,'R'] : 0) +
                    1 * ((n1,'S') in y_indexes ? y[n1,'S'] : 0) +
                    min(v_ub, exp(v_ub)) * ((n1,'E') in y_indexes ? y[n1,'E'] : 0) +
                    min(v_ub, log(v_ub)) * ((n1,'L') in y_indexes ? y[n1,'L'] : 0) +
                    c_ub * ((n1,'C') in y_indexes ? y[n1,'C'] : 0) +
                    pi * ((n1,'P') in y_indexes ? y[n1,'P'] : 0) +
                    sum(obs[i,o] * y[n1,o] for o in indvars)
            v_indvars_lb[i in 1:num_obs, n1 in nodes],
                v[i,n1] >=
                    v_lb * sum(y[n,o] for (n,o) in y_indexes if n==n1 && o in operB) +
                    max(v_lb, sqrt(bounds["sqrt_rhs_lb"])) * ((n1,'R') in y_indexes ? y[n1,'R'] : 0) +
                    max(v_lb, exp(bounds["exp_rhs_lb"])) * ((n1,'E') in y_indexes ? y[n1,'E'] : 0) +
                    max(v_lb, log(bounds["log_rhs_lb"])) * ((n1,'L') in y_indexes ? y[n1,'L'] : 0) +
                    c_lb * ((n1,'C') in y_indexes ? y[n1,'C'] : 0) +
                    pi * ((n1,'P') in y_indexes ? y[n1,'P'] : 0) +
                    sum(obs[i,o] * y[n1,o] for o in indvars)
        end)

        if 'C' in operA
            @constraints(model, begin
                v_cst_ub[i in 1:num_obs, n in nodes],
                    v[i,n] - c[n] <= (v_ub - c_lb) * (1 - y[n,'C'])
                v_cst_lb[i in 1:num_obs, n in nodes],
                    v[i,n] - c[n] >= (v_lb - c_ub) * (1 - y[n,'C'])
            end)
        end
    end

    if '+' in operA
        @constraints(model, begin
            v_plus_ub[i in 1:num_obs, n in nleaves],
                v[i,n] - (v[i,2*n] + v[i,2*n+1]) <= (v_ub - 2 * v_lb) * (1 - y[n,'+'])
            v_plus_lb[i in 1:num_obs, n in nleaves],
                v[i,n] - (v[i,2*n] + v[i,2*n+1]) >= (v_lb - 2 * v_ub) * (1 - y[n,'+'])
        end)
    end

    if '-' in operA
        @constraints(model, begin
            v_minus_ub[i in 1:num_obs, n in nleaves],
                v[i,n] - (v[i,2*n] - v[i,2*n+1]) <= (2 * v_ub - v_lb) * (1 - y[n,'-'])
            v_minus_lb[i in 1:num_obs, n in nleaves],
                v[i,n] - (v[i,2*n] - v[i,2*n+1]) >= (2 * v_lb - v_ub) * (1 - y[n,'-'])
        end)
    end

    if '*' in operA
        @constraints(model, begin
            v_mult_ub[i in 1:num_obs, n in nleaves],
                v[i,n] - (v[i,2*n] * v[i,2*n+1]) <=
                (v_ub - min(v_lb^2, v_ub^2, v_lb*v_ub)) * (1 - y[n,'*'])
            v_mult_lb[i in 1:num_obs, n in nleaves],
                v[i,n] - (v[i,2*n] * v[i,2*n+1]) >=
                (v_lb - max(v_lb^2, v_ub^2)) * (1 - y[n,'*'])
        end)
    end

    if 'D' in operA
        @constraints(model, begin
            v_div_ub[i in 1:num_obs, n in nleaves],
                v[i,n] * v[i,2*n+1] - v[i,2*n] <=
                (max(v_lb^2, v_ub^2) - v_lb) * (1 - y[n,'D'])
            v_div_lb[i in 1:num_obs, n in nleaves],
                v[i,n] * v[i,2*n+1] - v[i,2*n] >=
                (min(v_lb^2, v_ub^2, v_lb*v_ub) - v_ub) * (1 - y[n,'D'])
        end)
    end

    if 'R' in operA
        @constraints(model, begin
            v_sqrt_ub[i in 1:num_obs, n in nleaves],
                v[i,n] * v[i,n] - v[i,2*n+1] <=
                (max(v_lb^2, v_ub^2) - v_lb) * (1 - y[n,'R'])
            v_sqrt_lb[i in 1:num_obs, n in nleaves],
                v[i,n] * v[i,n] - v[i,2*n+1] >=
                (- v_ub) * (1 - y[n,'R'])
        end)
    end

    if 'S' in operA
        @constraints(model, begin
            v_asincos_ub[i in 1:num_obs, n in nleaves],
                v[i,n] * v[i,n] + v[i,2*n+1] * v[i,2*n+1] <= y[n,'S'] + 2 * max(v_lb^2, v_ub^2) * (1 - y[n,'S'])
            v_asincos_lb[i in 1:num_obs, n in nleaves],
                v[i,n] * v[i,n] + v[i,2*n+1] * v[i,2*n+1] >= y[n,'S']
        end)
    end

    if 'E' in operA
        @NLconstraints(model, begin
            v_exp_ub[i in 1:num_obs, n in nleaves],
                v[i,n] - exp(v[i,2*n+1]) <= v_ub * (1 - y[n,'E'])
            v_exp_lb[i in 1:num_obs, n in nleaves],
                v[i,n] - exp(v[i,2*n+1]) >= (v_lb - exp(v_ub)) * (1 - y[n,'E'])
        end)
    end

    if 'L' in operA
        @NLconstraints(model, begin
            v_log_ub[i in 1:num_obs, n in nleaves],
                exp(v[i,n]) - v[i,2*n+1] <= (exp(v_ub) - v_lb) * (1 - y[n,'L'])
            v_log_lb[i in 1:num_obs, n in nleaves],
                exp(v[i,n]) - v[i,2*n+1] >= (-v_ub) * (1 - y[n,'L'])
        end)
    end

    model
end

function add_redundancy_eliminating_constraints(model, formulation, y, c, y_indexes, nodes, nleaves, operA, operU)
    ### Nested Operations, e.g., no exp(log(x)) or log(exp(x))
    if issubset(Set("EL"), operA)
        @constraints(model, begin
            redun_explog[n in nleaves; 4*n+3 in nodes], y[n,'E'] + y[2*n+1,'L'] <= 1
            redun_logexp[n in nleaves; 4*n+3 in nodes], y[n,'L'] + y[2*n+1,'E'] <= 1
        end)
    end

    ### Constant Operations
    if 'C' in operA && !JuMP.is_integer(c[1])
        # no operation of two constants and no
        @constraint(model, redun_cst_oper[n in nleaves], y[2*n,'C'] + y[2*n+1,'C'] <= 1)

        # no x-cst, x/cst, and no unary operation with a constant
        if occursin("Coz", formulation)
            @info "Constrs to remove redundancy (Cozad version) $(formulation)"
            @constraints(model, begin
                redun_cst_unary[n in nleaves], y[2*n+1,'C'] + sum(y[nn,o] for (nn,o) in y_indexes if nn == n && o in operU) <= 1
                redun_cst_minus[n in nleaves], y[2*n+1,'C'] + sum(y[nn,o] for (nn,o) in y_indexes if nn == n && o == '-') <= 1
                redun_cst_div[n in nleaves], y[2*n+1,'C'] + sum(y[nn,o] for (nn,o) in y_indexes if nn == n && o == 'D') <= 1
            end)
        elseif occursin("Imp", formulation)
            @info "Constrs to remove redundancy (New version) $(formulation)"
            oper = intersect(operA, "+*")
            @constraint(model, redun_cst_rhs[n in nleaves], y[2*n+1,'C'] <= sum(y[n,o] for o in intersect(operA, "+*")))
        end
    end

    ### Association Properties
    if occursin("Imp", formulation)
        @info "Constrs to remove redundancy (New version 2) $(formulation)"
        nodes_constr = intersect(get_nodes_complete(nodes), get_nodes_grandparents(nodes))
        if issubset("+-", operA)
            @constraint(model, redun_pm[n in nodes_constr], y[n,'+'] + y[2*n+1,'-'] <= 1)
        end
        if issubset("*D", operA)
            @constraint(model, redun_pd[n in nodes_constr], y[n,'*'] + y[2*n+1,'D'] <= 1)
        end
    end

    model
end

function add_implication_cuts(model)
    model
end

function add_symmetry_breaking_constraints(model, y, v, nodes, nleaves, bounds)
    @info "Constrs to break symmetry"
    nodes_constr = intersect(get_nodes_complete(nodes), nleaves)
    @constraint(model, sym[n in nodes_constr], v[1,2*n] - v[1,2*n+1] >= (bounds["v_lb"] - bounds["v_ub"]) * (1 - y[n,'+'] - y[n,'*']))
    model
end

function get_MINLP_model(nodes, obs, operators, formulation, params)
    ### Declare useful parameters
    leaves		= get_leaves(nodes)
    nleaves		= get_nonleaves(nodes)
    num_obs 	= size(obs, 1)
    num_indvars = size(obs, 2) - 1
    indvars		= OrderedSet(1:num_indvars)
    operA		= union(operators, indvars)
    operB       = intersect(operators, Set("+-*D"))
    operU       = intersect(operators, Set("RSEL"))
    operBU      = union(operB, operU)
    operL       = union(intersect(operators, Set("CP")), indvars)
    y_indexes   = [(n,o) for n in nodes for o in operA if n in nleaves || o in operL]

    ### Create a model
    model = Model()
    bounds = get_bounds(obs, operators, params)
    model, y, c, v, e, df_vars = create_decision_variables(model, nodes, operA, y_indexes, num_obs, bounds, params)
    model = set_objective_function(model, e, num_obs)
    if get(params, "INT_NUM_CONSTANT", Inf) < length(nodes) / 2
        model = limit_number_of_constant(model, y, y_indexes, get(params, "INT_NUM_CONSTANT", Inf))
    end
    model = add_e_deining_constraints(model, v, e, obs, num_obs)
    model = add_tree_defining_constraints(model, formulation, y, y_indexes, nodes, nleaves, operBU, operB, operU, operL, indvars)
    model = add_value_defining_constraints(model, formulation, y, c, v, y_indexes, bounds, nodes, nleaves, obs, num_obs, operA, operB, indvars)
    if occursin("-R", formulation) || occursin("-F", formulation)
        model = add_redundancy_eliminating_constraints(model, formulation, y, c, y_indexes, nodes, nleaves, operA, operU)
    end
    if occursin("-S", formulation) || occursin("-F", formulation)
        model = add_symmetry_breaking_constraints(model, y, v, nodes, nleaves, bounds)
    end
    # if occursin("Imp", formulation)
    #     model = add_implication_cuts(model)
    # end

    model, df_vars, bounds
end

function save_solution(dir, model, df_vars, obs; save_df_sol=true, save_df_vars=false)
    num_results = JuMP.result_count(model)
    df_sol = DataFrame(index=Int[], formula=String[], objval=Float64[], treesize=Int[])
    if num_results > 0
        for i in 1:num_results
            if JuMP.has_values(model; result = i)
                df_vars.sol = JuMP.value.(JuMP.all_variables(model), result=i)
                ysol, csol = get_ysol_csol(df_vars)
                @info "ysol and csol", ysol, csol
                formula = get_formula(ysol, csol)[1]
                objval = compute_err_formula(obs, formula, "mse"; err2Inf=true)
                append!(df_sol, Dict("index"=>i, "formula"=>formula, "objval"=>objval, "treesize"=>length(ysol)))
            end
        end
        df_sol.time = repeat([JuMP.solve_time(model)], nrow(df_sol))
        df_sol.bnbnodes = repeat([JuMP.node_count(model)], nrow(df_sol))
    end

    ### Save the summary of the list of the solutions
    sort!(df_sol, ["objval"])
    if save_df_sol
        CSV.write(dir * "df_sol.csv", df_sol)
    end

    ### Save the best solution
    if num_results > 0
        df_vars.sol = JuMP.value.(JuMP.all_variables(model), result=df_sol[1,"index"])
        ysol, csol = get_ysol_csol(df_vars)
        if save_df_vars
            CSV.write(dir * "df_vars.csv", df_vars)
        end

        return df_sol[1,"objval"], ysol, csol, df_vars
    else
        return Inf, Dict(), Dict(), df_vars
    end
end

function MINLP(dir; datafile="data.txt", paramfile="params.txt")
    params = read_parameters(dir * paramfile)
    operators = params["OPERATORS"]
    maxdepth = params["INT_MAXDEPTH"]
    timelimit = params["DBL_TIMELIMIT"]
    formulation = params["FORMULATION"]
    loglevel = params["LOGLEVEL"]
    io = open_logger(dir * "log.log", loglevel)

    ### Parameters
    nodes = OrderedSet(1:(2^(maxdepth+1)-1))
    obs = DelimitedFiles.readdlm(dir * datafile)

    ### Build a MINLP model
    model, df_vars, bounds = get_MINLP_model(nodes, obs, operators, formulation, params)

    ### Set the MINLP solver
    set_optimizer(model, SCIP.Optimizer)
    set_optimizer_attribute(model, "limits/time", timelimit)

    ### Solve
    optimize!(model)

    ### Save the result
    save_solution(dir, model, df_vars, obs)

    close(io)
end

function change_c_variables(model, df_vars, params)

end


@testset "CM types" begin
  Qx, x = QQ["x"]
  K, a = NumberField(x^4 + 8*x^2 + 4)
  cK = Hecke._complex_embeddings(K)
  @test length(cK) == 4
  E = cK[1]
  @test number_field(E) === K
  @test place(E) isa InfPlc
  @test overlaps(conj(E)(a), conj(E(a)))
  @test cK[1] == cK[1]
  @test cK[1] != cK[2]

  Ec = conj(E)
  fl, c = Hecke.iscm_field(K)
  @test c * E == Ec

  sprint(show, E)
  sprint(show, MIME"text/plain"(), E)

  k, b = NumberField(x^2 + 16)
  ktoK = hom(k, K, a^3 + 6*a)
  z = rand(k, -10:10)
  @test overlaps(restrict(E, ktoK)(z), E(ktoK(z)))

  ct = cm_types(k)
  cembd = Hecke._complex_embeddings(k)
  C = cm_type(k, cembd[1:1])
  @test number_field(C) === k
  @test Hecke.embeddings(C) == Hecke._complex_embeddings(k)[1:1]
  @test_throws ArgumentError cm_type(k, Hecke._complex_embeddings(k))
  @test_throws ArgumentError cm_type(number_field(x - 1)[1], Hecke.NumFieldEmbedding[])
  @test_throws ArgumentError cm_type(K, Hecke._complex_embeddings(K)[1:2]) # they are conjugated
  fl, c = Hecke.iscm_field(k)
  @test c * C == cm_type(k, cembd[2:2])
  @test id_hom(k) * C == C

  K, a = cyclotomic_field(7)
  ct = Hecke.cm_types(K)
  @test count(isprimitive, ct) == 6

  Qx, x = QQ["x"]
  f = x^4 + 104x^2 + 796 # primitive non-galois quartic CM field
  K, a = number_field(f)
  ct = Hecke.cm_types(K)
  KK = Hecke.reflex(ct[1]).field
  @test isisomorphic(KK, number_field(x^4 + 52*x^2 + 477)[1])[1]
  @test isisomorphic(Hecke.reflex(Hecke.reflex(ct[1])).field, K)[1] # since primitive
end


# Test Bayesian Detctor
include("./lib/metadata.jl")
include("./2bayes_detector_replica.jl")
include("./lib/SignalGenerator.jl")
include("./lib/roc.jl")

using SignalGenerator
using BayesChpOrig
using Metadata
using Roc

function run_sig_gen()
    ret = genSig(5, 100.0, 1.0, 2.0, 0.5, 1)
    sig = ret["sig"]
    chps= ret["changes"]
    writedlm(outdir*"testsig.txt", sig)
    writedlm(outdir*"testsigchps.txt", chps)
    print("run1")
end

function run_detector()
    # genSig(n, mu, sigma, changeStep, sigStd, errRate)
    # ret = genSig(5, 100.0, 2.0, 1.0, 0.5, 1)
    # ret = genSig(5, 100.0, 2.0, 0.15, 0.75, 1)
    ret = genSig(10, 100.0, 10.0, 1.3, 0.75, 1)
    # ret=genSig(5, 100.0, 2.0, 1.15, 1.0, 1)
    sig = ret["sig"]
    chps= ret["changes"]
    #matrix_max_vals, matrix = bayes_detector(sig, lambda=150)
    matrix_max_vals, matrix = bayes_detector(sig, lambda = 120)
    detections = extract_changes(matrix_max_vals)

    max_delay = 30
    tpr = tpRate(detections, chps, max_delay)
    fpr = faRateScaled(detections, chps, max_delay, length(sig), tn_price_scale=0.1)

    println("... test_bd > TPrate = ", tpr, " / FArate = ", fpr)

    writedlm(outdir*"test_bd_detections.txt", detections)
    writedlm(outdir*"test_bd_signal.txt", sig)
    writedlm(outdir*"test_bd_changes.txt", chps)
    writedlm(outdir*"test_bd_matrix.txt", matrix, ",")
    writedlm(outdir*"test_bd_matrix_max_vals.txt", matrix_max_vals, ",")

    # println("... test_bd > changes:", chps)
    # println("... test_bd > detections:", detections)
end
run_detector()


"""
    module Geometry

This module defines basic geometrical objects.
"""
module Geometry

import Base: ==, length, in, iterate, getindex, lastindex, isempty, eltype, keys
import Base: union, setdiff, intersect, issubset

using DocStringExtensions
using StaticArrays
using LinearAlgebra
using ForwardDiff # for computing derivatives of parametric elements
using OrderedCollections

using WaveProp.Utils

include("point.jl")
include("hyperrectangle.jl")
include("referenceshapes.jl")
include("entities.jl")
include("domain.jl")
include("parametricbody.jl")

export
    # abstract types
    AbstractReferenceShape,
    AbstractEntity,
    AbstractParametricBody,
    # types
    ElementaryEntity,
    ParametricEntity,
    ParametricCurve,
    ClosedEntity,
    Circle,
    Kite,
    Domain,
    ReferenceLine,
    ReferenceTriangle,
    ReferenceTetrahedron,
    ReferenceSquare,
    HyperRectangle,
    SVector,
    # functions
    entities,
    key,
    remove,
    assertequaldim,
    boundary,
    skeleton,
    internal_boundary,
    external_boundary,
    domain,
    parametrization,
    jacobian,
    normal,
    ambient_dimension,
    geometric_dimension,
    bounding_box,
    center,
    diameter,
    radius,
    low_corner,
    high_corner,
    measure,
    number_of_nodes,
    vertices,
    line,
    # global variables
    TAGS,
    ENTITIES

end # module


@testsection "Bases" begin
    for T ∈ (Sym, Float64)
        @testsection "type $T" begin
            sq2 = √T(2)
            v = T[1 0 0; 0 1 0; 0 1 1]
            bv = Basis(v)
            @test eltype(bv) == T
            @test vecbasis(bv) == T[1 0 0; 0 1 0; 0 1 1]
            @test vecbasis(bv, :cont) == T[1 0 0; 0 1 -1; 0 0 1]
            @test metric(bv) == T[1 0 0; 0 2 1; 0 1 1]
            @test metric(bv, :cont) == bv.gⁱʲ == T[1 0 0; 0 1 -1; 0 -1 2]
            @test vecbasis(bv, :cont)' * vecbasis(bv, :cov) == I
            @test metric(bv, :cont)' * metric(bv, :cov) == I
            @test !isorthogonal(bv)

            bw = Basis(vecbasis(bv, :cont), :cont)
            @test eltype(bw) == T
            @test vecbasis(bw) == T[1 0 0; 0 1 0; 0 1 1]
            @test vecbasis(bw, :cont) == T[1 0 0; 0 1 -1; 0 0 1]
            @test metric(bw) == T[1 0 0; 0 2 1; 0 1 1]
            @test metric(bw, :cont) == T[1 0 0; 0 1 -1; 0 -1 2]
            @test vecbasis(bw, :cont)' * vecbasis(bw, :cov) == I
            @test metric(bw, :cont)' * metric(bw, :cov) == I
            @test !isorthogonal(bw)

            nb = normalize(bv)
            @test opequal(vecbasis(nb), vecbasis(nb, :cov))
            @test opequal(vecbasis(nb), T[1 0 0; 0 sq2/2 0; 0 sq2/2 1])
            θ, ϕ, ψ = T == Sym ? symbols("θ, ϕ, ψ", real = true) : rand(T, 3)
            br = Basis(θ, ϕ, ψ)
            @test br isa RotatedBasis
            @test opequal(metric(br), I)
            @test opequal(vecbasis(br, :cont), vecbasis(br, :cov))
            if T == Sym
                @test angles(br) == (θ = θ, ϕ = ϕ, ψ = ψ)
            else
                angbr = angles(br)
                @test angbr.θ ≈ θ && angbr.ϕ ≈ ϕ && angbr.ψ ≈ ψ
                angbr = angles(Array(br))
                @test angbr.θ ≈ θ && angbr.ϕ ≈ ϕ && angbr.ψ ≈ ψ
            end
            λ = T == Sym ? [symbols("λ$i", real = true) for i in 1:getdim(br)] : rand(getdim(br))
            e = vecbasis(br, :cov) .* λ'
            bo = Basis(e)
            @test bo isa OrthogonalBasis
            @test bo isa AllOrthogonalBasis

            br = Basis(θ)
            @test opequal(metric(br), I)
            @test opequal(vecbasis(br, :cont), vecbasis(br, :cov))
            if T == Sym
                @test angles(br) == (θ = θ,)
            else
                @test angles(br).θ ≈ θ
                @test angles(Array(br)).θ ≈ θ
            end
            λ = T == Sym ? [symbols("λ$i", real = true) for i in 1:getdim(br)] : rand(getdim(br))
            e = vecbasis(br, :cov) .* λ'
            bo = Basis(e)
            @test bo isa OrthogonalBasis
            @test bo isa AllOrthogonalBasis

            b2 = Basis(zero(T))
            @test b2 isa CanonicalBasis
            @test getdim(b2) == 2
            b3 = Basis(zero(T), zero(T), zero(T))
            @test b3 isa CanonicalBasis
            @test getdim(b3) == 3

            if T ≠ Sym
                v = rand(T, 3, 3)
                while det(v) ≈ 0
                    v = rand(T, 3, 3)
                end
                @test eltype(bv) == T
                @test vecbasis(bv, :cont)' * vecbasis(bv, :cov) == I
                @test metric(bv, :cont)' * metric(bv, :cov) == I
            end

        end
    end
end


using SeparatingAxisTheorem2D
using CoordinateTransformations
using StaticArrays
using Test

v = SVector(.5, .5)
w = SVector(1.5, 1.5)
B = AABB((0., 1.), (0., 1.))
L = LineSegment(v, w)
P = Polygon((0., 0.), (1.5, 0.), (0., 1.5))
C = Circle((1.4, 1.4), .2)
S = CompoundShape(B, P, C)

@test  intersecting(v, B)
@test  intersecting(v, P)
@test !intersecting(v, C)
@test  intersecting(v, S)
@test !intersecting(w, B)
@test !intersecting(w, P)
@test  intersecting(w, C)
@test  intersecting(w, S)
@test  intersecting(B, L)
@test  intersecting(B, P)
@test !intersecting(B, C)
@test  intersecting(L, P)
@test  intersecting(L, C)
@test !intersecting(P, C)


# Quadratic line minimization as suggested in original Beigi-Arias paper

function linmin1( pw, Vpot, α_t, psi, d, E, g )
    # compute directional derivative
    Δ = 2.0*real( sum( conj(d).*g ) )
    #
    psi_t = ortho_gram_schmidt( psi + α_t*d )
    E_trial = calc_Etot( pw, Vpot, psi_t )
    curvature = ( E_trial - ( E + α_t*Δ ) ) /α_t^2
    α = -Δ/(2*curvature)

    if abs(curvature) < eps()
        α = 0.0
    end

    return abs(α)   # need to use absolute value of α !!!
end


using Gadfly

set_default_plot_size(6inch, 3inch)

n = 10
plot(x=rand(n), y=rand(n), color=e.^rand(n),
     Scale.color_log(minvalue=1))


mutable struct Continuous_MountainCarEnv
    min_action::Float32
    max_action::Float32
    min_position::Float32
    max_position::Float32
    max_speed::Float32
    goal_position::Float32
    power::Float32

    state
    action_space::Box
    observation_space::Box
end

function Continuous_MountainCarEnv()
    min_action    =  -1f0
    max_action    =   1f0
    min_position  = -12f-1
    max_position  =   6f-1
    max_speed     =   7f-2
    goal_position =  45f-2 # was 0.5 in gym, 0.45 in Arnaud de Broissia's version
    power         =  15f-4
    action_space = Box(min_action, max_action, (1,), Float32)
    observation_space = Box([min_position, -max_speed], [max_position, max_speed], Float32)
    Continuous_MountainCarEnv(min_action, max_action, min_position, max_position,
                              max_speed, goal_position, power, nothing, action_space, observation_space)
end

function step!(env::Continuous_MountainCarEnv, action)
    @assert action ∈ env.action_space "Invalid action selected ($(action))"
    position, velocity = env.state[1:1], env.state[2:2]
    force = clamp.(action, env.min_action, env.max_action)

    v          = velocity .+ force * env.power .- 25f-4cos.(3f0position)
    velocity_  = clamp.(v, -env.max_speed, env.max_speed)
    x          = position .+ velocity_
    position_  = clamp.(x, env.min_position, env.max_position)
    if all(position_ .== env.min_position) && all(velocity_ .< 0)
        velocity_ = 0f0velocity_
    end

    done = all(position_ .≥ env.goal_position)

    r = [0f0]
    if done
        r = [1f2]
    end
    reward = r .- 1f-1action .^ 2

    env.state = vcat(position_, velocity_)
    return env.state, reward, done, Dict()
end

function reset!(env::Continuous_MountainCarEnv)
    env.state = param([2f-1rand(Float32) - 6f-1, 0f0])
end

show(io::IO, env::Continuous_MountainCarEnv) = print(io, "Continuous-MountainCarEnv")


#=

  Futoshiki problem in Julia ConstraintSolver.jl

  http://en.wikipedia.org/wiki/Futoshiki
  """
  The puzzle is played on a square grid, such as 5 x 5. The objective
  is to place the numbers 1 to 5 (or whatever the dimensions are) such
  that each row, and column contains each of the digits 1 to 5. Some
  digits may be given at the start. In addition, inequality
  constraints are also initially specifed between some of the squares,
  such that one must be higher or lower than its neighbour. These
  constraints must be honoured as the grid is filled out.
  """

  Model created by Hakan Kjellerstrand, hakank@gmail.com
  See also my Julia page: http://www.hakank.org/julia/

=#

using ConstraintSolver, JuMP
using Cbc, GLPK, Ipopt
const CS = ConstraintSolver
include("constraints_utils.jl")

function futoshiki(problem,print_solutions=true,all_solutions=true)

    cbc_optimizer = optimizer_with_attributes(Cbc.Optimizer, "logLevel" => 0)
    glpk_optimizer = optimizer_with_attributes(GLPK.Optimizer)
    ipopt_optimizer = optimizer_with_attributes(Ipopt.Optimizer)

    model = Model(optimizer_with_attributes(CS.Optimizer,   "all_solutions"=> all_solutions,
                                                            # "all_optimal_solutions"=>all_solutions,
                                                            "logging"=>[],

                                                            "traverse_strategy"=>:BFS,
                                                            # "traverse_strategy"=>:DFS,
                                                            # "traverse_strategy"=>:DBFS,

                                                            # "branch_split"=>:Smallest,
                                                            # "branch_split"=>:Biggest,
                                                            "branch_split"=>:InHalf,

                                                            # https://wikunia.github.io/ConstraintSolver.jl/stable/options/#branch_strategy-(:Auto)
                                                            "branch_strategy" => :IMPS, # default
                                                            # "branch_strategy" => :ABS, # Activity Based Search
                                                            # "activity.decay" => 0.999, # default 0.999
                                                            # "activity.max_probes" => 10, # default, 10
                                                            # "activity.max_confidence_deviation" => 20, # default 20

                                                            # "simplify"=>false,
                                                            # "simplify"=>true, # default

                                                            "time_limit"=>6,

                                                            # "backtrack" => false, # default true
                                                            # "backtrack_sorting" => false, # default true

                                                            # "lp_optimizer" => cbc_optimizer,
                                                            # "lp_optimizer" => glpk_optimizer,
                                                            # "lp_optimizer" => ipopt_optimizer,
                                        ))

    p = problem[:problem]
    less_than = problem[:less_than]
    n,_ = size(p)

    @variable(model, 1 <= x[1:n,1:n] <= n, Int)

    latin_square(model, x)

    # fill the data
    for i in 1:n, j in 1:n
        if p[i,j] > 0
            @constraint(model, x[i,j] == p[i,j])
        end
    end

    # less than hints
    for (i1,j1,i2,j2) in eachrow(less_than)
        @constraint(model, x[i1,j1] <= x[i2,j2])
    end

    # Solve the problem
    optimize!(model)

    status = JuMP.termination_status(model)
    # println("status:$status")
    if status == MOI.OPTIMAL
        num_sols = MOI.get(model, MOI.ResultCount())
        println("num_sols:$num_sols\n")
        if print_solutions
            for sol in 1:num_sols
                # println("solution #$sol")
                x_val = convert.(Integer,JuMP.value.(x; result=sol))
                # println("x:$x_val")
                display(x_val)

            end
        end
    else
        println("status:$status")
    end

    return status
end

futoshiki_problems = Dict(
  :1 => Dict(
    # Example from Tailor model futoshiki.param/futoshiki.param
    #
    # Solution:
    # 5 1 3 2 4
    # 1 4 2 5 3
    # 2 3 1 4 5
    # 3 5 4 1 2
    # 4 2 5 3 1
    #
    # Futoshiki instance, by Andras Salamon
    #
    :problem =>
       resize_matrix([[0,0,3,2,0],  # problem grid
        [0,0,0,0,0],
        [0,0,0,0,0],
        [0,0,0,0,0],
        [0,0,0,0,0]]),

    :less_than =>
       resize_matrix([[1,2,1,1], # [i1,j1, i2,j2] requires that values[i1,j1] < values[i2,j2]
        [1,4,1,5],
        [2,3,1,3],
        [3,3,2,3],
        [3,4,2,4],
        [2,5,3,5],
        [3,2,4,2],
        [4,4,4,3],
        [5,2,5,1],
        [5,4,5,3],
        [5,5,4,5]])
   ),

    # Example from http://en.wikipedia.org/wiki/Futoshiki
    # Solution:
    # 5 4 3 2 1
    # 4 3 1 5 2
    # 2 1 4 3 5
    # 3 5 2 1 4
    # 1 2 5 4 3
    #
:2 => Dict(
    :problem =>
       resize_matrix([[0,0,0,0,0],
        [4,0,0,0,2],
        [0,0,4,0,0],
        [0,0,0,0,4],
        [0,0,0,0,0]]),
    :less_than =>
       resize_matrix([[1,2, 1,1],
        [1,4, 1,3],
        [1,5, 1,4],
        [4,4, 4,5],
        [5,1, 5,2],
        [5,2, 5,3]])
),

    # From http://www.sudoku-puzzles.net/futoshiki06x6.html
   :3 => Dict(
    :problem =>
       resize_matrix([[5,0,2,0,1,0],
        [6,0,1,0,0,0],
        [0,0,3,0,0,0],
        [0,0,4,0,0,0],
        [0,0,0,4,2,5],
        [0,0,0,1,3,0]]),
    :less_than =>
       resize_matrix([[1,6, 2,6],
        [4,1, 3,1]])

    )

)

for p in sort(collect(keys(futoshiki_problems)))
    println("\nProblem $p")
    @time futoshiki(futoshiki_problems[p])
end


using naumer_Dimensionality_2022, Test
using LinearAlgebra: eigvals

@testset "Dynamic Programming" begin
    @testset "randomPSD" begin
        random_matrix = randomPSD(3, 1)

        # Check size is correct
        @test size(random_matrix) == (3,3)

        # Check for PSD
        @test minimum(eigvals(random_matrix)) >= 0

        # Check eigenvalue mean for appropriate scaling
        total = 0
        for i in 1:10000
            random_matrix = randomPSD(1,0.5)
            total += eigvals(random_matrix)[1]
        end
        total /= 10000
        @test isapprox(total, 2, atol=1e-1)
    end

    @testset "samplePSD" begin
        dataset = samplePSD(2,3,1)

        # Because this just call randomPSD, just check dimensionality
        @test size(dataset) == (3,3,2)
    end

    @testset "sampleStateSpace" begin
        # Minimal crash test for global RNG (With Burnin)
        system = LinearSystem([0.5 0; 0 0.5])
        trajectorySampleCount = 2
        timestepCount = 2
        burnin = 1
        τ = 1
        @test size(sampleStateSpace(system, trajectorySampleCount, timestepCount, burnin, τ)) == (2,trajectorySampleCount*timestepCount)
    end

    @testset "localAverage" begin
        # 2D Form
        targetPSD = [0 0; 0 0]
        targetState = [0, 0]
        psdSamples = zeros(2,2,2)
        stateSamples = zeros(2,2)
        stateSamples[:,1] = [1,0]
        stateSamples[:,2] = [0,1]
        values = [1,2,1,2]
        d_max = 2

        @test localAverage(targetPSD, targetState, psdSamples, stateSamples, values, d_max) ≈ 1.5

        # Reduction to nearest neighbor
        targetState = [0, 0.1]
        d_max = eps()
        @test localAverage(targetPSD, targetState, psdSamples, stateSamples, values, d_max) ≈ 1

        # Check for argument mismatch error
        @test_throws ArgumentError localAverage(targetPSD, targetState, psdSamples, stateSamples, values[1:2], d_max)

    end

    @testset "localAverageWeights" begin
        # 2D Form
        targetPSD = [0 0; 0 0]
        targetState = [0, 0]
        psdSamples = zeros(2,2,2)
        stateSamples = zeros(2,2)
        stateSamples[:,1] = [1,0]
        stateSamples[:,2] = [0,1]
        values = [1,2,1,2]
        d_max = 2
        @test localAverageWeights(targetPSD, targetState, psdSamples, stateSamples, d_max)[2] ≈ [0.25,0.25,0.25,0.25]
    end

    @testset "nearestNeighbor" begin
        # 1D form
        target = 1.1
        samples = [1.0,2.0]
        values = [-1,1]
        @test nearestNeighbor(target,samples,values) == -1
        @test nearestNeighbor(target+1,samples,values) == 1

        # 2D form
        target = [1.0,2.0]
        samples = [ 1.1 2.0;
                    2.0 1.1 ]
        values = [-1, 1]
        @test nearestNeighbor(target,samples,values) == -1

        # Matrix form
        target = [1.0 0.0;
                  0.0 1.0 ]
        samples = zeros(2,2,2)
        samples[:,:,1] = [ 1.2 0.1;
                          -0.1 0.9 ]
        values = [-1, 1]

        @test nearestNeighbor(target,samples,values) == -1

    end

    @testset "buildNearestNeighbor" begin
        # 1D test
        samples = -1:0.1:1
        values = buildNearestNeighbor(sin, samples)
        @test values == sin.(samples)

        # 2D test
        samples = [1 2;
                   2 1]
        f(x) = x[2] - x[1]
        values = buildNearestNeighbor(f,samples)
        @test values == [1,-1]
    end

    @testset "updateCRLB" begin

        # Identity Jacobian, Identity FIM
        crlb         = [1.0 0.0
                        0.0 1.0]

        action       = [1.0; 0]

        jacobian     = [1.0 0.0
                        0.0 1.0]
        σ² = 1.0

        new_info = updateCRLB(crlb, action, jacobian, σ²)
        @test new_info ≈ [0.5 0.0; 0.0 1.0]

        # Scale Second Eigenvalue, same setup
        jacobian     = [1.0 0.0
                        0.0 0.5]
        new_info = updateCRLB(crlb, action, jacobian, σ²)
        @test new_info ≈ [0.5 0.0; 0.0 0.25]
    end

    @testset "updateCRLB_naive" begin

        # Identity Jacobian, Identity FIM
        crlb         = [1.0 0.0
                        0.0 1.0]

        action       = [1.0; 0]

        jacobian     = [1.0 0.0
                        0.0 1.0]
        σ² = 1.0

        new_info = updateCRLB_naive(crlb, action, jacobian, σ²)
        @test new_info ≈ [0.5 0.0; 0.0 1.0]

        # Scale Second Eigenvalue, same setup
        jacobian     = [1.0 0.0
                        0.0 0.5]
        new_info = updateCRLB_naive(crlb, action, jacobian, σ²)
        @test new_info ≈ [0.5 0.0; 0.0 0.25]
    end

    @testset "optimalAction_NearestNeighbor" begin

        # Test without value function impact
        crlb         = [1.0 0.0;
                        0.0 5.0]

        actionSpace  = [[1.0, 0.0], [0.0, 1.0]]

        jacobian     = [1.0 0.0
                        0.0 1.0]
        σ² = 1.0

        samples = zeros(2,2,2)
        samples[:,:,1] = [0.5 0.0;
                          0.0 5.0 ]

        samples[:,:,2] = [1.0 0.0;
                          0.0 0.8 ]
        values =  [5.5,1.8]

        action = optimalAction_NearestNeighbor(crlb, actionSpace, samples, values, jacobian, σ²)
        @test action == 2

    end

    @testset "optimalAction_LocalAverage" begin

        # Test without value function impact
        crlb         = [1.0 0.0;
                        0.0 5.0]

        new_state = [0,0]

        actionSpace  = [[1.0, 0.0], [0.0, 1.0]]

        jacobian     = [1.0 0.0
                        0.0 1.0]
        σ² = 1.0

        psdsamples = zeros(2,2,2)
        psdsamples[:,:,1] = [0.5 0.0;
                             0.0 5.0 ]

        psdsamples[:,:,2] = [1.0 0.0;
                             0.0 0.8 ]

        statesamples = zeros(2,1)
        statesamples[:,1] = [0,0]

        values =  [5.5,1.8]

        d_max = 0.1

        action = optimalAction_LocalAverage(crlb, new_state, actionSpace, psdsamples, statesamples,  values, jacobian, σ², d_max)
        @test action == 2

    end

    @testset "valueUpdate_NearestNeighbor" begin
        crlb         = [1.0 0.0;
                        0.0 5.0]

        actionSpace  = [[1.0, 0.0], [0.0, 1.0]]

        jacobian     = [1.0 0.0
                        0.0 1.0]
        σ² = 1.0
        γ = 0.5

        samples = zeros(2,2,2)
        samples[:,:,1] = [0.5 0.0;
                          0.0 5.0 ]

        samples[:,:,2] = [1.0 0.0;
                          0.0 0.8 ]

        values =  [5.5,1.8]

        update = valueUpdate_NearestNeighbor(   crlb,
                                                γ,
                                                jacobian,
                                                actionSpace,
                                                samples, values,
                                                σ²)

        truth = 6.0 + 0.9
        @test update ≈ truth
    end

    @testset "valueIterate_NearestNeighbor" begin
        actionSpace  = [[1.0, 0.0], [0.0, 1.0]]

        jacobian     = [1.0 0.0
                        0.0 1.0]
        σ² = 1.0
        γ = 0.5

        samples = zeros(2,2,1)
        samples[:,:,1] = [1.0 0.0;
                          0.0 1.0 ]

        values =  [2.0]

        new_values = valueIterate_NearestNeighbor(γ, jacobian, actionSpace, samples, values, σ²)

        @test new_values[1] ≈ 3.0
    end

    @testset "ValueFunctionApproximation_NearestNeighbor_precompute!" begin
        # Basic test to ensure there are no crashes
        actionSpace = [[1.0,0], [0,1.0]]
        system = LinearSystem([0.5 0; 0 0.5])
        γ = 0.8
        τ = 0.25
        λ = 1
        psdSampleCount = 5
        trajectorySampleCount = 3
        timestepSampleCount = 2
        σ² = 1
        max_iterations = 2

        ValueFunctionApproximation_NearestNeighbor_precompute(system,
                                                              τ,
                                                              γ,
                                                              actionSpace,
                                                              λ,
                                                              psdSampleCount,
                                                              trajectorySampleCount,
                                                              timestepSampleCount,
                                                              σ²,
                                                              max_iterations)
        @test true
    end

    @testset "ValueFunctionApproximation_LocalAverage" begin
        # Basic test to ensure there are no crashes
        actionSpace = [[1.0,0], [0,1.0]]
        system = LinearSystem([0.5 0; 0 0.5])
        γ = 0.8
        τ = 0.25
        λ = 1
        psdSampleCount = 5
        trajectorySampleCount = 3
        timestepSampleCount = 2
        σ² = 1
        max_iterations = 2

        ValueFunctionApproximation_LocalAverage(    system,
                                                    τ,
                                                    γ,
                                                    actionSpace,
                                                    λ,
                                                    psdSampleCount,
                                                    trajectorySampleCount,
                                                    timestepSampleCount,
                                                    σ²,
                                                    max_iterations)
        @test true
    end

    @testset "ValueFunctionApproximation_LocalAverage_precompute" begin
        # Basic test to ensure there are no crashes
        actionSpace = [[1.0,0], [0,1.0]]
        system = LinearSystem([0.5 0; 0 0.5])
        γ = 0.8
        τ = 0.25
        λ = 1
        psdSampleCount = 5
        trajectorySampleCount = 3
        timestepSampleCount = 2
        σ² = 1
        max_iterations = 2

        ValueFunctionApproximation_LocalAverage_precompute(   system,
                                                              τ,
                                                              γ,
                                                              actionSpace,
                                                              λ,
                                                              psdSampleCount,
                                                              trajectorySampleCount,
                                                              timestepSampleCount,
                                                              σ²,
                                                              max_iterations)
        @test true
    end

    @testset "NearestNeighbor_OptimalPolicy" begin
        state = [1, 0]
        crlb = [1.0 0; 0 1.0]
        system = LinearSystem([0 0; 0 0])
        σ² = 0.1
        τ = 1.0
        values = [1, 2]
        psdSamples = zeros(2,2,2)
        psdSamples[:,:,1] = [.1 0; 0 1]
        psdSamples[:,:,2] = [1 0; 0 .1]
        stateSamples = zeros(2,1)
        stateSamples[:,1] = [1,0]

        a1 = [1.0, 0.0]
        a2 = [0.0, 1.0]
        actionSpace = [a1]

        # Test for single action
        @test NearestNeighbor_OptimalPolicy(state, crlb, system, σ², τ, values, psdSamples, stateSamples, actionSpace)[2] == 1

        # Test for actual decision
        actionSpace = [a1, a2]
        @test NearestNeighbor_OptimalPolicy(state, crlb, system, σ², τ, values, psdSamples, stateSamples, actionSpace)[2] == 1
        actionSpace = [a2, a1]
        @test NearestNeighbor_OptimalPolicy(state, crlb, system, σ², τ, values, psdSamples, stateSamples, actionSpace)[2] == 2
    end

    @testset "LocalAverage_OptimalPolicy" begin
        state = [1, 0]
        crlb = [1.0 0; 0 1.0]
        system = LinearSystem([0 0; 0 0])
        σ² = 0.1
        τ = 1.0
        values = [1, 2]
        psdSamples = zeros(2,2,2)
        psdSamples[:,:,1] = [.1 0; 0 1]
        psdSamples[:,:,2] = [1 0; 0 .1]
        stateSamples = zeros(2,1)
        stateSamples[:,1] = [1,0]

        d_max = 1

        a1 = [1.0, 0.0]
        a2 = [0.0, 1.0]
        actionSpace = [a1]

        # Test for single action
        @test LocalAverage_OptimalPolicy(state, crlb, system, σ², τ, values, psdSamples, stateSamples, actionSpace, d_max)[2] == 1

        # Test for actual decision
        actionSpace = [a1, a2]
        @test LocalAverage_OptimalPolicy(state, crlb, system, σ², τ, values, psdSamples, stateSamples, actionSpace, d_max)[2] == 1
        actionSpace = [a2, a1]
        @test LocalAverage_OptimalPolicy(state, crlb, system, σ², τ, values, psdSamples, stateSamples, actionSpace, d_max)[2] == 2
    end
end

@testset "Extended Kalman Filter" begin
    @testset "stateupdate_EKF" begin
        prediction = [1.0, 0.0]
        covariance = [1.0 0; 0 1.0]
        observation = 1
        action = [1.0, 0]
        σ² = 1

        @test stateupdate_EKF(prediction, covariance, observation, action, σ²) ≈ prediction

        observation = 2
        @test stateupdate_EKF(prediction, covariance, observation, action, σ²) ≈ [2.0,0.0]
    end
end

@testset "1D Approximation Method" begin
    @testset "measurementvalue_1DApprox" begin
        action = [1.0, 0.0]
        limitvector = [1.0, 0.0]
        crlb = [1.0 0; 0 1.0]
        σ² = 1

        @test measurementvalue_1DApprox(action, limitvector, crlb, σ²) ≈ 0.5
    end

    @testset "actiongradient_1DApprox" begin
        action = [1.0, 0.0]
        limitvector = [1.0, 0.0]
        crlb = [1.0 0; 0 1.0]
        σ² = 1

        @test actiongradient_1DApprox(action, limitvector, crlb, σ²) == [0,0]

        action = [0.5,0.5]
        action ./= sqrt(sum(abs2, action))
        decision = actiongradient_1DApprox(action, limitvector, crlb, σ²)
        decision ./= sqrt(sum(abs2, decision))
        truth = [1, -1] ./ sqrt(2)
        @test decision ≈ truth
    end

    @testset "exponentialmap_sphere" begin
        state = [1.0,0.0]
        grad = zeros(length(state))
        τ = 1
        @test exponentialmap_sphere(state, grad, τ) == state

        grad[1] = π
        @test exponentialmap_sphere(state, grad, τ) ≈ [-1.0, 0.0]
    end

    @testset "actiongradientstep_1DApprox" begin
        # For now, just test fixed point
        action = [1.0, 0.0]
        limitvector = [1.0, 0.0]
        crlb = [1.0 0; 0 1.0]
        σ² = 1
        τ = 1

        @test actiongradientstep_1DApprox(action, limitvector, crlb, σ², τ) == action
    end

    @testset "optimalaction_1DApprox" begin
        # Known limitvector version
        a1 = [1.0, 0.0]
        a2 = [0.0, 1.0]
        actionspace = [a1, a2]

        limitvector = [1.0, 0.0]
        crlb = [1.0 0; 0 1.0]
        σ² = 1

        @test optimalaction_1DApprox(actionspace, limitvector, crlb, σ²)[1] == 1

        # Given a system version
        system = LinearSystem([-0.9 0; 0 -0.1])
        state = [1,2]
        @test optimalaction_1DApprox(actionspace, system, state, crlb, σ², 10)[1] == 2
    end
end


module Normalize

using LinearAlgebra
using Statistics, StatsBase
using Random, Distributions

include("../src/util.jl")
using .Utility

# XXX: return coefficients (?)
function random_polynomial(order, scale)
    a = randn(order+1, order+1)
    return (u,v) -> sum(scale^(p₁+p₂)*a[p₁+1,p₂+1]*(u-.5)^p₁*(v-.5)^p₂ for p₁ ∈ 0:order for p₂ ∈ 0:order if (p₁ + p₂) ≤ order)
end

function rescale(x)
    x = x .- minimum(x,dims=2)
    x = x ./ mean(x,dims=2)

    return x
end

function random_surface(dim, len, order, scale)
    u, v = rand(len), rand(len)
    Λ = [ random_polynomial(order, scale) for d ∈ 1:dim ]
    w = hcat((λ.(u,v) for λ ∈ Λ)...)'
    w = rescale(w)

    return w, u, v
end

function random_mean(w, ngene)
    ncell = size(w,2)

    g = rand(Gamma(5,5), ngene) .+ 1
    c = rand(Gamma(5,5), ncell) .+ 1
    μ = Diagonal(g)*(rand(ngene,size(w,1))/size(w,1))*w*Diagonal(c)

    return μ, (gene=g, cell=c)
end

function random_sample(μ)
    ngene, ncell = size(μ)

    ϕ = rand(Gamma(10,10), ngene) .+ 1
    sample = Array{Int,2}(undef, ngene, ncell)
    for i in 1:ngene
        r = ϕ[i]
        for j in 1:ncell
            p = ϕ[i] / (μ[i,j] .+ ϕ[i])
            sample[i,j] = rand(NegativeBinomial(r,p))
        end
    end

    return sample, ϕ
end

function generate(ngene, ncell, rank; order=10)
    w, u, v   = random_surface(rank, ncell, order, 2)
    μ, scale  = random_mean(w, ngene)
    sample, ϕ = random_sample(μ)

    return (
        sample = sample,
        mean   = μ,
        ϕ      = ϕ,
        scale  = scale,
        w      = w,
        u      = u,
        v      = v,
    )
end

function main()
    model = generate(8000,2000,25; order=15)
    σ² = model.sample.*(model.sample .+ model.ϕ) ./ (1 .+ model.ϕ)
    u, v = sinkhorn(σ²)

    N = Diagonal(sqrt.(u))*model.sample*Diagonal(sqrt.(v))
    F = svd(N)

    return model, (norm=N, svd=F), u, v
end

end

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


"""
    SComplex{I,R}
"""
struct SComplex{R,I} <: Number
    SComplex(r::T, i::T) where T<:Real = new{r,i}()
end

@pure Base.real(::SComplex{R,I}) where {R,I} = R()
@pure Base.imag(::SComplex{R,I}) where {R,I} = I()
@pure Base.reim(::SComplex{R,I}) where {R,I} = (R(), I())


@doc raw"""
    CenteredMatrices{m,n,𝔽} <: AbstractDecoratorManifold{𝔽}

The manifold of $m × n$ real-valued or complex-valued matrices whose columns sum to zero, i.e.
````math
\bigl\{ p ∈ 𝔽^{m × n}\ \big|\ [1 … 1] * p = [0 … 0] \bigr\},
````
where $𝔽 ∈ \{ℝ,ℂ\}$.

# Constructor
    CenteredMatrices(m, n[, field=ℝ])

Generate the manifold of `m`-by-`n` (`field`-valued) matrices whose columns sum to zero.
"""
struct CenteredMatrices{M,N,𝔽} <: AbstractDecoratorManifold{𝔽} end

function CenteredMatrices(m::Int, n::Int, field::AbstractNumbers=ℝ)
    return CenteredMatrices{m,n,field}()
end

active_traits(f, ::CenteredMatrices, args...) = merge_traits(IsEmbeddedSubmanifold())

@doc raw"""
    check_point(M::CenteredMatrices{m,n,𝔽}, p; kwargs...)

Check whether the matrix is a valid point on the
[`CenteredMatrices`](@ref) `M`, i.e. is an `m`-by-`n` matrix whose columns sum to
zero.

The tolerance for the column sums of `p` can be set using `kwargs...`.
"""
function check_point(M::CenteredMatrices{m,n,𝔽}, p; kwargs...) where {m,n,𝔽}
    if !isapprox(sum(p, dims=1), zeros(1, n); kwargs...)
        return DomainError(
            p,
            string(
                "The point $(p) does not lie on $(M), since its columns do not sum to zero.",
            ),
        )
    end
    return nothing
end

"""
    check_vector(M::CenteredMatrices{m,n,𝔽}, p, X; kwargs... )

Check whether `X` is a tangent vector to manifold point `p` on the
[`CenteredMatrices`](@ref) `M`, i.e. that `X` is a matrix of size `(m, n)` whose columns
sum to zero and its values are from the correct [`AbstractNumbers`](https://juliamanifolds.github.io/ManifoldsBase.jl/stable/types.html#number-system).
The tolerance for the column sums of `p` and `X` can be set using `kwargs...`.
"""
function check_vector(M::CenteredMatrices{m,n,𝔽}, p, X; kwargs...) where {m,n,𝔽}
    if !isapprox(sum(X, dims=1), zeros(1, n); kwargs...)
        return DomainError(
            X,
            "The vector $(X) is not a tangent vector to $(p) on $(M), since its columns do not sum to zero.",
        )
    end
    return nothing
end

embed(::CenteredMatrices, p) = p
embed(::CenteredMatrices, p, X) = X

get_embedding(::CenteredMatrices{m,n,𝔽}) where {m,n,𝔽} = Euclidean(m, n; field=𝔽)

@doc raw"""
    manifold_dimension(M::CenteredMatrices{m,n,𝔽})

Return the manifold dimension of the [`CenteredMatrices`](@ref) `m`-by-`n` matrix `M` over the number system
`𝔽`, i.e.

````math
\dim(\mathcal M) = (m*n - n) \dim_ℝ 𝔽,
````
where $\dim_ℝ 𝔽$ is the [`real_dimension`](https://juliamanifolds.github.io/ManifoldsBase.jl/stable/types.html#ManifoldsBase.real_dimension-Tuple{ManifoldsBase.AbstractNumbers}) of `𝔽`.
"""
function manifold_dimension(::CenteredMatrices{m,n,𝔽}) where {m,n,𝔽}
    return (m * n - n) * real_dimension(𝔽)
end

@doc raw"""
    project(M::CenteredMatrices, p)

Projects `p` from the embedding onto the [`CenteredMatrices`](@ref) `M`, i.e.

````math
\operatorname{proj}_{\mathcal M}(p) = p - \begin{bmatrix}
1\\
⋮\\
1
\end{bmatrix} * [c_1 \dots c_n],
````
where $c_i = \frac{1}{m}\sum_{j=1}^m p_{j,i}$ for $i = 1, \dots, n$.
"""
project(::CenteredMatrices, ::Any)

project!(::CenteredMatrices, q, p) = copyto!(q, p .- mean(p, dims=1))

@doc raw"""
    project(M::CenteredMatrices, p, X)

Project the matrix `X` onto the tangent space at `p` on the [`CenteredMatrices`](@ref) `M`, i.e.

````math
\operatorname{proj}_p(X) = X - \begin{bmatrix}
1\\
⋮\\
1
\end{bmatrix} * [c_1 \dots c_n],
````
where $c_i = \frac{1}{m}\sum_{j=1}^m x_{j,i}$  for $i = 1, \dots, n$.
"""
project(::CenteredMatrices, ::Any, ::Any)

project!(::CenteredMatrices, Y, p, X) = (Y .= X .- mean(X, dims=1))

@generated representation_size(::CenteredMatrices{m,n,𝔽}) where {m,n,𝔽} = (m, n)

function Base.show(io::IO, ::CenteredMatrices{m,n,𝔽}) where {m,n,𝔽}
    return print(io, "CenteredMatrices($(m), $(n), $(𝔽))")
end