type L1RDASolver <: RegressionSolver end

function solve(model::RegressionModel,
			   method::RegERM, ::L1RDASolver,
			   X::AbstractMatrix, y::AbstractVector,
			   tolerance::Float64 = 1e-8, k::Int=1,
			   iterations::Int = 1000)
    (N,d) = size(X)
    (λ,ρ,γ) = method.params.λ, method.params.ρ, method.params.γ

    # define a derivative for a loss function
    function loss_grad(theta::Vector, i::Int)
        grad_model = gradient(model.f, X[i,:], theta)
        grad_loss = derivs(loss(method), values(model, X[i,:], theta), [y[i]])
        vec(broadcast(*, grad_loss',grad_model))
    end

    # Maintain current state in w, previous state in w_previous and initial dual average in g
    (w, w_previous, g) = copy(model.theta), copy(model.theta), zeros(d)

    for t=1:iterations
        λ_rda = λ+(ρ*γ)/sqrt(t)
        idx = int(ceil(N*rand(k)))

        eval = (X[idx,:]*w).*y[idx]
        grad = map(i->loss_grad(w,i), idx[find(eval .< 1)])

        # calculate dual average
        g = ((t-1)/t).*g - (1/t).*sum(grad)

        # find a close form solution
        w = -(sqrt(t)/γ).*(g - λ_rda.*sign(g))
        w[w .<= λ_rda] = 0

        # check the stopping criteria w.r.t. Tolerance
        if vecnorm(w - w_previous) < tolerance
            break
        end
    end
    -w
end


using JLD
thetas = load("MA2data.jld", "thetas")
Zs = load("MA2data.jld", "Zs")
# train is 900000, test is 100000
trainsize = 900000
testsize = 100000
# preprocess by OLS
x = [ones(size(Zs,1),1) Zs]
betas = x \ thetas
preprocess = x*betas
x = 0
errors = thetas - preprocess
# standardize the inputs using stds of training data
mZs = mean(Zs,1)
sZs = std(Zs,1)
Zs = (Zs .- mZs) ./ sZs
sErrors = std(errors,1)
errors = errors ./sErrors
Y = errors[1:trainsize,:]'
X = Zs[1:trainsize,:]'
YT = errors[trainsize+1:trainsize+testsize,:]'
XT = Zs[trainsize+1:trainsize+testsize,:]'
Zs = 0
errors = 0


#
# Get into the MGS reorthogonalization loop and see if it
# does its job.
#
function mgs_test(cond=1.e6)
(A, x0, b)=data_cook(cond)
V=zeros(3,20)
gout=kl_gmres(x0, b, matvec, V, 1.e-9; orth="mgs1",pdata=A)
gout2=kl_gmres(x0, b, matvec, V, 1.e-9; orth="mgs2",pdata=A)
del=gout.reshist-gout2.reshist
mgs2ok = (norm(del,Inf) > 1.e-12) && gout.idid && gout2.idid
mgs2ok || println("mgs_test fails")
return mgs2ok
#return(gout, gout2, mgs2ok)
end

function matvec(x,A)
return A*x
end

function data_cook(cond)
u1=[1,-2, 0]/sqrt(5.0)
u2=[0,0,1]
u3=[2,1,0]/sqrt(5.0)
U=[u1 u2 u3]
V=[u3 u1 u2]
D=diagm([1, cond, sqrt(cond)])
A=U * D * V'
xstar=ones(3,)
b=A*xstar
x0=[10.0,10.0,10.0]
return(A, x0, b)
end


using Rubin
using Tests
using Elliptic
using HypergeometricFunctions
using Polylogarithms
using SpecialFunctions

@test integrate((a+a*cos(e+f*x))^2*(-1c+c*cos(e+f*x))^-1*sec(e+f*x)^2, x) == :(-1*a^2*c^-1*f^-1*tan(e+f*x)+-3*a^2*c^-1*f^-1*arctanh(sin(e+f*x))+4*a^2*c^-1*f^-1*(1+-1*cos(e+f*x))^-1*sin(e+f*x))

test/4 Trig functions/4.2 Cosine/4.2.2.3 (g cos)^p (a+b cos)^m (c+d cos)^n.jl


using Oceananigans
using Oceananigans.Units
using Printf

@info "Building a grid..."

Ny = 128
Ly = 0.3

Nx = Int(Ny/2)
Lx = Ly/2

Nz = 2Ny # number of points in the vertical direction
Lz = Ly # domain depth

refinement = 1.5 # controls spacing near surface (higher means finer spaced)
stretching = 8   # controls rate of stretching at bottom
h(k) = (k - 1) / Nz
ζ₀(k) = 1 + (h(k) - 1) / refinement
Σ(k) = (1 - exp(-stretching * h(k))) / (1 - exp(-stretching))

arch = GPU()

grid = VerticallyStretchedRectilinearGrid(architecture = arch,
                                          size = (Nx, Ny, Nz),
                                          halo = (3, 3, 3),
                                          x = (0, Lx),
                                          y = (0, Ly),
                                          z_faces = k -> Lz * (ζ₀(k) * Σ(k) - 1),
                                          topology = (Periodic, Bounded, Bounded))
#=
grid = RegularRectilinearGrid(size = (Nx, Ny, Nz),
                              halo = (3, 3, 3),
                              x = (0, Lx),  # longer in streamwise direction
                              y = (0, Ly),  # wide enough to avoid finite-width effects?
                              z = (-Lz, 0), # not full depth, but deep enough?
                              topology = (Periodic, Bounded, Bounded))
=#

@show grid

#####
##### Parameters
#####

a = 0.000      # m, reference wave amplitude at t ≈ τᵃ
τ = 1e-4       # m² s⁻², surface stress
k = 2π / 0.02  # m⁻¹, constant wavenumber
ν = 1.05e-6    # m² s⁻¹, kinematic viscosity
κ = 1e-6       # tracer diffusivity

struct ConstantStokesShear{T}
    a :: T
    k :: T
    ω :: T
end

ConstantStokesShear(a, k; g = 9.81, T = 7.2e-5) =
    ConstantStokesShear{Float64}(a, k, sqrt(g * k + T * k^3))

(sh::ConstantStokesShear)(z, t) = 2 * sh.a^2 * sh.k^2 * sh.ω * exp(2 * sh.k * z)

∂z_uˢ = ConstantStokesShear(a, k)

# Calculations
ω = ∂z_uˢ.ω
ϵ = a * k
u★ = sqrt(abs(τ))
La = k * ν^(3/2) / (a * u★ * √(ω))

@info """

    Wave parameters | Values
    =============== | ======
                  a | $a
                  k | $k
             2π / k | $(2π / k)
                  ϵ | $ϵ
                 La | $La
               La⁻¹ | $(1 / La)
"""

#u_top_bc = FluxBoundaryCondition((x, y, t) -> - 1e-5 * sqrt(t))
u_top_bc = FluxBoundaryCondition(-τ)
u_bcs = FieldBoundaryConditions(top = u_top_bc)
boundary_conditions = (; u = u_bcs)

@info "Modeling..."

model = NonhydrostaticModel(architecture = arch,
                            advection = WENO5(),
                            timestepper = :RungeKutta3,
                            grid = grid,
                            tracers = :c,
                            boundary_conditions = boundary_conditions,
                            closure = IsotropicDiffusivity(ν=ν, κ=κ),
                            stokes_drift = UniformStokesDrift(∂z_uˢ=∂z_uˢ),
                            coriolis = nothing,
                            buoyancy = nothing)

Δz = try
    grid.Δz # for a regular grid
catch
    minimum(parent(grid.Δzᵃᵃᶜ)) # for a stretched grid
end

uᵢ(x, y, z) = 1e-9 * randn()
set!(model, u=uᵢ, v=uᵢ, w=uᵢ)

c = model.tracers.c
view(interior(c), :, :, grid.Nz) .= 1

@info "Revvving up a simulation..."

simulation = Simulation(model, Δt=0.01, stop_time=1minutes)

wizard = TimeStepWizard(cfl=0.5, max_Δt=1.0)
simulation.callbacks[:wizard] = Callback(wizard, IterationInterval(10))

function progress(sim)

    umax = maximum(abs, sim.model.velocities.u)
    vmax = maximum(abs, sim.model.velocities.v)
    wmax = maximum(abs, sim.model.velocities.w)

    t = time(sim)
    h = √(ν * t)
    Re = umax * h / ν

    @info @sprintf("Time: %s, iteration: %d, next Δt: %s, max|U|: (%.2e, %.2e, %.2e)  m s⁻¹, hk: %.2f, Re: %.1f",
                   prettytime(t),
                   iteration(sim),
                   prettytime(sim.Δt),
                   umax, vmax, wmax,
                   h * k,
                   Re)

    return nothing
end

simulation.callbacks[:progress] = Callback(progress, IterationInterval(10))

#####
##### Set up output
#####

prefix = @sprintf("simple_instability_%d_%d_%d_τ%.1e_a%.1e", grid.Nx, grid.Ny, grid.Nz, τ, a)

outputs = merge(model.velocities, model.tracers)

simulation.output_writers[:yz] = JLD2OutputWriter(model, outputs,
                                                    schedule = TimeInterval(0.1),
                                                    force = true,
                                                    prefix = prefix * "_yz",
                                                    field_slicer = FieldSlicer(i = round(Int, grid.Nx/2)))

simulation.output_writers[:xz] = JLD2OutputWriter(model, outputs,
                                                  schedule = TimeInterval(0.1),
                                                  force = true,
                                                  field_slicer = FieldSlicer(j = round(Int, grid.Ny/2)),
                                                  prefix = prefix * "_xz")

simulation.output_writers[:netcdf_xz] = NetCDFOutputWriter(model, outputs,
                                                           schedule = TimeInterval(0.1),
                                                           mode = "c",
                                                           field_slicer = FieldSlicer(j = round(Int, grid.Ny/2)),
                                                           filepath = prefix * "_xz.nc")

simulation.output_writers[:z] = NetCDFOutputWriter(model, outputs,
                                                   schedule = TimeInterval(0.01),
                                                   mode = "c",
                                                   field_slicer = FieldSlicer(i = round(Int, grid.Nx/2),
                                                                              j = round(Int, grid.Ny/2)),
                                                   filepath = prefix * "_z.nc")

C = AveragedField(model.tracers.c, dims=(1, 2))
U = AveragedField(model.velocities.u, dims=(1, 2))

simulation.output_writers[:averages] = NetCDFOutputWriter(model, (c=C, u=U),
                                                          schedule = TimeInterval(0.1),
                                                          mode = "c",
                                                          filepath = prefix * "_averages.nc")

@time run!(simulation)

@info "Simulation complete: $simulation. Output:"

for (name, writer) in simulation.output_writers
    absfilepath = abspath(writer.filepath)
    @info "OutputWriter $name, $absfilepath:\n $writer"
end


@polly function kernel_jacobi_2d(tsteps, A, B)
    n = size(A,1)

    for t = 1:tsteps
        for i = 2:(n-1), j = 2:(n-1)
            B[i,j] = 0.2 * (A[i,j] + A[i,j-1] + A[i,1+j] + A[1+i,j] + A[i-1,j])
        end
        for i = 2:(n-1), j = 2:(n-1)
            A[i,j] = 0.2 * (B[i,j] + B[i,j-1] + B[i,1+j] + B[1+i,j] + B[i-1,j])
        end
    end
end

let
    tsteps = 500
    n = 1300

    A = zeros(Float32, n, n)
    B = zeros(Float32, n, n)

    for i = 1:n, j = 1:n
        A[i,j] = (i*(j+2) + 2) / n
        B[i,j] = (i*(j+3) + 3) / n
    end

    SUITE["jacobi-2d"] = @benchmarkable kernel_jacobi_2d($tsteps, A, B) setup = (A = copy($A); B = copy($B))
end


#Mean functions
export ConstantMean, LinearMean, LinearDiffMean

abstract type MeanFunc end

#-------------
#Constant mean
#-------------
immutable ConstantMean <: MeanFunc
    m::Float64
end
(cm::ConstantMean)(x) = cm.m

#-----------
#Linear mean
#-----------

#Proportional to each parameter's absolute value
type LinearMean <: MeanFunc
    k::Vector{Float64}
end
(m::LinearMean)(x) = dot(m.k, x)

#----------------------
#Linear difference mean
#----------------------

#Mean function outputs proportional to the parameter at scl_idx and the vector
#norm of the difference between the parameters at msr_idx and ref_idx.
type LinearDiffMean <: MeanFunc
    k::Float64
    scl_idx::Int
    msr_idx::Vector{Int}
    ref_idx::Vector{Int}
end

function (m::LinearDiffMean)(x)
    return m.k * x[m.scl_idx] * norm(x[m.msr_idx] - x[m.ref_idx])
end


commands = readlines("input.txt")
x_position = 0
depth = 0
aim = 0

for command in commands
    local direction, dist = split(command)
    local distance = parse(Int64, dist)
    if direction == "forward"
        global x_position += distance
        global depth += distance * aim
    elseif direction == "down"
        global aim += distance
    elseif direction == "up"
        global aim -= distance
    else
        println("Unknown direction! :( ")
    end
end

println("x = ",x_position)
println("depth = ",depth)
println("x*depth = ",x_position*depth)


module TestMultivariateStats

using Test
using MLJBase
import MultivariateStats, Random
using MLJModels.MultivariateStats_
import Dates

using LinearAlgebra
import Random.seed!
seed!(1234)

@testset "Ridge" begin
    ## SYNTHETIC DATA TEST

    # Define some linear, noise-free, synthetic data:
    intercept = -42.0
    coefficients = Float64[1, 3, 7]
    n = 1000
    A = randn(n, 3)
    Xtable = table(A)
    y = A*coefficients

    # Train model on all data with no regularization and no
    # standardization of target:
    ridge = RidgeRegressor(lambda=0.0)

    fitresult, report, cache = fit(ridge, 0, Xtable, y)

    # Training error:
    yhat = predict(ridge, fitresult, Xtable)
    @test norm(yhat - y)/sqrt(n) < 1e-12

    # Get the true intercept?
    fr = fitted_params(ridge, fitresult)
    @test abs(fr.intercept) < 1e-10
    @test norm(fr.coefficients - coefficients) < 1e-10

    info_dict(ridge)
end

X, y = @load_crabs

@testset "PCA" begin
    X_array = matrix(X)
    pratio = 0.9999

    # MultivariateStats PCA
    pca_ms = MultivariateStats.fit(MultivariateStats.PCA,
                                   permutedims(X_array),
                                   pratio=pratio)
    Xtr_ms = permutedims(MultivariateStats.transform(pca_ms,
                                                     permutedims(X_array)))

    # MLJ PCA
    pca_mlj = PCA(pratio=pratio)
    fitresult, _, _ = fit(pca_mlj, 1, X)
    Xtr_mlj = matrix(transform(pca_mlj, fitresult, X))

    @test Xtr_mlj ≈ Xtr_ms
end

@testset "KernelPCA" begin
    X_array = matrix(X)

    # MultivariateStats KernelPCA
    kpca_ms = MultivariateStats.fit(MultivariateStats.KernelPCA
                                  , permutedims(X_array))
    Xtr_ms = permutedims(MultivariateStats.transform(kpca_ms,
                                                     permutedims(X_array)))

    # MLJ KernelPCA
    kpca_mlj = KernelPCA()
    fitresult, _, _ = fit(kpca_mlj, 1, X)
    Xtr_mlj = matrix(transform(kpca_mlj, fitresult, X))

    @test Xtr_mlj ≈ Xtr_ms
end

@testset "ICA" begin
    X_array = matrix(X)
    k = 5
    tolerance = 5.0

    # MultivariateStats ICA
    seed!(1234) # winit gets randomly initialised
    ica_ms = MultivariateStats.fit(MultivariateStats.ICA,
                                   permutedims(X_array),
                                   k;
                                   tol=tolerance)
    Xtr_ms = permutedims(MultivariateStats.transform(ica_ms,
                                                     permutedims(X_array)))

    # MLJ ICA
    seed!(1234) # winit gets randomly initialised
    ica_mlj = ICA(;k=k, tol=tolerance)
    fitresult, _, _ = fit(ica_mlj, 1, X)
    Xtr_mlj = matrix(transform(ica_mlj, fitresult, X))

    @test Xtr_mlj ≈ Xtr_ms
end

@testset "MulticlassLDA" begin
    Xfull, y = @load_smarket
    X = selectcols(Xfull, [:Lag1,:Lag2])

    train  = selectcols(Xfull, :Year) .< Dates.Date(2005)
    test   = .!train
    Xtrain = selectrows(X, train)
    ytrain = selectrows(y, train)
    Xtest  = selectrows(X, test)
    ytest  = selectrows(y, test)

    LDA_model = LDA()
    fitresult, = fit(LDA_model, 1, Xtrain, ytrain)
    class_means, projection_matrix = fitted_params(LDA_model, fitresult)

    preds = predict(LDA_model, fitresult, Xtest)

    mce = cross_entropy(preds, ytest) |> mean

    @test 0.685 ≤ mce ≤ 0.695

    @test round.(class_means', sigdigits = 3) == [0.0428 0.0339; -0.0395 -0.0313]

    d = info_dict(LDA)
    @test d[:input_scitype] == Table(Continuous)
    @test d[:target_scitype] == AbstractVector{<:Finite}
    @test d[:name] == "LDA"
end

@testset "MLDA-2" begin
    Random.seed!(1125)
    X1 = -2 .+ randn(100, 2)
    X2 = randn(100, 2)
    X3 = 2 .+ randn(100, 2)
    y1 = ones(100)
    y2 = 2ones(100)
    y3 = 3ones(100)
    X = vcat(X1, X2, X3)
    y = vcat(y1, y2, y3)
    p = Random.randperm(300)
    X = X[p, :]
    y = y[p]
    X = table(X)
    y = categorical(y)
    train, test = partition(eachindex(y), 0.7)
    Xtrain = selectrows(X, train)
    ytrain = selectrows(y, train)
    Xtest = selectrows(X, test)
    ytest = selectrows(y, test)

    lda_model = LDA()
    fitresult, = fit(lda_model, 1, Xtrain, ytrain)
    preds = predict_mode(lda_model, fitresult, Xtest)
    mcr = misclassification_rate(preds, ytest)
    @test mcr ≤ 0.15
end

@testset "BayesianMulticlassLDA" begin
    Xfull, y = @load_smarket
    X = selectcols(Xfull, [:Lag1,:Lag2])

    train  = selectcols(Xfull, :Year) .< Dates.Date(2005)
    test   = .!train
    Xtrain = selectrows(X, train)
    ytrain = selectrows(y, train)
    Xtest  = selectrows(X, test)
    ytest  = selectrows(y, test)

    BLDA_model = BayesianLDA()
    fitresult, = fit(BLDA_model, 1, Xtrain, ytrain)
    class_means, projection_matrix, priors = fitted_params(BLDA_model, fitresult)

    preds = predict(BLDA_model, fitresult, Xtest)

    mce = cross_entropy(preds, ytest) |> mean

    @test 0.685 ≤ mce ≤ 0.695

    @test round.(class_means', sigdigits = 3) ==
        [0.0428 0.0339; -0.0395 -0.0313]

    d = info_dict(BayesianLDA)
    @test d[:input_scitype] == Table(Continuous)
    @test d[:target_scitype] == AbstractVector{<:Finite}
    @test d[:name] == "BayesianLDA"
end

@testset "BayesianSubspaceLDA" begin
    X, y = @load_iris

    LDA_model = BayesianSubspaceLDA()

    fitresult, _, report = fit(LDA_model, 1, X, y)
    class_means,projection_matrix,prior_probabilities =
        fitted_params(LDA_model, fitresult)
    preds=predict(LDA_model, fitresult, X)
    predicted_class = predict_mode(LDA_model, fitresult, X)
    mcr = misclassification_rate(predicted_class, y)
    mce = cross_entropy(preds, y) |> mean

    @test mean(abs.(class_means' - [5.006 3.428 1.462 0.246;
                          5.936 2.770 4.260 1.326;
                          6.588 2.974 5.552 2.026])) < 0.01

    @test mean(abs.(projection_matrix ≈ [0.0675721   0.00271023;
  					   0.127666    0.177718;
 					  -0.180211   -0.0767255;
 				      -0.235382    0.231435])) < 0.05

    @test round.(prior_probabilities, sigdigits=7) ==
        [0.3333333, 0.3333333, 0.3333333]
    @test round.(mcr, sigdigits=1) == 0.02

    @test round.(report.explained_variance_ratio, digits=4) == [0.9915, 0.0085]
    @test 0.04 ≤ mce ≤ 0.045

    d = info_dict(BayesianSubspaceLDA)
    @test d[:input_scitype] == Table(Continuous)
    @test d[:target_scitype] == AbstractVector{<:Finite}
    @test d[:name] == "BayesianSubspaceLDA"
end

@testset "SubspaceLDA" begin
    Random.seed!(1125)
    X1 = -2 .+ randn(100, 2)
    X2 = randn(100, 2)
    X3 = 2 .+ randn(100, 2)
    y1 = ones(100)
    y2 = 2ones(100)
    y3 = 3ones(100)
    X = vcat(X1, X2, X3)
    y = vcat(y1, y2, y3)
    p = Random.randperm(300)
    X = X[p, :]
    y = y[p]
    X = table(X)
    y = categorical(y)
    train, test = partition(eachindex(y), 0.7)
    Xtrain = selectrows(X, train)
    ytrain = selectrows(y, train)
    Xtest = selectrows(X, test)
    ytest = selectrows(y, test)

    lda_model = SubspaceLDA()
    fitresult, = fit(lda_model, 1, Xtrain, ytrain)
    preds = predict_mode(lda_model, fitresult, Xtest)
    mcr = misclassification_rate(preds, ytest)
    @test mcr ≤ 0.15

    # MultivariateStats Linear Discriminant Analysis transform
     proj    = fitresult[1].projw
     projLDA = fitresult[1].projLDA
     proj =  proj * projLDA
     XWt = matrix(X) * proj
     tlda_ms = table(XWt, prototype=X)

    # MLJ Linear Discriminant Analysis transform
     tlda_mlj = transform(lda_model, fitresult, X)

     @test tlda_mlj == tlda_ms

     d = info_dict(SubspaceLDA)
     @test d[:input_scitype] == Table(Continuous)
     @test d[:target_scitype] == AbstractVector{<:Finite}
     @test d[:name] == "SubspaceLDA"
end

end
true


"""
Test minimize function with no constraints
"""

#push!(LOAD_PATH,"../src/")
using joptimise

function rosenbrock(x)
    f = (1 - x[1])^2 + 100*(x[2] - x[1]^2)^2
    return f
end

function rosenbrock!(g, x)
    # compute objective
    f = rosenbrock(x)
    g[1] = 0.0
    return f
end

# initial guess
x0 = [4.0; 4.0]
# bounds on variables
lx = [-5.0; -5.0]
ux = [5.0; 5.0]
# bounds on constriants
lg = [0.0]
ug = [0.0]
# number of constraints
ng = 1

## run minimizer with IPOPT
ip_options = Dict(
    "max_iter" => 2500,   # 1500 ~ 2500
    "tol" => 1e-6
)

xopt, fopt, info = minimize(rosenbrock!, x0, ng; lx=lx, ux=ux, lg=lg, ug=ug, solver="ipopt", options=ip_options, derivatives=ForwardAD());

println("Done with IPOPT!")
println(info)
println(xopt)

## run minimizer with SNOPT
sn_options = Dict(
    "Major feasibility tolerance" => 1.e-6,
    "Major optimality tolerance"  => 1.e-6,
    "Minor feasibility tolerance" => 1.e-6,
    "Major iterations limit" => 1000,
    "Major print level" => 1,
)

xopt, fopt, info = minimize(rosenbrock!, x0, ng; lx=lx, ux=ux, lg=lg, ug=ug, solver="snopt", options=sn_options, derivatives=ForwardAD());

println("Done with SNOPT!")
println(info)
println(xopt)


abstract type RandomSearch end
include("purerandomsearch.jl")
include("simulatedannealing.jl")
include("particleswarm.jl")


using Flipping
using LsqFit
using StatsPlots

##colors
c9090 = RGBA(108/255,218/255,165/255)#RGBA(181/255,94/255,50/255)
c9030 = RGBA(23/255,116/255,206/255)#RGBA(181/255,73/255,144/255)
c3030 = RGBA(181/255,94/255,51/255)#RGBA(111/255,108/255,166/255)
gr(grid=false,background_color = RGBA(1,1,1,1))
##
#probability that the state has change after n_fail consecutive failures
function rec_lr(n_fail, Prew, Psw; E1=0)
    #= estimate P(depleted) / P(rich); as a ratio not conditioned probability
    n_fail = number of pokes
    Prew = probability of reward
    Psw = probability of state transition
    E1 = first evidence need to be given been calculated recursively
    after a reward the evidence of P(depleted) / P(rich) = 0
    =#
    Es = zeros(n_fail)
    Es[1] = E1
    for i = 2:n_fail
        Es[i] = (Es[i-1]+Psw)/((1-Prew)*(1-Psw))
    end
    return Es
end
function rec_lr_with_uncertainty(n_fail, Prew, Psw, α; E1 = 0)
    #n_fail = numero di poke
    #Prew = probabilitá reward
    #Psw = probabilitá di transizione
    #E1= la prima evidenza deve essere assegnata perché la probabilitaá è calcolata in maniera ricorsiva
    #α = uncertainty
    Es = zeros(n_fail)
    Es[1] = E1
    for i = 2:n_fail
        Es[i] = (Es[i-1]+Psw)/((1-Prew)*(1-Psw))*α + (1-α)
    end
    return Es
end
#Probability of being in the correct side after n_fail consecutive failures, requires to compute rec_lr
function Pwrong(evidence_accumulation)
    # x = pdpl / prich; x è la stima ottenuta dalla funzione rec_lr
    # x*(1-pdpl)  = pdpl
    # x - x*pdpl = pdpl
    # x = pdpl + x*pdpl
    # x = pdpl(1+x)
    # x/(1+x) = pdpl
    evidence_accumulation ./ (1 .+ evidence_accumulation)
end

function Pwrong_with_uncertainty(n_fail, Prew, Psw, α; E1 = 0)
    evidence_accumulation = rec_lr(n_fail, Prew, Psw, E1, α)
    evidence_accumulation ./ (1 .+ evidence_accumulation)
end

function Pwrong(n_fail, Prew, Psw; E1 = 0)
    evidence_accumulation = rec_lr(n_fail, Prew, Psw; E1 = E1)
    evidence_accumulation ./ (1 .+ evidence_accumulation)
end

function Pcorrect(n_fail, Prew, Psw; E1 = 0)
    1 .- Pwrong(n_fail, Prew, Psw; E1 = E1)
end
@. expon(x,p) = p[1]*exp(x*p[2]) ## model for an exponent for fitting
inverse_exp(y,p) = (log(y/p[1]))*1/p[2] ## model for an inverse exponent

function fit_protocol(prot)
    xdata = 0:size(prot,1)-1
    ydata = prot #evidence accumulation of n_fail consecutive omissions ::array
    p0 = [0.5, 0.5] # initialization of parameters
    fit = curve_fit(expon, xdata, ydata, p0)
end

function fit_protocol(n_fail, Prew, Psw; E1 = 0, step = 0.1)
    ydata = Pcorrect(n_fail, Prew, Psw; E1 = E1) #evidence accumulation of n_fail consecutive omissions ::array
    xdata = 0:size(ydata,1)-1
    p0 = [0.5, 0.5] # initialization of parameters
    fit = curve_fit(expon, xdata, ydata, p0)
    x = 0.0:step:n_fail
    (collect(x),expon(x,fit.param))
end
##
c = fit_protocol(4,0.8,0.4)
c
##
plot(fit_protocol(4,0.8,0.4))
##
colorscheme = Dict("80/40" =>RGB(0/255, 158/255, 115/255),
    "60/30" => RGB(86/255, 180/255, 233/255),
    "40/20" => RGB(230/255, 159/255, 0/255))
##
gr(grid=false,background_color = RGBA(1,1,1,1), linewidth = 3)
plt = plot(0:1:9,Pcorrect(10, 0.4, 0.2),#plot(fit_protocol(5,0.8,0.4; step = 1),
    xlabel = "Consecutive failures",
    ylabel =  "Probability of current side high",
    label = "40/20",
    color = colorscheme["40/20"])

plot!(plt, 0:1:6,Pcorrect(7, 0.6, 0.3),
    label= "60/30",
    color = colorscheme["60/30"])

plot!(plt, 0:1:4,Pcorrect(5, 0.8, 0.4),
    label = "80/40",
    color = colorscheme["80/40"])

Plots.abline!(plt,0,0.15,
    label = "example threshold 15%",
    color = :black)

Plots.abline!(plt,0,0.05,label = "example threshold 5%",
    linestyle = :dash, color = :black)
xticks!(plt,0:1:9)

##
dir = "/Volumes/GoogleDrive/My Drive/Flipping/shared docs pietro dario/lab meeting 3-02-2020/inference qualitative observations"
f = joinpath(dir,"updated_inference.pdf")
savefig(plt,f)
##
plt = plot(fit_protocol(10,0.3,0.15),
    xticks = 0:10,
    title  = "Evidence accumulation",
    xlabel = "Evidence",
    ylabel =  "Latent state accumulation P(correct)",
    color = :black,
    legend = false,
    linewidth = 3,
    tickfont = font(:Bookman,12),
    background_color = RGBA(1,1,1,0)
    )
##

##
Plots.abline!(plt,0,0.15,linewidth = 3,
    color = :black,
    annotations=(5, 0.18, text("Latent threshold", :left)))
savefig(plt,"/home/beatriz/mainen.flipping.5ht@gmail.com/Flipping/Datasets/example2.png")
##
Plots.abline!(plt,0,0.05,linewidth = 3,
    color = :blue)
savefig(plt,"/home/beatriz/mainen.flipping.5ht@gmail.com/Flipping/Datasets/example3.png")
##
plot!(fit_protocol(10,0.2,0.1),
    linewidth = 3,
    color = :blue)
savefig(plt,"/home/beatriz/mainen.flipping.5ht@gmail.com/Flipping/Datasets/example4.png")


export For
using Random
import Base

"""
    For(f, base...)

`For` provides a convenient way to construct a `ProductMeasure`. There are
several options for the `base`. With Julia's `do` notation, this can look very
similar to a standard `for` loop, while maintaining semantics structure that's
easier to work with.

------------

# `For(f, base::Int...)`

When one or several `Int` values are passed for `base`, the result is treated as
depending on `CartesianIndices(base)`.

```
julia> For(3) do λ Exponential(λ) end |> marginals
3-element mappedarray(MeasureTheory.var"#17#18"{var"#15#16"}(var"#15#16"()), ::CartesianIndices{1, Tuple{Base.OneTo{Int64}}}) with eltype Exponential{(:λ,), Tuple{Int64}}:
 Exponential(λ = 1,)
 Exponential(λ = 2,)
 Exponential(λ = 3,)
```

```
julia> For(4,3) do μ,σ Normal(μ,σ) end |> marginals
4×3 mappedarray(MeasureTheory.var"#17#18"{var"#11#12"}(var"#11#12"()), ::CartesianIndices{2, Tuple{Base.OneTo{Int64}, Base.OneTo{Int64}}}) with eltype Normal{(:μ, :σ), Tuple{Int64, Int64}}:
 Normal(μ = 1, σ = 1)  Normal(μ = 1, σ = 2)  Normal(μ = 1, σ = 3)
 Normal(μ = 2, σ = 1)  Normal(μ = 2, σ = 2)  Normal(μ = 2, σ = 3)
 Normal(μ = 3, σ = 1)  Normal(μ = 3, σ = 2)  Normal(μ = 3, σ = 3)
 Normal(μ = 4, σ = 1)  Normal(μ = 4, σ = 2)  Normal(μ = 4, σ = 3)
```

-------

# `For(f, base::AbstractArray...)``

In this case, `base` behaves as if the arrays are `zip`ped together before
applying the map.

```
julia> For(randn(3)) do x Exponential(x) end |> marginals
3-element mappedarray(x->Main.Exponential(x), ::Vector{Float64}) with eltype Exponential{(:λ,), Tuple{Float64}}:
 Exponential(λ = -0.268256,)
 Exponential(λ = 1.53044,)
 Exponential(λ = -1.08839,)
```

```
julia> For(1:3, 1:3) do μ,σ Normal(μ,σ) end |> marginals
3-element mappedarray((:μ, :σ)->Main.Normal(μ, σ), ::UnitRange{Int64}, ::UnitRange{Int64}) with eltype Normal{(:μ, :σ), Tuple{Int64, Int64}}:
 Normal(μ = 1, σ = 1)
 Normal(μ = 2, σ = 2)
 Normal(μ = 3, σ = 3)
```

----

# `For(f, base::Base.Generator)`

For `Generator`s, the function maps over the values of the generator:

```
julia> For(eachrow(rand(4,2))) do x Normal(x[1], x[2]) end |> marginals |> collect
4-element Vector{Normal{(:μ, :σ), Tuple{Float64, Float64}}}:
 Normal(μ = 0.255024, σ = 0.570142)
 Normal(μ = 0.970706, σ = 0.0776745)
 Normal(μ = 0.731491, σ = 0.505837)
 Normal(μ = 0.563112, σ = 0.98307)
```

"""
For(f, dims...) = ProductMeasure(i -> f(i...), zip(dims...))

For(f, inds::AbstractArray) = ProductMeasure(f, inds)

For(f, n::Int) = ProductMeasure(f, 1:n)
For(f, dims::Int...) = ProductMeasure(i -> f(Tuple(i)...), CartesianIndices(dims))

function Base.eltype(d::ProductMeasure{F,I}) where {F,I<:AbstractArray}
    return eltype(d.f(first(d.pars)))
end

# """
#     indexstyle(a::AbstractArray, b::AbstractArray)

# Find the best IndexStyle that works for both `a` and `b`. This will return
# `IndexLinear` if both `a` and `b` support it; otherwise it will fall back on `IndexCartesian`.
# """
# function indexstyle(::A,::B)
#     if IndexStyle(A) == IndexStyle(B) == IndexLinear()
#         return IndexLinear()
#     end

#     return IndexCartesian()
# end

# function Base.rand(rng::AbstractRNG, μ::ForArray{D,N,T,F}) where {F,T<:AbstractArray,D,X}
#     s = size(μ.θ)
#     x = Array{X,length(s)}(undef, s...)
#     rand!(rng, x, μ)
# end

# function logdensity(μ::ForArray{D,N,T,F}, x)
#     getℓ(θⱼ, xⱼ) = logdensity(μ.f(θⱼ), xⱼ)
#     ℓ = mappedarray(getℓ, μ.θ, x)
#     _logdensity(μ, x, indexstyle(μ.θ, x), result_type)
# end

# function _logdensity(μ::ForArray{D,N,T,F}, x, ::IndexLinear, ::Type{R}) where {R<:AbstractFloat}
#     ℓ = zero(R)
#     μ.f(μ.θ)
# end

# function basemeasure(μ::ForArray{D,N,T,F}) where {F,T<:AbstractArray,D,X}

# ForGenerator


# ======================================
# Kronecker powers
# ======================================

using LinearAlgebra: checksquare

"""
    kron_pow(x::IA.Interval, pow::Int)

Given an interval x and an integer pow, compute `x^pow`.

### Input

- `x`   -- interval
- `pow` -- integer

### Output

An interval enclosure of `x^pow`.

### Examples

```julia
julia> [kron_pow(2 .. 3, i) for i in 1:3]
3-element Array{IntervalArithmetic.Interval{Float64},1}:
  [2, 3]
  [4, 9]
 [8, 27]
```
"""
function kron_pow(x::IA.Interval, pow::Int)
    return x^pow
end

"""
    kron_pow(x::Interval, pow::Int)

Given an interval x and an integer pow, compute `x^pow`.

### Input

- `x`   -- interval
- `pow` -- integer

### Output

An interval enclosure of `x^pow` as a LazySets `Interval`.
"""
function kron_pow(x::Interval, pow::Int)
    return Interval(kron_pow(x.dat, pow))
end

"""
    kron_pow(H::AbstractHyperrectangle, pow::Int)

Given hyperrectangular set `H` and an integer `pow`, compute the Kronecker power
`H^{⊗ pow}`.

### Input

- `H`   -- hyperrectangular set
- `pow` -- integer power

### Output

A hyperrectangle.

### Algorithm

We compute `H^{⊗ pow}` where `H` is a hyperrectangular set by working with `H`
as a product of intervals.

See also `kron_pow` which requires `DynamicPolynomials.jl`.
"""
function kron_pow(H::AbstractHyperrectangle, pow::Int, algorithm=nothing)
    if pow == 1
        return H
    end

    if isnothing(algorithm)
        if isdefined(@__MODULE__, :DynamicPolynomials)
            algorithm = "symbolic"
        else
            algorithm = "explicit"
        end
    end
    if algorithm == "explicit"
        return _kron_pow_explicit(H, pow)
    elseif algorithm == "symbolic"
        return _kron_pow_symbolic(H, pow)
    else
        throw(ArgumentError("algorithm \"$algorithm\" is not known"))
    end
end

function _kron_pow_explicit(H::AbstractHyperrectangle, pow::Int)
    x = [Interval(low(H, j), high(H, j)) for j in 1:dim(H)]
    r = reduce(kron, fill(x, pow))

    low_r = min.(r)
    high_r = max.(r)
    return Hyperrectangle(low=low_r, high=high_r)
end

"""
    kron_pow_stack(x::Union{<:Interval, <:IA.Interval}, pow::Int)

Return an array with the interval powers `[x, x^2, …, x^pow]`.

### Input

- `x`   -- interval
- `pow` -- integer power

### Output

A vector of elements of the same type as `x` such that the `i`-th element is the interval `x^i`.
"""
function kron_pow_stack(x::Union{<:Interval, <:IA.Interval}, pow::Int)
    return [kron_pow(x, i) for i in 1:pow]
end

"""
    kron_pow_stack(H::AbstractHyperrectangle, pow::Int)

Return the Cartesian product array ``H × H^{⊗2} × ⋯ × H^{⊗pow}`` where ``H`` is a
hyperrectangular set and ``H^{⊗ i}`` is the `i`-th Kronecker power of ``H``.

### Input

- `H`   -- hyperrectangular set
- `pow` -- integer power

### Output

A Cartesian product array of hyperrectangles.
"""
function kron_pow_stack(H::AbstractHyperrectangle, pow::Int)
    out = [kron_pow(H, i) for i in 1:pow]
    return CartesianProductArray(out)
end

# ------------------------------------------------------------
# Functionality that requires DynamicPolynomials.jl
# ------------------------------------------------------------

function load_kron_dynamicpolynomials()
return quote

function _kron_pow_symbolic(H::AbstractHyperrectangle{N}, pow::Int) where {N}
    n = dim(H)
    @polyvar x[1:n]
    B = convert(IntervalBox, H)
    dict = Dict((x[i] => i for i in 1:n)...)
    y = kron_pow(x, pow)

    out = Vector{IA.Interval{N}}(undef, length(y))
    for (i, p) in enumerate(y)
        aux = interval(1)
        for (xj, j) in powers(p)
            aux = aux * B[dict[xj]]^j
        end
        out[i] = aux
    end
    Bpow = IntervalBox(out)
    Hpow = convert(Hyperrectangle, Bpow)
    return Hpow
end

end end  # quote / load_kron_dynamicpolynomials()

# ------------------------------------------------------------
# Functionality that requires MultivariatePolynomials.jl
# ------------------------------------------------------------

function load_kron_multivariate()
return quote

import Base: findfirst, findall

"""
    kron_pow(x::Vector{<:AbstractVariable}, pow::Int)

Compute the higher order concrete Kronecker power: `x ⊗ x ⊗ ... ⊗ x`, `pow` times
for a vector of symbolic monomials.

### Input

- `x`   -- polynomial variable
- `pow` -- integer

### Output

Vector of multivariate monomial corresponding to `x^{⊗ pow}`.

### Examples

```julia
julia> using DynamicPolynomials

julia> @polyvar x[1:2]
(PolyVar{true}[x₁, x₂],)

julia> x
2-element Array{PolyVar{true},1}:
 x₁
 x₂

julia> kron_pow(x, 2)
4-element Array{Monomial{true},1}:
 x₁²
 x₁x₂
 x₁x₂
 x₂²
```
"""
function kron_pow(x::Vector{<:AbstractVariable}, pow::Int)
    @assert pow > 0 "expected positive power, got $pow"
    if pow == 1
        return x
    else
        return kron(x, kron_pow(x, pow-1))
    end
end

"""
    findfirst(y::Vector{<:AbstractMonomialLike}, x::AbstractMonomialLike)

Return the first position of the multivariate monomial ``x`` in the vector of monomials ``y``.

### Input

- `y` -- vector of multivariate monomials
- `x` -- multivariate monomials

### Output

An integer where each integer represents the index in the array `y`
corresponding to the first match, i.e. `y[i] == x` where `i` is the output integer,
or `nothing` if there is no match.

### Notes

Let ``x = (x_1, x_2, …, x_n)`` be given, and let ``y_i = x^{[i]}``.
Given the multi-index ``I = (i_1, i_2, …, i_n)``, this function returns the
first position of ``x^I`` in the array ``y`` (resp. all positions using `findall`).

### Examples

```julia
julia> using DynamicPolynomials

julia> @polyvar x[1:2]
(PolyVar{true}[x₁, x₂],)

julia> y = kron_pow(x, 2)
4-element Array{Monomial{true},1}:
 x₁²
 x₁x₂
 x₁x₂
 x₂²

julia> findfirst(y, x[1]*x[2])
2
```
"""
function findfirst(y::Vector{<:AbstractMonomialLike}, x::AbstractMonomialLike)
    pow = exponents(x)
    @assert sum(pow) == sum(exponents(first(y))) "power indices don't match the power in the lifted vector"
    xv = variables(x)
    _findfirst(y, xv, pow)
end

"""
    findall(y::Vector{<:AbstractMonomialLike}, x::AbstractMonomialLike)

Return all positions of the multivariate monomial ``x`` in the vector of monomials ``y``.

### Input

- `y` -- vector of multivariate monomials
- `x` -- multivariate monomials

### Output

A vector of integers where each integer represents the index in the array `y`
corresponding to a match, i.e. `y[i] == x` for all elements `i` in the output array.

### Notes

Let ``x = (x_1, x_2, …, x_n)`` be given, and let ``y_i = x^{[i]}``.
Given the multi-index ``I = (i_1, i_2, …, i_n)``, this function returns all
positions of ``x^I`` in the array ``y`` (resp. the first position using `findfirst`).

### Examples

```julia
julia> using DynamicPolynomials

julia> @polyvar x[1:2]
(PolyVar{true}[x₁, x₂],)

julia> y = kron_pow(x, 2)
4-element Array{Monomial{true},1}:
 x₁²
 x₁x₂
 x₁x₂
 x₂²

julia> findall(y, x[1]*x[2])
2-element Array{Int64,1}:
 2
 3
```
"""
function findall(y::Vector{<:AbstractMonomialLike}, x::AbstractMonomialLike)
    pow = exponents(x)
    @assert sum(pow) == sum(exponents(first(y))) "power indices don't match the power in the lifted vector"
    xv = variables(x)
    _findall(y, xv, pow)
end

function _findfirst(y, x, pow)
    yp = powers.(y)
    t = ntuple(i -> (x[i], pow[i]), length(pow))
    idx = findfirst(pi -> pi == t, yp)
end

function _findall(y, x, pow)
    yp = powers.(y)
    t = ntuple(i -> (x[i], pow[i]), length(pow))
    idx = findall(pi -> pi == t, yp)
end

end end  # quote / load_kron_multivariate()


# License for this file: MIT (expat)
# Copyright 2017-2018, DLR Institute of System Dynamics and Control

"""
    module ModiaMath.Frames

This module contains functions for **frames** that is coordinate systems in 3D.
The orientation of a frame is described either with a 3x3 **rotation matrix**
or with a **quaternion vector** and its origin is described with a **Vector3D**:

- `const `[`ModiaMath.RotationMatrix`](@ref) = SMatrix{3,3,Float64,9}`:
  Type of a Rotation matrix to rotate from a frame 1 into a frame 2.

- `const `[`ModiaMath.Quaternion`](@ref) = SVector{4,Float64}`:
  Type of a Quaternion vector to rotate from a frame 1 into a frame 2.

- `const `[`ModiaMath.Vector3D`](@ref) = SVector{3,Float64}`:
  Type of a vector in 3D space (e.g. position vector of the origin of a frame).

The following constants are defined

- `const `[`ModiaMath.NullRotation`](@ref):
  RotationMatrix with no rotation from a frame 1 into a frame 2.

- `const `[`ModiaMath.NullQuaternion`](@ref):
  Quaternion vector with no rotation from a frame 1 into a frame 2.

- `const `[`ModiaMath.ZeroVector3D`](@ref):
  Vector3D with only zero elements.

If an angle is given as an argument to one of the functions below, it might be a
number (interpreted as having unit `rad`) or a number with a unit
(for example: `using Unitful; angle = 90u"°"`).

# Constructors for a RotationMatrix R

The following functions return a [`ModiaMath.RotationMatrix`](@ref)` R`
to rotate a frame 1 into a frame 2.

| Function                                           | Description                                |
|:---------------------------------------------------|:-------------------------------------------|
| [`ModiaMath.rot1`](@ref)(angle)                    | Rotate around `angle` along x-axis         |
| [`ModiaMath.rot2`](@ref)(angle)                    | Rotate around `angle` along y-axis         |
| [`ModiaMath.rot3`](@ref)(angle)                    | Rotate around `angle` along z-axis         |
| [`ModiaMath.rot123`](@ref)(angle1, angle2, angle3) | Rotate around angles along x,y,z-axes      |
| [`ModiaMath.rot_e`](@ref)(e, angle)                | Rotate around `angle` along unit vector `e`|
| [`ModiaMath.rot_nxy`](@ref)(nx, ny)                | `nx`/`ny` are in x/y-direction of frame 2  |
| [`ModiaMath.from_q`](@ref)(q)                      | Return `R` from Quaternion `q`             |

# Constructors for a Quaternion q

The following functions return a [`ModiaMath.Quaternion`](@ref)` q`
to rotate a frame 1 into a frame 2.
Since `q` and `-q` define the same rotation the constructor functions have
a keyword argument `q_guess::Quaternion = NullQuaternion`.
From the two possible solutions `q` the one is returned that is closer to `q_guess`.

| Function                                            | Description                                |
|:----------------------------------------------------|:-------------------------------------------|
| [`ModiaMath.qrot1`](@ref)(angle)                    | Rotate around `angle` along x-axis         |
| [`ModiaMath.qrot2`](@ref)(angle)                    | Rotate around `angle` along y-axis         |
| [`ModiaMath.qrot3`](@ref)(angle)                    | Rotate around `angle` along z-axis         |
| [`ModiaMath.qrot123`](@ref)(angle1, angle2, angle3) | Rotate around angles along x,y,z-axes      |
| [`ModiaMath.qrot_e`](@ref)(e, angle)                | Rotate around `angle` along unit vector `e`|
| [`ModiaMath.qrot_nxy`](@ref)(nx, ny)                | `nx`/`ny` are in x/y-direction of frame 2  |
| [`ModiaMath.from_R`](@ref)(R)                       | Return `q` from RotationMatrix `R`         |

# Operations on Frames

The following functions provide operations on frames. The orientation of a frame is
defined with argument `Rq` meaning it can be either a
[`ModiaMath.RotationMatrix`](@ref) ` R` or a
[`ModiaMath.Quaternion`](@ref) ` q` (to rotate a frame 1 into a frame 2).

| Function                                         | Description                                   |
|:-------------------------------------------------|:----------------------------------------------|
| [`ModiaMath.resolve1`](@ref)(Rq, v2)             | Transform vector `v` from frame 2 to frame 1  |
| [`ModiaMath.resolve2`](@ref)(Rq, v1)             | Transform vector `v` from frame 1 to frame 2  |
| [`ModiaMath.absoluteRotation`](@ref)(Rq01, Rq12) | Return rotation 0->2 from rot. 0->1 and 1->2  |
| [`ModiaMath.relativeRotation`](@ref)(Rq01, Rq02) | Return rotation 1->2 from rot. 0->1 and 0->2  |
| [`ModiaMath.inverseRotation`](@ref)(Rq01)        | Return rotation 1->0 from rot, 0->1           |
| [`ModiaMath.planarRotationAngle`](@ref)(e,v1,v2) | Return angle of planar rotation along `e`     |
| [`ModiaMath.eAxis`](@ref)(axis)                  | Return unit vector `e` in direction of `axis` |
| [`ModiaMath.skew`](@ref)(v)                      | Return skew-symmetric matrix of vector v      |

# Interpolation of Frames

Given a set of frames by a vector `r` of position vectors to their origins and
and an optional vector `q` of Quaternions of their absolute orientations, then
the following functions interpolate linearly in these frames:

| Function                                          | Description                                    |
|:--------------------------------------------------|:-----------------------------------------------|
| [`ModiaMath.Path`](@ref)(r,q)                     | Return path defined by a vector of frames      |
| [`ModiaMath.t_pathEnd`](@ref)(path)               | Return path parameter `t_end` of last frame    |
| [`ModiaMath.interpolate`](@ref)(path,t)           | Return `(rt,qt)` of Path at path parameter `t` |
| [`ModiaMath.interpolate_r`](@ref)(path,t)         | Return `rt` of Path at path parameter `t`      |

# Examples

```julia
using ModiaMath
using Unitful

# R1,R2,R3 are the same RotationMatrices
R1 = ModiaMath.rot1(pi/2)
R2 = ModiaMath.rot1(90u"°")
R3 = ModiaMath.rot_e([1,0,0], pi/2)
```

# Main developer

[Martin Otter](https://rmc.dlr.de/sr/de/staff/martin.otter/),
[DLR - Institute of System Dynamics and Control](https://www.dlr.de/sr/en)

*The functions of this module are mostly a mapping of some of the functions
of the Modelica Standard Library from Modelica
([Modelica.Mechanics.MultiBody.Frames](https://doc.modelica.org/help/Modelica_Mechanics_MultiBody_Frames.html#Modelica.Mechanics.MultiBody.Frames))
to Julia (taking advantage of Julia features
such as multiple dispatch and unit package Unitful)*.
"""
module Frames

export RotationMatrix, NullRotation  , assertRotationMatrix
export Quaternion    , NullQuaternion, assertQuaternion
export Vector3D      , ZeroVector3D

export  rot1,  rot2,  rot3,  rot123,  rot_e,  rot_nxy, from_q
export qrot1, qrot2, qrot3, qrot123, qrot_e, qrot_nxy, from_R

export resolve1, resolve2, absoluteRotation, relativeRotation, inverseRotation
export planarRotationAngle, eAxis

export Path, t_pathEnd, interpolate, interpolate_r
export skew

using StaticArrays
import ModiaMath

@static if VERSION >= v"0.7.0-DEV.2005"
    using LinearAlgebra
end

include("vector3D.jl")
include("rotationMatrix.jl")
include("quaternion.jl")
include("interpolation.jl")

end


using TopologyPreprocessing
using Test
using LinearAlgebra

# using MatrixOrganization

@testset "MatrixProcessing.jl -> basics" begin
    # Ints
    positive_matrix = [1 2 3;
                       4 5 6]
    negative_matrix = [1 2 -3;
                       -4 5 6]
    @test shift_to_non_negative(positive_matrix) == positive_matrix
    @test isempty(findall(x->x<0, shift_to_non_negative(negative_matrix)))

    # Floats
    positive_matrix2 = [1. 2 3; 4 5 6]
    negative_matrix2 = [1 2 -3; -4 5 6]
    @test shift_to_non_negative(positive_matrix2) == positive_matrix2
    @test isempty(findall(x->x<0, shift_to_non_negative(negative_matrix2)))

    positive_3d_matrix = rand(2,3,4)
    negative_3d_matrix = rand(2,3,4).*2 .-1
    @test shift_to_non_negative(positive_3d_matrix) == positive_3d_matrix
    @test isempty(findall(x->x<0, shift_to_non_negative(negative_3d_matrix)))

    # =====================

    @test findmin(normalize_to_01(negative_matrix))[1] == 0
    @test findmax(normalize_to_01(negative_matrix))[1] == 1

    powers_of_2 = [0 ; [2^k for k in 0:2:8]]
    powers_of_3 = [0 ; [3^k for k in 0:2:8]]
    @test findmin(normalize_to_01(powers_of_2, use_factor=true))[1] == 0
    @test findmax(normalize_to_01(powers_of_2, use_factor=true))[1] == 1

    @test findmin(normalize_to_01(powers_of_3, use_factor=true))[1] == 0
    @test findmax(normalize_to_01(powers_of_3, use_factor=true))[1] > 1

    @test findmin(normalize_to_01(powers_of_3, use_factor=true, norm_factor=3^8))[1] == 0
    @test findmax(normalize_to_01(powers_of_3, use_factor=true, norm_factor=3^8))[1] == 1

    # =====================

    square_matrix = [1 2 3;
                     4 5 6;
                     7 8 9]
    @test issymmetric(diagonal_symmetrize(square_matrix))
    @test LinearAlgebra.checksquare(diagonal_symmetrize(square_matrix)) == 3

    @test issymmetric(diagonal_symmetrize(positive_matrix))
    @test LinearAlgebra.checksquare(diagonal_symmetrize(positive_matrix)) == 2

    @test diagonal_symmetrize(square_matrix)[end,1] == square_matrix[1,end]
    @test diagonal_symmetrize(square_matrix)[2,1] == square_matrix[1,2]

    @test diagonal_symmetrize(square_matrix, below_over_upper=true)[1,end] == square_matrix[end,1]
    @test diagonal_symmetrize(square_matrix, below_over_upper=true)[1,2] == square_matrix[2,1]
end

@testset "MatrixProcessing.jl -> distances and indices" begin
    # Ints
    positive_matrix = [1 2 3;4 5 6]
    positive_matrix2 = [1. 2 3; 4 5 6]
    positive_3d_matrix = rand(2,3,4)

    negative_matrix = [1 2 -3;-4 5 6]
    negative_matrix2 = [1 2 -3; -4 5 6]
    negative_3d_matrix = rand(2,3,4).*2 .-1

    negative_6d_matrix = rand(2,3,4,5,6,7).*2 .-1

    powers_of_2 = [0 ; [2^k for k in 0:2:8]]
    powers_of_3 = [0 ; [3^k for k in 0:2:8]]

    square_matrix = [1 2 3;
                     4 5 6;
                     7 8 9]

    # ========================================================

    @test length(unique(group_distances(square_matrix,1))) == 1
    @test length(unique(group_distances(square_matrix,2))) == 2
    @test length(unique(group_distances(square_matrix,9))) == 9
    @test_throws DomainError group_distances(square_matrix,20)

    @test length(unique(group_distances(positive_matrix,1))) == 1
    @test length(unique(group_distances(positive_matrix,2))) == 2
    @test length(unique(group_distances(positive_matrix,6))) == 6
    @test_throws DomainError group_distances(positive_matrix,20)

    @test length(unique(group_distances(positive_3d_matrix,2))) == 2
    @test length(unique(group_distances(negative_3d_matrix,2))) == 2
    @test length(unique(group_distances(negative_6d_matrix,2))) == 2

    group_distances(square_matrix,2)

    # ========================================================

    @test length(generate_indices(3))==3^2
    @test length(generate_indices(223))==223^2

    n=3
    @test length(generate_indices(n, symmetry_order=true, include_diagonal=false)) == ((n^2 - n) ÷ 2)
    n=9
    @test length(generate_indices(n, symmetry_order=true, include_diagonal=false)) == ((n^2 - n) ÷ 2)
    index_matrix1 = generate_indices(n, symmetry_order=true, include_diagonal=false)
    @test length(findall(x-> x == CartesianIndex(n,n),index_matrix1)) == 0
    @test length(findall(x-> x == CartesianIndex(n-2,n-1),index_matrix1)) == 1
    @test length(findall(x-> x == CartesianIndex(n-1,n-2),index_matrix1)) == 0

    n=223
    @test length(generate_indices(n, symmetry_order=true, include_diagonal=false)) == ((n^2 - n) ÷ 2)
    @test length(generate_indices(n, symmetry_order=true, include_diagonal=true)) == ((n^2 - n) ÷ 2 + n)

    n=3
    index_matrix2 = generate_indices(n, symmetry_order=true, include_diagonal=true)
    @test length(index_matrix2) == (((n-1)*n)÷2 +n)
    @test findmax(index_matrix2)[1] == CartesianIndex(n,n)
    @test length(findall(x-> x == CartesianIndex(1,n),index_matrix2)) == 1
    @test length(findall(x-> x == CartesianIndex(1,1),index_matrix2)) == 1
    @test length(findall(x-> x == CartesianIndex(n,n),index_matrix2)) == 1

    n=9
    @test length(generate_indices(n, symmetry_order=true, include_diagonal=true)) == (((n-1)*n)÷2 +n)
    n=223
    @test length(generate_indices(n, symmetry_order=true, include_diagonal=true)) == (((n-1)*n)÷2 +n)

    n=9
    @test length(generate_indices((n,n), symmetry_order=true, include_diagonal=true)) == (((n-1)*n)÷2 +n)
    n=223
    @test length(generate_indices((n,n), symmetry_order=true, include_diagonal=true)) == (((n-1)*n)÷2 +n)
end

# @testset "MatrixProcessing.jl -> arrays of arrays" begin
#     # Need to be taken from Bettis
#     # arr_of_arrs
#
#     # reduce_arrs_to_min_len(arr_of_arrs)
#
# end

@testset "MatrixProcessing.jl -> matrix ordering helping functions" begin
    let testing_matrix0 = Array{Float64,2}(undef, 2, 3),
        testing_matrix1 = [1 2 3; 4 5 6; 7 8 9],
        testing_matrix2 = ones((2,3,4))
        testing_matrix3 = [1 2 3; 4 5 6]
        testing_matrix4 = [1, 4, 7, 2, 5, 8, 3, 6, 9]

        @test arr_to_vec(testing_matrix0) isa Vector
        @test length(testing_matrix0) == length(arr_to_vec(testing_matrix0))

        @test arr_to_vec(testing_matrix1) isa Vector
        @test length(testing_matrix1) == length(arr_to_vec(testing_matrix1))
        @test arr_to_vec(testing_matrix1) == [1, 4, 7, 2, 5, 8, 3, 6, 9]

        @test arr_to_vec(testing_matrix2) isa Vector
        @test length(testing_matrix2) == length(arr_to_vec(testing_matrix2))

        @test arr_to_vec(testing_matrix3) isa Vector
        @test length(testing_matrix3) == length(arr_to_vec(testing_matrix3))
        @test arr_to_vec(testing_matrix3) == [1, 4, 2, 5, 3, 6]

        @test arr_to_vec(testing_matrix4) isa Vector
        @test length(testing_matrix4) == length(arr_to_vec(testing_matrix4))
        @test arr_to_vec(testing_matrix4) == [1, 4, 7, 2, 5, 8, 3, 6, 9]
    end

    let testing_matrix1 = CartesianIndices((2,2)),
        testing_matrix2 = CartesianIndices((9,1))

        @test cartesianInd_to_vec(testing_matrix1) isa Vector
        @test length(cartesianInd_to_vec(testing_matrix1)) == length(testing_matrix1)

        @test cartesianInd_to_vec(testing_matrix2) isa Vector
        @test length(cartesianInd_to_vec(testing_matrix2)) == length(testing_matrix2)
    end

    let vals_matrix1a = [1 2 3; 4 5 6],
        vals_matrix1b = [6 5 4; 3 2 1],
        vals_matrix1c = [4 5 6; 1 2 3]
        let index_matrix1a = [CartesianIndex(1,1), CartesianIndex(1,2), CartesianIndex(1,3), CartesianIndex(2,1), CartesianIndex(2,2), CartesianIndex(2,3)]
            @test length(sort_indices_by_values(vals_matrix1a, index_matrix1a)) == length(vals_matrix1a)
            @test sort_indices_by_values(vals_matrix1a, index_matrix1a) == collect(1:6)

            @test length(sort_indices_by_values(vals_matrix1b, index_matrix1a)) == length(vals_matrix1b)
            @test sort_indices_by_values(vals_matrix1b, index_matrix1a) == collect(6:-1:1)

            @test length(sort_indices_by_values(vals_matrix1c, index_matrix1a)) == length(vals_matrix1c)
            @test sort_indices_by_values(vals_matrix1c, index_matrix1a) == [4, 5, 6, 1, 2, 3]
        end

        let index_matrix1b = CartesianIndices((2,3))
            @test_throws TypeError sort_indices_by_values(vals_matrix1a, index_matrix1b)
        end
    end

    let vals_matrix2 = [1 2 3; 4 5 6; 7 8 9],
        index_matrix2a = CartesianIndices((3,3)),
        index_matrix2b =   [CartesianIndex(1,1) CartesianIndex(1,2) CartesianIndex(1,3);
                            CartesianIndex(2,1) CartesianIndex(2,2) CartesianIndex(2,3);
                            CartesianIndex(3,1) CartesianIndex(3,2) CartesianIndex(3,3)],
        index_matrix2c =   [CartesianIndex(1,1), CartesianIndex(1,2), CartesianIndex(1,3),
                            CartesianIndex(2,1), CartesianIndex(2,2), CartesianIndex(2,3),
                            CartesianIndex(3,1), CartesianIndex(3,2), CartesianIndex(3,3)]

        @test_throws TypeError sort_indices_by_values(vals_matrix2, index_matrix2a)
        @test_throws TypeError sort_indices_by_values(vals_matrix2, index_matrix2b)

        @test sort_indices_by_values(vals_matrix2, index_matrix2c) isa Vector
        @test sort_indices_by_values(vals_matrix2, index_matrix2c) == 1:9
        @test length(sort_indices_by_values(vals_matrix2, index_matrix2c)) == length(vals_matrix2)
    end

    let vals_matrix3 = [1, 4, 7, 2, 5, 8, 3, 6, 9],
        index_matrix3a = CartesianIndices((9,1)),
        index_matrix3b = CartesianIndices((9,)),
        index_matrix3c = [1, 4, 7, 2, 5, 8, 3, 6, 9]

        @test_throws TypeError sort_indices_by_values(vals_matrix3, index_matrix3a)
        @test_throws TypeError sort_indices_by_values(vals_matrix3, index_matrix3b)

        @test sort_indices_by_values(vals_matrix3, index_matrix3c) isa Vector
        @test sort_indices_by_values(vals_matrix3, index_matrix3c) == 1:9
        @test length(sort_indices_by_values(vals_matrix3, index_matrix3c)) == length(vals_matrix3)
    end

    let target_coords1 = CartesianIndex(2,3),
        target_value = -20

        let some_matrix = [1 2 3; 4 5 6; 7 8 9]

            set_values!(some_matrix, target_coords1, target_value; do_symmetry=false)
            @test some_matrix[target_coords1] == target_value
        end

        let some_matrix = [1 2 3; 4 5 6; 7 8 9]

            another_matrix = set_values!(some_matrix, target_coords1, target_value; do_symmetry=false)
            @test some_matrix[target_coords1] == target_value
            @test another_matrix[target_coords1] == target_value
            @test another_matrix === some_matrix
        end

        let some_matrix = [1 2 3; 4 5 6; 7 8 9]

            another_matrix = set_values!(some_matrix, target_coords1, target_value; do_symmetry=true)
            @test some_matrix[target_coords1] == target_value
            @test some_matrix[target_coords1[1],target_coords1[2]] == target_value
            @test some_matrix[target_coords1[2],target_coords1[1]] == target_value
            @test another_matrix === some_matrix
        end

        let some_matrix = [1 2 3; 4 5 6; 7 8 9],
            some_matrix2 = [1 2 3; 4 5 6; 7 8 9],
            target_coords2 = CartesianIndex(8,9)

            @test_throws BoundsError set_values!(some_matrix, target_coords2, target_value; do_symmetry=false)
            @test some_matrix == some_matrix2
        end

        let some_matrix = [1 2 3; 4 5 6; 7 8 9],
            target_coords3 = CartesianIndices((2,2))

            @test_throws MethodError set_values!(some_matrix, target_coords3, target_value; do_symmetry=false)
        end
    end
end

@testset "MatrixProcessing.jl -> matrix ordering" begin

    """
        check_for_min_val_position(input_matrix::Matrix; force_symmetry=false, assign_same_values=false)

    Takes an 'input_matrix' and checks if its ordered form has the minimum value
    in the same position as the minimum value of 'input_matrix'.

    """
    function check_for_min_val_position(input_matrix::Matrix; force_symmetry=false, assign_same_values=false)
        # check if min val in uniquely value matrix is in the same position
        ordered_matrix = get_ordered_matrix(input_matrix, force_symmetry=force_symmetry, assign_same_values=assign_same_values)
        if issymmetric(input_matrix) || force_symmetry
            off_diag_ind = generate_indices(size(input_matrix),include_diagonal=false)
        else
            off_diag_ind = generate_indices(size(input_matrix),include_diagonal=true)
        end
        min_val = findmin(input_matrix[off_diag_ind])[1]
        # min_orig = findmin(input_matrix[off_diag_ind])[2]
        all_min_input = findall(x->x==min_val,input_matrix[off_diag_ind])
        all_min_ordered = findall(x->x==1,ordered_matrix[off_diag_ind])
        return off_diag_ind[all_min_input] == off_diag_ind[all_min_ordered]
    end

    # ==
    let power_matrix = zeros(Int, (2,3,4))
        for k = 1:4
            power_matrix[:,:,k] = reshape([(k+1)^n for n =1:6], (2,3))
        end

        ordered_power_matrix = copy(power_matrix)
        ordered_power_matrix[:, :, 1] =[1  6  12;
                                        3  8  13]
        ordered_power_matrix[:, :, 2] =[2  11  17;
                                            7  15  20]
        ordered_power_matrix[:, :, 3] = [4  14  21;
                                            9  18  23]
        ordered_power_matrix[:, :, 4] =[5  16  22;
                                        10  19  24]

        @test !isempty(get_ordered_matrix(power_matrix) == ordered_power_matrix)
        @test get_ordered_matrix(power_matrix) == ordered_power_matrix

        # @test_throws MethodError get_ordered_matrix(power_matrix)
    end

    # ==
    let square_matrix1 = [1 2 3; 4 5 6; 7 8 9],
        square_matrix2 = [1 4 7; 2 5 8; 3 6 9]
        @test get_ordered_matrix(10 .*square_matrix1) == (square_matrix1 )
        @test get_ordered_matrix(10 .*square_matrix2) == (square_matrix2 )
        @test sum(get_ordered_matrix(square_matrix1) .== square_matrix1 ) == 9
        @test sum(get_ordered_matrix(square_matrix2) .== square_matrix2 ) == 9
        @test get_ordered_matrix(10 .*square_matrix1, force_symmetry=true) == get_ordered_matrix(square_matrix1, force_symmetry=true)
        @test get_ordered_matrix(10 .*square_matrix2, force_symmetry=true) == get_ordered_matrix(square_matrix2, force_symmetry=true)

        # check if min val in uniquely value matrix is in the same position
        let input_matrix = 10square_matrix1
            @test check_for_min_val_position(input_matrix)
        end
    end

    # ==
    let square_matrix_same_vals1 = [1 2 3; 3 4 5; 6 7 8],
        square_matrix_same_vals2 = [1 3 6; 2 4 7; 3 5 8]

        @test get_ordered_matrix(10 .*square_matrix_same_vals1, assign_same_values=true) == (square_matrix_same_vals1 )
        @test get_ordered_matrix(10 .*square_matrix_same_vals2, assign_same_values=true) == (square_matrix_same_vals2 )

        # forcing symmetry test
        some_ord_mat = get_ordered_matrix(10 .*square_matrix_same_vals1, force_symmetry=true, assign_same_values=false)
        # remove 1, because 0 is not off diagonal
        @test length(unique(some_ord_mat))-1 == (size(square_matrix_same_vals1,1)*(size(square_matrix_same_vals1,1)-1))/2
    end

    # ==
    let square_matrix_same_vals3 = [1 2 3; 3 4 3; 5 6 7],
        square_matrix_same_vals4 = [1 3 3; 2 4 6; 3 3 7]
        @test get_ordered_matrix(10 .*square_matrix_same_vals3, force_symmetry=true, assign_same_values=true) ==
                            get_ordered_matrix(square_matrix_same_vals3, force_symmetry=true, assign_same_values=true)
        @test get_ordered_matrix(10 .*square_matrix_same_vals4, force_symmetry=true, assign_same_values=true) ==
                            get_ordered_matrix(square_matrix_same_vals4, force_symmetry=true, assign_same_values=true)
    end

    # ==================
    # Tests on symmetric matrices
    let test_mat_b1 = [ 1  1  1  4  5  9 ;
                        1  1  2  3  6  10;
                        1  2  1  3  7  11;
                        4  3  3  1  8  12;
                        5  6  7  8  1  13;
                        9  10 11 12 13 1 ;]

        # no same values
        let test_mat_b1_ord1 = [  0   1   2   6   7  11;
                                  1   0   3   4   8  12;
                                  2   3   0   5   9  13;
                                  6   4   5   0  10  14;
                                  7   8   9  10   0  15;
                                 11  12  13  14  15   0],
            test_mat_b1_indices = generate_indices(size(test_mat_b1))

            ord_mat_b1_1 = get_ordered_matrix(10 .*test_mat_b1, force_symmetry=false, assign_same_values=false)
            @test issymmetric(ord_mat_b1_1)
            @test ord_mat_b1_1[test_mat_b1_indices] == test_mat_b1_ord1[test_mat_b1_indices]

            ord_mat_b1_2 = get_ordered_matrix(10 .*test_mat_b1, force_symmetry=true, assign_same_values=false)
            @test issymmetric(ord_mat_b1_2)
            @test ord_mat_b1_2[test_mat_b1_indices] == test_mat_b1_ord1[test_mat_b1_indices]
        end

        # assign same values
        let test_mat_b1_ord2 = [ 0  1  1  4  5  9 ;
                                 1  0  2  3  6  10;
                                 1  2  0  3  7  11;
                                 4  3  3  0  8  12;
                                 5  6  7  8  0  13;
                                 9 10  11 12 13 0 ;]

            let ord_mat_b1_3 = get_ordered_matrix(10 .*test_mat_b1, force_symmetry=false, assign_same_values=true)
                @test issymmetric(ord_mat_b1_3)
                # Removed, because ordered diagonal is all 1: @test ord_mat_b1_3 == (test_mat_b1 )
                @test ord_mat_b1_3 == test_mat_b1_ord2
            end

            let ord_mat_b1_4 = get_ordered_matrix(10 .*test_mat_b1, force_symmetry=true, assign_same_values=true)
                @test issymmetric(ord_mat_b1_4)
                # Removed, because ordered diagonal is all 1: @test ord_mat_b1_4 == (test_mat_b1 )
                @test ord_mat_b1_4 == test_mat_b1_ord2
            end
        end

        let input_matrix = 10 .*test_mat_b1
            @test check_for_min_val_position(input_matrix; force_symmetry=false, assign_same_values=true)
        end
    end

    # ==
    # Non-symmetric matrix test
    let test_mat_b2 = [ 1  1  3  4  5  9 ;
                        1  1  2  3  6  10;
                        14 2  1  3  7  11;
                        4  15 3  1  8  12;
                        5  5  7  8  1  13;
                        9  10 11 12 13 1 ;]

        let ord_mat_b2_1 = get_ordered_matrix(10 .*test_mat_b2, force_symmetry=false, assign_same_values=false)
            @test !issymmetric(ord_mat_b2_1)
            @test findall(x->x<=8,ord_mat_b2_1) == findall(x->x==1,test_mat_b2) # all values with one are first 8 values used for ordering
            @test length(unique(ord_mat_b2_1)) == length(test_mat_b2) #check if all values are unique
        end

        let test_mat_b2_ord1 = [0  1  6  4  7  11;
                                1  0  2  5  8  12;
                                6  2  0  3  9  13;
                                4  5  3  0 10  14;
                                7  8  9 10  0  15;
                                11 12 13 14 15 0 ;]

            test_mat_b2_indices = generate_indices(size(test_mat_b2))
            filter!(x->x!=CartesianIndex(1,3), test_mat_b2_indices)
            filter!(x->x!=CartesianIndex(1,4), test_mat_b2_indices)
            filter!(x->x!=CartesianIndex(2,4), test_mat_b2_indices)
            filter!(x->x!=CartesianIndex(3,4), test_mat_b2_indices)
            filter!(x->x!=CartesianIndex(3,1), test_mat_b2_indices)
            filter!(x->x!=CartesianIndex(4,1), test_mat_b2_indices)
            filter!(x->x!=CartesianIndex(4,2), test_mat_b2_indices)
            filter!(x->x!=CartesianIndex(4,3), test_mat_b2_indices)

            ord_mat_b2_2 = get_ordered_matrix(10 .*test_mat_b2, force_symmetry=true, assign_same_values=false)
            @test issymmetric(ord_mat_b2_2)
            @test ord_mat_b2_2[test_mat_b2_indices] == test_mat_b2_ord1[test_mat_b2_indices]
            # forcing symmetry test:
            @test length(unique(ord_mat_b2_2))-1 == (size(test_mat_b2,1)*(size(test_mat_b2,1)-1))/2
        end

        # TODO to fix tests, add 1 to all values off diagonal for the input matrix
        let test_mat_b2_ord2 = [0  0  2  3  4  8;
                                0  0  1  2  5  9;
                                13 1  0  2  6  10;
                                3  14 2  0  7  11;
                                4  4  6  7  0  12;
                                8  9  10 11 12 0 ;]
            ord_mat_b2_3 = get_ordered_matrix(10 .*test_mat_b2, force_symmetry=false, assign_same_values=true)
            @test_skip !issymmetric(ord_mat_b2_3)
            @test_skip  ord_mat_b2_3 == test_mat_b2_ord2
        end

        # TODO to fix tests, add 1 to all values off diagonal for the input matrix
        let test_mat_b2_ord3 = [0  0  2   3   4  8
                                0  0  1   2   5  9
                                2  1  0   2   6  10
                                3  2  2   0   7  11
                                4  5  6   7   0  12
                                8  9  10  11  12 0 ;]
            ord_mat_b2_4 = get_ordered_matrix(10 .*test_mat_b2, force_symmetry=true, assign_same_values=true)
            @test_skip issymmetric(ord_mat_b2_4)
            @test_skip ord_mat_b2_4 == test_mat_b2_ord3
        end
    end

    # ==
    let test_mat_b3 = [ 1  1  3  4  5  9 ;
                        1  1  2  3  6  10;
                        3  2  1  3  7  11;
                        4  3  3  1  8  12;
                        5  6  7  8  1  13;
                        9  10 11 12 13 1 ;]

        let test_mat_b3_ord1 = [ 0  0  5  3  6  10;
                                    0  0  1  4  7  11;
                                    5  1  0  2  8  12;
                                    3  4  2  0  9  13;
                                    6  7  8  9  0  14;
                                    10 11 12 13 14 0 ;],
          test_mat_b3_indices = generate_indices(size(test_mat_b3))
            filter!(x->x!=CartesianIndex(1,3), test_mat_b3_indices)
            filter!(x->x!=CartesianIndex(1,4), test_mat_b3_indices)
            filter!(x->x!=CartesianIndex(2,4), test_mat_b3_indices)
            filter!(x->x!=CartesianIndex(3,4), test_mat_b3_indices)
            filter!(x->x!=CartesianIndex(3,1), test_mat_b3_indices)
            filter!(x->x!=CartesianIndex(4,1), test_mat_b3_indices)
            filter!(x->x!=CartesianIndex(4,2), test_mat_b3_indices)
            filter!(x->x!=CartesianIndex(4,3), test_mat_b3_indices)

            ord_mat_b3_1 = get_ordered_matrix(10 .*test_mat_b3, force_symmetry=false, assign_same_values=false)
            @test issymmetric(ord_mat_b3_1)
            @test_skip ord_mat_b3_1[test_mat_b3_indices] == test_mat_b3_ord1[test_mat_b3_indices]

            ord_mat_b3_2 = get_ordered_matrix(10 .*test_mat_b3, force_symmetry=true, assign_same_values=false)
            @test issymmetric(ord_mat_b3_2)
            @test_skip ord_mat_b3_2[test_mat_b3_indices] == test_mat_b3_ord1[test_mat_b3_indices]
        end

        # TODO remove tests that do not add anything new for testing and are just another similar case
        let test_mat_b3_ord2 = [ 0  0  2  3  4  8 ;
                                 0  0  1  2  5  9 ;
                                 2  1  0  2  6  10;
                                 3  2  2  0  7  11;
                                 4  5  6  7  0  12;
                                 8  9  10 11 12 0 ;]

            ord_mat_b3_3 = get_ordered_matrix(10 .*test_mat_b3, force_symmetry=false, assign_same_values=true)
            @test issymmetric(ord_mat_b3_3)
            @test_skip ord_mat_b3_3 == (test_mat_b3 .-1)
            @test_skip ord_mat_b3_3 == test_mat_b3_ord2

            ord_mat_b3_4 = get_ordered_matrix(10 .*test_mat_b3, force_symmetry=true, assign_same_values=true)
            @test issymmetric(ord_mat_b3_4)
            @test_skip ord_mat_b3_4 == (test_mat_b3 .-1)
            @test_skip ord_mat_b3_4 == test_mat_b3_ord2
        end

        let input_matrix = 10 .*test_mat_b3
            @test check_for_min_val_position(input_matrix; force_symmetry=false, assign_same_values=true)
        end
    end

    # ==
    let test_mat_b4 = [  1  1  41 4  5  9  13 17 25 33;
                        1  1  2  42 6  10 14 18 26 34;
                        41 2  1  3  7  11 15 19 27 35;
                        4  42 3  1  8  12 16 20 28 36;
                        5  6  7  8  1  21 43 24 29 37;
                        9  10 11 12 21 1  22 44 30 38;
                        13 14 15 16 43 22 1  23 31 39;
                        17 18 19 20 24 44 23 1  32 40;
                        25 26 27 28 29 30 31 32 1  45;
                        33 34 35 36 37 38 39 40 45 1;]

        @test_skip get_ordered_matrix(10 .*test_mat_b4, force_symmetry=false, assign_same_values=false) == (test_mat_b4 .-1)
        @test_skip get_ordered_matrix(10 .*test_mat_b4, force_symmetry=false, assign_same_values=true) == (test_mat_b4 .-1)

        let ord_mat = get_ordered_matrix(10 .*test_mat_b4, force_symmetry=true, assign_same_values=false)
            @test issymmetric(ord_mat)
            @test_skip ord_mat == (test_mat_b4 .-1)
        end

        let ord_mat = get_ordered_matrix(10 .*test_mat_b4, force_symmetry=true, assign_same_values=true)
            @test issymmetric(ord_mat)
            @test_skip  ord_mat == (test_mat_b4 .-1)
        end

        let input_matrix = 10test_mat_b4
            @test check_for_min_val_position(input_matrix)
        end
    end

    # ==
    let test_mat_b5 = -[1  1  3  4  5  9 ;
                        1  1  2  3  6  10;
                        14 2  1  3  7  11;
                        4  15 3  1  8  12;
                        5  5  7  8  1  13;
                        9  10 11 12 13 1 ;]
        let ord_mat_b5_1 = get_ordered_matrix(10 .*test_mat_b5, force_symmetry=false, assign_same_values=false)
            @test !issymmetric(ord_mat_b5_1)
            @test_skip findall(x->x>=28,ord_mat_b5_1) == findall(x->x==-1,test_mat_b5) # all values with one are first 8 values used for ordering
            @test length(unique(ord_mat_b5_1)) == length(test_mat_b5) #check if all values are unique
        end

        let test_mat_b5_ord1 = [ 0  14  9 11  8  4
                                 14 0  13 10  7  3
                                 9  13 0  12  6  2
                                 11 10 12 0   5  1
                                 8  7  6  5   0  0
                                 4  3  2  1   0  0 ;]

            test_mat_b5_indices = generate_indices(size(test_mat_b5))
            filter!(x->x!=CartesianIndex(1,3), test_mat_b5_indices)
            filter!(x->x!=CartesianIndex(1,4), test_mat_b5_indices)
            filter!(x->x!=CartesianIndex(2,4), test_mat_b5_indices)
            filter!(x->x!=CartesianIndex(3,4), test_mat_b5_indices)
            filter!(x->x!=CartesianIndex(3,1), test_mat_b5_indices)
            filter!(x->x!=CartesianIndex(4,1), test_mat_b5_indices)
            filter!(x->x!=CartesianIndex(4,2), test_mat_b5_indices)
            filter!(x->x!=CartesianIndex(4,3), test_mat_b5_indices)

            ord_mat_b5_2 = get_ordered_matrix(10 .*test_mat_b5, force_symmetry=true, assign_same_values=false)

            @test issymmetric(ord_mat_b5_2)
            @test_skip ord_mat_b5_2[test_mat_b5_indices] == test_mat_b5_ord1[test_mat_b5_indices]
            # forcing symmetry test:
            @test length(unique(ord_mat_b5_2))-1 == (size(test_mat_b5,1)*(size(test_mat_b5,1)-1))/2
        end

        let test_mat_b5_ord2 = [14 14 12 11 10 6;
                                14 14 13 12 9  5;
                                1  13 14 12 8  4;
                                11 0  12 14 7  3;
                                10 10 8  7  14 2;
                                6  5  4  3  2  14],
            ord_mat_b5_3 = get_ordered_matrix(10 .*test_mat_b5, force_symmetry=false, assign_same_values=true)

            @test !issymmetric(ord_mat_b5_3)
            @test_skip ord_mat_b5_3 == test_mat_b5_ord2
        end

        let test_mat_b5_ord3 = [0  12 10 9  8  4;
                                12 0  11 10 7  3;
                                10 11 0  10 6  2;
                                9  10 10 0  5  1;
                                8  7  6  5  0  0;
                                4  3  2  1  0  0]
            ord_mat_b5_4 = get_ordered_matrix(10 .*test_mat_b5, force_symmetry=true, assign_same_values=true)
            @test issymmetric(ord_mat_b5_4)
            @test_skip ord_mat_b5_4 == test_mat_b5_ord3
        end
    end
    # ==================
end


function metropolis_sampler(;target_pdf, known_pdf, known_sampler, starting_state, n_iter, n_burn = 100, n_thinning = 100)

    function transition_probability(target_pdf, known_pdf, new_state, previous_state)

        """metropolis_independant_sampler transition probability, see eq 1 in https://arxiv.org/abs/1705.00686"""
        min(1, target_pdf(new_state) * known_pdf(previous_state)/(target_pdf(previous_state) * known_pdf(new_state)))
    end

    function do_with_probability(p)

        """returns true with probability p, false with (1-p)"""

        rand() < p ? true : false

    end

    function update_state_parameters(target_pdf, known_pdf, new_state, previous_state)

        if do_with_probability(transition_probability(target_pdf, known_pdf, new_state, previous_state))

            return new_state
        else
            return previous_state
        end

    end

    function metropolis_iteration!(target_pdf, known_pdf, known_sampler, new_state, previous_state)
        new_state[:] = known_sampler()[:]
        new_state[:] = update_state_parameters(target_pdf, known_pdf, new_state, previous_state)[:]
        previous_state[:] = new_state[:]
    end

    samples = Array{eltype(starting_state)}(undef, n_iter, size(starting_state,1))

    previous_state = similar(starting_state) #x
    new_state = similar(previous_state)
    previous_state = starting_state

    if n_burn == n_thinning
        @showprogress for iter in 1:n_iter
            for i in 1:n_thinning
                metropolis_iteration!(target_pdf, known_pdf, known_sampler, new_state, previous_state)
            end
            samples[iter,:] = new_state[:]
        end
    else
        throw(Exception("n_burn != n_thinning not implemented"))
    end

    samples

end


using Clustering
using Test
using CodecZlib

@testset "hclust() (hierarchical clustering)" begin

@testset "param checks" begin
    D = rand(5, 5)
    Dsym = D + D'
    Dnan = copy(Dsym)
    Dnan[1, 3] = Dnan[3, 1] = NaN
    @testset "hclust()" begin
        @test_throws ArgumentError hclust(Dsym, linkage=:typo)
        @test_throws ArgumentError hclust(D, linkage=:single)
        @test_throws ArgumentError hclust(Dnan, linkage=:single)
    end
    hclu = @inferred(hclust(Dsym, linkage=:single))
    @test hclu isa Clustering.Hclust
    @testset "cutree()" begin
        @test_throws ArgumentError cutree(hclu)
        @test_throws ArgumentError cutree(hclu, k=-1)
        @test_throws ArgumentError cutree(hclu, k=0)
        @test cutree(hclust(Dsym), k=10) isa Vector{Int}
        @test cutree(hclust(fill(0.0, 0, 0)), k=0) == Int[]
        @test cutree(hclust(fill(0.0, 0, 0)), k=1) == Int[]
    end
end

@testset "R hclust() generated examples" begin
# load the examples
hclu_examples_filename = joinpath(@__DIR__, "data", "hclust_generated_examples.jl.gz")
Base.include_string(@__MODULE__,
                    open(io -> read(io, String), GzipDecompressorStream, hclu_examples_filename),
                    hclu_examples_filename)

# test to make sure many random examples match R's implementation
@testset "example #$i (linkage=:$(example["linkage"]), n=$(size(example["D"], 1)))" for
        (i, example) in enumerate(examples)

    linkage = example["linkage"]
    hclu = @inferred(hclust(example["D"], linkage=linkage))
    @test hclu isa Clustering.Hclust
    @test Clustering.nnodes(hclu) == size(example["D"], 1)
    @test Clustering.nmerges(hclu) == Clustering.nnodes(hclu)-1
    @test Clustering.height(hclu) ≈ maximum(example["height"]) atol=1e-5
    @test hclu.merges == example["merge"]
    @test hclu.heights ≈ example["height"] atol=1e-5
    @test hclu.order == example["order"]

    # compare hclust_nn_lw() (the default) and hclust_nn() (slower) methods
    if linkage ∈ [:complete, :average]
        @testset "hclust_nn()" begin
            hclu2 = Hclust(Clustering.hclust_nn(example["D"], linkage == :complete ? Clustering.slicemaximum : Clustering.slicemean),
                           linkage)
            @test hclu2.merges == hclu.merges
            @test hclu2.heights ≈ hclu.heights atol=1e-5
            @test hclu2.order == hclu.order
        end
    end

    local cut_k = example["cut_k"]
    local cut_h = example["cut_h"]
    if cut_h !== nothing
        # due to small arithmetic differences between R and Julia heights might be slightly different
        # find the matching height
        cut_h_r = cut_h
        cut_h_ix = findmin(abs.(hclu.heights .- cut_h_r))[2]
        cut_h = hclu.heights[cut_h_ix]
        @assert isapprox(cut_h, cut_h_r, atol=1e-6) "h=$cut_h ≈ h_R=$cut_h_r"
    end
    @testset "cutree(hclu, k=$(repr(cut_k)), h=$(repr(cut_h)))" begin
        cutt = @inferred(cutree(hclu, k=cut_k, h=cut_h))
        @test cutt isa Vector{Int}
        @test length(cutt) == Clustering.nnodes(hclu)
        @test cutt == example["cutree"]
    end
end
end

@testset "hclust_n3()" begin
    # no thorough testing (it's O(N³)), just test one example
    example_n3 = examples[10]
    hclu_n3 = @inferred(Clustering.hclust_n3(example_n3["D"], Clustering.slicemaximum))
    @test hclu_n3.mleft == example_n3["merge"][:, 1]
    @test hclu_n3.mright == example_n3["merge"][:, 2]
    @test hclu_n3.heights ≈ example_n3["height"] atol=1e-5
end

local hclust_linkages = [:single, :average, :complete]

@testset "hclust(0×0 matrix, linkage=$linkage)" for linkage in hclust_linkages
    hclu = hclust(fill(0.0, 0, 0), linkage=linkage)
    @test Clustering.nnodes(hclu) == 0
    @test Clustering.nmerges(hclu) == 0
    @test Clustering.height(hclu) == -Inf
    cut1 = @inferred(cutree(hclu, h=1))
    @test cut1 == Int[]
end

@testset "hclust([$dist] 1×1 matrix, linkage=$linkage)" for
    linkage in hclust_linkages, dist in [-Inf, 0, 1, 2, Inf]
    hclu = hclust(fill(dist, 1, 1), linkage=linkage)
    @test Clustering.nnodes(hclu) == 1
    @test Clustering.nmerges(hclu) == 0
    @test Clustering.height(hclu) == -Inf
    cut1 = @inferred(cutree(hclu, h=1))
    @test cut1 == [1]
end

@testset "hclust(linkage=$linkage) when data contains an isolated point (#109)" for linkage in hclust_linkages
    # point #2 is isolated: distances to all the other points are Inf
    mdist = [
        0.0 Inf 0.1  Inf  Inf  Inf  Inf  Inf  Inf  Inf;
        Inf 0.0 Inf  Inf  Inf  Inf  Inf  Inf  Inf  Inf;
        0.1 Inf 0.0  0.11 Inf  Inf  Inf  Inf  Inf  Inf;
        Inf Inf 0.11 0.0  0.86 0.86 Inf  Inf  Inf  Inf;
        Inf Inf Inf  0.86 0.0  0.67 0.72 1.93 Inf  Inf;
        Inf Inf Inf  0.86 0.67 0.0  0.72 Inf  Inf  Inf;
        Inf Inf Inf  Inf  0.72 0.72 0.0  Inf  0.72 Inf;
        Inf Inf Inf  Inf  1.93 Inf  Inf  0.0  3.91 Inf;
        Inf Inf Inf  Inf  Inf  Inf  0.72 3.91 0.0  0.52;
        Inf Inf Inf  Inf  Inf  Inf  Inf  Inf  0.52 0.0]

    hclu = hclust(mdist, linkage=linkage)
    @test Clustering.nnodes(hclu) == 10
    @test Clustering.nmerges(hclu) == 9
    @test Clustering.height(hclu) == Inf
    cut1 = @inferred(cutree(hclu, h=1))
    @test cut1 isa Vector{Int}
    @test length(cut1) == Clustering.nnodes(hclu)
    if linkage == :single
        @test cut1 == [1, 2, 1, 1, 1, 1, 1, 3, 1, 1]
    end
end

@testset "cutree(hclust, h=h) when the height of all subtrees greater than h (#141)" begin
    A = [0.0 0.7; 0.7 0.0]
    hA = hclust(A, linkage=:average)
    @test cutree(hA, h=0.5) == [1, 2]
    @test cutree(hA, h=0.7) == [1, 1]
    @test cutree(hA, h=0.9) == [1, 1]
end

end


using Images

function loadDataFromFiles!(fileDir,fileList,Dataset;hist_eq=false)
  for sampleIdx = 1:size(Dataset,2)
    imFile = joinpath(fileDir,fileList[sampleIdx])
    im = load(imFile)
    ImgData = convert(Array{Float64,2},im.data)
    if hist_eq
      ImgData = histogramEqualization(ImgData)
    end
    Dataset[:,sampleIdx] = ImgData[:]
  end
end

function histogramEqualization(im)
  # println("=================EQ====================")
  L  = 256   # Intesnity levels
  Lm = 255   # L-1
  N = length(im)
  inc = 1./N;

  # First, convert everything back to integers [0,255]
  im = floor(im.*255.0)
  # println(im)
  UIim = convert(Array{UInt8,2},im)

  # Next, we want the PMF
  P = zeros(L,1)
  for i=1:N
    P[UIim[i]+1] += inc;
  end

  gim = copy(im)
  # Now apply equalizaitons
  for i=1:N
    gim[i] = Lm.*sum(P[1:(UIim[i]+1)])
  end

  return gim
end

function traindata(;hist_eq=false)
  Pixels  = 19*19

  # Load Face Data
  FaceDataDir = joinpath(dirname(@__FILE__),"..","data","train","face")
  FaceFiles   = readdir(FaceDataDir)
  FaceDataset = Array(Float64,Pixels,length(FaceFiles))
  loadDataFromFiles!(FaceDataDir,FaceFiles,FaceDataset;hist_eq=hist_eq)

  # Load Non-Face Data
  NonFaceDataDir = joinpath(dirname(@__FILE__),"..","data","train","non-face")
  NonFaceFiles = readdir(NonFaceDataDir)
  NonFaceDataset = Array(Float64,Pixels,length(NonFaceFiles))
  loadDataFromFiles!(NonFaceDataDir,NonFaceFiles,NonFaceDataset;hist_eq=hist_eq)

  return FaceDataset, NonFaceDataset
end

function testdata(;hist_eq=false)
  Pixels  = 19*19

  # Load Face Data
  FaceDataDir = joinpath(dirname(@__FILE__),"..","data","test","face")
  FaceFiles   = readdir(FaceDataDir)
  FaceDataset = Array(Float64,Pixels,length(FaceFiles))
  loadDataFromFiles!(FaceDataDir,FaceFiles,FaceDataset;hist_eq=hist_eq)

  # Load Non-Face Data
  NonFaceDataDir = joinpath(dirname(@__FILE__),"..","data","test","non-face")
  NonFaceFiles = readdir(NonFaceDataDir)
  NonFaceDataset = Array(Float64,Pixels,length(NonFaceFiles))
  loadDataFromFiles!(NonFaceDataDir,NonFaceFiles,NonFaceDataset;hist_eq=hist_eq)

  return FaceDataset, NonFaceDataset
end


# Parents: use findall in adjacency matrix, instead of `Graphs.innerneigbhors` method:
children(mat::SparseMatrixCSC, node::Int64) = findall(x -> x != 0, mat[node, :])
children(net::AbstractSimpleWeightedGraph, node::Int64) = Graphs.inneighbors(net, node)

# Parents: use findall in adjacency matrix, instead of `Graphs.outneigbhors` method:
parents(mat::SparseMatrixCSC, node::Int64) = findall(x -> x != 0, mat[:, node])
parents(net::AbstractSimpleWeightedGraph, node::Int64) = Graphs.outneighbors(net, node)

# Parents: use findall in adjacency matrix, instead of `Graphs.all_neigbhors` method:
neighbors(mat::SparseMatrixCSC, node::Int64) = union(children(mat, node), parents(mat, node))
neighbors(net::AbstractSimpleWeightedGraph, node::Int64) = Graphs.all_neighbors(net, node)

function toplogicalsort(net::AbstractSimpleWeightedGraph; method = "graphs")
    if Graphs.is_directed(net)
        if method == "dfs"
            L, acyclic = _topologicalsort_dfs(net)
        elseif method == "kahn"
            L, acyclic = _topologicalsort_kahn(net)
        else
            try # Todo: ooh, that's not valid in all error cases!
                acyclic = true
                L = Graphs.topological_sort_by_dfs(net)
            catch e
                L = []
                acyclic = false
            end
        end
        return L, acyclic
    else
        print("undirected graph or cyclic graph. No topological order")
        return [], false
    end
end

function _topologicalsort_kahn(g::AbstractSimpleWeightedGraph)

    # if there exists a topological order of directed graph g, then it is acyclic
    acyclic = true

    adj = g.weights
    # get set $S$ of all nodes with no incoming edge
    n = nv(g)
    S = Int64[]
    num_parents = zeros(Int, 1, n) # no incoming EDGELABELSIZE
    for i = 1:n
        num_parents[i] = length(WeightedNetwork.parents(adj, i))
        if num_parents[i] == 0
            push!(S, i)
        end
    end

    L = [] # return value

    while !isempty(S)
        node = pop!(S)
        push!(L, node)
        childs = WeightedNetwork.children(adj, node)
        for kid in childs
            adj[node, kid] = 0
            if length(WeightedNetwork.parents(adj, kid)) == 0
                push!(S, kid)
            end
        end
    end

    if sum(adj) == 0.0 || sum(adj) == 0
        acyclic = true
        return L, acyclic
    else
        acyclic = false
        print("No tolpological order. This is not a DAG. The direct graph is cyclic!")
        return [], acyclic
    end
end

function _topologicalsort_dfs(g::AbstractSimpleWeightedGraph)

    # if there exists a topological order of directed graph g, then it is acyclic
    acyclic = true

    mat = g.weights
    L = Int64[]
    n = nv(g)
    marked = zeros(Int, 1, n)

    # choose node with most children:
    unmarked_node = marked[1]

    function visit(node)
        if marked[node] == 2 # 2 = permanent_mark
            return
        elseif marked[node] == 1 # 1 = temporary mark
            acyclic = false # then this unacyclic and not a DAG.
            print("No tolpological order. This is not a DAG. The direct graph is cyclic!")
            return [], acyclic
        end

        marked[node] = 1

        childs = WeightedNetwork.children(g, node)

        for kid in childs
            visit(kid)
        end

        marked[node] = 2
        push!(L, node)
    end

    i = 1
    while !isnothing(unmarked_node)
        unmarked_node = findfirst(isequal(0), marked)
        if !isnothing(unmarked_node)
            visit(unmarked_node[2])
        end
    end

    return reverse(L), acyclic
end

# 8-May-1998	 4:44 PM	ATC	Created under MATLAB 5.1.0.421
# ATC = Ali Taylan Cemgil,
# SNN - University of Nijmegen, Department of Medical Physics and Biophysics
# modified for Julia, Röttgermann, 12/2021
# Test-only
function _toposort(g::AbstractSimpleWeightedGraph)
    adj = g.weights
    N = nv(g)
    indeg = sum(g, 2)
    outdeg = sum(g, 1)

    seq = Int64[]

    for i = 1:N
        idx = findall(x -> x == 0, indeg)

        if isempty(idx)
            seq = []
            break
        end

        max, idx_max = findmax(outdeg[idx])
        indx = idx[idx_max]
        push!(seq, indx)

        indeg[indx] = -1
        idx = findall(x -> x > 0, adj[indx, :])
        !isempty(idx) ? indeg[idx] = indeg[idx] .- 1 : nothing
        @show seq
    end

    return vec(seq)'

end

function is_tree(mat::SparseMatrixCSC)
    # zusammenhängend
    error("Is not implemented yet!")
    #
end

is_tree(g::SimpleWeightedDiGraph) = is_tree(g.weights)
is_tree(g::SimpleWeightedGraph) = is_tree(g.weights)

function is_symmetric(M::SparseMatrixCSC)
    sz = size(M)

    if sz[1] != sz[2]
        return false
    else

        if M == M'
            return true
        else
            return false
        end
    end
end

is_symmetric(g::AbstractSimpleWeightedGraph) = is_symmetric(g.weights)


@testset "DataMatrix" begin
    a = @data [1.0, 2.0, 3.0]
    v_a = [1.0, 2.0, 3.0]

    b = @data eye(3, 3)
    m_b = eye(3, 3)

    #
    # Transposes
    #

    @test all(a' .== v_a')
    @test all(a'' .== v_a'')
    @test all(b' .== m_b')
    @test all(b'' .== m_b'')

    #
    # DataVector * DataMatrix
    #

    # TODO: Get indexing for b[1, :] to work
    # @test all(a * b[1, :] .== v_a * m_b[1, :])

    #
    # DataMatrix * DataVector
    #

    @test all(b * a .== m_b * v_a)
    @test all(convert(Array, b * a) .== m_b * v_a)

    #
    # DataMatrix * DataMatrix
    #

    @test all(b * b .== m_b * m_b)

    #
    # DataVector * DataMatrix w/ missings
    #

    b[1, 1] = missing
    res = a * b[1:1, :]
    @test all(ismissing.(res[:, 1]))
    @test all(.!(ismissing.(res[:, 2])))
    @test all(.!(ismissing.(res[:, 3])))
    res = a * b[2:2, :]
    @test all(.!(ismissing.(res)))

    #
    # DataMatrix w missings * DataVector
    #

    res = b * a
    @test ismissing.(res[1])
    @test .!(ismissing.(res[2]))
    @test .!(ismissing.(res[3]))

    #
    # DataMatrix * DataMatrix
    #

    res = b * b
    # 3x3 Float64 DataMatrix:
    #  missing   missing   missing
    #  missing  1.0  0.0
    #  missing  0.0  1.0
    @test ismissing.(res[1, 1])
    @test ismissing.(res[1, 2])
    @test ismissing.(res[1, 3])
    @test ismissing.(res[2, 1])
    @test .!(ismissing.(res[2, 2]))
    @test .!(ismissing.(res[2, 3]))
    @test ismissing.(res[3, 1])
    @test .!(ismissing.(res[3, 2]))
    @test .!(ismissing.(res[3, 3]))

    res = b * @data eye(3)
    # 3x3 Float64 DataMatrix:
    #   missing   missing   missing
    #  0.0  1.0  0.0
    #  0.0  0.0  1.0
    @test ismissing.(res[1, 1])
    @test ismissing.(res[1, 2])
    @test ismissing.(res[1, 3])
    @test .!(ismissing.(res[2, 1]))
    @test .!(ismissing.(res[2, 2]))
    @test .!(ismissing.(res[2, 3]))
    @test .!(ismissing.(res[3, 1]))
    @test .!(ismissing.(res[3, 2]))
    @test .!(ismissing.(res[3, 3]))

    res = (@data eye(3)) * b
    # julia> dataeye(3) * b
    # 3x3 Float64 DataMatrix:
    #  missing  0.0  0.0
    #  missing  1.0  0.0
    #  missing  0.0  1.0
    @test ismissing.(res[1, 1])
    @test .!(ismissing.(res[1, 2]))
    @test .!(ismissing.(res[1, 3]))
    @test ismissing.(res[2, 1])
    @test .!(ismissing.(res[2, 2]))
    @test .!(ismissing.(res[2, 3]))
    @test ismissing.(res[3, 1])
    @test .!(ismissing.(res[3, 2]))
    @test .!(ismissing.(res[3, 3]))

    # Test row operations
    dm = @data eye(6, 2)
    # rowmeans(dm)

    # Test column operations
    dm = @data eye(6, 2)
    # colmeans(dm)

    # Test linear algebra
    du, dd, dv = svd((@data eye(3, 3)))
    u, d, v = svd(eye(3, 3))
    @test all(du .== u)
    @test all(dd .== d)
    @test all(dv .== v)

    # Test elementary functions
    dm = -(@data eye(5, 5))
    @test all(abs(dm) .== eye(5, 5))
end


type FiveStepFunctionArms <: BanditEnvironment
	shift_time::Int64
end

function arms(env::FiveStepFunctionArms, trial::Int64)
  if trial > env.shift_time
    [
      BernoulliArm(0.1),
      BernoulliArm(0.2),
      BernoulliArm(0.3),
      BernoulliArm(0.4),
      BernoulliArm(0.5)
    ]
  else
    [
      BernoulliArm(0.5),
      BernoulliArm(0.4),
      BernoulliArm(0.3),
      BernoulliArm(0.2),
      BernoulliArm(0.1)
    ]
  end
end

function n_arms(env::FiveStepFunctionArms)
  5
end


module Polygons

using Compat
using Compat.Statistics: mean
using Compat: reverse
import Base: length, show, isinf

export Polygon,vertex,interiorangle,isinpoly,naca4

"""
    Polygon(x::Vector{Float64}, y::Vector{Float64})

A polygon defined by its vertices, which must be provided in
counter-clockwise order.

# Example

```jldoctest
julia> p = Polygon([-1.0,0.2,1.0,-1.0],[-1.0,-1.0,0.5,1.0])
Polygon with 4 vertices at
             (-1.0,-1.0) (0.2,-1.0) (1.0,0.5) (-1.0,1.0)
             interior angles/π = [0.5, 0.656, 0.422, 0.422]
```
"""
struct Polygon
  vert :: Vector{ComplexF64}
  angle :: Vector{Float64}

  Polygon(vert,angle) = abs(vert[end]-vert[1])<eps() ? new(vert[1:end-1],angle) : new(vert,angle)
end

Polygon(x::T,y::T,angle) where T = Polygon(x+im*y,angle)

Polygon(x::T,y::T) where T = Polygon(x+im*y,interiorangle(x+im*y))

"""
    Polygon(w::Vector{Complex128})

Sets up a polygon with the coordinates of the vertices specified
with complex vector `w`. As usual, these must be supplied in
counter-clockwise order.

# Example

```jldoctest
julia> p = Polygon([-1.0-1.0im,0.2-1.0im,1.0+0.5im,-1.0+1.0im])
Polygon with 4 vertices at
             (-1.0,-1.0) (0.2,-1.0) (1.0,0.5) (-1.0,1.0)
             interior angles/π = [0.5, 0.656, 0.422, 0.422]
```
"""
Polygon(w::Vector{ComplexF64}) = Polygon(w,interiorangle(w))

"""
    vertex(p::Polygon) -> Vector{Complex128}

Returns the vector of vertices of the polygon `p`, in complex form.

# Example

```jldoctest
julia> p = Polygon([-1.0,0.2,1.0,-1.0],[-1.0,-1.0,0.5,1.0]);

julia> vertex(p)
4-element Array{Complex{Float64},1}:
 -1.0-1.0im
  0.2-1.0im
  1.0+0.5im
 -1.0+1.0im
```
"""
vertex(p::Polygon) = p.vert

"""
    isinf(p::Polygon) -> Bool

Returns `true` if any vertex in polygon `p` is at infinity.

# Example

```jldoctest
julia> p = Polygon([-1.0,0.2,1.0,-1.0],[-1.0,-1.0,0.5,1.0]);

julia> isinf(p)
false
```
"""
Base.isinf(p::Polygon) = any(isinf.(vertex(p)))

"""
    length(p::Polygon) -> Integer

Returns the number of vertices of the polygon `p`.

# Example

```jldoctest
julia> p = Polygon([-1.0,0.2,1.0,-1.0],[-1.0,-1.0,0.5,1.0]);

julia> length(p)
4
```
"""
Base.length(p::Polygon) = length(vertex(p))

"""
    interiorangle(p::Polygon) -> Vector{Float64}

Returns the vector of interior angles (divided by \$\\pi\$) of the polygon `p`.

# Example

```jldoctest
julia> p = Polygon([-1.0,0.2,1.0,-1.0],[-1.0,-1.0,0.5,1.0]);

julia> interiorangle(p)
4-element Array{Float64,1}:
 0.5
 0.655958
 0.422021
 0.422021
```
"""
interiorangle(p::Polygon) = length(p.angle) != 0 ? p.angle : interiorangle(p.vertex)

function interiorangle(w::Vector{ComplexF64})
  if length(w)==0
    return []
  end

  atinf = isinf.(w)
  mask = .~(atinf .| circshift(atinf,-1) .| circshift(atinf,1))

  dw = w - circshift(w,1)
  dwshift = circshift(dw,-1)
  beta = fill(NaN,length(w))
  beta[mask] = mod.(angle.( -dw[mask].*conj.(dwshift[mask]) )/π,2)

  mods = abs.(beta .+ 1) .< 1e-12
  beta[mods] = fill!(similar(beta[mods]), 1)

  return beta

end

function isinpoly(z::ComplexF64,w::Vector{ComplexF64},beta::Vector{Float64},tol)

  index = 0.0

  scale = mean(abs.(diff(circshift(w,-1))))
  if ~any(scale > eps())
        return
  end
  w = w/scale
  z = z/scale
  d = w .- z
  d[abs.(d) .< eps()] .= eps()
  ang = angle.(circshift(d,-1)./d)/π
  tangents = sign.(circshift(w,-1)-w)

  # Search for repeated points and fix these tangents
  for p = findall( tangents .== 0 )
    v = [w[p+1:end];w]
    tangents[p] = sign(v[findfirst(v.!=w[p])]-w[p])
  end

  # Points close to an edge
  onbdy = abs.(imag.(d./tangents)) .< 10*tol
  onvtx = abs.(d) .< tol
  onbdy = onbdy .& ( (abs.(ang) .> 0.9) .| onvtx .| circshift(onvtx,-1) )

  if ~any(onbdy)
    index = round(sum(ang)/2)
  else
    S = sum(ang[.~onbdy])
    b = beta[onvtx]
    augment = sum(onbdy) - sum(onvtx) - sum(b)

    index = round(augment*sign(S) + S)/2
  end
  return index==1

end

isinpoly(z,w,beta) = isinpoly(z,w,beta,eps())

"""
    isinpoly(z::Complex128,p::Polygon) -> Bool

Returns `true` or `false` depending on whether `z` is inside
or outside polygon `p`.

# Example

```jldoctest
julia> p = Polygon([-1.0,0.2,1.0,-1.0],[-1.0,-1.0,0.5,1.0]);

julia> isinpoly(0.0+0.0im,p)
true

julia> isinpoly(1.0+2.0im,p)
false
```
"""
isinpoly(z,p::Polygon) = isinpoly(z,p::Polygon,eps())

"""
    isinpoly(z::Complex128,p::Polygon,tol::Float64) -> Bool

Returns `true` if `z` is inside or within distance `tol` of polygon `p`.

# Example

```jldoctest
julia> p = Polygon([-1.0,0.2,1.0,-1.0],[-1.0,-1.0,0.5,1.0]);

julia> isinpoly(-1.01+0.0im,p)
false

julia> isinpoly(-1.01+0.0im,p,1e-2)
true
```
"""
isinpoly(z,p::Polygon,tol) = isinpoly(z,p.vert,p.angle,tol)

winding(z,x...) = float.(isinpoly(z,x...))

#=  some specific shape families =#

"""
    naca4(cam,pos,t[;np=20][,Zc=0.0+0.0im][,len=1.0]) -> Vector{Complex128}

Generates the vertices of a NACA 4-digit airfoil of chord length 1. The
relative camber is specified by `cam`, the position of
maximum camber (as fraction of chord) by `pos`, and the relative thickness
by `t`.

The optional parameter `np` specifies the number of points on the upper
or lower surface. The optional parameter `Zc` specifies the mean position of
the vertices (which is set to the origin by default). The optional parameter
`len` specifies the chord length.

# Example

```jldoctest
julia> w = naca4(0.0,0.0,0.12);

julia> p = Polygon(w);
```
"""
function naca4(cam::Number,pos::Number,t::Number;np=20,Zc=0.0+0.0im,len=1.0)

# Here, cam is the fractional camber, pos is the fractional chordwise position
# of max camber, and t is the fractional thickness.

npan = 2*np-2

# Trailing edge bunching
an = 1.5
anp = an+1
x = zeros(np)

θ = zero(x)
yc = zero(x)

for j = 1:np
    frac = Float64((j-1)/(np-1))
    x[j] = 1 - anp*frac*(1-frac)^an-(1-frac)^anp;
    if x[j] < pos
        yc[j] = cam/pos^2*(2*pos*x[j]-x[j]^2)
        if pos > 0
            θ[j] = atan(2*cam/pos*(1-x[j]/pos))
        end
    else
        yc[j] = cam/(1-pos)^2*((1-2*pos)+2*pos*x[j]-x[j]^2)
        if pos > 0
            θ[j] = atan(2*cam*pos/(1-pos)^2*(1-x[j]/pos))
        end
    end
end

xu = zero(x)
yu = xu
xl = xu
yl = yu

yt = t/0.20*(0.29690*sqrt.(x)-0.12600*x-0.35160*x.^2+0.28430*x.^3-0.10150*x.^4)

xu = x-yt.*sin.(θ)
yu = yc+yt.*cos.(θ)

xl = x+yt.*sin.(θ)
yl = yc-yt.*cos.(θ)

rt = 1.1019*t^2;
θ0 = 0
if pos > 0
    θ0 = atan(2*cam/pos)
end
# Center of leading edge radius
xrc = rt*cos(θ0)
yrc = rt*sin(θ0)
θle = collect(0:π/50:2π)
xlec = xrc .+ rt*cos.(θle)
ylec = yrc .+ rt*sin.(θle)

# Assemble data
coords = [xu yu xl yl x yc]
cole = [xlec ylec]

# Close the trailing edge
xpanold = [0.5*(xl[np]+xu[np]); reverse(xl[2:np-1], dims = 1); xu[1:np-1]]
ypanold = [0.5*(yl[np]+yu[np]); reverse(yl[2:np-1], dims = 1); yu[1:np-1]]

xpan = zeros(npan)
ypan = zeros(npan)
for ipan = 1:npan
    if ipan < npan
        xpan1 = xpanold[ipan]
        ypan1 = ypanold[ipan]
        xpan2 = xpanold[ipan+1]
        ypan2 = ypanold[ipan+1]
    else
        xpan1 = xpanold[npan]
        ypan1 = ypanold[npan]
        xpan2 = xpanold[1]
        ypan2 = ypanold[1]
    end
    xpan[ipan] = 0.5*(xpan1+xpan2)
    ypan[ipan] = 0.5*(ypan1+ypan2)
end
w = ComplexF64[1;reverse(xpan, dims = 1)+im*reverse(ypan,dims = 1)]*len
w .-= sum(w)/length(w)
return w .+ Zc

end

function show(io::IO, p::Polygon)
    println(io, "Polygon with $(length(p.vert)) vertices at")
    print(io,   "             ")
    for i = 1:length(p.vert)
        print(io,"($(real(p.vert[i])),$(imag(p.vert[i]))) ")
    end
    println(io)
    println(io, "             interior angles/π = $(round.(p.angle, digits=3))")
end

end