using SpecialFunctions: erfc
using StaticArrays
using LinearAlgebra

const Mat3{T} = SMatrix{3, 3, T, 9} where T
const Vec3{T} = SVector{3, T} where T

"""
Compute the electrostatic interaction energy per unit cell between point
charges in a uniform background of compensating charge to yield net
neutrality. the `lattice` and `recip_lattice` should contain the
lattice and reciprocal lattice vectors as columns. `charges` and
`positions` are the point charges and their positions (as an array of
arrays) in fractional coordinates. If `forces` is not nothing, minus the derivatives
of the energy with respect to `positions` is computed.
"""
function energy_ewald_handwritten(lattice, charges, positions; η=nothing, forces=nothing)
    T = eltype(lattice)

    for i=1:3
        if norm(lattice[:,i]) == 0
            ## TODO should something more clever be done here? For now
            ## we assume that we are not interested in the Ewald
            ## energy of non-3D systems
            return T(0)
        end
    end
    energy_ewald_handwritten(lattice, T(2π) * inv(lattice'), charges, positions; η=η, forces=forces)
end

function energy_ewald_handwritten(lattice, recip_lattice, charges, positions; η=nothing, forces=nothing)
    T = eltype(lattice)
    @assert T == eltype(recip_lattice)
    @assert length(charges) == length(positions)
    if η === nothing
        # Balance between reciprocal summation and real-space summation
        # with a slight bias towards reciprocal summation
        η = sqrt(sqrt(T(1.69) * norm(recip_lattice ./ 2T(π)) / norm(lattice))) / 2
    end
    if forces !== nothing
        @assert size(forces) == size(positions)
        forces_real = copy(forces)
        forces_recip = copy(forces)
    end

    #
    # Numerical cutoffs
    #
    # The largest argument to the exp(-x) function to obtain a numerically
    # meaningful contribution. The +5 is for safety.
    max_exponent = -log(eps(T)) + 5

    # The largest argument to the erfc function for various precisions.
    # To get an idea:
    #   erfc(5) ≈ 1e-12,  erfc(8) ≈ 1e-29,  erfc(10) ≈ 2e-45,  erfc(14) ≈ 3e-87
    max_erfc_arg = get(
        Dict(Float32 => 5, Float64 => 8, BigFloat => 14),
        T,
        something(findfirst(arg -> 100 * erfc(arg) < eps(T), 1:100), 100) # fallback for not yet implemented cutoffs
    )

    #
    # Reciprocal space sum
    #
    # Initialize reciprocal sum with correction term for charge neutrality
    sum_recip = - (sum(charges)^2 / 4η^2)

    # Function to return the indices corresponding
    # to a particular shell
    # TODO switch to an O(N) implementation
    function shell_indices(ish)
        (Vec3(i,j,k) for i in -ish:ish for j in -ish:ish for k in -ish:ish
         if maximum(abs.((i,j,k))) == ish)
        # Iterators.flatten(
        #     (
        #         # 8 sides
        #         (Vec3(i,j,k) for i in -(ish-1):(ish-1), j in -(ish-1):(ish-1), k in (-ish,ish)),
        #         (Vec3(i,j,k) for j in -(ish-1):(ish-1), k in -(ish-1):(ish-1), i in (-ish,ish)),
        #         (Vec3(i,j,k) for k in -(ish-1):(ish-1), i in -(ish-1):(ish-1), j in (-ish,ish)),
        #         # 12 edges
        #         (Vec3(i,j,k) for i in (-ish,ish), j in (-ish,ish), k in -(ish-1):(ish-1)),
        #         (Vec3(i,j,k) for j in (-ish,ish), k in (-ish,ish), i in -(ish-1):(ish-1)),
        #         (Vec3(i,j,k) for k in (-ish,ish), i in (-ish,ish), j in -(ish-1):(ish-1)),
        #         # 8 corners
        #         (Vec3(i,j,k) for i in (-ish,ish), j in (-ish,ish), k in (-ish,ish))
        #     )
        # )
    end

    # Loop over reciprocal-space shells
    gsh = 1 # Exclude G == 0
    any_term_contributes = true
    while any_term_contributes
        any_term_contributes = false

        # Compute G vectors and moduli squared for this shell patch
        for G in shell_indices(gsh)
            Gsq = sum(abs2, recip_lattice * G)

            # Check if the Gaussian exponent is small enough
            # for this term to contribute to the reciprocal sum
            exponent = Gsq / 4η^2
            if exponent > max_exponent
                continue
            end

            cos_strucfac = sum(Z * cos(2T(π) * dot(r, G)) for (r, Z) in zip(positions, charges))
            sin_strucfac = sum(Z * sin(2T(π) * dot(r, G)) for (r, Z) in zip(positions, charges))
            sum_strucfac = cos_strucfac^2 + sin_strucfac^2

            any_term_contributes = true
            sum_recip += sum_strucfac * exp(-exponent) / Gsq

            if forces !== nothing
                for (ir, r) in enumerate(positions)
                    Z = charges[ir]
                    dc = -Z*2T(π)*G*sin(2T(π) * dot(r, G))
                    ds = +Z*2T(π)*G*cos(2T(π) * dot(r, G))
                    dsum = 2cos_strucfac*dc + 2sin_strucfac*ds
                    forces_recip[ir] -= dsum * exp(-exponent)/Gsq
                end
            end
        end
        gsh += 1
    end
    # Amend sum_recip by proper scaling factors:
    sum_recip *= 4T(π) / abs(det(lattice))
    if forces !== nothing
        forces_recip .*= 4T(π) / abs(det(lattice))
    end

    #
    # Real-space sum
    #
    # Initialize real-space sum with correction term for uniform background
    sum_real = -2η / sqrt(T(π)) * sum(Z -> Z^2, charges)

    # Loop over real-space shells
    rsh = 0 # Include R = 0
    any_term_contributes = true
    while any_term_contributes || rsh <= 1
        any_term_contributes = false

        # Loop over R vectors for this shell patch
        for R in shell_indices(rsh)
            for i = 1:length(positions), j = 1:length(positions)
                ti = positions[i]
                Zi = charges[i]
                tj = positions[j]
                Zj = charges[j]

                # Avoid self-interaction
                if rsh == 0 && ti == tj
                    continue
                end

                v = lattice * (ti - tj - R)
                dist = norm(v)

                # erfc decays very quickly, so cut off at some point
                if η * dist > max_erfc_arg
                    continue
                end

                any_term_contributes = true
                energy_contribution = Zi * Zj * erfc(η * dist) / dist
                sum_real += energy_contribution
                if forces !== nothing
                    # grad = ForwardDiff.gradient(r -> (dist=norm(lattice * (r - tj - R)); Zi * Zj * erfc(η * dist) / dist), ti)
                    ddist = Zi * Zj * η * (-2exp(-(η*dist)^2) / sqrt(T(π)))
                    ddist += -energy_contribution
                    grad = lattice'*((ddist / dist^2) * v)
                    forces_real[i] += -grad
                    forces_real[j] += grad
                end
            end # i,j
        end # R
        rsh += 1
    end
    energy = (sum_recip + sum_real) / 2  # Divide by 2 (because of double counting)
    if forces !== nothing
        forces .= (forces_recip .+ forces_real) ./ 2
    end
    energy
end


# Networks found by Bose-Nelson algorithm, except for N=16
# http://pages.ripco.net/~jgamble/nw.html
# http://www.cs.brandeis.edu/~hugues/sorting_networks.html
# https://github.com/JeffreySarnoff/SortingNetworks.jl
sorting_network_parameters = Dict(
    2 => (((1,2),),),
    4 => (((1,2), (3,4)), ((1,3), (2,4)), ((2,3),)),
    8 => (
        ((1, 2), (3, 4), (5, 6), (7, 8)), ((1, 3), (2, 4), (5, 7), (6, 8)),
        ((2, 3), (6, 7), (1, 5), (4, 8)), ((2, 6), (3, 7)),
        ((2, 5), (4, 7)), ((3, 5), (4, 6)), ((4, 5),)
    ),
    16 => (
        ((4, 11), (12, 15), (5, 14), (3, 13), (1, 7), (9, 10), (2,8), (6, 16)),
        ((1, 2), (3, 5), (7, 8), (13, 14), (4, 6), (9, 12), (11, 16), (10, 15)),
        ((1, 4), (7, 11), (14, 15), (2, 6), (8, 16), (3, 9), (10, 13), (5, 12)),
        ((1, 3), (8, 14), (15, 16), (2, 5), (6, 12), (4, 9), (11, 13), (7, 10)),
        ((2, 3), (4, 7), (8, 9), (12, 14), (6, 10), (13, 15), (5, 11)),
        ((3, 7), (12, 13), (2, 4), (6, 8), (9, 10), (14, 15)),
        ((3, 4), (5, 7), (11, 12), (13, 14)),
        ((9, 11), (5, 6), (7, 8), (10, 12)),
        ((4, 5), (6, 7), (8, 9), (10, 11), (12, 13)),
        ((7, 8), (9, 10))
    ),
    32 => (
        ((1, 2), (3, 4), (5, 6), (7, 8), (9, 10), (11, 12), (13, 14), (15, 16), (17, 18), (19, 20), (21, 22), (23, 24),
         (25, 26), (27, 28), (29, 30), (31, 32)),
        ((1, 3), (2, 4), (5, 7), (6, 8), (9, 11), (10, 12), (13, 15), (14, 16), (17, 19), (18, 20), (21, 23), (22, 24),
         (25, 27), (26, 28), (29, 31), (30, 32)),
        ((2, 3), (6, 7), (1, 5), (4, 8), (10, 11), (14, 15), (9, 13), (12, 16), (18, 19), (22, 23), (17, 21), (20, 24),
         (26, 27), (30, 31), (25, 29), (28, 32)),
        ((2, 6), (3, 7), (10, 14), (11, 15), (1, 9), (8, 16), (18, 22), (19, 23), (26, 30), (27, 31), (17, 25),
         (24, 32)),
        ((2, 5), (4, 7), (10, 13), (12, 15), (18, 21), (20, 23), (26, 29), (28, 31), (1, 17), (16, 32)),
        ((3, 5), (4, 6), (11, 13), (12, 14), (2, 10), (7, 15), (19, 21), (20, 22), (27, 29), (28, 30), (18, 26), (23, 31)),
        ((4, 5), (12, 13), (2, 9), (3, 11), (6, 14), (8, 15), (20, 21), (28, 29), (18, 25), (19, 27), (22, 30), (24, 31)),
        ((4, 12), (3, 9), (5, 13), (8, 14), (20, 28), (19, 25), (21, 29), (24, 30), (2, 18), (15, 31)),
        ((4, 11), (6, 13), (20, 27), (22, 29), (2, 17), (3, 19), (14, 30), (16, 31)),
        ((4, 10), (7, 13), (20, 26), (23, 29), (3, 17), (16, 30)),
        ((4, 9), (8, 13), (6, 10), (7, 11), (20, 25), (24, 29), (22, 26), (23, 27)),
        ((5, 9), (8, 12), (21, 25), (24, 28), (4, 20), (13, 29)),
        ((6, 9), (8, 11), (22, 25), (24, 27), (4, 19), (5, 21), (12, 28), (14, 29)),
        ((7, 9), (8, 10), (23, 25), (24, 26), (4, 18), (6, 22), (11, 27), (15, 29)),
        ((8, 9), (24, 25), (4, 17), (6, 21), (7, 23), (10, 26), (12, 27), (16, 29)),
        ((8, 24), (7, 21), (5, 17), (6, 18), (9, 25), (12, 26), (15, 27), (16, 28)),
        ((8, 23), (6, 17), (7, 19), (10, 25), (14, 26), (16, 27)), ((8, 22), (7, 17), (11, 25), (16, 26)),
        ((8, 21), (12, 25)), ((8, 20), (13, 25)), ((8, 19), (14, 25), (12, 20), (13, 21)),
        ((8, 18), (15, 25), (11, 19), (14, 22)), ((8, 17), (16, 25), (10, 18), (12, 19), (14, 21), (15, 23)),
        ((9, 17), (12, 18), (16, 24), (15, 21)), ((10, 17), (16, 23), (14, 18), (15, 19)), ((11, 17), (16, 22)),
        ((12, 17), (16, 21)), ((13, 17), (16, 20)), ((14, 17), (16, 19)), ((15, 17), (16, 18)), ((16, 17),)
    )
)


using Test
# Toy dataset.
# Format: each row is an example.
# The last column is the label name, a string.
# The first two columns are features.
# Feel free to play with it by adding more features & examples.
# Interesting note: I've written this so the 2nd and 5th examples
# have the same features, but different labels - so we can see how the
# tree handles this case.

"""
Find the unique values for a column in a dataset.
Could make the result Union{Set{Real}, Set{AbstractString}}???
"""
function unique_vals(rows::Data, col::Int)::Set{Feature}
    Set([row[col] for row in rows])
end

training_data = Data([
    ["Green", 3, "Apple"],
    ["Yellow", 3, "Apple"],
    ["Red", 1, "Grape"],
    ["Red", 1, "Grape"],
    ["Yellow", 3, "Lemon"],
])

"""
    Column labels.
    These are used only to print the tree.
"""
const header = ["color", "diameter", "label"]
Base.show(io::IO, question::Question) = CART.show(io, question, header)
#######
# Demo:
# Let"s look at some example to understand how Gini Impurity works.
#
# First, we"ll look at a dataset with no mixing.
@testset "Gini" begin
    no_mixing = Data([["Apple"], ["Apple"]])
    # this will return 0
    @test gini(no_mixing) == 0.0
    # Now, we"ll look at dataset with a 50:50 apples:oranges ratio
    some_mixing = Data([["Apple"], ["Orange"]])
    # this will return 0.5 - meaning, there"s a 50% chance of misclassifying
    # a random example we draw from the dataset.
    @test gini(some_mixing) == 0.5

    # Now, we'll look at a dataset with many different labels
    lots_of_mixing = Data([["Apple"],
                           ["Orange"],
                           ["Grape"],
                           ["Grapefruit"],
                           ["Blueberry"]])
    # This will return 0.8
    @test gini(lots_of_mixing) ≈ 0.8
end
#######

#######
# Demo:
# Calculate the uncertainy of our training data.
@testset "Information gain" begin
    current_uncertainty = gini(training_data)
    @test current_uncertainty ≈ 0.6399999999999999

    # How much information do we gain by partioning on "Green"?
    true_rows, false_rows = partition(training_data, Question(1, "Green"))
    gain = info_gain(true_rows, false_rows, current_uncertainty)
    @test gain ≈ 0.1399999999999999

    # What about if we partioned on "Red" instead?
    true_rows, false_rows = partition(training_data, Question(1,"Red"))
    gain = info_gain(true_rows, false_rows, current_uncertainty)
    @test gain ≈ 0.37333333333333324
end

@testset "Best split" begin
end

@testset "Print tree" begin
    my_tree = build_tree(training_data)
    print_tree(my_tree)
end


function mincut(graph, source, sink, interdiction, ::Val{:raw})
    m = model(; kind=:cut)
    n = size(graph, 1)

    d = domain(0:n)

    separator = n + 1 # value that separate the two sides of the cut

    # Add variables:
    foreach(_ -> variable!(m, d), 0:n)

    # Extract error function from usual_constraint
    e1 = (x; param=nothing, dom_size=n + 1) -> error_f(
        usual_constraints[:ordered])(x; param, dom_size
    )
    e2 = (x; param=nothing, dom_size=n + 1) -> error_f(
        usual_constraints[:all_different])(x; param, dom_size
    )

    # Add constraint
    constraint!(m, e1, [source, separator, sink])
    constraint!(m, e2, 1:(n + 1))

    # Add objective
    objective!(m, (x...) -> o_mincut(graph, x...; interdiction))

    return m
end

function mincut(graph, source, sink, interdiction, ::Val{:JuMP})
    m = JuMP.Model(CBLS.Optimizer)
    n = size(graph, 1)
    separator = n + 1

    @variable(m, 0 ≤ X[1:separator] ≤ n, Int)

    @constraint(m, [X[source], X[separator], X[sink]] in Ordered())
    @constraint(m, X in AllDifferent())

    obj(x...) = o_mincut(graph, x...; interdiction)
    @objective(m, Min, ScalarFunction(obj))

    return m, X
end

"""
    mincut(graph; source, sink, interdiction =0, modeler = :JuMP)

Compute the minimum cut of a graph.

# Arguments:
- `graph`: Any matrix <: AbstractMatrix that describes the capacities of the graph
- `source`: Id of the source node; must be set
- `sink`: Id of the sink node; must be set
- `interdiction`: indicates the number of forbidden links
- `modeler`: Default to `:JuMP`.
"""
function mincut(graph; source, sink, interdiction =0, modeler = :JuMP)
    return mincut(graph, source, sink, interdiction, Val(modeler))
end


using RandomGeometricGraphs, LightGraphs, Test, Statistics

@testset "Boundaries" begin
    nodes = 100
    @test random_geometric_graph(nodes,0) |> edges |> length == 0
    @test random_geometric_graph(nodes,2) |> edges |> length ≈ nodes * (nodes - 1) / 2 atol=0.6
end

@testset "Equivalence to LightGraphs" begin
    nodes = 100
    dimensions = 2
    radius = 0.2
    lg_edges = [euclidean_graph(nodes,dimensions,cutoff=radius)[1] |> edges |> length for _ in 1:1000]
    rgg_edges = [random_geometric_graph(nodes,radius;dim=dimensions) |> edges |> length for _ in 1:1000]

    # This is an abuse of statistics
    @test mean(lg_edges) ≈ mean(rgg_edges) atol=std(lg_edges)
end


include("attributes.jl")
include("mechanical-problems.jl")
include("chemical-problems.jl")
include("power-electronics-problems.jl")
include("power-system-problems.jl")
include("process-design-and-synthesis.jl")

# Mechanical problems
export pressure_vessel,
vibrating_platform,
two_bar_Truss_design_problems,
weldan_beam_design,
disc_brake_design,
speed_reducer_design,
gear_train_design,
car_side_impact_design,
four_bar_plane_truss,
two_bar_plane_truss,
water_tesource_management,
simply_supported_I_beam_design,
gear_box_design,
multiple_disk_clutch_brake_design,
spring_design,
cantilever_beam_design,
bulk_carriers_design,
front_rail_design,
multi_product_batch_plant,
hydro_static_thrust_bearing_design,
crash_energy_management_for_high_speed_train

# chemical problems
export haverlys_pooling_problem,
reactor_network_design,
heat_exchanger_network_design

# Process Design and Synthesis
export process_synthesis,
process_sythesis_and_design,
process_flow_sheeting_problem,
two_reactor_problem,
process_synthesis_problem

# Power electronics Problem
export SOPWM_for_3_level_invereters,
SOPWM_for_5_level_Inverters,
SOPWM_for_7_level_inverters,
SOPWM_for_9_level_inverters,
SOPWM_for_11_level_inverters,
SOPWM_for_13_level_inverters

export sizing_of_single_phase_distributed_generation_1,
 sizing_of_single_phase_distributed_generation_2,
 sizing_of_single_phase_distributed_generation_3,
 sizing_of_single_phase_distributed_generation_4,
 active_and_reactive_power_loss,
 optimal_power_flow_1,
 optimal_power_flow_2,
 optimal_power_flow_3,
 optimal_power_flow_4,
 optimal_power_flow_5,
 optimal_power_flow_6,
 minimization_of_active_power_loss_in_islanded_microgrids_1,
 minimization_of_active_power_loss_in_islanded_microgrids_2,
 minimization_of_active_power_loss_in_islanded_microgrids_3,
 power_distribution_system_planning

export get_RW_MOP_problem


#=
For details, please visit:
https://computationalmindset.com/en/neural-networks/experiments-with-neural-odes-in-julia.html#exp3
=#

using Plots
using Flux
using DiffEqFlux

tbegin = 0.0
tend = 4.0
datasize = 51
dataset_in = range(tbegin, tend, length=datasize)
dataset_out = log.(1 .+ dataset_in) .- 3 * sqrt.(dataset_in)

function neural_network(data_dim)
    fc = FastChain(FastDense(data_dim, 64, swish),
                  FastDense(64, 32, swish),
                  FastDense(32, data_dim))
end

nn = neural_network(1)
theta = initial_params(nn)

predict(t, p) = nn(t', p)'

loss(p) = begin
  yhat = predict(dataset_in, p)
  l = Flux.mse(yhat, dataset_out)
end

learning_rate=1e-2
opt = ADAMW(learning_rate)
epochs = 500

function cb_train(theta, loss)
    println("Loss: ", loss)
    false
end

res_train = DiffEqFlux.sciml_train(
    loss, theta, opt,
    maxiters = epochs,
    cb = cb_train)

y_pred = predict(dataset_in, res_train.minimizer)

pl = plot(
    dataset_in,
    dataset_out,
    linewidth=2, ls=:dash,
    title="MLP by FastChain without Neural ODEs",
    xaxis="t",
    label="original y(t)",
    legend=:topright)
display(pl)

pl = plot!(
    dataset_in,
    y_pred,
    linewidth=1,
    label="predicted y(t)")
display(pl)


# # Porting Yao with Zygote
using Zygote, YaoExtensions, Yao
using Random; Random.seed!(2)

# first include this patch file in your working directory.
include("zygote_patch.jl")

# ## setup a circuit
# wrap differentiable nodes with `Diff`, please try `print(c)` to see the marked nodes with `[∂]`.
c = variational_circuit(5) |> autodiff(:BP)
dispatch!(c, :random)

# ## define a loss function
h = rand_hermitian(1<<5)
function loss2(reg::ArrayReg, params::Vector{T}, circuit) where T
    reg = dispatch_diff_apply!(copy(reg), circuit, params)
    v = vec(reg.state)
    v' * h *v |> real
end

# ## obtaining the gradient
reg = rand_state(5)
params = rand!(parameters_of_diff(c))
grad_input, grad_params, _ = Zygote.gradient(loss2, reg, params, c)

# ## Check gradients
using Test

# to check gradients of input wave function, we can introduce a one-hot vector
# perturbation to get numerical finite-difference gradient.
η = 1e-5
perturb_input_real = η*product_state(bit"00111")
perturb_input_imag = im*η*product_state(bit"00111")

num_gradient = (loss2(reg + perturb_input_real, params, c) - loss2(reg - perturb_input_real, params, c)) / 2η +
                im*(loss2(reg + perturb_input_imag, params, c) - loss2(reg - perturb_input_imag, params, c)) / 2η
exact_gradient = select(grad_input, bit"00111").state[]
@test num_gradient ≈ exact_gradient

# check the gradient of circuit parameters
index = 1
perturb_params = (p=zeros(length(params)); p[index]=1; p)
loss_pos = loss2(reg, params+η*perturb_params, c)
loss_neg = loss2(reg, params-η*perturb_params, c)
@test grad_params[index] ≈ (loss_pos - loss_neg) / 2η


# process-based-bank-model.jl
# SimLynx.jl Process-Based Discrete-Event Simulation Example

using SimLynx
SimLynx.greet()

using Distributions: Exponential, Uniform
using Random

const N_TELLERS = 2
const MEAN_INTERARRIVAL_TIME = 4.0
const MIN_SERVICE_TIME = 2.0
const MAX_SERVICE_TIME = 10.0

tellers = nothing

"Process to generate n customers arriving into the system."
@process generator(n::Integer) begin
    dist = Exponential(MEAN_INTERARRIVAL_TIME)
    for i = 1:n
        @schedule now customer(i)
        # We don't want to wait after the last customer, which could extend the
        # simulation time past the last customer leaving.
        if i < n
            work(rand(dist))
        end
    end
end

"The ith customer into the system."
@process customer(i::Integer) begin
    dist = Uniform(MIN_SERVICE_TIME, MAX_SERVICE_TIME)
    @with_resource tellers begin
        work(rand(dist))
    end
end

"Run the simulation for n customers."
function run_simulation(n::Integer)
    @assert n > 0
    println("SimLynx.jl Process-Based Discrete-Event Simulation Example")
    println("$N_TELLERS Teller, Single Queue Bank Model")
    println("  Inter-arrival time = Exponential($MEAN_INTERARRIVAL_TIME)")
    println("  Service time = Uniform($MIN_SERVICE_TIME, $MAX_SERVICE_TIME)")
    println("  Number of customers = $n")
    @simulation begin
        global tellers = Resource(N_TELLERS, "tellers")
        @schedule at 0.0 generator(n)
        start_simulation()
        println("--- Simulation results after $(current_time()) minutes ---")
        print_stats(tellers.allocated, title="Teller Allocation Statistics")
        print_stats(tellers.queue_length, title="Teller Queue Length Statistics")
        plot_history(tellers.queue_length, title="Teller Queue Length History")
        print_stats(tellers.wait, title="Teller Queue Wait Time Statistics")
        plot_history(tellers.wait, title="Teller Queue Wait Time History")
    end
end

run_simulation(100)


##
#   File:    qo1.jl
#
#   Purpose: Demonstrate how to solve a quadratic
#            optimization problem using the MOSEK Python API.
##

using Mosek

printstream(msg::String) = print(msg)

# Define a stream printer to grab output from MOSEK

bkc   = [ MSK_BK_LO ]
blc   = [ 1.0 ]
buc   = [ Inf ]

bkx   = [ MSK_BK_LO, MSK_BK_LO, MSK_BK_LO ]
blx   = [ 0.0,  0.0, 0.0 ]
bux   = [ Inf,  Inf, Inf ]
numvar = length(bkx)
numcon = length(bkc)

c     = [ 0.0, -1.0, 0.0 ]
A     = sparse( [ 1, 1, 1 ],
                [ 1, 2, 3 ],
                [ 1.0, 1.0, 1.0 ],
                numcon, numvar )

maketask() do task
    putstreamfunc(task,MSK_STREAM_LOG,printstream)

    # Append 'numcon' empty constraints.
    # The constraints will initially have no bounds.
    appendcons(task,numcon)

    # Append 'numvar' variables.
    # The variables will initially be fixed at zero (x=0).
    appendvars(task,numvar)

    # Set the linear term c_j in the objective.
    putclist(task,[1:numvar;],c)

    # Set the bounds on variable j
    # blx[j] <= x_j <= bux[j]
    putvarboundslice(task,1,numvar+1,bkx,blx,bux)

    putacolslice(task,1,numvar+1,
                 A.colptr[1:numvar],A.colptr[2:numvar+1],
                 A.rowval,A.nzval)

    # Set up and input quadratic objective
    qsubi = [ 1,   2,    3,   3   ]
    qsubj = [ 1,   2,    1,   3   ]
    qval  = [ 2.0, 0.2, -1.0, 2.0 ]

    putqobj(task,qsubi,qsubj,qval)

    putobjsense(task,MSK_OBJECTIVE_SENSE_MINIMIZE)

    # Optimize
    optimize(task)
    # Print a summary containing information
    # about the solution for debugging purposes
    solutionsummary(task,MSK_STREAM_MSG)

    prosta = getprosta(task,MSK_SOL_ITR)
    solsta = getsolsta(task,MSK_SOL_ITR)

    if solsta in     [ MSK_SOL_STA_OPTIMAL,
                      MSK_SOL_STA_NEAR_OPTIMAL ]
        xx = getxx(task,MSK_SOL_ITR)
        println("Optimal solution:")
        print(xx)
    elseif solsta in [ MSK_SOL_STA_DUAL_INFEAS_CER,
                      MSK_SOL_STA_PRIM_INFEAS_CER,
                      MSK_SOL_STA_NEAR_DUAL_INFEAS_CER,
                      MSK_SOL_STA_NEAR_PRIM_INFEAS_CER ]
        println("Primal or dual infeasibility certificate found.\n")
    elseif solsta == MSK_SOL_STA_UNKNOWN
        println("Unknown solution status")
    else
        @printf("Other solution status (%d)\n",solsta)
    end

end


using NCDatasets
using Formatting
using ArgParse
using Statistics
using JSON

include("constants.jl")

function mreplace(x)
    return replace(x, missing=>NaN)
end

function calΨ!(
    ilev  :: Array{Float64, 1},
    lat   :: Array{Float64, 1},
    v     :: Array{Float64, 2},
    Ψ :: Array{Float64, 2},
)

    Ny, Nz = size(v)
    Nz_bnd = length(ilev)
    dp = ilev[1:end-1] - ilev[2:end]

    Ψ[:, Nz_bnd] .= 0.0
    for j=1:Ny
        for k=Nz_bnd-1:-1:1
            Ψ[j, k] = Ψ[j, k+1] + dp[k] * v[j, k]
        end

        Ψ[j, :] .*= 2π * Re * cos(lat[j] * π/180.0) / g
    end

end

function parse_commandline()

    s = ArgParseSettings()
    @add_arg_table s begin
        "--input-data-file-prefix"
            help = "Input data filename prefix including folder and path until the timestamp. File extension `nc` is assumed."
            arg_type = String
            required = true

        "--output-data-file-prefix"
            help = "Output data filename prefix including folder and path until the timestamp. File extension `nc` is assumed."
            arg_type = String
            required = true

        "--domain-file"
            help = "Domain file."
            arg_type = String
            required = true

        "--ilev-varname"
            help = "Variable name of pressure level. Assume in unit of hPa."
            arg_type = String
            default = "ilev"

        "--V-varname"
            help = "Variable name of meridional velocity"
            arg_type = String
            default = "V"

        "--beg-year"
            help = "Year of begin."
            arg_type = Int64
            required = true

        "--end-year"
            help = "Year of end."
            arg_type = Int64
            required = true

        "--overwrite"
            action = :store_true

    end

    return parse_args(ARGS, s)
end

parsed = parse_commandline()
print(json(parsed, 4))

Dataset(parsed["domain-file"], "r") do ds
    global lon = replace(ds["xc"][:, 1], missing=>NaN)
    global lat = replace(ds["yc"][1, :], missing=>NaN)
end

Ψ   = nothing
ilev = nothing
for y=parsed["beg-year"]:parsed["end-year"], m=1:12

    i_filename = format("{:s}{:04d}-{:02d}.nc", parsed["input-data-file-prefix"], y, m)
    o_filename = format("{:s}{:04d}-{:02d}.nc", parsed["output-data-file-prefix"], y, m)

    if !isfile(o_filename) || parsed["overwrite"]

        Dataset(i_filename, "r") do ds

            v = convert(Array{Float64}, nomissing(ds[parsed["V-varname"]][:, :, 1]))

            if ilev == nothing
                global Ny, Nz = size(v)
                global ilev = ( ds[parsed["ilev-varname"]][:] |> mreplace ) * 100.0

                global Ψ = zeros(Float64, Ny, length(ilev))
            end

            calΨ!(ilev, lat, v, Ψ)

            Dataset(o_filename, "c") do ds
                defDim(ds, "Ny", Ny)
                defDim(ds, "ilev", length(ilev))
                defDim(ds, "time", Inf)
                for (varname, vardata, vardim, attrib) in [
                    ("psi", Ψ, ("Ny", "ilev"), Dict()),
                    ("ilev", ilev, ("ilev",), Dict()),
                ]

                    var = defVar(ds, varname, Float64, vardim)
                    var.attrib["_FillValue"] = 1e20

                    var = ds[varname]

                    for (k, v) in attrib
                        var.attrib[k] = v
                    end

                    rng = []
                    for i in 1:length(vardim)-1
                        push!(rng, Colon())
                    end
                    push!(rng, 1:size(vardata)[end])
                    var[rng...] = vardata

                end
            end
            println("Output: ", o_filename)
        end

    end

end


function best_psis(psis::Vector)
  k = length(psis)
  idx = sortperm(psis, rev=true)    # sort, get the indices

  max_id = 1
  max_val = psis[idx[1]]
  for j=2:k
    v = ( sum(psis[idx[1:j]]) + j - 1 ) / j
    if v >= max_val
      max_val = v
      max_id = j
    else
      break
    end
  end

  psi_id = idx[1:max_id]

  return psi_id::Vector{Int}, max_val::Float64
end

function calc_const(psi_list::Vector, psi_id::Vector)
  j = length(psi_id)
  ret = ( sum(psi_list[psi_id]) + j - 1 ) / j
  return ret::Float64
end

function calc_cconst(psi_id::Vector)
  n_psi = length(psi_id)
  ret = (n_psi - 1 ) / n_psi
  return ret::Float64
end

## get featrure indeces for class i
function idi(m::Integer, i::Integer)
  return ((i-1)*m+1 : i*m)::UnitRange{Int64}
end

function psi_list(w::Vector, X::Matrix, y::Vector, i::Integer, c::Integer, idmi::Vector)
  psis = zeros(c)
  yi = y[i]
  for j=1:c
    if j != yi
      v1 = dot(w[idmi[j]], view(X, :, i))
      v2 = dot(w[idmi[yi]], -view(X, :, i))
      psis[j] = v1 + v2
    end
  end

  return psis::Vector{Float64}
end

## no y. for prediction
function psi_list(w::Vector, X::Matrix, i::Integer, c::Integer, idmi::Vector)
  psis = zeros(c)
  for j=1:c
    val = dot(w[idmi[j]], view(X, :, i))
    psis[j] = val
  end

  return psis::Vector{Float64}
end

# in terms of dual variable (useful for kernel methods)
function psi_list_dual(alpha::Vector, LPsi::Matrix, i::Integer, c::Integer)
  LPsi_i = view(LPsi, (i-1)*c+1 : i*c, :)
  psis = -(LPsi_i * alpha)
  return psis::Vector{Float64}
end

# in terms of dual variable (useful for kernel methods) dec
function psi_list_dual(sLa::Vector, i::Integer, c::Integer)
  psis = -(sLa[(i-1)*c+1 : i*c])
  return psis::Vector{Float64}
end

function eta_list(w::Vector, X::Matrix, i::Integer, c::Integer, idmi::Vector)
  etas = zeros(c)
  for j=1:c
    etas[j] = dot(w[idmi[j]], view(X, :, i))
  end

  return etas::Vector{Float64}
end

function calc_dot(key::Tuple{Integer,Vector,Integer,Vector}, K_ij::Float64, y::Vector)

  i = key[1]
  j = key[3]
  psi_i = key[2]
  psi_j = key[4]
  li = length(psi_i)
  lj = length(psi_j)
  yi = y[i]
  yj = y[j]

  mult = 0.0

  inii = yi in psi_i
  inij = yi in psi_j
  inji = yj in psi_i
  injj = yj in psi_j

  if yi == yj
    mult += (li - round(Int, inii)) * (lj - round(Int, injj))
  else
    if inij
      mult -= (li - round(Int, inii))
    end
    if inji
      mult -= (lj - round(Int, injj))
    end
  end

  ii = 1
  ij = 1
  while ii <= li && ij <= lj
    if psi_i[ii] > psi_j[ij]
      ij += 1
    elseif psi_i[ii] < psi_j[ij]
      ii += 1
    else  # equal
      if psi_i[ii] != yi && psi_j[ij] != yj && psi_i[ii] != yj && psi_j[ij] != yi
        mult += 1
      end
      ii += 1
      ij += 1
    end
  end

  d = (mult * K_ij) / (li * lj)

  return d::Float64
end

function calc_dot(key::Tuple{Integer,Vector,Integer,Vector}, K::Matrix, y::Vector)
  i = key[1]
  j = key[3]
  K_ij = K[i,j]

  d = calc_dot(key, K_ij, y)

  return d::Float64
end

function calc_dconst(key::Tuple{Integer,Vector}, X::Matrix, y::Vector, n_c::Integer, idmi::Vector)

    i = key[1]
    psi_i = key[2]
    li = length(psi_i)
    yi = y[i]

    xi = view(X, :, i)

    m = length(xi)
    dc = zeros(m * n_c)

    inii = yi in psi_i
    dc[idmi[yi]] = - ( (li - round(Int, inii)) * xi ) / li

    for ii = 1:li
      if psi_i[ii] != yi
        dc[idmi[psi_i[ii]]] = xi / li
      end
    end

    return dc::Vector{Float64}
end

# calc Lambda * Phi for kernel prediction
function calc_dotlphi(key::Tuple{Integer,Vector,Integer}, K_ij::Float64, y::Vector, n_c::Integer)
  i = key[1]
  psi_i = key[2]
  li = length(psi_i)
  yi = y[i]

  mults = zeros(n_c)
  for ii in psi_i
    if ii == yi
      mults[ii] = -(li - 1.0) / li
    else
      mults[ii] = 1.0 / li
    end
  end

  if !(yi in psi_i)
    mults[yi] = -1.0
  end

  return (mults * K_ij)::Vector{Float64}
end

function calc_dotlphi(key::Tuple{Integer,Vector,Integer}, K::Matrix, y::Vector, n_c::Integer)
  i = key[1]
  j = key[3]
  K_ij = K[i,j]
  return calc_dotlphi(key, K_ij, y, n_c)::Vector{Float64}
end


# This file simulates a closed system

using Causal

# Construct the model
@defmodel model begin
    @nodes begin
        gen = FunctionGenerator(readout=sin)
        adder = Adder(signs=(+,-))
        ds = ContinuousLinearSystem()
        writer = Writer(input=Inport(2))
    end
    @branches begin
        gen[1]      =>  adder[1]
        adder[1]    =>  ds[1]
        ds[1]       =>  adder[2]
        gen[1]      =>  writer[1]
        ds[1]       =>  writer[2]
    end
end

# Simualate the model
sim = simulate!(model, 0., 0.01, 10.)

# Read and plot data
t, x = read(getnode(model, :writer).component)
using Plots
plot(t, x[:, 1], label="r(t)", xlabel="t", lw=3)
plot!(t, x[:, 2], label="y(t)", xlabel="t", lw=3)
plot!(t, 6 / 5 * exp.(-2t) + 1 / 5 * (2 * sin.(t) - cos.(t)), label="Analytical Solution", lw=3)
# fileanme = "readme_example.svg"
# path = joinpath(@__DIR__, "../docs/src/assets/ReadMePlot/")
# savefig(joinpath(path, fileanme))


using LPVSpectral
using Test, LinearAlgebra, Statistics, Random
using Plots, DSP
Random.seed!(0)
# write your own tests here
function generate_signal(f,w,N, modphase=false)
    x = sort(10rand(N)) # Sample points
    v = range(0, stop=1, length=N) # Scheduling variable

    # generate output signal
    # phase_matrix
    dependence_matrix = Float64[f[(i-1)%length(f)+1].(v) for v in v, i in eachindex(w)] # N x nw
    frequency_matrix = [cos(w*x -0.5modphase*(dependence_matrix[i,j])) for (i,x) in enumerate(x), (j,w) in enumerate(w)] # N x nw
    y = sum(dependence_matrix.*frequency_matrix,dims=2)[:] # Sum over all frequencies
    y += 0.1randn(size(y))
    y,v,x,frequency_matrix, dependence_matrix
end

@testset "LPV methods" begin

    N      = 500 # Number of training data points
    f      = [v->2v^2, v->2/(5v+1), v->3exp(-10*(v-0.5)^2),] # Functional dependences on the scheduling variable
    w      = 2pi*[2,10,20] # Frequency vector
    w_test = 2π*collect(2:2:25)
    Y,V,X,frequency_matrix, dependence_matrix = generate_signal(f,w,N,true)

    λ      = 0.02 # Regularization parmater
    normal = true # Use normalized basis functions
    Nv     = 50 # Number of basis functions

    se = ls_spectral_lpv(Y,X,V,w_test,Nv; λ = λ, normalize = normal) # Perform LPV spectral estimation
    windowpsd = ls_windowpsd_lpv(Y,X,V,w_test,Nv; λ = λ, normalize = normal)

    spectrum_lpv   = psd(se)

    si = sortperm(spectrum_lpv[:],rev=true)
    @test Set(si[1:3]) == Set([1,5,10])

    si = sortperm(windowpsd[:],rev=true)
    @test Set(si[1:3]) == Set([1,5,10])
end

@testset "detrend" begin
    tre = [1,2,3]
    @test detrend(tre) == [-1,0,1]
    @test tre == [1,2,3] # Shouldn't have changed
    detrend!(tre)
    @test tre == [-1,0,1] # Should have changed
end

## Test ComplexNormal
@testset "Complex Normal" begin
    n  = 10
    n2 = 5
    a  = randn(n)
    A  = randn(n,n)
    A  = A'A + I
    b  = randn(n2)
    B  = randn(n2,n2)
    B  = B'B + I
    X  = randn(n,n2)
    Y  = randn(n,n2)
    cn = ComplexNormal(X,Y)
    @test size(cn.m) == (n2,)
    @test size(cn.Γ) == (n2,n2)
    @test size(cn.C) == (n2,n2)
    # affine_transform(cn, B, b)
    pdf(cn,b)
    pdf(cn,im*b)

    @test isa(cn_Vxx(A,A), Cholesky)
    @test isa(cn_fVxx(A,A), Matrix)
    @test issymmetric(cn_fVxx(A,A))
    cn_V(A,0.1A)

    cn  = ComplexNormal(a,A)
    cn  = ComplexNormal(im*b,A)

    A   = randn(4,4);
    A   = A'A
    x   = randn(1000,3)
    y   = randn(1000,3)
    cn  = ComplexNormal(x,y)
    z   = rand(cn,1000000);
    cn2 = ComplexNormal(z)
    @test norm(Matrix(cn.Γ)-Matrix(cn2.Γ)) < 0.01
    @test norm(Matrix(cn.C)-Matrix(cn2.C)) < 0.01

end

@testset "ls methods" begin
    T = 1000
    t = 0:T-1
    f = LPVSpectral.default_freqs(t)
    @test f[1] == 0
    @test f[end] == 0.5-1/length(t)/2
    @test length(f) == T÷2

    # f2 = LPVSpectral.default_freqs(t,10)
    # @test f2[1] == 0
    # @test f2[end] == 0.5-1/length(t)/2
    # @test length(f2) == T÷2÷10

    @test LPVSpectral.check_freq(f) == 1
    @test_throws ArgumentError LPVSpectral.check_freq([1,0,2])
    A, z = LPVSpectral.get_fourier_regressor(t,f)
    @test size(A) == (T,2length(f)-1)

    Base.isapprox(t1::Tuple{Float64,Int64}, t2::Tuple{Float64,Int64}; atol) = all(t -> isapprox(t[1],t[2],atol=atol), zip(t1,t2))
    y = sin.(t)
    x,_ = ls_spectral(y,t)
    @test findmax(abs.(x)) ≈ (0.9999773730281, 160) atol=0.001

    W = ones(length(y))
    x,_ = ls_spectral(y,t,f,W)
    @test findmax(abs.(x)) ≈ (0.999977373027, 160) atol=0.001

    x,_ = tls_spectral(y,t)
    @test findmax(abs.(x)) ≈ (0.9999777508878254, 160) atol=0.001

    x,_ = ls_windowpsd(y,t; nw=20)
    @test findmax(abs.(x)) ≈ (1.000783557456378, 160) atol=0.001

    x,_ = ls_windowpsd(y,t)
    @test findmax(abs.(x)) ≈ (1.0011490769234443, 160) atol=0.001

    x,_ = ls_windowcsd(y,y,t)
    @test findmax(abs.(x)) ≈ (1.0011490769234443, 160) atol=0.001

    x,_ = ls_cohere(y,y,t)
    @test findmax(abs.(x))[1] ≈ 1 atol=0.001

    x,_ = ls_cohere(y,y .+ 10randn.(),t)
    @test mean(abs.(x)) ≈ 0.5 atol = 0.15

end

@testset "plots" begin

    y = randn(1000)
    plot(periodogram(y))
    plot(periodogram(filtfilt(ones(4), [4], y)))
    plot(welch_pgram(y))

end

include("test_lasso.jl")


#!/usr/bin/env julia
################################################################################
#
#  Copyright (c) 2015 Wojciech Migda
#  All rights reserved
#  Distributed under the terms of the Apache 2.0 license
#
################################################################################
#
#  Filename: MFCC.jl
#
#  Decription:
#       Mel Frequency Cepstral Coefficients
#
#  Authors:
#       Wojciech Migda
#
################################################################################
#
#  History:
#  --------
#  Date         Who  Ticket     Description
#  ----------   ---  ---------  ------------------------------------------------
#  2015-09-20   wm              Initial version
#
################################################################################

module MFCC

export mfcc

"""
MFCC code ported from https://github.com/jameslyons/python_speech_features
The MIT License (MIT)
"""

function hz2mel(hz::Float64)
    """Convert a value in Hertz to Mels

    :param hz: a value in Hz. This can also be a numpy array, conversion proceeds element-wise.
    :returns: a value in Mels. If an array was passed in, an identical sized array is returned.
    """
    return 2595. * log10(1. + hz / 700.0)
end

function mel2hz(mel::Float64)
    """Convert a value in Mels to Hertz

    :param mel: a value in Mels. This can also be a numpy array, conversion proceeds element-wise.
    :returns: a value in Hertz. If an array was passed in, an identical sized array is returned.
    """
    return 700. * (10. ^ (mel / 2595.0) - 1.)
end

function preemphasis(signal::Vector{Float64}, coeff=0.95)
    """perform preemphasis on the input signal.

    :param signal: The signal to filter.
    :param coeff: The preemphasis coefficient. 0 is no filter, default is 0.95.
    :returns: the filtered signal.
    """
    return vcat(signal[1], signal[2:end] - signal[1:end - 1] * coeff)
end

function preemphasis!(signal::Vector{Float64}, coeff=0.95)
    """perform preemphasis on the input signal.

    :param signal: The signal to filter.
    :param coeff: The preemphasis coefficient. 0 is no filter, default is 0.95.
    """
    #@inbounds for i = -length(signal):-2 signal[-i] -= signal[-i - 1] * coeff end
    #"""
    @inbounds begin
        m = signal[1]
        @simd for i = 2:length(signal)
            n = signal[i] - m * coeff
            m = signal[i]
            signal[i] = n
        end
    end
    #"""
end

function lifter(cepstra::Array{Float64, 2}, L=22.)
    """Apply a cepstral lifter the the matrix of cepstra. This has the effect of increasing the
    magnitude of the high frequency DCT coeffs.

    :param cepstra: the matrix of mel-cepstra, will be numframes * numcep in size.
    :param L: the liftering coefficient to use. Default is 22. L <= 0 disables lifter.
    """
    if L > 0
        ncoeff = size(cepstra, 2)
        lift = 1. + (L / 2.) * sin(pi * [0:ncoeff - 1] / L)
        return cepstra .* lift'
    else
        # values of L <= 0, do nothing
        return cepstra
    end
end

function lifter!(cepstra::Array{Float64, 2}, L=22.)
    """Apply a cepstral lifter the the matrix of cepstra. This has the effect of increasing the
    magnitude of the high frequency DCT coeffs.

    :param cepstra: the matrix of mel-cepstra, will be numframes * numcep in size.
    :param L: the liftering coefficient to use. Default is 22. L <= 0 disables lifter.
    """
    if L > 0
        ncoeff = size(cepstra, 2)
        for j = 1:size(cepstra, 2)
            lift = 1. + (L / 2.) * sin(pi * (j - 1) / L)
            for i = 1:size(cepstra, 1)
                @inbounds cepstra[i, j] *= lift
            end
        end
    end
end

function magspec(frames::Array{Float64, 2}, NFFT)
    """Compute the magnitude spectrum of each frame in frames. If frames is an NxD matrix, output will be NxNFFT.

    :param frames: the array of frames. Each row is a frame.
    :param NFFT: the FFT length to use. If NFFT > frame_len, the frames are zero-padded.
    :returns: If frames is an NxD matrix, output will be NxNFFT. Each row will be the magnitude spectrum of the corresponding frame.
    """
    padded_frames = hcat(frames, repeat(zeros(size(frames, 1)), outer=[1, NFFT - size(frames, 2) % NFFT]))
    complex_spec = rfft(padded_frames, 2)
    #TODO magspec2 -> abs2
    return abs(complex_spec)
end

function magspec2(frames::Array{Float64, 2}, NFFT)
    """Compute the squared magnitude spectrum of each frame in frames. If frames is an NxD matrix, output will be NxNFFT.

    :param frames: the array of frames. Each row is a frame.
    :param NFFT: the FFT length to use. If NFFT > frame_len, the frames are zero-padded.
    :returns: If frames is an NxD matrix, output will be NxNFFT. Each row will be the magnitude spectrum of the corresponding frame.
    """
    #padded_frames = hcat(frames, repeat(zeros(size(frames, 1)), outer=[1, NFFT - size(frames, 2) % NFFT]))
    @time padded_frames = reshape(resize!(vec(frames), size(frames, 1) * NFFT), size(frames, 1), NFFT)

    return abs2(rfft(padded_frames, 2))
end

function powmagspec(frames::Array{Float64, 2}, NFFT)
    """Compute the squared magnitude spectrum of each frame in frames. If frames is an NxD matrix, output will be NxNFFT.

    :param frames: the array of frames. Each row is a frame.
    :param NFFT: the FFT length to use. If NFFT > frame_len, the frames are zero-padded.
    :returns: If frames is an NxD matrix, output will be NxNFFT. Each row will be the magnitude spectrum of the corresponding frame.
    """
    #padded_frames = hcat(frames, repeat(zeros(size(frames, 1)), outer=[1, NFFT - size(frames, 2) % NFFT]))
    padded_vec = resize!(vec(frames), size(frames, 1) * NFFT)
    padded_vec[length(frames) + 1:end] = 0.
    padded_frames = reshape(padded_vec, size(frames, 1), NFFT)

    return abs2(rfft(padded_frames, 2)) / NFFT
end

function magspec2!(spec::Array{Float64, 2}, frames::Array{Float64, 2}, NFFT)
    """Compute the squared magnitude spectrum of each frame in frames. If frames is an NxD matrix, output will be NxNFFT.

    :param frames: the array of frames. Each row is a frame.
    :param NFFT: the FFT length to use. If NFFT > frame_len, the frames are zero-padded.
    :returns: If frames is an NxD matrix, output will be NxNFFT. Each row will be the magnitude spectrum of the corresponding frame.
    """
    padded_frames = hcat(frames, repeat(zeros(size(frames, 1)), outer=[1, NFFT - size(frames, 2) % NFFT]))
    copy!(spec, abs2(rfft(padded_frames, 2)))
end

function powspec(frames::Array{Float64, 2}, NFFT)
    """Compute the power spectrum of each frame in frames. If frames is an NxD matrix, output will be NxNFFT.

    :param frames: the array of frames. Each row is a frame.
    :param NFFT: the FFT length to use. If NFFT > frame_len, the frames are zero-padded.
    :returns: If frames is an NxD matrix, output will be NxNFFT. Each row will be the power spectrum of the corresponding frame.
    """
    return 1.0 / NFFT * magspec(frames, NFFT) .^ 2
end

function powspec2(frames::Array{Float64, 2}, NFFT)
    """Compute the power spectrum of each frame in frames. If frames is an NxD matrix, output will be NxNFFT.

    :param frames: the array of frames. Each row is a frame.
    :param NFFT: the FFT length to use. If NFFT > frame_len, the frames are zero-padded.
    :returns: If frames is an NxD matrix, output will be NxNFFT. Each row will be the power spectrum of the corresponding frame.
    """
    return 1.0 / NFFT * magspec2(frames, NFFT)
end

function powspec2!(pspec::Array{Float64, 2}, frames, NFFT)
    """Compute the power spectrum of each frame in frames. If frames is an NxD matrix, output will be NxNFFT.

    :param frames: the array of frames. Each row is a frame.
    :param NFFT: the FFT length to use. If NFFT > frame_len, the frames are zero-padded.
    :returns: If frames is an NxD matrix, output will be NxNFFT. Each row will be the power spectrum of the corresponding frame.
    """
    magspec2!(pspec, frames, NFFT)
    pspec /= float(NFFT)
end

function framesig(sig::Vector{Float64}, frame_len, frame_step, winfunc=ones)
    """Frame a signal into overlapping frames.

    :param sig: the audio signal to frame.
    :param frame_len: length of each frame measured in samples.
    :param frame_step: number of samples after the start of the previous frame that the next frame should begin.
    :param winfunc: the analysis window to apply to each frame. By default no window is applied.
    :returns: an array of frames. Size is NUMFRAMES by frame_len.
    """
    const slen = size(sig, 1)
    frame_len = int(round(frame_len))
    frame_step = int(round(frame_step))
    if slen <= frame_len
        numframes = 1
    else
        numframes = 1 + int(ceil((1.0 * slen - frame_len) / frame_step))
    end

    const padlen = int((numframes - 1) * frame_step + frame_len)

    # TODO, check which is faster
    #padsignal = vcat(sig, zeros(padlen - slen))
    resize!(sig, padlen)
    sig[padlen + 1:end] = 0

    frames = Array(Float64, numframes, frame_len)
    @simd for q = 1:frame_len
        i = q
        @simd for p = 0:numframes - 1
            @inbounds frames[p + 1, q] = sig[i]
            i += frame_step
        end
    end

    if winfunc != ones
        win = repeat(winfunc(frame_len)', outer=[numframes, 1])
        return frames .* win
    else
        return frames
    end
end

function get_filterbanks(nfilt=20, nfft=512, samplerate=16000, lowfreq=0, highfreq=samplerate / 2)
    """Compute a Mel-filterbank. The filters are stored in the rows, the columns correspond
    to fft bins. The filters are returned as an array of size nfilt * (nfft/2 + 1)

    :param nfilt: the number of filters in the filterbank, default 20.
    :param nfft: the FFT size. Default is 512.
    :param samplerate: the samplerate of the signal we are working with. Affects mel spacing.
    :param lowfreq: lowest band edge of mel filters, default 0 Hz
    :param highfreq: highest band edge of mel filters, default samplerate/2
    :returns: A numpy array of size nfilt * (nfft/2 + 1) containing filterbank. Each row holds 1 filter.
    """
    @assert highfreq <= (samplerate / 2) "highfreq is greater than samplerate/2"

    # compute points evenly spaced in mels
    lowmel = hz2mel(float(lowfreq))
    highmel = hz2mel(float(highfreq))

    melpoints = linspace(lowmel, highmel, nfilt + 2)
    # our points are in Hz, but we use fft bins, so we have to convert
    #  from Hz to fft bin number
    bin = int(floor((nfft + 1) * map(mel2hz, melpoints) / samplerate)) + 1

    fbank = zeros(div(nfft, 2) + 1, nfilt)
    @inbounds (
        @simd for j in 1:nfilt
            @simd for i in int(bin[j]):int(bin[j + 1])
                fbank[i, j] = (i - bin[j]) / (bin[j + 1] - bin[j])
            end
        end
    )
    @inbounds (
        @simd for j in 1:nfilt
            @simd for i in int(bin[j + 1]):int(bin[j + 2])
                fbank[i, j] = (bin[j + 2] - i) / (bin[j + 2] - bin[j + 1])
            end
        end
    )
    return fbank
end

function fbank(signal::Vector{Float64}, samplerate=16000, winlen=0.025, winstep=0.01,
    nfilt=26, nfft=512, lowfreq=0, highfreq=samplerate / 2, preemph=0.97)
    """Compute Mel-filterbank energy features from an audio signal.

    :param signal: the audio signal from which to compute features. Should be an N*1 array
    :param samplerate: the samplerate of the signal we are working with.
    :param winlen: the length of the analysis window in seconds. Default is 0.025s (25 milliseconds)
    :param winstep: the step between successive windows in seconds. Default is 0.01s (10 milliseconds)
    :param nfilt: the number of filters in the filterbank, default 26.
    :param nfft: the FFT size. Default is 512.
    :param lowfreq: lowest band edge of mel filters. In Hz, default is 0.
    :param highfreq: highest band edge of mel filters. In Hz, default is samplerate/2
    :param preemph: apply preemphasis filter with preemph as coefficient. 0 is no filter. Default is 0.97.
    :returns: 2 values. The first is a numpy array of size (NUMFRAMES by nfilt) containing features. Each row holds 1 feature vector. The
        second return value is the energy in each frame (total energy, unwindowed)
    """
    preemphasis!(signal, preemph)

    frames = framesig(signal, winlen * samplerate, winstep * samplerate)
    pspec = powmagspec(frames, nfft)

    energy = sum(pspec, 2) # this stores the total energy in each frame
    @inbounds energy[find(x -> x == 0, energy)] = eps(eltype(energy)) # if energy is zero, we get problems with log

    fb = get_filterbanks(nfilt, nfft, samplerate, lowfreq, highfreq)
    feat = pspec * fb # compute the filterbank energies
    @inbounds feat[find(x -> x == 0, feat)] = eps(eltype(feat)) # if feat is zero, we get problems with log

    return feat, energy
end

################################################################################

function mfcc(signal::Vector{Float64};
              samplerate=16000., winlen=0.025, winstep=0.01, numcep=13,
              nfilt=26, nfft=512, lowfreq=0, highfreq=samplerate/2, preemph=0.97, ceplifter=22, appendEnergy=true)
    """Compute MFCC features from an audio signal.

    :param signal: the audio signal from which to compute features. Should be an N*1 array
    :param samplerate: the samplerate of the signal we are working with.
    :param winlen: the length of the analysis window in seconds. Default is 0.025s (25 milliseconds)
    :param winstep: the step between successive windows in seconds. Default is 0.01s (10 milliseconds)
    :param numcep: the number of cepstrum to return, default 13
    :param nfilt: the number of filters in the filterbank, default 26.
    :param nfft: the FFT size. Default is 512.
    :param lowfreq: lowest band edge of mel filters. In Hz, default is 0.
    :param highfreq: highest band edge of mel filters. In Hz, default is samplerate/2
    :param preemph: apply preemphasis filter with preemph as coefficient. 0 is no filter. Default is 0.97.
    :param ceplifter: apply a lifter to final cepstral coefficients. 0 is no lifter. Default is 22.
    :param appendEnergy: if this is true, the zeroth cepstral coefficient is replaced with the log of the total frame energy.
    :returns: A numpy array of size (NUMFRAMES by numcep) containing features. Each row holds 1 feature vector.
    """
    feat, energy = fbank(signal, samplerate, winlen, winstep, nfilt, nfft, lowfreq, highfreq, preemph)

    feat = log(feat)
    dct!(feat, 2)
    feat = feat[:, 1:numcep]

    lifter!(feat, ceplifter)
    if appendEnergy
        feat[:,1] = log(energy) # replace first cepstral coefficient with log of frame energy
    end

    return feat
end

end # module MFCC


# This is an implementation of the CMSA-ES as described in the paper:
#  H. Beyer and B. Sendhoff, "Covariance Matrix Adaptation Revisited –
#  the CMSA Evolution Strategy", 2008
#

#type CMSA_ES <: PopulationOptimizer
#  decomposer::Function    # Function that decomposes covar matrix into sqrtC
#
#  # Constants that we save for performance reasons
#  n::Int64
#  mu::Int64
#  lambda::Int64
#  tau::Float64
#  tau_c::Float64
#
#  # State
#  sigma::Float64            # Current mutation strength, sigma
#  sigmas::Array{Float64,2}  # Sigma per individual in last generated population
#  C::Array{Float64,2}       # Covariance matrix
#  sqrtC::Array{Float64,2}   # sqrtC' * sqrtC == C, i.e. some composition
#
#  CMSA_ES(n, mu, lambda, decomposeFunc = chol) = begin
#
#    tau = 1 / sqrt(2*n)           # Equation (1) on page 5 in Beyer2008
#    tau_c = 1 + n * (n + 1) / 2   # Equation (2) on page 5 in Beyer2008
#
#    sigma = 1.0                   # Mutation strength, sigma (self-adapted)
#    sigmas = sigma * ones(1, n)
#    C = eye(n, n)                 # Covariance matrix
#
#    new(n, mu, lambda, decomposeFunc, tau, tau_c,
#      sigma, sigmas, C, decomposeFunc(C))
#
#  end
#end
#
#
#
## Get a set of new individuals to be ranked based on fitness.
#function ask(o::CMSA_ES)
#  o.sigmas = o.sigma * exp(o.tau * randn(1, o.lambda))
#  o.s = randn(o.n, o.n) * o.sqrtC
#end

function normalize_utilities(utilities)
  utilities / sum(utilities)
end

function linear_utilities(mu, lambda)
  normalize_utilities(hcat(ones(1, mu), zeros(1, lambda-mu)))
end

function log_utilities(mu, lambda)
  normalize_utilities(hcat((log(mu+1) - log(1:mu))', zeros(1, lambda-mu)))
end

# This weights both positive and negative ones
function active_linear_utilities(mu, lambda)
  pos = ones(1, mu) / mu
  rest = lambda-mu
  neg = (-1) * ones(1, rest) / (2*rest)
  2 * hcat(pos, neg)'
end

# Optimize the n-dimensional objective func with a (mu,lambda) CMSA-ES.
function cmsa_es(p;
  trace = true,

  # Stopping criteria related
  max_seconds = 4*numdims(p), max_evals_per_dim = 1e7,
  ftol = 1e-7, xtol = 1e-10, stol = 1e-10,
  max_rounds_without_improvement = 500,

  # Starting points, will be random unless specified
  xmean = false,

  # Algorithm specific params:
  covarMatrixSampler = CholeskyCovarSampler,
  utilitiesFunc = log_utilities,
  lambda = 4*numdims(p),
  mu = int(max(ceil(lambda/rand(4:20)), 1.0)),
  tau = 1 / sqrt(2*numdims(p)), # Equation (1) on page 5 in Beyer2008
  tau_c = 1 + numdims(p) * (numdims(p) + 1) / (2 * mu),  # Equation (2) on page 5 in Beyer2008
  sigma = 0.05*rand(1:8)*minimum(diameters(search_space(p))),
  decompose_covar_prob = 0.4,
  max_successes_before_increasing = 3,
  max_failures_before_decreasing = 3

  )

  N = numdims(p)
  ss = search_space(p)
  max_evals = max_evals_per_dim * N
  a = 1 - 1 / tau_c
  C = covarMatrixSampler(N)
  utilities = utilitiesFunc(mu, lambda)
  xbest = xmean = rand_individual(ss)   # Current best mean value.
  fbest = eval1(xbest, p)
  num_fevals = 1
  archive = BlackBoxOptim.TopListArchive(N, 10)
  add_candidate!(archive, fbest, xbest[:], num_fevals)
  fevals_last_best = num_fevals
  next_print_covar = 100
  termination_reason = "?"

  num_successes = 0
  num_failures = 0

  start_time = time()

  # Now lets optimize! Ensure we run at least one iteration.
  while(true)

    if (time() - start_time) > max_seconds
      termination_reason = "Exceeded time budget"
      break
    end

    if num_fevals > max_evals
      termination_reason = "Exceeded function eval budget"
      break
    end

    # Decompose only with a given probability => saves time
    if rand() <= decompose_covar_prob
      decompose!(C)
    end

    # Generate new population
    sigmas = sigma * exp( tau * randn(1, lambda) )  # 1*lambda
    s = multivariate_normal_sample(C, N, lambda)
    z = broadcast(*, sigmas, s)                     # n*lambda
    xs = repmat(xmean, 1, lambda) + z               # n*lambda

    # Evaluate fitness
    fitnesses = eval_fitnesses(p, xs, lambda)
    num_fevals += lambda

    # Check if best new fitness is best ever and print some info if tracing.
    indbest = indmin(fitnesses)
    fbest_new = fitnesses[indbest]

    if fbest_new < fbest
      num_successes += 1 # From Solis-Wets
      num_failures = 0   # From Solis-Wets
      xbest = xs[:, indbest]
      fbest = fbest_new
      add_candidate!(archive, fbest, xbest, num_fevals)
      fevals_last_best = num_fevals

      if fitness_is_within_ftol(p, ftol, fbest)
        termination_reason = "Within ftol"
        break
      end

      if trace
        println("$(num_fevals): Best fitness = $(fbest)")
        if num_fevals > next_print_covar
          next_print_covar = num_fevals + 100
          println("covar summary: ", sumstats(C.C, (x) -> @sprintf("%.2e", x)))
          println("sigma: ", sigma)
        end
      end
    else
      if (num_fevals - fevals_last_best) > max_rounds_without_improvement * lambda
        termination_reason = "Max rounds without improvement reached"
        break
      end
      num_failures += 1 # From Solis-Wets
      num_successes = 0   # From Solis-Wets
    end

    # Assign weights to the best individuals according to the utilities vector.
    weights = assign_weights(lambda, fitnesses, utilities)

    # Calculate new mean value based on weighted version of steps in population.
    xmean += (z * weights)

    # Update the covariance matrix
    uc = zeros(Float64, N, N)
    for(i in 1:lambda)
      if weights[i] > 0.0
        se = s[:,i]
        uc += weights[i] * (se * se')
      end
    end
    update_covariance_matrix!(C, uc, a)
    #ws = broadcast(*, weights', s)
    #update_covariance_matrix!(C, (ws * ws'), covar_learning_rate)

    # Adapt sigma for next round
    sigma = sigmas * weights

    # From solis-wets
    if (max_successes_before_increasing >= 1.0) && (num_successes >= max_successes_before_increasing)
      sigma = sigma * 2
      num_successes = 0
    elseif (max_failures_before_decreasing >= 1.0) && (num_failures >= max_failures_before_decreasing)
      sigma = sigma / 2
      num_failures = 0
    end

    # Terminate if sigma very small
    #println("sigma = $(sigma[1,1])")
    if sigma[1,1] < stol
      termination_reason = "Sigma too small"
      break
    end
  end

  return xbest, fbest, num_fevals, termination_reason, archive
end

function restart_cmsa_es(p;
  trace = true,

  # Stopping criteria related
  max_seconds = 2*numdims(p), max_evals_per_dim = 1e7,
  ftol = 1e-7, xtol = 1e-10, stol = 1e-10,

  # Starting points, will be random unless specified
  xmean = false,

  # Start values for algorithm specific params:
  lambda = 4*numdims(p),
  lambda_divisor = 7,
  covarMatrixSampler = CholeskyCovarSampler,
  utilitiesFunc = log_utilities,
  tau = 1 / sqrt(2*numdims(p)), # Equation (1) on page 5 in Beyer2008
  covar_learning_rate = 0.88,

  )

  now = start_time = time()

  total_fbest = Inf
  xb = total_fevals = termination_reason = 0
  num_runs = 0

  while(now < start_time + max_seconds)
    time_left = max_seconds - (now - start_time)

    num_runs += 1

    x, fbest, num_fevals, termination_reason, archive = cmsa_es(p;
      trace = trace, max_seconds = time_left,
      ftol = ftol, xtol = xtol, stol = stol,
      max_evals_per_dim = max_evals_per_dim,
      max_rounds_without_improvement = 1000,
      lambda = lambda,
      mu = int(maximum([ceil(lambda/lambda_divisor), 2.0])),
      covarMatrixSampler = CholeskyCovarSampler, utilitiesFunc = log_utilities,
      tau = tau, covar_learning_rate = covar_learning_rate
    );

    lambda += (4*numdims(p))
    lambda_divisor = 5 + rand(1:15)
    sigma = 0.10*rand(1:4)*minimum(diameters(search_space(p)))
    covar_learning_rate = 0.85+0.01*rand(1:4)

    total_fevals += num_fevals

    if fbest < total_fbest
      total_fbest = fbest
      xb = x
      println(num_runs, ": New global best = $(total_fbest)")

      if fitness_is_within_ftol(p, ftol, fbest)
        termination_reason = "Within ftol"
        break
      end
    end

    now = time()
  end

  if time() > start_time + max_seconds
    termination_reason = "Exceeded time budget (elasped = $(time()-start_time))"
  end

  return xb, total_fbest, total_fevals, termination_reason, archive
end

function eval_fitnesses(problem, xs, lambda = size(xs, 2))
  fitnesses = zeros(lambda)
  for(i in 1:lambda)
    fitnesses[i] = eval1(xs[:,i], problem)
  end
  fitnesses
end

function assign_weights(n, fitnesses, utilities; minimize = true)
  us_ordered = zeros(n, 1)
  perms = sortperm(fitnesses, rev = !minimize)
  for(i in 1:n)
    us_ordered[perms[i]] = utilities[i]
  end
  us_ordered
end

# We create different types of Covariance matrix samplers based on different
# decompositions.
abstract CovarianceMatrixSampler

function update_covariance_matrix!(cms::CovarianceMatrixSampler, delta, a)
  C = a * cms.C + (1 - a) * delta
  cms.C = triu(C) + triu(C,1)' # Ensure C is symmetric. Should not be needed, investigate...
end

type EigenCovarSampler <: CovarianceMatrixSampler
  C::Array{Float64,2}
  B::Array{Float64,2}
  diagD::Array{Float64,1}

  EigenCovarSampler(n) = begin
    new(eye(n,n), eye(n,n), ones(n))
  end
end

function decompose!(cms::EigenCovarSampler)
  try
    EV, B = eig(cms.C)
    cms.B = B
    cms.diagD = sqrt(EV)
  catch
    # We don't update if there is some problem
  end
end

function multivariate_normal_sample(cms::CovarianceMatrixSampler, n, m)
  cms.B * (cms.diagD .* randn(n, m))
end

type CholeskyCovarSampler <: CovarianceMatrixSampler
  C::Array{Float64,2}
  sqrtC::Array{Float64,2}

  CholeskyCovarSampler(n) = begin
    new(eye(n,n), eye(n,n))
  end
end

function decompose!(cms::CholeskyCovarSampler)
  try
    cms.sqrtC = chol(cms.C)'
  catch error
    # We don't update if there is some problem
  end
end

function multivariate_normal_sample(cms::CholeskyCovarSampler, n, m)
  cms.sqrtC * randn(n, m)
end

# Block-wise Covar sampling selects a block of coordinates for which the
# matrix is decomposed but then discards the covariance for the coordinates
# not in the selected block and only uses their (individual) variances.
# Since this is simpler to relate to the Eigen decomposition we start with
# that version.


include("../utils.jl")
include("group_utils.jl")

@testset "Translation action" begin
    M = Euclidean(2, 3)
    G = TranslationGroup(2, 3)
    A = TranslationAction(Euclidean(2, 3), G)

    @test repr(A) == "TranslationAction($(repr(M)), $(repr(G)), LeftAction())"
    @test repr(switch_direction(A)) ==
          "TranslationAction($(repr(M)), $(repr(G)), RightAction())"

    types_a = [Matrix{Float64}]

    types_m = [Matrix{Float64}]

    @test g_manifold(A) == M
    @test base_group(A) == G
    @test base_manifold(G) == M

    for (T_A, T_M) in zip(types_a, types_m)
        a_pts =
            convert.(
                T_A,
                [
                    [0.0 1.0 2.0; 2.0 4.0 1.0],
                    [-1.0 0.0 2.0; -2.0 3.0 3.0],
                    [1.0 1.0 2.0; 3.0 2.0 1.0],
                ],
            )
        m_pts =
            convert.(
                T_M,
                [
                    [0.0 1.0 2.0; 2.0 4.0 1.0],
                    [-1.0 0.0 2.0; -2.0 3.0 3.0],
                    [1.0 1.0 2.0; 3.0 2.0 1.0],
                ],
            )
        v_pts =
            convert.(
                T_M,
                [
                    [0.0 1.0 2.0; 2.0 4.0 1.0],
                    [-1.0 0.0 2.0; -2.0 3.0 3.0],
                    [1.0 1.0 2.0; 3.0 2.0 1.0],
                ],
            )

        atol = if eltype(T_M) == Float32
            1e-7
        else
            1e-15
        end
        @test apply_diff(A, a_pts[1], m_pts[1], v_pts[1]) === v_pts[1]
        @test inverse_apply_diff(A, a_pts[1], m_pts[1], v_pts[1]) === v_pts[1]
        test_action(
            A,
            a_pts,
            m_pts,
            v_pts;
            test_optimal_alignment=false,
            test_diff=true,
            atol=atol,
        )
    end
end


"""
    B,Z = bkz(H,K)

Block Korkine-Zolotarev (BKZ) reduction of the lattice (matrix) H with
blocksize K.

BKZ was originaly described in "A hierarchy of polynomial time basis
reduction algorithms" by Claus Peter Schnorr, Theoretical Computer Science,
1987. This function is not well tested and could be wrong or otherwise
different from Schnorr's BKZ.
"""
function bkz(H::AbstractArray{Td,2},K) where {Td<:Number}
    m,n = size(H)
    Ti= LLLplus.getIntType(Td)

    B,Z = lll(H)
    Ti = eltype(Z)
    k = 1
    clean = zeros(n)
    norms = zeros(n) .+ Inf
    prevclean = false
    while !all(clean .==2)
        kx = k:min(k+K-1,n)
        Bk = view(B,:,kx)
        Zk = view(Z,:,kx)
        if kx[end]==n
            Bh,Zh = hkz(Bk)
            bhnorms = diag(Bh'*Bh)
            ih = sortperm(diag(Bh'*Bh))
            bhnorms = bhnorms[ih]
            Bh = Bh[:,ih]
            Zh = Zh[:,ih]
            if all(bhnorms .>= norms[kx])
                clean[kx] .= prevclean ? 2 : 1
            else
                clean[kx] .= 0
                Bk[:,:] = Bh
                Zh[:,:] = Zh
                norms[kx] .=bhnorms
            end
            k=1
            prevclean = all(clean .>=0)
        else
            u = LLLplus.svpu(Bk)
            b = Bk*u
            bnrm = b'*b
            z = Zk*u
            kk =findfirst(u .!=zero(Ti))
            if kk==1
                if (abs(u[1])==one(Ti) && all(u[2:end] .==zero(Ti))) || bnrm>=norms[k]
                    clean[k]= prevclean ? 2 : 1
                else
                    clean[k]=0
                    B[:,k] = b
                    Z[:,k] = z
                    norms[k] = bnrm
                end
            else
                if bnrm>=norms[k]
                    clean[k]= prevclean ? 2 : 1
                else
                    clean[k]=0
                    B[:,[k,k+kk-1]] = [b B[:,k]]
                    Z[:,[k,k+kk-1]] = [z Z[:,k]]
                    norms[[k,k+kk-1]] = [bnrm norms[k]]
                end
            end
            k=k+1
        end
    end
    return B,Z
end


@recipe function f(
    sol::ODESolution,
    H::OpenQuantumBase.AbstractHamiltonian,
    lvl::Integer,
    t_axis=sol.t,
)
    y = instantaneous_population(sol, H, lvl=lvl, t_axis=t_axis)
    lab = ["\$E_$x\$" for x in (0:lvl - 1)']
    label --> lab
    (t_axis, y)
end

@recipe function f(
    sol::ODESolution,
    H::OpenQuantumBase.AbstractHamiltonian,
    lvl::AbstractArray{T,1},
    t_axis=sol.t,
) where {T <: Integer}
    y = instantaneous_population(sol, H, lvl=maximum(lvl) + 1, t_axis=t_axis)
    lab = ["\$E_$x\$" for x in (lvl)']
    label --> lab
    (t_axis, y[:, lvl .+ 1])
end

@recipe function f(
    sol::ODESolution,
    lvl::AbstractArray{T,1},
    s_axis=sol.t,
) where {T <: Integer}
    y = []
    if ndims(sol.u[1]) == 1
        for x in s_axis
            push!(y, abs2.(sol(x)[lvl .+ 1]))
        end
    elseif ndims(sol.u[1]) == 2
        for x in s_axis
            push!(y, real.(diag(sol(x))[lvl .+ 1]))
        end
    else
        throw(ArgumentError("Solution needs to be state vector or density matrix."))
    end

    y = hcat(y...)'
    (s_axis, y[:, lvl .+ 1])
end


using LinearAlgebra: eigvecs
using SparseArrays: SparseVector, sparsevec

"""
    ShortestPathGraphKernel <: AbstractGraphKernel

A graph kernel that compares two graphs `g` and `g'` by comparing all pairs
of vertices `(u, v)` of the `g` and `(u', v')` of `g'` if their shortest distance
is smaller than `tol`. In that case, the vertices `u` and `u'`, as well as `v`, `v'` are
compared with `vertex_kernel`.

# Keywords
- `tol=0.0`: Only pairs of vertices where the shortest distance is at most `tol` are
    compared.
- `dist_key=:`: The key for the edge values to compute the shortest distance with. Can be either
    an `Integer` or a `Symbol` for a key to a specific edge value, `nothing` to use a default distance
    of `1` for each edge, or `:` in which case the default edge weight for that graph type
    is used.
- `vertex_kernel=ConstVertexKernel(1.0)`: The kernel used to compare two vertices.

# References
[Borgwardt, K. M., & Kriegel, H. P.: Shortest-path kernels on graphs](https://www.dbs.ifi.lmu.de/~borgward/papers/BorKri05.pdf)
"""
struct ShortestPathGraphKernel{VK <: AbstractVertexKernel} <: AbstractGraphKernel

    tol::Float64
    vertex_kernel::VK
    dist_key::Union{Int, Symbol, Colon, Nothing}
end

function ShortestPathGraphKernel(;tol=0.0, vertex_kernel=ConstVertexKernel(1.0), dist_key=Colon())

    return ShortestPathGraphKernel(tol, vertex_kernel, dist_key)
end

function preprocessed_form(kernel::ShortestPathGraphKernel, g::AbstractGraph)

    dists = _make_dists(g, kernel.dist_key)

    ds = map(t -> t.dist, dists)
    us =  map(t -> t.u, dists)
    vs =  map(t -> t.v, dists)

    return (g=g, ds=ds, us=us, vs=vs)
end

function apply_preprocessed(kernel::ShortestPathGraphKernel, pre1, pre2)

    g1, ds1, us1, vs1 = pre1
    g2, ds2, us2, vs2 = pre2

    # TODO there might be some issues with unsigned types here
    ε = kernel.tol
    vertex_kernel = kernel.vertex_kernel

    result = 0.0

    len1 = length(ds1)
    len2 = length(ds2)

    i2 = 1
    @inbounds for i1 in Base.OneTo(length(ds1))
        d1 = ds1[i1]
        while i2 <= len2 && d1 > (ds2[i2] + ε)
            i2 += 1
        end
        j2 = i2
        while j2 <= len2 && ds2[j2] <= (d1 + ε)
            result += vertex_kernel(g1, us1[i1], g2, us2[j2])
            result += vertex_kernel(g1, vs1[i1], g2, vs2[j2])
            j2 += 1
        end
    end

    return result

end

function _make_dists(g, dist_key)

    dists = if dist_key === Colon()
        floyd_warshall_shortest_paths(g).dists
    elseif dist_key === nothing
        floyd_warshall_shortest_paths(g, Graphs.DefaultDistance(nv(g))).dists
    else
        floyd_warshall_shortest_paths(g, weights(g, dist_key)).dists
    end

    verts = vertices(g)
    tm = typemax(eltype(dists))
    dists_list = [(dist=dists[u, v], u=u, v=v) for u in verts for v in verts if u != v && dists[u, v] != tm]
    sort!(dists_list, by=t->t.dist)
    return dists_list
end


using DynMultiply
using Base.Test

# first test set
A = rand(1000, 1000)
B = rand(1000, 1000)
v = rand(1000)

@test dynmultiply(A, B, v) ≈ A * B * v
@test dynmultiplystr(A, B, v) == "(A1*(A2*A3))"
@test dynmultiply_plan(A, B, v).data == [0,0,0,1,1,1]

# second test set
A = rand(1000, 1)
B = rand(1, 1000)
C = rand(1000, 50)
D = rand(50, 1)

@test dynmultiply(A, B, C, D) ≈ A * B * C * D
@test dynmultiplystr(A, B, C, D) == "(A1*((A2*A3)*A4))"
@test dynmultiply_plan(A, B, C, D).data == [0,0,0,0,1,1,1,1,2,1]


# Unit tests for Distances

function test_metricity(dist, x, y, z)
    @testset "Test metricity of $(typeof(dist))" begin
        dxy = evaluate(dist, x, y)
        dxz = evaluate(dist, x, z)
        dyz = evaluate(dist, y, z)
        if isa(dist, PreMetric)
            # Unfortunately small non-zero numbers (~10^-16) are appearing
            # in our tests due to accumulating floating point rounding errors.
            # We either need to allow small errors in our tests or change the
            # way we do accumulations...
            @test evaluate(dist, x, x) + one(eltype(x)) ≈ one(eltype(x))
            @test evaluate(dist, y, y) + one(eltype(y)) ≈ one(eltype(y))
            @test evaluate(dist, z, z) + one(eltype(z)) ≈ one(eltype(z))
            @test dxy ≥ zero(eltype(x))
            @test dxz ≥ zero(eltype(x))
            @test dyz ≥ zero(eltype(x))
        end
        if isa(dist, SemiMetric)
            @test dxy ≈ evaluate(dist, y, x)
            @test dxz ≈ evaluate(dist, z, x)
            @test dyz ≈ evaluate(dist, y, z)
        else # Not symmetric, so more PreMetric tests
            @test evaluate(dist, y, x) ≥ zero(eltype(x))
            @test evaluate(dist, z, x) ≥ zero(eltype(x))
            @test evaluate(dist, z, y) ≥ zero(eltype(x))
        end
        if isa(dist, Metric)
            # Again we have small rounding errors in accumulations
            @test dxz ≤ dxy + dyz || dxz ≈ dxy + dyz
            dyx = evaluate(dist, y, x)
            @test dyz ≤ dyx + dxz || dyz ≈ dyx + dxz
            dzy = evaluate(dist, z, y)
            @test dxy ≤ dxz + dzy || dxy ≈ dxz + dzy
        end
    end
end

@testset "PreMetric, SemiMetric, Metric on $T" for T in (Float64, F64)
    Random.seed!(123)
    n = 100
    x = rand(T, n)
    y = rand(T, n)
    z = rand(T, n)

    test_metricity(SqEuclidean(), x, y, z)
    test_metricity(Euclidean(), x, y, z)
    test_metricity(Cityblock(), x, y, z)
    test_metricity(TotalVariation(), x, y, z)
    test_metricity(Chebyshev(), x, y, z)
    test_metricity(Minkowski(2.5), x, y, z)

    test_metricity(CosineDist(), x, y, z)
    test_metricity(CorrDist(), x, y, z)

    test_metricity(ChiSqDist(), x, y, z)

    test_metricity(Jaccard(), x, y, z)
    test_metricity(SpanNormDist(), x, y, z)

    test_metricity(BhattacharyyaDist(), x, y, z)
    test_metricity(HellingerDist(), x, y, z)
    test_metricity(Bregman(x -> sqeuclidean(x, zero(x)), x -> 2*x), x, y, z);

    x₁ = rand(T, 2)
    x₂ = rand(T, 2)
    x₃ = rand(T, 2)

    test_metricity(Haversine(6371.0), x₁, x₂, x₃)

    k = rand(1:3, n)
    l = rand(1:3, n)
    m = rand(1:3, n)

    test_metricity(Hamming(), k, l, m)

    a = rand(Bool, n)
    b = rand(Bool, n)
    c = rand(Bool, n)

    test_metricity(RogersTanimoto(), a, b, c)
    test_metricity(BrayCurtis(), a, b, c)
    test_metricity(Jaccard(), a, b, c)

    w = rand(T, n)

    test_metricity(PeriodicEuclidean(w), x, y, z)
    test_metricity(WeightedSqEuclidean(w), x, y, z)
    test_metricity(WeightedEuclidean(w), x, y, z)
    test_metricity(WeightedCityblock(w), x, y, z)
    test_metricity(WeightedMinkowski(w, 2.5), x, y, z)
    test_metricity(WeightedHamming(w), a, b, c)

    Q = rand(T, n, n)
    Q = Q * Q'  # make sure Q is positive-definite

    test_metricity(SqMahalanobis(Q), x, y, z)
    test_metricity(Mahalanobis(Q), x, y, z)

    p = rand(T, n)
    q = rand(T, n)
    r = rand(T, n)
    p[p .< median(p)] .= 0
    p /= sum(p)
    q /= sum(q)
    r /= sum(r)

    test_metricity(KLDivergence(), p, q, r)
    test_metricity(RenyiDivergence(0.0), p, q, r)
    test_metricity(RenyiDivergence(1.0), p, q, r)
    test_metricity(RenyiDivergence(Inf), p, q, r)
    test_metricity(RenyiDivergence(0.5), p, q, r)
    test_metricity(RenyiDivergence(2), p, q, r)
    test_metricity(RenyiDivergence(10), p, q, r)
    test_metricity(JSDivergence(), p, q, r)
end

@testset "individual metrics" begin
    a = 1
    b = 2
    @test sqeuclidean(a, b) == 1.0

    @test euclidean(a, b) == 1.0
    @test cityblock(a, b) == 1.0
    @test totalvariation(a, b) == 0.5
    @test chebyshev(a, b) == 1.0
    @test minkowski(a, b, 2) == 1.0
    @test hamming(a, b) == 1
    @test peuclidean(a, b, 0.5) == 0
    @test peuclidean(a, b, 2) == 1.0

    bt = [true, false, true]
    bf = [false, true, true]
    @test rogerstanimoto(bt, bf) == 4.0 / 5.0
    @test braycurtis(bt, bf) == 0.5

    for T in (Float64, F64)

        for (_x, _y) in (([4.0, 5.0, 6.0, 7.0], [3.0, 9.0, 8.0, 1.0]),
                         ([4.0, 5.0, 6.0, 7.0], [3. 8.; 9. 1.0]))
            x, y = T.(_x), T.(_y)
            @test sqeuclidean(x, y) == 57.0
            @test euclidean(x, y) == sqrt(57.0)
            @test jaccard(x, y) == 13.0 / 28
            @test cityblock(x, y) == 13.0
            @test totalvariation(x, y) == 6.5
            @test chebyshev(x, y) == 6.0
            @test braycurtis(x, y) == 1.0 - (30.0 / 43.0)
            @test minkowski(x, y, 2) == sqrt(57.0)
            @test peuclidean(x, y, fill(10.0, 4)) == sqrt(37)
            @test peuclidean(x - vec(y), zero(y), fill(10.0, 4)) == peuclidean(x, y, fill(10.0, 4))
            @test peuclidean(x, y, [10.0, 10.0, 10.0, Inf]) == sqrt(57)
            @test_throws DimensionMismatch cosine_dist(1.0:2, 1.0:3)
            @test cosine_dist(x, y) ≈ (1.0 - 112. / sqrt(19530.0))
            x_int, y_int = Int64.(x), Int64.(y)
            @test cosine_dist(x_int, y_int) == (1.0 - 112.0 / sqrt(19530.0))
            @test corr_dist(x, y) ≈ cosine_dist(x .- mean(x), vec(y) .- mean(y))
            @test chisq_dist(x, y) == sum((x - vec(y)).^2 ./ (x + vec(y)))
            @test spannorm_dist(x, y) == maximum(x - vec(y)) - minimum(x - vec(y))

            @test gkl_divergence(x, y) ≈ sum(i -> x[i] * log(x[i] / y[i]) - x[i] + y[i], 1:length(x))

            @test meanad(x, y) ≈ mean(Float64[abs(x[i] - y[i]) for i in 1:length(x)])
            @test msd(x, y) ≈ mean(Float64[abs2(x[i] - y[i]) for i in 1:length(x)])
            @test rmsd(x, y) ≈ sqrt(msd(x, y))
            @test nrmsd(x, y) ≈ sqrt(msd(x, y)) / (maximum(x) - minimum(x))

            w = ones(4)
            @test sqeuclidean(x, y) ≈ wsqeuclidean(x, y, w)

            w = rand(Float64, size(x))
            @test wsqeuclidean(x, y, w) ≈ dot((x - vec(y)).^2, w)
            @test weuclidean(x, y, w) == sqrt(wsqeuclidean(x, y, w))
            @test wcityblock(x, y, w) ≈ dot(abs.(x - vec(y)), w)
            @test wminkowski(x, y, w, 2) ≈ weuclidean(x, y, w)
        end

        # Test ChiSq doesn't give NaN at zero
        @test chisq_dist([0.0], [0.0]) == 0.0

        # Test weighted Hamming distances with even weights
        a = T.([1.0, 2.0, 1.0, 3.0, 2.0, 1.0])
        b = T.([1.0, 3.0, 0.0, 2.0, 2.0, 0.0])
        w = rand(T, size(a))

        @test whamming(a, a, w) === T(0.0)
        @test whamming(a, b, w) === sum((a .!= b) .* w)

        # Minimal test of Jaccard - test return type stability.
        @inferred evaluate(Jaccard(), rand(T, 3), rand(T, 3))
        @inferred evaluate(Jaccard(), [1, 2, 3], [1, 2, 3])
        @inferred evaluate(Jaccard(), [true, false, true], [false, true, true])

        # Test Bray-Curtis. Should be 1 if no elements are shared, 0 if all are the same
        @test braycurtis([1,0,3],[0,1,0]) == 1.0
        @test braycurtis(rand(10), zeros(10)) == 1.0
        @test braycurtis([1,0],[1,0]) == 0.0

        # Test KL, Renyi and JS divergences
        r = rand(T, 12)
        p = copy(r)
        p[p .< median(p)] .= 0.0
        scale = sum(p) / sum(r)
        r /= sum(r)
        p /= sum(p)
        q = rand(T, 12)
        q /= sum(q)

        klv = 0.0
        for i = 1:length(p)
            if p[i] > 0
                klv += p[i] * log(p[i] / q[i])
            end
        end
        @test kl_divergence(p, q) ≈ klv
        @test typeof(kl_divergence(p, q)) == T

        @test renyi_divergence(p, r, 0) ≈ -log(scale)
        @test renyi_divergence(p, r, 1) ≈ -log(scale)
        @test renyi_divergence(p, r, 10) ≈ -log(scale)
        @test renyi_divergence(p, r, rand()) ≈ -log(scale)
        @test renyi_divergence(p, r, Inf) ≈ -log(scale)
        @test isinf(renyi_divergence([0.0, 0.5, 0.5], [0.0, 1.0, 0.0], Inf))
        @test renyi_divergence([0.0, 1.0, 0.0], [0.0, 0.5, 0.5], Inf) ≈ log(2.0)
        @test renyi_divergence(p, q, 1) ≈ kl_divergence(p, q)

        pm = (p + q) / 2
        jsv = kl_divergence(p, pm) / 2 + kl_divergence(q, pm) / 2
        @test js_divergence(p, q) ≈ jsv
    end
end # testset

@testset "NaN behavior" begin
    a = [NaN, 0]; b = [0, NaN]
    @test isnan(chebyshev(a, b)) == isnan(maximum(a - b))
    a = [NaN, 0]; b = [0, 1]
    @test isnan(chebyshev(a, b)) == isnan(maximum(a - b))
    @test isnan(renyi_divergence([0.5, 0.0, 0.5], [0.5, 0.5, NaN], 2))
end #testset

@testset "empty vector" begin
    for T in (Float64, F64)
        a = T[]
        b = T[]
        @test sqeuclidean(a, b) == 0.0
        @test isa(sqeuclidean(a, b), T)
        @test euclidean(a, b) == 0.0
        @test isa(euclidean(a, b), T)
        @test cityblock(a, b) == 0.0
        @test isa(cityblock(a, b), T)
        @test totalvariation(a, b) == 0.0
        @test isa(totalvariation(a, b), T)
        @test chebyshev(a, b) == 0.0
        @test isa(chebyshev(a, b), T)
        @test braycurtis(a, b) == 0.0
        @test isa(braycurtis(a, b), T)
        @test minkowski(a, b, 2) == 0.0
        @test isa(minkowski(a, b, 2), T)
        @test hamming(a, b) == 0.0
        @test isa(hamming(a, b), Int)
        @test renyi_divergence(a, b, 1.0) == 0.0
        @test isa(renyi_divergence(a, b, 2.0), T)
        @test braycurtis(a, b) == 0.0
        @test isa(braycurtis(a, b), T)

        w = T[]
        @test isa(whamming(a, b, w), T)
        @test peuclidean(a, b, w) == 0.0
        @test isa(peuclidean(a, b, w), T)
    end
end # testset

@testset "DimensionMismatch throwing" begin
    a = [1, 0]; b = [2]
    @test_throws DimensionMismatch sqeuclidean(a, b)
    a = [1, 0]; b = [2.0] ; w = [3.0]
    @test_throws DimensionMismatch wsqeuclidean(a, b, w)
    @test_throws DimensionMismatch peuclidean(a, b, w)
    a = [1, 0]; b = [2.0, 4.0] ; w = [3.0]
    @test_throws DimensionMismatch wsqeuclidean(a, b, w)
    @test_throws DimensionMismatch peuclidean(a, b, w)
    p = [0.5, 0.5]; q = [0.3, 0.3, 0.4]
    @test_throws DimensionMismatch bhattacharyya(p, q)
    @test_throws DimensionMismatch hellinger(q, p)
    Q = rand(length(p), length(p))
    Q = Q * Q'  # make sure Q is positive-definite
    @test_throws DimensionMismatch mahalanobis(p, q, Q)
    @test_throws DimensionMismatch mahalanobis(q, q, Q)
    mat23 = [0.3 0.2 0.0; 0.1 0.0 0.4]
    mat22 = [0.3 0.2; 0.1 0.4]
    @test_throws DimensionMismatch colwise!(mat23, Euclidean(), mat23, mat23)
    @test_throws DimensionMismatch colwise!(mat23, Euclidean(), mat23, q)
    @test_throws DimensionMismatch colwise!(mat23, Euclidean(), mat23, mat22)
    @test_throws DimensionMismatch colwise!(mat23, Bregman(x -> sqeuclidean(x, zero(x)), x -> 2*x), mat23, mat22)
    @test_throws DimensionMismatch evaluate(Bregman(x -> sqeuclidean(x, zero(x)), x -> 2*x), [1, 2, 3], [1, 2])
    @test_throws DimensionMismatch evaluate(Bregman(x -> sqeuclidean(x, zero(x)), x -> [1, 2]), [1, 2, 3], [1, 2, 3])
end # testset

@testset "mahalanobis" begin
    for T in (Float64, F64)
        x, y = T.([4.0, 5.0, 6.0, 7.0]), T.([3.0, 9.0, 8.0, 1.0])
        a = T.([1.0, 2.0, 1.0, 3.0, 2.0, 1.0])
        b = T.([1.0, 3.0, 0.0, 2.0, 2.0, 0.0])

        Q = rand(T, length(x), length(x))
        Q = Q * Q'  # make sure Q is positive-definite
        @test sqmahalanobis(x, y, Q) ≈ dot(x - y, Q * (x - y))
        @test eltype(sqmahalanobis(x, y, Q)) == T
        @test mahalanobis(x, y, Q) == sqrt(sqmahalanobis(x, y, Q))
        @test eltype(mahalanobis(x, y, Q)) == T
    end
end #testset

@testset "haversine" begin
    for T in (Float64, F64)
        @test haversine([-180.,0.], [180.,0.], 1.) ≈ 0 atol=1e-10
        @test haversine([0.,-90.],  [0.,90.],  1.) ≈ π atol=1e-10
        @test haversine((-180.,0.), (180.,0.), 1.) ≈ 0 atol=1e-10
        @test haversine((0.,-90.),  (0.,90.),  1.) ≈ π atol=1e-10
        @test haversine((1.,-15.625), (-179.,15.625), 6371.) ≈ 20015. atol=1e0
        @test_throws ArgumentError haversine([0.,-90., 0.25], [0.,90.], 1.)
    end
end

@testset "bhattacharyya / hellinger" begin
    for T in (Float64, F64)
        x, y = T.([4.0, 5.0, 6.0, 7.0]), T.([3.0, 9.0, 8.0, 1.0])
        a = T.([1.0, 2.0, 1.0, 3.0, 2.0, 1.0])
        b = T.([1.0, 3.0, 0.0, 2.0, 2.0, 0.0])
        p = rand(T, 12)
        p[p .< median(p)] .= 0.0
        q = rand(T, 12)

        # Bhattacharyya and Hellinger distances are defined for discrete
        # probability distributions so to calculate the expected values
        # we need to normalize vectors.
        px = x ./ sum(x)
        py = y ./ sum(y)
        expected_bc_x_y = sum(sqrt.(px .* py))
        @test Distances.bhattacharyya_coeff(x, y) ≈ expected_bc_x_y
        @test bhattacharyya(x, y) ≈ (-log(expected_bc_x_y))
        @test hellinger(x, y) ≈ sqrt(1 - expected_bc_x_y)

        pa = a ./ sum(a)
        pb = b ./ sum(b)
        expected_bc_a_b = sum(sqrt.(pa .* pb))
        @test Distances.bhattacharyya_coeff(a, b) ≈ expected_bc_a_b
        @test bhattacharyya(a, b) ≈ (-log(expected_bc_a_b))
        @test hellinger(a, b) ≈ sqrt(1 - expected_bc_a_b)

        pp = p ./ sum(p)
        pq = q ./ sum(q)
        expected_bc_p_q = sum(sqrt.(pp .* pq))
        @test Distances.bhattacharyya_coeff(p, q) ≈ expected_bc_p_q
        @test bhattacharyya(p, q) ≈ (-log(expected_bc_p_q))
        @test hellinger(p, q) ≈ sqrt(1 - expected_bc_p_q)

        # Ensure it is semimetric
        @test bhattacharyya(x, y) ≈ bhattacharyya(y, x)
    end
end #testset

function test_colwise(dist, x, y, T)
    @testset "Colwise test for $(typeof(dist))" begin
        n = size(x, 2)
        r1 = zeros(T, n)
        r2 = zeros(T, n)
        r3 = zeros(T, n)
        for j = 1:n
            r1[j] = evaluate(dist, x[:, j], y[:, j])
            r2[j] = evaluate(dist, x[:, 1], y[:, j])
            r3[j] = evaluate(dist, x[:, j], y[:, 1])
        end
        # ≈ and all( .≈ ) seem to behave slightly differently for F64
        @test all(colwise(dist, x, y) .≈ r1)
        @test all(colwise(dist, x[:, 1], y) .≈ r2)
        @test all(colwise(dist, x, y[:, 1]) .≈ r3)
    end
end

@testset "column-wise metrics on $T" for T in (Float64, F64)
    m = 5
    n = 8
    X = rand(T, m, n)
    Y = rand(T, m, n)
    A = rand(1:3, m, n)
    B = rand(1:3, m, n)

    P = rand(T, m, n)
    Q = rand(T, m, n)
    # Make sure not to remove all of the non-zeros from any column
    for i in 1:n
        P[P[:, i] .< median(P[:, i]) / 2, i] .= 0.0
    end

    test_colwise(SqEuclidean(), X, Y, T)
    test_colwise(Euclidean(), X, Y, T)
    test_colwise(Cityblock(), X, Y, T)
    test_colwise(TotalVariation(), X, Y, T)
    test_colwise(Chebyshev(), X, Y, T)
    test_colwise(Minkowski(2.5), X, Y, T)
    test_colwise(Hamming(), A, B, T)
    test_colwise(Bregman(x -> sqeuclidean(x, zero(x)), x -> 2*x), X, Y, T);

    test_colwise(CosineDist(), X, Y, T)
    test_colwise(CorrDist(), X, Y, T)

    test_colwise(ChiSqDist(), X, Y, T)
    test_colwise(KLDivergence(), P, Q, T)
    test_colwise(RenyiDivergence(0.0), P, Q, T)
    test_colwise(RenyiDivergence(1.0), P, Q, T)
    test_colwise(RenyiDivergence(Inf), P, Q, T)
    test_colwise(RenyiDivergence(0.5), P, Q, T)
    test_colwise(RenyiDivergence(2), P, Q, T)
    test_colwise(RenyiDivergence(10), P, Q, T)
    test_colwise(JSDivergence(), P, Q, T)
    test_colwise(SpanNormDist(), X, Y, T)

    test_colwise(BhattacharyyaDist(), X, Y, T)
    test_colwise(HellingerDist(), X, Y, T)
    test_colwise(BrayCurtis(), X, Y, T)

    w = rand(T, m)

    test_colwise(WeightedSqEuclidean(w), X, Y, T)
    test_colwise(WeightedEuclidean(w), X, Y, T)
    test_colwise(WeightedCityblock(w), X, Y, T)
    test_colwise(WeightedMinkowski(w, 2.5), X, Y, T)
    test_colwise(WeightedHamming(w), A, B, T)
    test_colwise(PeriodicEuclidean(w), X, Y, T)

    Q = rand(T, m, m)
    Q = Q * Q'  # make sure Q is positive-definite

    test_colwise(SqMahalanobis(Q), X, Y, T)
    test_colwise(Mahalanobis(Q), X, Y, T)
end

function test_pairwise(dist, x, y, T)
    @testset "Pairwise test for $(typeof(dist))" begin
        nx = size(x, 2)
        ny = size(y, 2)
        rxy = zeros(T, nx, ny)
        rxx = zeros(T, nx, nx)
        for j = 1:ny, i = 1:nx
            rxy[i, j] = evaluate(dist, x[:, i], y[:, j])
        end
        for j = 1:nx, i = 1:nx
            rxx[i, j] = evaluate(dist, x[:, i], x[:, j])
        end
        # As earlier, we have small rounding errors in accumulations
        @test pairwise(dist, x, y) ≈ rxy
        @test pairwise(dist, x) ≈ rxx
        @test pairwise(dist, x, y, dims=2) ≈ rxy
        @test pairwise(dist, x, dims=2) ≈ rxx
        @test pairwise(dist, permutedims(x), permutedims(y), dims=1) ≈ rxy
        @test pairwise(dist, permutedims(x), dims=1) ≈ rxx
    end
end

@testset "pairwise metrics on $T" for T in (Float64, F64)
    m = 5
    n = 8
    nx = 6
    ny = 8

    X = rand(T, m, nx)
    Y = rand(T, m, ny)
    A = rand(1:3, m, nx)
    B = rand(1:3, m, ny)

    P = rand(T, m, nx)
    Q = rand(T, m, ny)

    test_pairwise(SqEuclidean(), X, Y, T)
    test_pairwise(Euclidean(), X, Y, T)
    test_pairwise(Cityblock(), X, Y, T)
    test_pairwise(TotalVariation(), X, Y, T)
    test_pairwise(Chebyshev(), X, Y, T)
    test_pairwise(Minkowski(2.5), X, Y, T)
    test_pairwise(Hamming(), A, B, T)

    test_pairwise(CosineDist(), X, Y, T)
    test_pairwise(CorrDist(), X, Y, T)

    test_pairwise(ChiSqDist(), X, Y, T)
    test_pairwise(KLDivergence(), P, Q, T)
    test_pairwise(RenyiDivergence(0.0), P, Q, T)
    test_pairwise(RenyiDivergence(1.0), P, Q, T)
    test_pairwise(RenyiDivergence(Inf), P, Q, T)
    test_pairwise(RenyiDivergence(0.5), P, Q, T)
    test_pairwise(RenyiDivergence(2), P, Q, T)
    test_pairwise(JSDivergence(), P, Q, T)

    test_pairwise(BhattacharyyaDist(), X, Y, T)
    test_pairwise(HellingerDist(), X, Y, T)
    test_pairwise(BrayCurtis(), X, Y, T)
    test_pairwise(Bregman(x -> sqeuclidean(x, zero(x)), x -> 2*x), X, Y, T)

    w = rand(m)

    test_pairwise(WeightedSqEuclidean(w), X, Y, T)
    test_pairwise(WeightedEuclidean(w), X, Y, T)
    test_pairwise(WeightedCityblock(w), X, Y, T)
    test_pairwise(WeightedMinkowski(w, 2.5), X, Y, T)
    test_pairwise(WeightedHamming(w), A, B, T)
    test_pairwise(PeriodicEuclidean(w), X, Y, T)

    Q = rand(m, m)
    Q = Q * Q'  # make sure Q is positive-definite

    test_pairwise(SqMahalanobis(Q), X, Y, T)
    test_pairwise(Mahalanobis(Q), X, Y, T)
end

@testset "Euclidean precision" begin
    X = [0.1 0.2; 0.3 0.4; -0.1 -0.1]
    pd = pairwise(Euclidean(1e-12), X, X)
    @test pd[1, 1] == 0
    @test pd[2, 2] == 0
    pd = pairwise(Euclidean(1e-12), X)
    @test pd[1, 1] == 0
    @test pd[2, 2] == 0
    pd = pairwise(SqEuclidean(1e-12), X, X)
    @test pd[1, 1] == 0
    @test pd[2, 2] == 0
    pd = pairwise(SqEuclidean(1e-12), X)
    @test pd[1, 1] == 0
    @test pd[2, 2] == 0
end

@testset "Bregman Divergence" begin
    # Some basic tests.
    @test_throws ArgumentError bregman(x -> x, x -> 2*x, [1, 2, 3], [1, 2, 3])
    # Test if Bregman() correctly implements the gkl divergence between two random vectors.
    F(p) = LinearAlgebra.dot(p, log.(p));
    ∇(p) = map(x -> log(x) + 1, p)
    testDist = Bregman(F, ∇)
    p = rand(4)
    q = rand(4)
    p = p/sum(p);
    q = q/sum(q);
    @test evaluate(testDist, p, q) ≈ gkl_divergence(p, q)
    # Test if Bregman() correctly implements the squared euclidean dist. between them.
    @test bregman(x -> norm(x)^2, x -> 2*x, p, q) ≈ sqeuclidean(p, q)
    # Test if Bregman() correctly implements the IS distance.
    F(p) = -1 * sum(log.(p))
    ∇(p) = map(x -> -1 * x^(-1), p)
    function ISdist(p::AbstractVector, q::AbstractVector)
        return sum([p[i]/q[i] - log(p[i]/q[i]) - 1 for i in 1:length(p)])
    end
    @test bregman(F, ∇, p, q) ≈ ISdist(p, q)
end

#=
@testset "zero allocation colwise!" begin
    d = Euclidean()
    a = rand(2, 41)
    b = rand(2, 41)
    z = zeros(41)
    colwise!(z, d, a, b)
    # This fails when bounds checking is enforced
    bounds = Base.JLOptions().check_bounds
    if bounds == 0
        @test (@allocated colwise!(z, d, a, b)) == 0
    else
        @test_broken (@allocated colwise!(z, d, a, b)) == 0
    end
end
=#


@testset "TrapezoidalPoints" begin
    @test_throws Exception TrapezoidalPoints{Int}(4)
    @test eltype(@inferred(TrapezoidalPoints(4))) == Float64

    for T in Floats
        @test eltype(@inferred(TrapezoidalPoints{T}(3))) == T
        @test typeof(@inferred(TrapezoidalPoints{T}(3)[1])) == T

        @test TrapezoidalPoints{T}(0) ≈ T[]
        @test TrapezoidalPoints{T}(1) ≈ T[0]  atol = eps(real(T))
        @test TrapezoidalPoints{T}(2) ≈ T[-1,1]
        @test TrapezoidalPoints{T}(3) ≈ T[-1,0,1]   atol = eps(real(T))
        @test TrapezoidalPoints{T}(4) ≈ T[-1,-T(1)/3,T(1)/3,1]
        @test TrapezoidalPoints{T}(5) ≈ T[-1,-T(1)/2,0,T(1)/2,1]   atol = eps(real(T))
    end
end


using StaticArrays, Test
using OrdinaryDiffEq

u0 = fill(zero(MVector{2,Float64}), 2)
u0[1] = ones(MVector{2,Float64}) .+ 1
f = (du,u,p,t) -> du .= u
ode = ODEProblem(f, u0, (0.,1.))
sol = solve(ode, Euler(), dt=1.e-2)
sol = solve(ode, Tsit5())

u0 = fill(zero(SVector{2,Float64}), 2) .+ 1
u0[1] = ones(SVector{2,Float64}) .+ 1
ode = ODEProblem(f, u0, (0.,1.))
sol = solve(ode, Euler(), dt=1.e-2)
sol = solve(ode, Tsit5())

sol = solve(ode, SSPRK22(), dt=1.e-2)

u0 = ones(MVector{2,Float64})
ode = ODEProblem(f, u0, (0.,1.))
sol = solve(ode, Euler(), dt=1.e-2)
sol = solve(ode, Tsit5(), dt=1.e-2)

u0 = ones(SVector{2,Float64})
f = (u,p,t) -> u
ode = ODEProblem(f, u0, (0.,1.))
sol = solve(ode, Euler(), dt=1.e-2)
sol = solve(ode, ImplicitEuler(nlsolve=NLAnderson()), dt=1.e-2)
sol = solve(ode, Tsit5(), dt=1.e-2)