using Parameters, QuantEcon, LinearAlgebra, Interpolations, Random
using Optim, Roots

Household = @with_kw (r = 0.01,
                      R = 1 + r,
                      α = 0.33,
                      δ = 0.05,
                      w = (1-α)*(α/(r+δ))^(α/(1-α)),
                      σ = 1.0,
                      κ = 4.00,
                      ν = 0.25,
                      β = 0.96,
                      MC = MarkovChain([0.9 0.1; 0.1 0.9], [0.95; 1.05]),
                      Π = MC.p,
                      z_vals = MC.state_values,
                      z_size = length(z_vals),
                      b = 0.0, # Borrowing limit/Minimum level of assets
                      r_ss = (1/β-1), # SS interest rate
                      a_max = ((r_ss+δ)/α)^(1/(α-1))*3, # 3 times SS capital
                      a_size = 200,
                      a_vals = collect(range(b, a_max, length = a_size)),
                      u = σ == 1 ? (x, h=0.0) -> log(x) -  κ*(h^1+ν)/(1+ν) : (x, h=0.0) -> (x^(1 - σ) - 1) / (1 - σ) - κ*(h^1+ν)/(1+ν),
                      duc = c -> c^(-σ),
                      duh = h -> κ*h^(ν),
                      duc_inv = x -> x^(-1/σ),
                      duh_inv = x -> (x/κ)^(1/ν))

function BruteForce_Mtcy(V::AbstractArray,
                            am::NamedTuple;
                            return_policy::Bool=false)

    @unpack R, r, w, b, σ, β, Π, z_vals, a_size, a_vals, z_size, duc, duh, duc_inv, duh_inv, u = am
    U = fill(-Inf, a_size, z_size, a_size)
    V_new = similar(V)
    g_a = zeros(Int, a_size, z_size)

    for i_a in 1:a_size
        for i_z in 1:z_size
            if i_a == 1
                a_new_min = 1
            else
                a_new_min = g_a[i_a-1, i_z]
            end
            a_new_max = min(a_new_min + 15, a_size)
            for i_a_new in a_new_min:a_new_max
                c = w * z_vals[i_z] + R * a_vals[i_a] - a_vals[i_a_new]
                if c > 0.0
                    U[i_a, i_z, i_a_new] = u(c) + β*dot(Π[i_z,:],V[i_a_new,:])
                end
            end
            V_new[i_a, i_z], g_a[i_a, i_z] = findmax(U[i_a, i_z, :])
        end
    end
    if return_policy == true
        return V_new, g_a
    else
        return V_new
    end
end

function BruteForce_Std(V::AbstractArray,
                            am::NamedTuple;
                            return_policy::Bool=false)

    @unpack R, r, w, b, σ, β, Π, z_vals, a_size, a_vals, z_size, duc, duh, duc_inv, duh_inv, u = am
    U = fill(-Inf, a_size, z_size, a_size)
    V_new = similar(V)
    g_a = similar(V)

    for i_a in 1:a_size
        for i_z in 1:z_size
            for i_a_new in 1:a_size
                c = w * z_vals[i_z] + R * a_vals[i_a] - a_vals[i_a_new]
                if c > 0
                    U[i_a, i_z, i_a_new] = u(c) + β*dot(Π[i_z,:],V[i_a_new,:])
                end
            end
            V_new[i_a, i_z], g_a[i_a, i_z] = findmax(U[i_a, i_z, :])
        end
    end
    if return_policy == true
        return V_new, g_a
    else
        return V_new
    end
end

function BellmanOperator(V::AbstractArray,
                            am::NamedTuple;
                            return_policy::Bool=false)

    @unpack R, r, w, b, σ, β, Π, z_vals, a_size, a_vals, z_size, duc, duh, duc_inv, duh_inv, u = am
    vf = interp(a_vals, V)
    z_idx = 1:z_size
    V_new = similar(V)
    g_a = similar(V)
    g_c = similar(V)

    for (i_a, a) in enumerate(a_vals)
        for (i_z, z) in enumerate(z_vals)
            function obj(x)
                EV = dot(Π[i_z, :], vf(R * a + w * z - x, z_idx)) # compute expectation
                return  u(x) +  β * EV
            end
            res = Optim.maximize(obj, 0.0, R * a + w * z + b)
            Optim.converged(res) || error("Didn't converge") # important to check
            g_c[i_a, i_z] = Optim.maximizer(res)
            g_a[i_a, i_z] = R * a + w * z - g_c[i_a, i_z]
            V_new[i_a, i_z] = Optim.maximum(res)
        end
    end
    if return_policy
        V_new, g_a, g_c
    else
        V_new
    end
end

function ColemanOperator(g::AbstractArray,
                            am::NamedTuple;
                            return_policy::Bool=false)
    @unpack R, r, w, b, σ, β, Π, z_vals, a_size, a_vals, z_size, duc, duh, duc_inv, duh_inv, u = am
    gf = interp(a_vals, g)
    z_idx = 1:z_size
    V_new = similar(g)
    g_a = similar(g)
    g_c = similar(g)
    opt_lb = 1e-8
    for (i_a, a) in enumerate(a_vals)
        for (i_z, z) in enumerate(z_vals)
            function obj(x)
                cps = gf.(R * a + z*w - x, z_idx) # c' for each z'
                expectation = dot(duc.(cps), Π[i_z, :])
                return abs(duc(x) - max(R*β * expectation, duc(R * a + z*w + b)))
            end
            opt_ub = R*a + z*w + b  # addresses issue #8 on github
            res = optimize(obj, min(opt_lb, opt_ub - 1e-2), opt_ub,
                           method = Optim.Brent())
            g_c[i_a, i_z] = Optim.minimizer(res)
            g_a[i_a, i_z] = R * a + z*w - g_c[i_a, i_z]
        end
    end
    V_new = u.(g_c)
    if return_policy
        V_new, g_a, g_c
    else
        g_c
    end
end

function EGMOperator(Kg_f::Function,
                            am::NamedTuple;
                            return_policy::Bool=false)

    @unpack R, r, w, b, σ, β, Π, z_vals, a_size, a_vals, z_size, duc, duh, duc_inv, duh_inv, u = am
    g_a = zeros(Float64, a_size, z_size)
    g_c_nb = zeros(Float64, a_size, z_size)
    g_c_new = zeros(Float64, a_size, z_size)

    for (i_z, z) in enumerate(z_vals)
        for (i_a, a_new) in enumerate(a_vals)
            expectation = dot(duc.(g_c[i_a, :]), Π[i_z, :])
            g_c_nb[i_a, i_z] = duc_inv(R*β * expectation)
            g_a[i_a, i_z] = (g_c_nb[i_a, i_z] + a_new - w*z) /R
        end
        # Update policy function and return
        g_c_new[:, i_z] = LinearInterpolation(g_a[:,i_z], g_c_nb[:,i_z], extrapolation_bc=Line())(a_vals)
    end
    for (i_a, a_new) in enumerate(a_vals)
        for (i_z, z) in enumerate(z_vals)
            if a_new < g_a[1, i_z]
                g_c_new[i_a, i_z] = a_new*R + duh_inv(w*z*duc_inv(c))*w*z + b
            end
        end
    end
    return g_c_new
end

function EGMOperator_Labor(g_c::AbstractArray,
                            am::NamedTuple;
                            return_policy::Bool=false)

    @unpack R, r, w, b, σ, β, Π, z_vals, a_size, a_vals, z_size, duc, duh, duc_inv, duh_inv, u = am
    # Initialize arrays
    g_a = zeros(Float64, a_size, z_size)
    g_c_nb = zeros(Float64, a_size, z_size)
    g_c_new = zeros(Float64, a_size, z_size)

    # Iterate over different combinations of (a',z)
    for (i_z, z) in enumerate(z_vals)
        for (i_a, a_new) in enumerate(a_vals)
            # Solve for C(a',z) in the no binding case
            expectation = dot(duc.(g_c[i_a, :]), Π[i_z, :])
            g_c_nb[i_a, i_z] = duc_inv(R*β * expectation)
            # Solve for H(c,z)
            h = duh_inv(w*z*duc_inv(g_c_nb[i_a, i_z]))
            # Solve for A(a',z)
            g_a[i_a, i_z] = (g_c_nb[i_a, i_z] + a_new - h*w*z) /R
        end
        # Update to c1(a',z')
        g_c_new[:, i_z] = LinearInterpolation(g_a[:,i_z], g_c_nb[:,i_z], extrapolation_bc=Line())(a_vals)
    end

    # Iterate over different combinations of (a',z)
    for (i_a, a_new) in enumerate(a_vals)
        for (i_z, z) in enumerate(z_vals)
            # Check if it is a binding borrowing constraint case
            if a_new < g_a[1, i_z]
                # If so, solve for C(a',z) using a root finder and use the result value in c1(a',z')
                cf = c -> abs(a_new*R + duh_inv(w*z*duc_inv(c))*w*z + b - c)
                res = optimize(cf, 0.0, g_c_new[i_a, i_z])
                g_c_new[i_a, i_z] = Optim.minimizer(res)
            end
        end
    end
    # Return c1(a',z')
    return g_c_new
end

function initialize(am::NamedTuple)
    # simplify names, set up arrays
    @unpack R, r, w, b, σ, β, Π, z_vals, a_size, a_vals, z_size, duc, duh, duc_inv, duh_inv, u = am
    V0 = zeros(Float64, a_size, z_size)
    c0 = zeros(Float64, a_size, z_size)
    a0 = [a_vals a_vals]
    h0 = zeros(Float64, a_size, z_size)

    # Populate V and c
    for (i_z, z) in enumerate(z_vals)
        for (i_a, a) in enumerate(a_vals)
            c_max = R * a + w*z + b
            c0[i_a, i_z] = c_max
            V0[i_a, i_z] = u(c_max) / (1 - β)
        end
    end
    return V0, a0, c0, h0
end

function SimulationAssets(am::NamedTuple,
                    g_c::AbstractArray,
                    T::Int = 100000;
                    labor::Bool = false,
                    verbose::Bool = false)

    @unpack Π, r, w, α, δ, a_vals, z_vals, R, duc_inv, duh_inv = am  # Simplify names

    cf = interp(a_vals, g_c)

    a_sim = zeros(Float64, T + 1)
    h_sim = zeros(Float64, T)
    z_sim = simulate(MarkovChain(Π), T)
    for t in 1:T
        i_z = z_sim[t]
        z = z_vals[i_z]
        if labor
            h_sim[t] = duh_inv(w*z*duc_inv(cf(a_sim[t], i_z)))
            a_sim[t+1] = R * a_sim[t] + h_sim[t]*w*z - cf(a_sim[t], i_z)
        else
            a_sim[t+1] = R * a_sim[t] + w*z - cf(a_sim[t], i_z)
        end
    end
    # Burn 1000 first periods
    a_sim = a_sim[1000:end]

    # Compute Aggregates from market clearin
    K = mean(a_sim)
    L = mean(z_sim.*h_sim)
    r_new = α*(K/L)^(α-1) - δ
    return a_sim, K, L, r_new
end


using SequGen
using Unitful, Test

@testset "Constant" begin
  t = 1:10
  for c in (5, 1.23, u"1s", u"1K")
    s = Constant(c)
    @test sample(s, t) == fill(c, 10)
  end
end

@testset "Trend" begin
  for Δt in (1, 1.0, 1.12u"s", 2.5u"hr"),
    t₀ in (zero(Δt), oneunit(Δt)),
    Δx in (3, 3.14, 3.4u"m", 7.5u"K")

    t = collect(range(t₀, t₀ + Δt, length = 101))
    s = Trend(t₀, Δx / Δt)
    x = collect(range(zero(Δx), Δx, length = 101))
    @test sample(s, t) ≈ x
  end
end


using DelimitedFiles
using Printf
using Combinatorics

# Import Data
nums = readdlm("./input_a.dat")

# Function to check all possible sums of `arr` in groups of `n` that equal `target`
function check_sums(arr, n, target)
    pairs = collect(combinations(arr, n))
    for p in pairs
        if sum(p) == target
            @printf("Combination of %d values that results in %d\n", n, target)
            println(p)
            @printf("Sum: %i\tProduct: %i\n\n", sum(p), prod(p))
            return p
        end
    end
end

check_sums(nums, 2, 2020) # Check 2 values that sum to 2020
check_sums(nums, 3, 2020) # Check 3 values that sum to 2020


using Statistics
using NCDatasets

using ArgParse
using JSON

include("LinearRegression.jl")

correlation = (x1, x2) -> x1' * x2 / (sum(x1.^2)*sum(x2.^2)).^0.5

function parse_commandline()

    s = ArgParseSettings()
    @add_arg_table s begin

        "--data-file"
            help = "Atm data file. New variable will be appended."
            arg_type = String
            required = true

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

        "--output-file"
            help = "Output file."
            arg_type = String
            required = true

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

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

    end

    return parse_args(ARGS, s)
end

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

output_file = parsed["output-file"]

Dataset(parsed["data-file"], "r") do ds

    dims = size(ds["PSL"])

    beg_t = (parsed["beg-year"] - 1) * 12 + 1
    end_t = (parsed["end-year"] - 1) * 12 + 12

    if length(dims) == 3 # 2D + time
        rng = (:,:,beg_t:end_t)
    elseif length(dims) == 4 # 2D + time
        rng = (:,:,1,beg_t:end_t)
    else
        ErrorException("Unknown dimension") |> throw
    end

    global PSL  = replace(ds["PSL"][rng...], missing=>NaN)

    global (Nx, Ny, Nt) = size(PSL)

    if mod(Nt, 12) != 0

        PSL = PSL[:, :, 1:12*floor(Integer, Nt/12.0)]
        (Nx, Ny, Nt) = size(PSL)

        #ErrorException("Time record is not multiple of 12") |> throw
    end

    global nyears = Int64(Nt / 12)
end

Dataset(parsed["domain-file"], "r") do ds
    global mask = replace(ds["mask"], missing=>NaN)
end

# MM   = Monthly Mean
# A    = Anomaly
# YYC  = Year-to-Year Correlation

PSLMM   = zeros(Float64, size(PSL)[1:2]..., 12)
PSLA    = zeros(Float64, size(PSL)...)
PSLAYYC = zeros(Float64, size(PSL)[1:2]..., 12)
PSLAVAR = zeros(Float64, size(PSL)[1:2]..., 12)

x = collect(Float64, 1:Nt)
for i=1:Nx, j=1:Ny

    d = detrend(x, view(PSL, i, j, :))

    PSLMM[i, j, :] = mean( reshape(d, 12, :), dims=2 )[:, 1]
    PSLA[i, j, :]  = d - repeat( PSLMM[i, j, :], outer=nyears)

    for m = 1:12
        d_yy = view(PSLA, i, j, m:12:(m+nyears*12-1))
        PSLAYYC[i, j, m] = correlation(d_yy[1:end-1], d_yy[2:end])
        PSLAVAR[i, j, m] = std(d_yy)
    end

end

Dataset(output_file, "c") do ds

    defDim(ds, "months", 12)
    defDim(ds, "time", Inf)
    defDim(ds, "Nx", Nx)
    defDim(ds, "Ny", Ny)

    for (varname, vardata, vardim, attrib) in [
        ("PSLMM",  PSLMM, ("Nx", "Ny", "months"), Dict()),
        ("PSLA",    PSLA, ("Nx", "Ny", "time"),   Dict()),
        ("PSLAYYC",  PSLAYYC, ("Nx", "Ny", "months"), Dict()),
        ("PSLAVAR",  PSLAVAR, ("Nx", "Ny", "months"), Dict()),
    ]

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

        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


include("euler/euler.jl")
using .Grid: max_path_sum
using BenchmarkTools
BenchmarkTools.DEFAULT_PARAMETERS.samples = 100

function compute(path::String)::Int
    triangle = map(x -> parse.(Int, split(x, " ")), split(read(path, String), "\n"))
    return max_path_sum(triangle)
end

compute("problems/0067/p067_triangle.txt")

@benchmark compute("problems/0067/p067_triangle.txt")


using ModelingToolkit
using Test

@parameters t
@variables x(t) y(t) z(t)

null_op = 0*t
@test isequal(simplify(null_op), 0)

one_op = 1*t
@test isequal(simplify(one_op), t)

identity_op = Operation(identity,[x])
@test isequal(simplify(identity_op), x)

minus_op = -x
@test isequal(simplify(minus_op), -1x)
simplify(minus_op)

@variables x

@test simplified_expr(expand_derivatives(Differential(x)((x-2)^2))) == :(2 * (-2 + x))
@test simplified_expr(expand_derivatives(Differential(x)((x-2)^3))) == :(3 * (-2 + x)^2)
@test simplified_expr(simplify(x+2+3)) == :(5 + x)

d1 = Differential(x)((-2 + x)^2)
d2 = Differential(x)(d1)
d3 = Differential(x)(d2)

@test simplified_expr(expand_derivatives(d3)) == :(0)
@test simplified_expr(simplify(x^0)) == :(1)


module Ising

export IsingModel, update!, simulate!

"""
Ising model with periodic boundary conditions and arbitrary number of dimensions.

To simplify the calculations, the units are set in way that J=1 and kB=1
(the boltzman constant). So, the energy is in units of J and the temperature
is in units of J/kB. β can be thought of as J/(kB*T) or 1/T in this unit system.

`dims` in the constructor takes the number of particles in each dimension,
so that the number of elements of `dims` is equal to the dimensionality of the system.

The simulation here utilizes the Metropolis Monte Carlo method. Also, a checkerboard
coloring pattern is used for flipping many spins at the same time
"""
mutable struct IsingModel{N} <: Integer
    spins::Array{Int8, N}
    energy::Float64
    magnetization::Float64

    β::Float64
    mass::Float64
    dims::NTuple{N, Integer}

    # pre-calculated values for possible state changes
    _probs::Vector{Float64}

    # spins broken down into a checkerboard pattern. spins of each color
    # do not interact with each other, so they can be flipped simultatiously
    _black::SubArray{Int8, 1, Vector{Int8}, Tuple{Vector{Int64}}, false} # even index sum
    _white::SubArray{Int8, 1, Vector{Int8}, Tuple{Vector{Int64}}, false} # odd index sum

    _blackneighbors::SubArray{Int8, 2, Vector{Int8}, Tuple{Matrix{Int64}}, false}
    _whiteneighbors::SubArray{Int8, 2, Vector{Int8}, Tuple{Matrix{Int64}}, false}

    function IsingModel(β, dims::Integer...)
        spins = rand([Int8(-1), Int8(1)], dims...)

        mass = prod(dims)
        indices = CartesianIndices(spins)
        linindices = LinearIndices(spins)
        neighbors = _get_neighbors(indices, linindices, dims, mass)

        energy = -sum(reshape(spins, mass, 1) .* sum(spins[neighbors], dims=2)) / 2
        magnetization = sum(spins) / mass
        probs = exp.(-4β .* collect(-length(dims):length(dims)))

        isblack = sum.(Tuple.(indices)) .% 2 .== 0
        blackindices = linindices[findall(isblack)]
        whiteindices = linindices[findall(map(!, isblack))]

        black = @view spins[blackindices]
        white = @view spins[whiteindices]
        blackneighbors = @view spins[neighbors[blackindices, :]]
        whiteneighbors = @view spins[neighbors[whiteindices, :]]

        return new{length(dims)}(spins, energy, magnetization, β, mass, dims,
            probs, black, white, blackneighbors, whiteneighbors)
    end
end

"""
get the indices of the neighbors of every point in `indices`
"""
@inline function _get_neighbors(indices::CartesianIndices, linindices::LinearIndices,
        dims::Tuple{Vararg{Integer}}, mass::Integer)
    neighborindices = Matrix{Int}(undef, mass, 0)
    for dim in 1:length(dims)
        basevec = zeros(Integer, length(dims))
        basevec[dim] = 1
        step = CartesianIndex(basevec...)
        basevec[dim] = dims[dim]
        modulus = CartesianIndex(basevec...)

        # collecting the CartesianIndices into a vector to enable setindex
        nextneighbors = collect(indices .+ step)
        prevneighbors = collect(indices .- step)
        # referencing fixes broadcasting issues
        selectdim(prevneighbors, dim, 1) .+= Ref(modulus)
        selectdim(nextneighbors, dim, dims[dim]) .-= Ref(modulus)
        neighborindices = hcat(
            neighborindices,
            reshape(linindices[prevneighbors], mass, 1),
            reshape(linindices[nextneighbors], mass, 1))
    end
    return neighborindices
end

"""
Update an `IsingModel` for one time step, where all spins are updated once.
The operations are vectorized on the two "black" and "white" indices.
"""
function update!(ising::IsingModel)
    blackΔEs = vec(2 * ising._black .* sum(ising._blackneighbors, dims=2))
    blackprobs = ising._probs[blackΔEs .÷ 4 .+ (length(ising.dims) + 1)]
    flipblack = rand(length(blackprobs)) .<= blackprobs
    ising._black[flipblack] *= -1

    whiteΔEs = vec(2 * ising._white .* sum(ising._whiteneighbors, dims=2))
    whiteprobs = ising._probs[whiteΔEs .÷ 4 .+ (length(ising.dims) + 1)]
    flipwhite = rand(length(whiteprobs)) .<= whiteprobs
    ising._white[flipwhite] *= -1

    ising.magnetization += 2 / ising.mass * (
        sum(ising._black[flipblack]) + sum(ising._white[flipwhite]))
    ising.energy += sum(blackΔEs[flipblack]) + sum(whiteΔEs[flipwhite])
end

"""
Update an `IsingModel` for `time` steps. If `fullstats` is `false`, calculate and return
the mean and mean squared energy and magnetization (in the order ⟨E⟩, ⟨m⟩, ⟨E²⟩, ⟨m²⟩)
over the simulation time period. If `fullstats` is `true`, record and return the energy
and magnetization for every time step of the simulation (in the order Es, ms)

note: m is averaged over the absolute value
"""
function simulate!(isingmodel::IsingModel, time::Integer; fullstats::Bool=false)
    if fullstats
        energy = Vector{Float64}(undef, time)
        magnetization = Vector{Float64}(undef, time)

        for t in 1:time
            update!(isingmodel)
            energy[t] = isingmodel.energy
            magnetization[t] = isingmodel.magnetization
        end

        return energy, magnetization
    else
        Ē, m̄, Ē², m̄² = zeros(4)

        for _ in 1:time
            update!(isingmodel)
            Ē += isingmodel.energy / time
            m̄ += abs.(isingmodel.magnetization) / time
            Ē² += isingmodel.energy^2 / time
            m̄² += isingmodel.magnetization^2 / time
        end

        return Ē, m̄, Ē², m̄²
    end
end

using Statistics: mean, var

"""
    Calculate the correlation function of an `IsingModel` along the dimension `dim`,
    taking into account the periodic boundary conditions.
"""
function correlation(ising::IsingModel; dim=1)
    variance = var(ising.spins, mean=ising.magnetization)
    if variance == 0
        return zeros(ceil(Integer, size(ising.spins, dim) / 2))
    end

    correlations = Vector{Float64}(undef, ceil(Integer, size(ising.spins, dim) / 2))
    shiftvec = zeros(length(size(ising.spins)))
    for shift in 1:length(correlations)
        shiftvec[dim] = shift
        correlations[shift] = mean(ising.spins .* circshift(ising.spins, shiftvec))
    end

    # demean and normalize
    correlations .-= ising.magnetization ^ 2
    correlations ./= variance

    return correlations
end

"""
    Calculate the correlation length of an `IsingModel`, taking into account the periodic
    boundary conditions. Correlation data smaller than `thresh` is discarded (this is
    considered as "noise"). To counteract cases where the correlation length in one
    direction diverges, in these cases the correlation length is calculated for two
    directions and the minimum of the two correlation lengths is returned.
"""
function corlen(ising::IsingModel; thresh::Float64 = 0.1)
    correlations = correlation(ising, dim=1)
    threshindex = findfirst(x -> abs(x) <= thresh, correlations)

    if isnothing(threshindex)
        threshindex = length(correlations)

        # newcorr = correlation(ising, dim=2)
        # newthresh = findfirst(x -> abs(x) <= thresh, newcorr)
        # if isnothing(newthresh)
        #     newthresh = length(newcorr)
        # elseif newthresh < 3
        #     return 0.0
        # end

        # negativeξ⁻¹ = ([collect(1:threshindex) ones(threshindex)]
        #     \ log.(abs.(correlations[1:threshindex])))[1]
        # ξ = isnan(negativeξ⁻¹) ? 0.0 : 1/negativeξ⁻¹
        # if ξ < 0.0
        #     return 0.0
        # end

        # newnegativeξ⁻¹ = ([collect(1:newthresh) ones(newthresh)]
        #     \ log.(abs.(newcorr[1:newthresh])))[1]
        # newξ = isnan(newnegativeξ⁻¹) ? 0.0 : 1/newnegativeξ⁻¹
        # if newξ < 0.0
        #     return 0.0
        # end

        # return min(ξ, newξ)
    elseif threshindex < 3
        return 0.0
    end

    negativeξ⁻¹ = ([collect(1:threshindex) ones(threshindex)]
        \ log.(abs.(correlations[1:threshindex])))[1]
    if isnan(negativeξ⁻¹)
        return 0.0
    end

    ξ = -1/negativeξ⁻¹
    if ξ < 0
        return 0.0
    end

    return ξ
end

end


"""
    mes = meshmeasures(p, t)

compute areas of every element
"""
function meshareas(p, t)
    d12 = p[t[:, 2], :] - p[t[:, 1], :]
    d13 = p[t[:, 3], :] - p[t[:, 1], :]
    mes = (d12[:, 1] .* d13[:, 2] .- d12[:, 2] .* d13[:, 1]) / 2.
    return reshape(abs.(mes), length(mes))
end
#~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~-------------------------
"""
   IK, JK, SK, IF, SF = assembKM_vemKsource(pv, cellsb, rhs, boundary_condition,
                                            meshboundary)

Assemble stiffness and source terms of virtual element approach
"""
function assembKM_vemKsource(pv, cellsb, rhs, boundary_condition,
                             meshboundary)
    # type of input (for overloaded differentiation calls)
    typein = typeof(pv[1])
    n_dofs, n_polys = size(pv, 1), 3 # method has 1 degree of freedom per vertex
    # stiffness coefficients and source coefficients
    IK, JK, IF = Int64[], Int64[], Int64[]
    SK, SF = if typein == Float64
        Float64[], Float64[]
    else
        Any[], Any[]
    end
    # impose an ordering on the linear polynomials
    linear_polynomials = [[0,0], [1,0], [0,1]]
    # a utility function for wrapping around a vector
    mod_wrap(x, a) = mod(x - 1, a) + 1
    for el_id = 1:length(cellsb)
        #println(cellsb[el_id])
        vert_ids = cellsb[el_id]   # IDs of the vertices of this element
        verts = pv[vert_ids, :]    # coordinates of the vertices of this element
        n_sides = length(vert_ids) # start computing the geometric information
        area_components = verts[:, 1] .* [verts[2:end, 2]..., verts[1, 2]] .-
                         [verts[2:end, 1]..., verts[1, 1]] .* verts[:, 2]
        area = 0.5 * abs(sum(area_components))
        centroid = sum((verts .+ vcat(verts[2:end, :], verts[1, :]')) .*
                        repeat(area_components, 1, 2), dims = 1) / (6 * area)
        diameter = 0 # compute the diameter by looking at every pair of vertices
        for i = 1:(n_sides - 1)
            for j = (i + 1):n_sides
                diameter = max(diameter, norm(verts[i, :] .- verts[j, :]))
            end
        end
        D = zeros(typein, n_sides, n_polys)
        D[:, 1] .= 1
        B = zeros(typein, n_polys, n_sides)
        B[1, :] .= 1 / n_sides
        for vertex_id = 1:n_sides
            vert = verts[vertex_id, :] # this vertex and its neighbours
            prev = verts[mod_wrap(vertex_id - 1, n_sides), :]
            next = verts[mod_wrap(vertex_id + 1, n_sides), :]
            # average of normals on edges
            vertex_normal = [next[2] - prev[2], prev[1] - next[1]]
            # only need to loop over non-constant polynomials
            for poly_id = 2:n_polys
                poly_degree = linear_polynomials[poly_id]
                # gradient of a linear polynomial is constant
                monomial_grad = poly_degree / diameter
                D[vertex_id, poly_id] = dot(vert .- centroid[:],
                                            poly_degree) / diameter
                B[poly_id, vertex_id] = 0.5 * dot(monomial_grad, vertex_normal)
            end
        end
        # compute the local Ritz projector to polynomials
        projector = (B * D) \ B
        stabilising_term = (eye(n_sides) .- D * projector)' *
                           (eye(n_sides) .- D * projector)
        G = B * D
        G[1, :] .= 0
        local_stiffness = projector' * G * projector + stabilising_term
        # copy local to global
        # K[vert_ids, vert_ids] += local_stiffness
        # F[vert_ids] .+= rhs(centroid) * area / n_sides
        append!(IK, repeat(vert_ids, length(vert_ids)))
        append!(JK, repeat(vert_ids', length(vert_ids))[:])
        append!(SK, local_stiffness[:])
        append!(IF, vert_ids)
        append!(SF, repeat(rhs(centroid) * area / n_sides, length(vert_ids)))
    end
    # add boundary contribution in the source term
    # old term F -= K[:, meshboundary] * boundary_vals
    boundary_vals = boundary_condition(pv[meshboundary, :])
    Jinmesh = indexin(JK, meshboundary)
    Jboundary = findall(Jinmesh .!= nothing)
    for l ∈ Jboundary
        push!(IF, IK[l])
        push!(SF, - SK[l] * boundary_vals[Jinmesh[l]])
    end
    return IK, JK, SK, IF, SF
end
#~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~-------------------------
"""
   I, J = indKM_sparse(t)

non zero indices in stiffness and mass P1 sparse matrices
"""
function indKM_sparse(t)
    it1, it2, it3 = t[:, 1], t[:, 2], t[:, 3]
    I = vcat(it1, it2, it3, it2, it3, it1, it1, it2, it3)
    J = vcat(it2, it3, it1, it1, it2, it3, it1, it2, it3)
    return I, J
end
#~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~-------------------------
"""
     K, M = assembKM_P12D(p, t)

assemble the P1 stifness and mass matrices.
"""
function assembKM_P12D(p, t)
    np, dim, nt = size(p, 1), size(p, 2), size(t, 1)
    # compute ar add define arK and arM (not used ?)
    ar = meshareas(p, t)
    arK, arM = ar, ar
    it1, it2, it3 = t[:, 1], t[:, 2], t[:, 3]
    g1x, g1y, g2x, g2y, g3x, g3y = pdetrg(p, t, ar)
    c3 = (g1x .* g2x .+ g1y .* g2y) .* arK
    c1 = (g2x .* g3x .+ g2y .* g3y) .* arK
    c2 = (g3x .* g1x .+ g3y .* g1y) .* arK
    SK = vcat(c3, c1, c2, c3, c1, c2, -c2 - c3, -c3 - c1, -c1 - c2)
    # mass matrix
    aod = arM / 12. # off diagonal element
    ad = 2 * aod    # diagonal element
    SM = vcat(aod, aod, aod, aod, aod, aod, ad, ad, ad)
    # matrices
    I, J = indKM_sparse(t)
    K = sparse(I, J, SK, np, np)
    M = sparse(I, J, SM, np, np)
    return K, M
end
#-------------------------------------------------------------------------------
"""
   g1x, g1y, g2x, g2y, g3x, g3y = pdetrg(p, t, ar)

returns the  the gradient components of the triangle base functions
"""
function pdetrg(p, t, ar)
    # corner point indices
    a1, a2, a3 =  t[:, 1], t[:, 2], t[:, 3]
    # triangle sides
    r23x =  p[a3, 1] .- p[a2, 1]
    r23y =  p[a3, 2] .- p[a2, 2]
    r31x =  p[a1, 1] .- p[a3, 1]
    r31y =  p[a1, 2] .- p[a3, 2]
    r12x =  p[a2, 1] .- p[a1, 1]
    r12y =  p[a2, 2] .- p[a1, 2]
    g1x =  - 0.5 .* r23y ./ ar
    g1y =    0.5 .* r23x ./ ar
    g2x =  - 0.5 .* r31y ./ ar
    g2y =    0.5 .* r31x ./ ar
    g3x =  - 0.5 .* r12y ./ ar
    g3y =    0.5 .* r12x ./ ar
    return g1x, g1y, g2x, g2y, g3x, g3y
end
#-------------------------------------------------------------------------------
"""
    be = btri(T::Array{Int64, 2})

finds boundary edges
"""
function btri(T::Array{Int64, 2})
    # form all faces, non - duplicates are surface triangles
    faces = vcat(T[:, [1, 2]], T[:, [1, 3]], T[:, [2, 3]])
    faces = sort(faces, dims = 2)
    _, ix, jx = uniquerows_int(faces)
    ee = collect(1:(maximum(jx) + 1)) .- 1e-8
    result = fit(Histogram, jx, (1:(maximum(jx) + 1)) .- 1e-8,
                 closed = :right)
    vec = result.weights
    qx = findall(x -> x != 0, vec .- 2)
    return faces[ix[qx], :]
end
#-------------------------------------------------------------------------------
function uniquerows_int(x::Matrix{Int64})
    Vx = [x[k, :] for k = 1:size(x, 1)]
    jj = groupslices(Vx)
    ix, jx = firstinds(jj), similar(jj)
    for k = 1:length(ix)
        jx[jj .== ix[k]] .= k
    end
    Vy = Vx[ix]
    y = [Vy[k][l] for k = 1:length(Vy), l = 1:length(Vy[1])]
    return y, ix, jx
end


### Compute empirical covariance matrix recursively

covariance(lastC::Real, k::Integer, x::Real, lastmean::Real, secondlastmean::Real) =
  ((k-1)*lastC+abs2(x)-(k+1)*abs2(lastmean)+k*abs2(secondlastmean))/k

function covariance!(
  C::RealMatrix,
  lastC::RealMatrix,
  k::Integer,
  x::RealVector,
  lastmean::RealVector,
  secondlastmean::RealVector
)
  C[:, :] = (k-1)*lastC
  BLAS.ger!(1.0, x, x, C)
  BLAS.ger!(1.0, -(k+1)*lastmean, lastmean, C)
  BLAS.ger!(1.0, k*secondlastmean, secondlastmean, C)
  C[:, :] /= k
end


using ITensors,
  ITensorGPU,
  LinearAlgebra, # For tr()
  Combinatorics, # For permutations()
  Random,
  CUDA,
  Test

@testset "cuITensor $T Contractions" for T in (Float64, ComplexF64)
  mi, mj, mk, ml, ma = 2, 3, 4, 5, 6, 7
  i = Index(mi, "i")
  j = Index(mj, "j")
  k = Index(mk, "k")
  l = Index(ml, "l")
  a = Index(ma, "a")
  @testset "Test contract cuITensors" begin
    A = cuITensor(randomITensor(T))
    B = cuITensor(randomITensor(T))
    Ai = cuITensor(randomITensor(T, i))
    Bi = cuITensor(randomITensor(T, i))
    Aj = cuITensor(randomITensor(T, j))
    Aij = cuITensor(randomITensor(T, i, j))
    Aji = cuITensor(randomITensor(T, j, i))
    Bij = cuITensor(randomITensor(T, i, j))
    Aik = cuITensor(randomITensor(T, i, k))
    Ajk = cuITensor(randomITensor(T, j, k))
    Ajl = cuITensor(randomITensor(T, j, l))
    Akl = cuITensor(randomITensor(T, k, l))
    Aijk = cuITensor(randomITensor(T, i, j, k))
    Ajkl = cuITensor(randomITensor(T, j, k, l))
    Aikl = cuITensor(randomITensor(T, i, k, l))
    Akla = cuITensor(randomITensor(T, k, l, a))
    Aijkl = cuITensor(randomITensor(T, i, j, k, l))
    @testset "Test contract cuITensor (Scalar*Scalar -> Scalar)" begin
      C = A * B
      @test scalar(C) ≈ scalar(A) * scalar(B)
      C = cuITensor(T(2.0)) * cuITensor(T(2.0))
      @test scalar(C) ≈ T(4.0)
    end
    @testset "Test contract cuITensor (Scalar*Vector -> Vector)" begin
      C = A * Ai
      @test cpu(C) ≈ scalar(A) * cpu(Ai)
      C = cuITensor(T(2.0)) * Ai
      @test cpu(C) ≈ T(2.0) * cpu(Ai)
      C = Ai * cuITensor(T(2.0))
      @test cpu(C) ≈ T(2.0) * cpu(Ai)
    end
    @testset "Test contract cuITensor (Vector*Scalar -> Vector)" begin
      C = Aj * A
      @test cpu(C) ≈ scalar(A) * cpu(Aj)
    end
    @testset "Test contract cuITensors (Vectorᵀ*Vector -> Scalar)" begin
      C = Ai * Bi
      Ccpu = cpu(Ai) * cpu(Bi)
      @test scalar(Ccpu) ≈ scalar(C)
    end
    @testset "Test contract cuITensors (Vector*Vectorᵀ -> Matrix)" begin
      C = Ai * Aj
      Ccpu = cpu(Ai) * cpu(Aj)
      @test Ccpu ≈ cpu(permute(C, i, j))
    end
    @testset "Test contract cuITensors (Matrix*Scalar -> Matrix)" begin
      Aij = permute(Aij, i, j)
      C = Aij * A
      @test cpu(permute(C, i, j)) ≈ scalar(A) * cpu(Aij)
    end
    @testset "Test contract cuITensors (Matrixᵀ*Vector -> Vector)" begin
      cAij = permute(copy(Aij), j, i)
      Ccpu = cpu(Aij) * cpu(Aj)
      C = cAij * Aj
      @test Ccpu ≈ cpu(C)
    end
    @testset "Test contract cuITensors (Matrix*Vector -> Vector)" begin
      cpAij = permute(copy(Aij), i, j)
      Ccpu = cpu(cpAij) * cpu(Aj)
      C = copy(cpAij) * copy(Aj)
      @test Ccpu ≈ cpu(C)
    end
    @testset "Test contract cuITensors (Vector*Matrix -> Vector)" begin
      Aij = permute(Aij, i, j)
      C = Ai * Aij
      Ccpu = cpu(Ai) * cpu(Aij)
      @test Ccpu ≈ cpu(C)
    end
    @testset "Test contract cuITensors (Vector*Matrixᵀ -> Vector)" begin
      C = Ai * Aji
      Ccpu = cpu(Ai) * cpu(Aji)
      @test Ccpu ≈ cpu(C)
    end
    @testset "Test contract cuITensors (Matrix*Matrix -> Scalar)" begin
      Aij = permute(Aij, i, j)
      Bij = permute(Bij, i, j)
      C = Aij * Bij
      Ccpu = cpu(Aij) * cpu(Bij)
      @test scalar(Ccpu) ≈ scalar(C)
    end
    @testset "Test contract cuITensors (Matrix*Matrix -> Matrix)" begin
      Aij = permute(Aij, i, j)
      Ajk = permute(Ajk, j, k)
      C = Aij * Ajk
      Ccpu = cpu(Aij) * cpu(Ajk)
      @test Ccpu ≈ cpu(C)
    end
    @testset "Test contract cuITensors (Matrixᵀ*Matrix -> Matrix)" begin
      Aij = permute(Aij, j, i)
      Ajk = permute(Ajk, j, k)
      C = Aij * Ajk
      Ccpu = cpu(Aij) * cpu(Ajk)
      @test Ccpu ≈ cpu(C)
    end
    @testset "Test contract cuITensors (Matrix*Matrixᵀ -> Matrix)" begin
      Aij = permute(Aij, i, j)
      Ajk = permute(Ajk, k, j)
      C = Aij * Ajk
      Ccpu = cpu(Aij) * cpu(Ajk)
      @test Ccpu ≈ cpu(C)
    end
    @testset "Test contract cuITensors (Matrixᵀ*Matrixᵀ -> Matrix)" begin
      Aij = permute(Aij, j, i)
      Ajk = permute(Ajk, k, j)
      C = Aij * Ajk
      Ccpu = cpu(Aij) * cpu(Ajk)
      @test Ccpu ≈ cpu(C)
    end
    @testset "Test contract cuITensors (3-Tensor*Scalar -> 3-Tensor)" begin
      Aijk = permute(Aijk, i, j, k)
      C = Aijk * A
      @test cpu(permute(C, i, j, k)) ≈ scalar(A) * cpu(Aijk)
    end
    @testset "Test contract cuITensors (3-Tensor*Vector -> Matrix)" begin
      cAijk = permute(copy(Aijk), i, j, k)
      C = cAijk * Ai
      Ccpu = cpu(cAijk) * cpu(Ai)
      @test Ccpu ≈ cpu(C)
    end
    @testset "Test contract cuITensors (Vector*3-Tensor -> Matrix)" begin
      Aijk = permute(Aijk, i, j, k)
      C = Aj * Aijk
      Ccpu = cpu(Aj) * cpu(Aijk)
      @test Ccpu ≈ cpu(C)
    end
    @testset "Test contract cuITensors (3-Tensor*Matrix -> Vector)" begin
      Aijk = permute(Aijk, i, j, k)
      Aik = permute(Aik, i, k)
      C = Aijk * Aik
      Ccpu = cpu(Aijk) * cpu(Aik)
      @test Ccpu ≈ cpu(C)
    end
    @testset "Test contract cuITensors (3-Tensor*Matrix -> 3-Tensor)" begin
      Aijk = permute(Aijk, i, j, k)
      Ajl = permute(Ajl, j, l)
      C = Aijk * Ajl
      Ccpu = cpu(Aijk) * cpu(Ajl)
      @test Ccpu ≈ cpu(permute(C, i, k, l))
    end
    @testset "Test contract cuITensors (Matrix*3-Tensor -> 3-Tensor)" begin
      Aijk = permute(Aijk, i, j, k)
      Akl = permute(Akl, k, l)
      C = Akl * Aijk
      Ccpu = cpu(Aijk) * cpu(Akl)
      @test Ccpu ≈ cpu(permute(C, l, i, j))
    end
    @testset "Test contract cuITensors (3-Tensor*3-Tensor -> 3-Tensor)" begin
      Aijk = permute(Aijk, i, j, k)
      Ajkl = permute(Ajkl, j, k, l)
      C = Aijk * Ajkl
      Ccpu = cpu(Aijk) * cpu(Ajkl)
      @test Ccpu ≈ cpu(permute(C, i, l))
    end
    @testset "Test contract cuITensors (3-Tensor*3-Tensor -> 3-Tensor)" begin
      for inds_ijk in permutations([i, j, k]), inds_jkl in permutations([j, k, l])
        Aijk = permute(Aijk, inds_ijk...)
        Ajkl = permute(Ajkl, inds_jkl...)
        C = Ajkl * Aijk
        Ccpu = cpu(Ajkl) * cpu(Aijk)
        @test Ccpu ≈ cpu(C)
      end
    end
    @testset "Test contract cuITensors (4-Tensor*3-Tensor -> 3-Tensor)" begin
      for inds_ijkl in permutations([i, j, k, l]), inds_kla in permutations([k, l, a])
        Aijkl = permute(Aijkl, inds_ijkl...)
        Akla = permute(Akla, inds_kla...)
        C = Akla * Aijkl
        Ccpu = cpu(Akla) * cpu(Aijkl)
        @test Ccpu ≈ cpu(C)
      end
    end
    @testset "Test contract cuITensors (4-Tensor*3-Tensor -> 1-Tensor)" begin
      for inds_ijkl in permutations([i, j, k, l]), inds_jkl in permutations([j, k, l])
        Aijkl = permute(Aijkl, inds_ijkl...)
        Ajkl = permute(Ajkl, inds_jkl...)
        C = Ajkl * Aijkl
        Ccpu = cpu(Ajkl) * cpu(Aijkl)
        @test Ccpu ≈ cpu(C)
      end
    end
    @testset "Test supersized contract cuITensors (14-Tensor*14-Tensor -> 14-Tensor)" begin
      a_only_inds = [Index(2) for ii in 1:7]
      b_only_inds = [Index(2) for ii in 1:7]
      shared_inds = [Index(2) for ii in 1:7]
      A = randomITensor(IndexSet(vcat(a_only_inds, shared_inds)))
      B = randomITensor(IndexSet(vcat(b_only_inds, shared_inds)))
      cA = cuITensor(A)
      cB = cuITensor(B)
      inds_a = vcat(a_only_inds, shared_inds)
      inds_b = vcat(b_only_inds, shared_inds)
      cA_ = permute(cA, inds_a...)
      cB_ = permute(cB, inds_b...)
      @disable_warn_order begin
        C = cA_ * cB_
        Ccpu = cpu(cA_) * cpu(cB_)
      end
      @test Ccpu ≈ cpu(C)
      for shuffles in 1:1 # too many permutations to test all
        inds_a = shuffle(vcat(a_only_inds, shared_inds))
        inds_b = shuffle(vcat(b_only_inds, shared_inds))
        cA_ = permute(cA, inds_a...)
        cB_ = permute(cB, inds_b...)
        @disable_warn_order begin
          C = cA_ * cB_
          Ccpu = cpu(cA_) * cpu(cB_)
        end
        @test Ccpu ≈ cpu(C)
      end
    end
  end # End contraction testset
end


mutable struct parameters
      H::Float64
     Ly::Float64
     f₀::Float64
      g::Float64
     Lj::Float64
     Uj::Float64
     Ny::Int64
     dk::Float64
     kₘ::Float64
     Fr::Float64
     Ro::Float64
end

function parameters(;
    H  = 500,
    Ly = 2e5,
    f₀ = 1e-4,
    g  = 9.81,
    Lj = 1e4,
    Uj = 0.1,
    Ny = 100,
    dk = 5e-6,
    kₘ = 2e-4,
    Fr = 0.0,
    Ro = 0.0
    )
    return parameters(H, Ly, f₀ , g, Lj, Uj, Ny, dk, kₘ, Fr, Ro)
end


using ITensors

function local_boltzmann_weight(m::String, args...; kwargs...)
  return local_boltzmann_weight(Model(m), args...; kwargs...)
end

f(λ₊, λ₋) = [
  (λ₊ + λ₋)/2 (λ₊ - λ₋)/2
  (λ₊ - λ₋)/2 (λ₊ + λ₋)/2
]

function sqrt_bond_matrix(; β::Real, J::Real=1.0)
  # Alternative method
  #Q = [exp(β * J) exp(-β * J); exp(-β * J) exp(β * J)]
  #return √Q
  λ₊ = √(exp(β * J) + exp(-β * J))
  λ₋ = √(exp(β * J) - exp(-β * J))
  return f(λ₊, λ₋)
end

# The local Boltzmann weight for the Ising
# model in 1 dimension
function local_boltzmann_weight(
  ::Model"ising", ::Val{1}; β::Real, J::Real=1.0, sz::Bool=false
)
  d = 2 # local dimension of the Ising local Boltzmann factor
  s, s′ = Index.((d, d))
  T = ITensor(s, s′)
  for i in 1:d
    T[i, i] = 1.0
  end
  if sz
    T[1, 1] = -T[1, 1]
  end
  s̃, s̃′ = sim.((s, s′))
  T̃ = T * δ(s, s̃) * δ(s′, s̃′)
  sqrtQ = sqrt_bond_matrix(; β=β, J=J)
  @show sqrtQ
  X = itensor(vec(sqrtQ), s̃, s)
  X′ = itensor(vec(sqrtQ), s̃′, s′)
  return array(permute(T̃ * X′ * X, s, s′))
end

# The local Boltzmann weight for the Ising
# model in 2 dimensions
function local_boltzmann_weight(
  ::Model"ising", ::Val{2}; β::Real, J::Real=1.0, sz::Bool=false
)
  d = 2 # local dimension of the Ising local Boltzmann factor
  sₕ, sₕ′ = Index.((d, d))
  sᵥ, sᵥ′ = Index.((d, d))
  @assert dim(sₕ) == dim(sᵥ)
  d = dim(sₕ)
  T = ITensor(sₕ, sₕ′, sᵥ, sᵥ′)
  for i in 1:d
    T[i, i, i, i] = 1.0
  end
  if sz
    T[1, 1, 1, 1] = -T[1, 1, 1, 1]
  end
  s̃ₕ, s̃ₕ′, s̃ᵥ, s̃ᵥ′ = sim.((sₕ, sₕ′, sᵥ, sᵥ′))
  T̃ = T * δ(sₕ, s̃ₕ) * δ(sₕ′, s̃ₕ′) * δ(sᵥ, s̃ᵥ) * δ(sᵥ′, s̃ᵥ′)
  X = sqrt_bond_matrix(; β=β, J=J)
  Xₕ = itensor(vec(X), s̃ₕ, sₕ)
  Xₕ′ = itensor(vec(X), s̃ₕ′, sₕ′)
  Xᵥ = itensor(vec(X), s̃ᵥ, sᵥ)
  Xᵥ′ = itensor(vec(X), s̃ᵥ′, sᵥ′)
  return array(permute(T̃ * Xₕ′ * Xᵥ′ * Xₕ * Xᵥ, sₕ, sₕ′, sᵥ, sᵥ′))
end

function mpo_itensor(
  m::Model, pair_sₕ::Pair{<:Index,<:Index}, pair_sᵥ::Pair{<:Index,<:Index}; kwargs...
)
  sₕ, sₕ′ = pair_sₕ
  sᵥ, sᵥ′ = pair_sᵥ
  return itensor(mpo_array(m; kwargs...), sₕ, sₕ′, sᵥ, sᵥ′)
end

critical_point(::Model"ising") = 0.5 * log(√2 + 1)

function free_energy(::Model"ising"; β::Real, J::Real=1.0)
  k = β * J
  c = cosh(2 * k)
  s = sinh(2 * k)
  xmin = 0.0
  xmax = π
  integrand(x) = log(c^2 + √(s^4 + 1 - 2 * s^2 * cos(x)))
  integral, err = quadgk(integrand, xmin, xmax)::Tuple{Float64,Float64}
  return -(log(2) + integral / π) / (2 * β)
end

function magnetization(::Model"ising"; β::Real)
  β > βc && return (1 - sinh(2 * β)^(-4))^(1 / 8)
  return 0.0
end


"""
    x,y = squaredalitz(k,σs,ms)

    Calculate normalized values in square coordinates, -1 ≤ x,y ≤ 1
"""
function squaredalitz(k,σs,ms::ThreeBodyMasses)
    cθ = cosθij(k,σs,ms^2)
    cθ = cθ ≥ 1.0 ? 1.0 : (cθ ≤ -1.0 ? -1.0 : cθ)
    y = acos(cθ) / π
    mlimsk = sqrt.(lims(k, ms))
    xn = 2*(sqrt(σs[k]) - mlimsk[1]) / (mlimsk[2]-mlimsk[1])-1
    x = acos(xn) / π
    return (x, y)
end

squaredalitz1(σs,ms) = squaredalitz(1,σs,ms)
squaredalitz2(σs,ms) = squaredalitz(2,σs,ms)
squaredalitz3(σs,ms) = squaredalitz(3,σs,ms)

"""
    invsquaredalitz(k,x,y,ms)

    Inverse transformation from square coordinates
"""
function invsquaredalitz(k,x,y,ms::ThreeBodyMasses)
    mlimsk = sqrt.(lims(k, ms))
    #
    _mk = mlimsk[1] + (mlimsk[2]-mlimsk[1]) * (cos(x*π) + 1) / 2
    _σk = _mk^2
    _cosθij = cos(y*π)
    #
    k == 1 && return Invariants(ms, σ1 = _σk,  σ3 = σ3of1(_cosθij,_σk,ms^2))
    k == 2 && return Invariants(ms, σ2 = _σk,  σ1 = σ1of2(_cosθij,_σk,ms^2))
    k != 3 && error("k should be 1,2, or 3, while it is $k")
    return Invariants(ms, σ3 = _σk,  σ2 = σ2of3(_cosθij,_σk,ms^2))
end

invsquaredalitz1(x,y,ms) = invsquaredalitz(1,x,y,ms)
invsquaredalitz2(x,y,ms) = invsquaredalitz(2,x,y,ms)
invsquaredalitz3(x,y,ms) = invsquaredalitz(3,x,y,ms)

"""
    jacobian_squaredalitz(k,σs,ms)

calculates jacobian of transformation to square coordinates
"""
function jacobean_squaredalitz(k,σs,ms::ThreeBodyMasses)
    (i,j) = ij_from_k(k)
    #
    cθ = cosθij(k,σs,ms^2)
    cθ = cθ ≥ 1.0 ? 1.0 : (cθ ≤ -1.0 ? -1.0 : cθ)
    dydcθ = 1 / sqrt(1-cθ^2) / π
    #
    m = sqrt(σs[k])
    sqrtlimsk = sqrt.(lims(k, ms))
    xn = 2*(m - sqrtlimsk[1]) / (sqrtlimsk[2]-sqrtlimsk[1]) - 1
    dxdm = 2 / (sqrtlimsk[2]-sqrtlimsk[1]) / sqrt(1-xn^2) / π
    ρρm = sqrt(Kallen(σs[k],ms[i]^2,ms[j]^2)*Kallen(ms[4]^2,σs[k],ms[k]^2)) / (σs[k]*ms[4]^2) * m
    return ρρm / dydcθ / dxdm #
end

jacobean_squaredalitz1(σs,ms) = jacobean_squaredalitz(1,σs,ms)
jacobean_squaredalitz2(σs,ms) = jacobean_squaredalitz(2,σs,ms)
jacobean_squaredalitz3(σs,ms) = jacobean_squaredalitz(3,σs,ms)


using BasisFunctions: hasderivative, hasantiderivative

P = 80

@testset "$(rpad("Lagrange polynomials",P))" begin

    ξ = [-1.0, +1.0]
    b = Lagrange(ξ)
    blob = LagrangeLobattoLegendre(2)
    bgau = LagrangeGaussLegendre(2)
    blch = LagrangeLobattoChebyshev(2)
    bgch = LagrangeGaussChebyshev(2)

    BasisFunctions.Test.test_generic_dict_interface(LagrangeGaussLegendre(10))

    @testset "$(rpad("Basic functionality",P))" begin
        @test nodes(b) == ξ
        @test nnodes(b) == 2
        @test degree(b) == 1
        @test support(b) == LagrangeInterval{eltype(ξ)}()
        @test nodes(blob) == nodes(b)
        @test nodes(blch) == nodes(b)
        @test nodes(bgau) == [-0.5773502691896258, 0.5773502691896258]
        @test nodes(bgch) == [-0.7071067811865475, 0.7071067811865475]
        @test nodes(blob) == nodes(similar(blob, eltype(ξ), 2))
        @test nodes(bgau) == nodes(similar(bgau, eltype(ξ), 2))
    end

    @testset "$(rpad("Basis evaluation",P))" begin
        @test b[1](-1.0) == 1.0
        @test b[1]( 0.0) == 0.5
        @test b[1](+1.0) == 0.0

        @test b[2](-1.0) == 0.0
        @test b[2]( 0.0) == 0.5
        @test b[2](+1.0) == 1.0
    end

    @testset "$(rpad("Basis derivative",P))" begin
        @test eval_element_derivative(b, 1, -1.0) == -0.5
        @test eval_element_derivative(b, 1,  0.0) == -0.5
        @test eval_element_derivative(b, 1, +1.0) == -0.5

        @test eval_element_derivative(b, 2, -1.0) == +0.5
        @test eval_element_derivative(b, 2,  0.0) == +0.5
        @test eval_element_derivative(b, 2, +1.0) == +0.5
    end

    @testset "$(rpad("Basis antiderivative",P))" begin
        # @test moment(b, 1) == 1.0
        # @test moment(b, 2) == 1.0

        @test eval_element_antiderivative(b, 1, -1.0) == 0.0
        @test eval_element_antiderivative(b, 1,  0.0) == 0.75
        @test eval_element_antiderivative(b, 1, +1.0) == 1.0

        @test eval_element_antiderivative(b, 2, -1.0) == 0.0
        @test eval_element_antiderivative(b, 2,  0.0) == 0.25
        @test eval_element_antiderivative(b, 2, +1.0) == 1.0
    end

end


@testset "Testing branch length distance LS score" begin

# on a tree:
net = readTopology("(A,(B,(C,D)));")
net2 = readTopology("(((A,B),C),D);")
divergenceTimes = Dict(1 => 0.0, 2 => 0.0, 3 => 0.0, 4 => 0.0, -4 => 10.0, -3 => 15.0, -2 => 25.0)
setBLfromDivergenceTimes!(net2, divergenceTimes)
dnaDistances = pairwiseTaxonDistanceMatrix(net2)
distanceMismatch = pairwiseDistanceLSscore(net, divergenceTimes, dnaDistances)
@test distanceMismatch==130.7669683062202

# on a network:
net = readTopology("(((A,(B)#H1:::0.9),(C,#H1:::0.1)),D);")
net2 = readTopology("(((D,(C)#H1:::0.9,(B,#H1:::0.1))),A);")
divergenceTimes = Dict(1 => 0.0, 2 => 0.0, 4 => 0.0, 5 => 0.0, -4 => 10.0, -3 => 15.0, -2 => 25.0, -6 => 7.0, 3 => 5.0)
setBLfromDivergenceTimes!(net2, divergenceTimes)
dnaDistances = pairwiseTaxonDistanceMatrix(net2)
distanceMismatch = pairwiseDistanceLSscore(net, divergenceTimes, dnaDistances)
@test distanceMismatch==63.75829357816911

end


struct StdSquare{T}
    half_side::T
end

struct StdCircle{T}
    radius::T
end

function get_projection(square::StdSquare, position::AbstractVector)
    half_side = square.half_side
    return clamp.(position, -half_side, half_side)
end

function is_colliding(square::StdSquare, circle::StdCircle, position::AbstractVector)
    projection = get_projection(square, position)
    vec = position .- projection
    radius = circle.radius
    return sum(vec .^ 2) < radius * radius
end

function is_player_colliding(obstacle_map, player_position_wu, player_radius_wu::T) where {T}
    height_tile_map_tu, width_tile_map_tu = size(obstacle_map)

    square = StdSquare(convert(T, 0.5))
    circle = StdCircle(player_radius_wu)

    i_player_position_tu = wu_to_tu(player_position_wu[1])
    j_player_position_tu = wu_to_tu(player_position_wu[2])

    for j in j_player_position_tu - 1 : j_player_position_tu + 1
        for i in i_player_position_tu - 1 : i_player_position_tu + 1
            tile_position_wu = similar(player_position_wu)
            tile_position_wu[1] = i - convert(T, 0.5)
            tile_position_wu[2] = j - convert(T, 0.5)
            if obstacle_map[i, j] && is_colliding(square, circle, player_position_wu .- tile_position_wu)
                return true
            end
        end
    end

    return false
end


export sinint, cosint

function sinint{T<:BigFloat}(z::T;n::Integer=2^6)
    Tπ=convert(T,π);dDE=Tπ/2;ga=one(T);b2=one(T);h=log(Tπ*dDE*ga*n/b2)/ga/n
    phi(t::T) = log(exp(Tπ/2*sinh(t))+one(T))
    phip(t::T) = Tπ/2*cosh(t)/(one(T)+exp(-Tπ/2*sinh(t)))
    phiinv(x::T) = asinh(2/Tπ*log(expm1(x)))
    val1,val2=zero(T),zero(T)
    for j=-n:n
        jh=j*h
        phijh = phi(jh)
        temp = exp(-phijh)*phip(jh)/(phijh^2+z^2)*(one(T)-(-1)^(one(T)*j)*exp(-phiinv(abs(z))))
        val1 += temp*z
        val2 += temp*z*phijh
    end
    retval = Tπ/2 - h*(val1*cos(z)+val2*sinc(z/Tπ))
    return z == zero(T) ? zero(T) : z >= zero(T) ? retval : retval-Tπ
end

function Base.sinc{T<:Number}(n::Integer,x::T)
#This program computes the nth sinc differentiation matrices.
    val = zero(T)
    if n ≥ 0
        if x == 0
            val = isodd(n) ? zero(T) : (-1)^(n/2)*convert(T,π)^n/(n+1)
        else
            sp,cp,πk,xk,den = sinpi(x),cospi(x),one(T)/π,one(T),1
            for k=0:n
                m = mod(k,4)
                if m == 0
                    val -= sp*πk*xk/den
                elseif m == 1
                    val += cp*πk*xk/den
                elseif m == 2
                    val += sp*πk*xk/den
                else
                    val -= cp*πk*xk/den
                end
                πk,xk,den = πk*π,xk*x,den*(k+1)
            end
            val *= iseven(n) ? -den/(n+1)/xk : den/(n+1)/xk
        end
    elseif n == -1
        val = 0.5+sinint(π*x)/π
    end
    return val
end
Base.sinc{T<:Number}(n::T,x::T) = sinc(rount(Int,n),x)


"""
    CategoricalLikelihood

Categorical likelihood is to be used if we assume that the
uncertainity associated with the data follows a Categorical distribution.
```math
    p(y|f_1, f_2, \\dots, f_{n-1}) = Categorical(y | softmax(f_1, f_2, \\dots, f_{n-1}, 0))
```
On calling, this would return a Categorical distribution with `f_i`
probability of `i` category.
"""
struct CategoricalLikelihood end

(l::CategoricalLikelihood)(f::AbstractVector{<:Real}) = Categorical(softmax(vcat(f, 0)))

(l::CategoricalLikelihood)(fs::AbstractVector) = Product(Categorical.(softmax.(vcat.(fs, 0))))


function tdc(T::SymTridiagonal)

    # Tridiagonal divide-and-conquer using arrowhead matrices
    # COMPUTES: eigenvectors U and eigenvalues E of a SymTridiagonal matrix T,
    # T = U*diagm(E)*U'
    # RETURNS: U, E

    n=length(T.dv)
    # set the tolerances for eig
    tols=[1e2,1e2,1e2,1e2,1e2]
    if n==1
        U=[1.0]
        E=T.dv[1]
    elseif n==2
        U,E,rest=eig(SymArrow([T.dv[1]],[T.ev[1]],T.dv[2],2),tols)
    else
        k=round(Int,floor(n/2))
        T1=SymTridiagonal(T.dv[1:k],T.ev[1:k-1])
        T2=SymTridiagonal(T.dv[k+2:end],T.ev[k+2:end])
        U1,E1=tdc(T1)
        U2,E2=tdc(T2)
        # form arrowhead
        D=[E1;E2]
        z=[U1'[:,end]*T.ev[k]; U2'[:,1]*T.ev[k+1]]
        a=T.dv[k+1]
        A=SymArrow(D,z,a,k+1)
        U,E,rest=eig(A,tols)
        U[1:k,:]=U1*U[1:k,:]
        U[k+2:end,:]=U2*U[k+2:end,:]
    end
    return U,E
end

# roots
function rootsah(pol::Union{Poly{Float32},Poly{Float64},Poly{Int32},Poly{Int64}}, D::Vector{Float64})

    # COMPUTES: the roots of polynomials with all distinct real roots.
    # The computation is forward stable. The program uses SymArrow (arrowhead) companion matrix and
    # corresponding eig routine
    # D are barycentric coordinates - elements od D must interpolate the roots of P,
    # for example
    #              D=roots(polyder(pol))
    # RETURNS: roots E

    T = Float64

    # tols = [1e2,1e2,1e2,1e2,1e2] or similar is the vector of tolerances for eig
    tols=[1e2,1e2,1e2,1e2,1e2]

    Dm=map(T,D)
    Dm=sort(Dm,rev=true)
    D=sort(D,rev=true)

    p=map(T,[pol[i] for i=0:1:length(pol)-1])
    p=p[end:-1:1]

    # Compute z of the arrowhead
    # First we compute values s=p(D) we use Horner sheme with Double
    n=length(p)-1
    pD=map(Double,p)
    DD=map(Double,Dm)
    oneD=Double(one(T))
    s=Array(Double{T},n-1)
    for i=1:n-1
        s[i]=oneD
    end
    s=pD[1]*s
    for i=2:n+1
        r=s.*DD
        s=r+pD[i]
    end

    # Compute t's
    t=Array(Double{T},n-1)
    for j=1:n-1
        h=oneD
        for i=1:j-1
            g=DD[j]-DD[i]
            h=h*g
        end
        for i=j+1:n-1
            g=DD[j]-DD[i]
            h=h*g
        end
        t[j]=h
    end
    # Compute alphaD
    alphaD=pD[2]/pD[1]
    for i=1:n-1
        alphaD+=DD[i]
    end
    # alphaD=-alphaD
    alphaD=Double(-alphaD.hi,-alphaD.lo)
    #  Compute z
    zD=Array(Double{T},n-1)
    for i=1:n-1
        #         zD[i]=sqrt((-s[i])/(t[i]*pD[1]))
        zD[i]=sqrt(Double(-s[i].hi,-s[i].lo)/(t[i]*pD[1]))
    end

    z=Array(Float64,n-1)
    for i=1:n-1
        z[i]=Float64(zD[i].lo+zD[i].hi)
    end
    alpha=Float64(alphaD.hi+alphaD.lo)
    A=SymArrow(D,z,alpha,n)

    E=zeros(Float64,n)
    Qout=zeros(Int64,n)
    for k=1:n
        E[k], Qout[k] = eig( A,zD,alphaD,k,tols )
    end

    E, Qout
end

#------------------------
#--------  Functions

#-------- Inverses

function inv{T}(A::SymArrow{T},zD::Vector{Double{T}}, alphaD::Double{T},i::Int64,tols::Vector{Float64})

    # COMPUTES: inverse of a SymArrow matrix A, inv(A-A.D[i]*I) which is again SymArrow
    # uses higher precision to compute top of the arrow element accurately, if
    # needed.
    # tols=[tolb,tolz] are tolerances, usually [1e3, 10*n]
    # [0.0,0.0] forces DoubleDouble, [1e50,1e50] would never use it
    # RETURNS:  SymArrow(D,z,b,i), Kb, Kz, Qout
    # Kb - condition Kb, Kz - condition Kz, Qout = 1 / 0 - double was / was not used

    n=length(A.D)
    D=Array(T,n)
    z=Array(T,n)
    wz=1/A.z[i]
    shift=A.D[i]

    for k=1:i-1
        D[k]=one(T)/(A.D[k]-shift)
        z[k]=-A.z[k]*D[k]*wz
    end
    for k=i+1:n
        D[k-1]=one(T)/(A.D[k]-shift)
        z[k-1]=-A.z[k]*D[k-1]*wz
    end

    # D=[1./(A.D[1:i-1]-shift),1./(A.D[i+1:end]-shift)]
    # wz=1/A.z[i]
    # z=-[A.z[1:i-1], A.z[i+1:end]].*D1*wz1
    a=A.a-A.D[i]

    # compute the sum in a plain loop
    P=zero(T)
    Q=zero(T)
    Qout=0

    nn=n-1
    for k=1:i-1
        D[k]>0.0 ? P+=A.z[k]^2*D[k] : Q+=A.z[k]^2*D[k]
    end
    for k=i:nn
        D[k]>0.0 ? P+=A.z[k+1]^2*D[k] : Q+=A.z[k+1]^2*D[k]
    end
    a<0 ? P=P-a : Q=Q-a

    Kb=(P-Q)/abs(P+Q)
    Kz=maxabs(A.z)*abs(wz)

    if Kb<tols[1] ||  Kz<tols[2]
        b=(P+Q)*wz*wz
    elseif Kb==Inf
       b=0.0
    else  # recompute in Double
        Qout=1
        AD=map(Double{T},A.D)
        shiftd=AD[i]
        #shiftd=map(Double,A.D[i])
        Dd=[Double{T}[AD[k]-shiftd for k=1:i-1];
            Double{T}[AD[k]-shiftd for
                            k=i+1:length(A.D)]]
        wzd=zD[i] # Double(A.z[i])
        ad=alphaD-shiftd

        Pd,Qd=map(Double{T},(0.0,0.0))

        for k=1:i-1
            Dd[k].hi>0.0 ? Pd+=zD[k]^2/Dd[k] :
            Qd+=zD[k]^2/Dd[k]
        end

        for k=i:nn
            Dd[k].hi>0.0 ? Pd+=zD[k+1]^2/Dd[k] :
            Qd+=zD[k+1]^2/Dd[k]
        end

        ad.hi<0 ?   Pd=Pd-ad : Qd=Qd-ad

        bd=(Pd+Qd)/(wzd*wzd)
        b=Float64(bd.hi+bd.lo)
    end

    if i<A.i
        # SymArrow([D[1:A.i-2],zero(T),D[A.i-1:end]],[z[1:A.i-2],wz,z[A.i-1:end]],b,i),Kb,Kz,Qout
        # This reduces memory allocation (4 vectors) and is 30-40% faster
        for k=n-1:-1:A.i-1
            D[k+1]=D[k]
            z[k+1]=z[k]
        end
        D[A.i-1]=zero(T)
        z[A.i-1]=wz
        SymArrow(D,z,b,i),Kb,Kz,Qout
    else
        # SymArrow([D[1:A.i-1],zero(T),D[A.i:end]],[z[1:A.i-1],wz,z[A.i:end]],b,i+1),Kb,Kz,Qout
        for k=n-1:-1:A.i
            D[k+1]=D[k]
            z[k+1]=z[k]
        end
        D[A.i]=zero(T)
        z[A.i]=wz
        SymArrow(D,z,b,i+1),Kb,Kz,Qout
    end

end # inv

function inv{T}(A::SymArrow{T}, zD::Vector{Double{T}}, alphaD::Double{T}, shift::Float64, tolr::Float64)

    # COMPUTES: inverse of the shifted SymArrow A, inv(A-shift*I) which is SymDPR1
    # uses DoubleDouble to compute rho accurately, if needed.
    # tolr is tolerance, usually 1e3,  0.0 forces Double, 1e50 would never use it
    # RETURNS: SymDPR1(D,u,rho), Krho, Qout
    # Krho - condition Krho, Qout = 1 / 0 - Double was / was not used

    n=length(A.D)
    D=Array(T,n+1)
    u=Array(T,n+1)

    # Ds=A.D-shift

    # D=[1./Ds[1:A.i-1];0.0;1./Ds[A.i:end]];
    # u=[A.z[1:A.i-1]./Ds[1:A.i-1];-1.0;A.z[A.i:end]./Ds[A.i:end]];

    for k=1:A.i-1
        D[k]=one(T)/(A.D[k]-shift)
        u[k]=A.z[k]*D[k]
    end
    D[A.i]=zero(T)
    u[A.i]=-one(T)
    for k=A.i:n
        D[k+1]=one(T)/(A.D[k]-shift)
        u[k+1]=A.z[k]*D[k+1]
    end

    # compute rho and Krho
    #--- compute the sum in a plain loop

    a=A.a-shift
    P=zero(T)
    Q=zero(T)
    Qout=0

    for k=1:A.i-1
        D[k]>0.0 ? P=P+A.z[k]^2*D[k] : Q=Q+A.z[k]^2*D[k]
    end
    for k=A.i:n
        D[k+1]>0.0 ? P=P+A.z[k]^2*D[k+1] : Q=Q+A.z[k]^2*D[k+1]
    end

    a<0 ? P=P-a : Q=Q-a

    # Condition of rho
    Krho=(P-Q)/abs(P+Q)

    if Krho<tolr
        rho=-1.0/(P+Q)

    else  # recompute in Double
        Qout=1
        Pd,Qd=map(Double{T},(0.0,0.0))
        AD=map(Double{T},A.D)
        shiftd=map(Double{T},shift)
        ad=alphaD-shiftd
        #   Dd=[Double(A.D[k])-shiftd for k=1:length(A.D) ]
        #   Dd=map(Double,A.D)-shiftd
        for k=1:A.i-1
            D[k]>0.0 ? Pd=Pd+zD[k]^2/(AD[k]-shiftd) : Qd=Qd+zD[k]^2/(AD[k]-shiftd)
        end
        for k=A.i:n
            D[k+1]>0.0 ? Pd=Pd+zD[k]^2/(AD[k]-shiftd) : Qd=Qd+zD[k]^2/(A.D[k]-shiftd)
        end

        # for k=1:length(A.D)
        #    Dd[k].hi>0.0 ? Pd=Pd+Double(A.z[k])^2/Dd[k] : Qd=Qd+Double(A.z[k])^2/Dd[k]
        # end
        ad.hi+ad.lo<0 ? Pd=Pd-ad : Qd=Qd-ad
        r=map(Double{T},1.0)/(Pd+Qd)
        rho=Float64(-(r.hi+r.lo))
    end

    # returns the following
    SymDPR1(D,u,rho), Krho, Qout

end # inv

function inv{T}(A::SymArrow{T}, zD::Vector{Double{T}}, alphaD::Double{T}, shift::Double{Float64})

    # COMPUTES: inverse of the shifted SymArrow A, inv(A-shift*I), which is a SymDPR1
    # here shift is Double so it uses Double to compute everything
    # RETURNS: SymDPR1(D1,u1,rho), Qout
    # Qout = 1 on exit meaning Double was used

    n=length(A.D)
    D=Array(Double,n+1)
    u=Array(Double,n+1)
    shiftd=map(Double{T},shift)

    # D[1:n]=map(Double,A.D)-shift
    for k=1:n
        D[k]=map(Double{T},A.D[k])-shiftd
    end

    u[1:n]=zD
    a=alphaD-shiftd
    i=A.i

    oned=Double(one(T))
    zerod=Double(zero(T))

    for k=1:i-1
        u[k]=u[k]/D[k]
        D[k]=oned/D[k]
    end

    for k=i:n
        u[k+1]=u[k]/D[k]
        D[k+1]=oned/D[k+1]
    end

    D[i]=zerod
    u[i]=-oned

    # compute rho and Krho
    #--- compute the sum in a plain loop

    P,Q=zerod,zerod
    Qout=1

    for k=1:i-1
        # D[k].hi+D[k].lo > 0.0  # .lo cannot change sign, I think
        # D[k].hi > 0.0 ? P=P+Double(A.z[k])^2*D[k] : Q=Q+Double(A.z[k])^2*D[k]
        D[k].hi > 0.0 ? P=P+zD[k]*u[k] : Q=Q+zD[k]*u[k]
    end
    for k=i:n
        D[k+1].hi >0.0 ? P=P+zD[k]*u[k+1] : Q=Q+zD[k]*u[k+1]
    end

    a.hi+a.lo<0.0  ? P=P-a : Q=Q-a
    r=oned/(P+Q)

    rho=Float64(-(r.hi+r.lo))

    # returns the following
    D1=Array(T,n+1)
    u1=Array(T,n+1)
    for k=1:n+1
        D1[k]=Float64(D[k].hi+D[k].lo)
    end
    for k=1:n+1
        u1[k]=Float64(u[k].hi+u[k].lo)
    end

    SymDPR1(D1,u1,rho), Qout

end # inv

function  eig{T}( A::SymArrow{T},zD::Vector{Double{T}},alphaD::Double{T},k::Int64,tols::Vector{Float64})

    # COMPUTES: k-th eigenvalue of an ordered irreducible SymArrow
    # A = [diag (D) z; z' alpha]
    # Specially designed to be used in the polynomial rootfinder rootsah !!!!
    # tols=[tolb,tolz,tolnu,tolrho,tollambda] = [1e3,10.0*n,1e3,1e3,1e3]
    # RETURNS: lambda
    # lambda - k-th eigenvalue
    # Kb, Kz, Knu, Krho - condition numbers
    # Qout = 1 / 0 - Double was / was not used

    # If Qout>0, quad precision was used
    # i was the shift index

    # Set the dimension
    n = length(A.D) + 1

    # Set all conditions initially to zero
    Kb,Kz,Knu,Krho=0.0,0.0,0.0,0.0
    Qout=0
    v=zeros(T,n)
    # Kz is former kappa_nu

    # Determine the shift sigma, the shift index i, and whether lambda
    # is on the left or the right side of the nearest pole
    # Exterior eigenvalues (k = 1 or k = n):

    if k == 1
        sigma,i,side = A.D[1],1,'R'
    elseif k==n
        sigma,i,side = A.D[n-1],n-1,'L'
    else
        # Interior eigenvalues (k in (2,...n-1) ):
        # We need to compute this in Double
        Dd=map(Double,A.D)
        # Dtemp = Dd-Dd[k]
        # atemp = alphaD-Dd[k]
        # middle = Dtemp[k-1]/2.0
        # Fmiddle = (atemp-middle)-sum(A.z.^2./(Dtemp-middle))
        middle=(Dd[k-1]-Dd[k])/Double(2.0)
        Fmiddle=Double(zero(T))
        for l=1:n-1
            Fmiddle+=zD[l]^2/((Dd[l]-Dd[k])-middle)
        end
        Fmiddle=((alphaD-Dd[k])-middle)-Fmiddle
        sigma,i,side = Fmiddle.hi+Fmiddle.lo < 0.0 ? (A.D[k],k,'R') : (A.D[k-1],k-1,'L')
    end

    # Compute the inverse of the shifted matrix, A_i^(-1), Kb and Kz

    Ainv,Kb,Kz,Qout = inv(A,zD,alphaD,i,tols[1:2])

    # Compute the eigenvalue of the inverse shifted matrix

    nu = bisect( Ainv,side )

    #  nu=fastzero([invD1; 0; invD2], [w1;wz;w2], b, side); # implement later
    #  [nu-nu1]/nu, [nu-nueva]/nu, pause  # and compare

    if abs(nu)==Inf
        # this is nonstandard
        # Deflation in aheig (nu=Inf)
        v[i]=1.0
        lambda=sigma
    else
        # standard case, full computation
        # nu1 is the F- or 1-norm of the inverse of the shifted matrix
        # nu10=maximum([sum(abs(Ainv.z))+abs(Ainv.a), maximum(abs(Ainv.D)+abs(Ainv.z))])
        nu1=0.0
        for l=1:n-1
            nu1=max(nu1,abs(Ainv.D[l])+abs(Ainv.z[l]))
        end
        nu1=max(sumabs(Ainv.z)+abs(Ainv.a), nu1)

        Knu=nu1/abs(nu)
        if Knu<tols[3]
            # Accuracy is fine, compute the eigenvector
            mu = 1.0/nu
            # v=[ A.z./((A.D-sigma)-mu);-1.0]
            for l=1:n-1
            	v[l] = A.z[l]/((A.D[l]-sigma)-mu)
            end
            v[n]=-1.0
            lambda, v = mu+sigma, v/norm(v)
        else
            # Remedies according to Remark 3 - we shift between original
            # eigenvalues and compute DPR1 matrix
            # 1/nu1+sigma, 1/nu+sigma
            # println(k, " Remedy 3 ")
            nu = side=='R' ? abs(nu) : -abs(nu)
            nu1=-sign(nu)*nu1
            sigma1=(nu1+nu)/(2.0*nu*nu1)+sigma

            if findfirst(A.D-sigma1,0.0)>0 # we came back with a pole
                # recompute sigmav more accurately according with dekker
                sigmav=(Double(nu1)+Double(nu))/(Double(2.0)*Double(nu)*Double(nu1))+Double(sigma)
                # Compute the inverse of the shifted arrowhead (DPR1)
                Ainv,Qout1=inv(A,sigmav) # Ainv is Float64
                nu1=bisect(Ainv,side)
                mu1 = 1.0/nu1
                D0=map(Double,A.D)-sigmav
                D1=[D0[l].hi+D0[l].lo for l=1:length(D0)]
                if findfirst(D1-mu1,0.0)>0
                    ind=find(D1-mu1==0.0)
                    v=zeros(n)
                    v[ind]=1.0
                else
                    v=[ A.z./(D1-mu1);-1.0]
                end
                # Shift the eigenvalue back in Double
                lam = Double(1.0)/Double(nu1)+sigmav
                # Return this
                lambda,v = lam.hi+lam.lo, v/norm(v)
            else
                # Compute the inverse of the shifted arrowhead (DPR1)
                Ainv, Krho,Qout1=Arrowhead.inv(A,sigma1,tols[4]) # Ainv is Float64
                # Compute the eigenvalue by bisect for DPR1
	            # Note: instead of bisection could use dlaed4 (need a wrapper) but
	            # it is not faster. There norm(u)==1
                nu1= bisect(Ainv,side)
                mu1=1.0/nu1
                # standard v
                v=[ A.z./((A.D-sigma1)-mu1);-1.0]
                # Return this - shift the eigenvalue back and normalize the vector
                lambda, v = mu1+sigma1, v/norm(v)
            end
            Qout==1 && (Qout=Qout+2)
        end

        # Remedy according to Remark 1 - we recompute the the eigenvalue
        # near zero from the inverse of the original matrix (a DPR1 matrix).
        if (abs(A.D[i])+abs(1.0/nu))/abs(lambda)>tols[5]

            if k==1 && A.D[1]<0.0 || k==n && A.D[n-1]>0.0 || i<n-1 && side=='L' && sign(A.D[i])+sign(A.D[i+1])==0 ||
                side=='R' && sign(A.D[i])+sign(A.D[i-1])==0
                # println(k, " Remedy 1 ")
                # Compute the inverse of the original arrowhead (DPR1)
                Ainv,Krho,Qout1 = inv(A,0.0,tols[4]) # Ainv is Float64
                Qout==1 && (Qout=Qout+4)
                if abs(Ainv.r)==Inf
                    lambda=0.0
                else
                    # Here we do not need bisection. We compute the Rayleigh
                    # quotient by using already computed vectors which is
                    # componentwise accurate
                    nu1=sum(v.^2.*Ainv.D)+Ainv.r*sum(v.*Ainv.u)^2;
                    lambda=1.0/nu1
                end
            end

        end
    end

    # Return this
    lambda, Qout

end # eig (k)


using PlotlyJS, Random, DSP, Dierckx, FileIO, JLD2

# This script makes the trace data to compare a step in the tine shifts

#ricker wavelet function
"""
ricker(omega)
returns Ricker wavelet

input:

    peak frequency   = 20 Hz
    sampling time    = 0.001 seconds
    number of points = 100;

output:
    s: ricker wavelet with length = number of points
    t: time
"""
function ricker(;f::Float64=20.0, n::Int64=100, dt::Float64=0.001, nargout::Bool=true)
    # Create the wavelet and shift in time if needed
    T = dt*(n-1);
    t = collect(0:dt:T);
    t0 = 1/f;
    tau = t.-t0;

    s = (1.0 .- tau.*tau.*f.^2*π.^2).*exp.(-tau.^2*π.^2*f.^2);

    if nargout
        trace = scatter(;x=t, y=s, mode="lines")
        layout = Layout(
            xaxis=attr(title="Time"),
            yaxis=attr(title="u"),
        )
        p=plot(trace, layout)
        display(p)
    end

    return (s, t)
end

# make Ricker wavelet

dt     = 0.05; # s
f      = 1.0; # Hz
w,tw = ricker(f=f, n=41, dt=dt, nargout=false);

## make random reflectivity sequence and convolve with Ricker

npts = 601; # length of time series

rng = MersenneTwister(20190603)
f    = randn(rng, Float64, npts);
u0   = conv(f,w)[1:npts]; # make waveform time series

## make time varying stretched + Noise
dvV = - 0.01 #dv over V [%] = dt over T, therfore dt = dvV * T

tvec  = collect(( 0:npts - 1 ) .* dt);
st   = tvec .* dvV
# add random noise (5 % of maximum st)
rng = MersenneTwister(20190603);
noise = randn(rng, Float64, npts);
noise = 0.05*maximum(abs.(st)).*(noise./maximum(abs.(st)))
st = st .+ noise;

tvec2 = tvec .+ st;

spl = Spline1D(tvec, u0)
u1 = spl(tvec2);

st = st;

@save "./StretchedData.jld2" dt u0 u1 st

if nargout

    #time shift curve

    function lineplot1()
        tr = scatter(;x=tvec, y=st, mode="lines", name="Timeshift")
        layout = Layout(
            xaxis=attr(title="Time"),
            yaxis=attr(title="Timeshift"),
        )
        plot(tr, layout)
    end

    function lineplot2()
        tr1 = scatter(;x=tvec, y=u0, mode="lines+markers", name="u0")
        tr2 = scatter(;x=tvec, y=u1, mode="lines+markers", name="u1")
        layout = Layout(
            xaxis=attr(title="Time"),
            yaxis=attr(title="Signal"),
        )
        plot([tr1, tr2], layout)
    end

    p = [lineplot1(), lineplot2()]

end


using Test

using Distributed
using Random
using SymbolicUtils

addprocs(2)

@testset "Basic test" begin
    @everywhere include("src/SymbolicRegression.jl")
    @everywhere using .SymbolicRegression

    Random.seed!(0)

    X = randn(Float32, 100, 2)
    y = 2 * cos.(X[:, 2]) + X[:, 1] .^ 2 .- 2

    options = SymbolicRegression.Options(
        binary_operators=(+, *),
        unary_operators=(cos,),
        npopulations=3,
        maxsize=15
    )

    niterations = 5
    hallOfFame = RunSR(X, y, niterations, options)
    dominating = calculateParetoFrontier(X, y, hallOfFame, options)

    @test dominating[end].score < 1e-3
    @test length(stringTree(dominating[end].tree, options)) > 10
end

rmprocs()


using DistCtrl4DistMan
using LinearAlgebra
using Test

# @testset "DistCtrl4DistMan.jl" begin
# end

@testset "ObjectAgentModule.jl" begin
    include("test_data.jl")

    @testset "Test - MAG" begin# Used types
        F = Float64;
        U = UInt64;

        # Load test data
        Gxy_ref, MAG_test_currents, MAG_fi_ref, MAG_F_ref = test_data_MAG();

        n = U(8);
        dx = F(25e-3);
        aa = DistCtrl4DistMan.ActuatorArray(n, n, dx, dx, :coil);
        ball_pos = ( (n-1)*dx/2 - dx/2, (n-1)*dx/2 + 1.3*dx );

        # calcMAGForce() and fi() test values
        fi_test = DistCtrl4DistMan.ObjectAgentModule.mag_fi(MAG_test_currents);
        F_MAG = DistCtrl4DistMan.calcMAGForce(aa, ball_pos, fi_test);

        aL, a_used = DistCtrl4DistMan.genActList(aa, (ball_pos[1], ball_pos[2], F(0)), F(2.5*dx), F(3.5*dx));
        oa_mag = DistCtrl4DistMan.ObjectAgent_MAG("Agent_MAG", ball_pos, F_MAG, aa, aL, a_used, F(1));

        @testset "Test the function generating the Gxy matrix" begin
            @test DistCtrl4DistMan.ObjectAgentModule.genGxy(aa, ball_pos) ≈ Gxy_ref
        end
        @testset "Test the function fi()" begin
            @test fi_test ≈ MAG_fi_ref
        end
        @testset "Test the function calcMagForce()" begin
            @test all(F_MAG .≈ MAG_F_ref)
        end
    end

    @testset "Test - DEP" begin
        # Used types
        F = Float64;
        U = UInt64;

        # Load test data
        Gamma_ref, Lambda_x, Lambda_y, Lambda_z, DEP_test_phases, DEP_F_ref = test_data_DEP();
        n = U(5);
        dx = F(100e-6);
        aa = DistCtrl4DistMan.ActuatorArray(n, n, dx, dx/2, :square);
        center_pos = ( (n-1)*dx/2, (n-1)*dx/2, F(100e-6));

        F_DEP = DistCtrl4DistMan.calcDEPForce(aa, center_pos, DEP_test_phases);
        aL, a_used = DistCtrl4DistMan.genActList(aa, center_pos, F(350e-6), F(350e-6));
        oa_dep = DistCtrl4DistMan.ObjectAgent_DEP("Agent1", center_pos, F_DEP, aa, aL, a_used);

        @testset "Test the function generating the Gamma and Lambda_a marices" begin
            Gamma_test, Lambda_a = DistCtrl4DistMan.ObjectAgentModule.genC_DEP(aa, center_pos);
            @test Gamma_ref ≈ Gamma_test
            @test Lambda_a[1] ≈ Lambda_x
            @test Lambda_a[2] ≈ Lambda_y
            @test Lambda_a[3] ≈ Lambda_z
        end
        @testset "Test the function calculating the DEP force" begin
            @test all(F_DEP .≈ DEP_F_ref)
        end
        @testset "Test the cost function" begin # The cost function (oa_dep,fvk) should be zero since we set oa_dep.vk_r and oa_dep.vk_i to values which should result in the desired force
            # Calculate the cost function (oa_dep.fvk)
            oa_dep.vk_r = cos.(DEP_test_phases'[:]);
            oa_dep.vk_i = sin.(DEP_test_phases'[:]);
            DistCtrl4DistMan.costFun!(oa_dep);

            @test all( abs.(oa_dep.fvk) .< [1,1,1].*1e-6 )
        end
    end

    @testset "Test - ACU" begin
        # Used types
        F = Float64;
        U = UInt64;

        # Load test data
        C1ref, C2ref, th1, th2, objVal1_ref, objVal2_ref, D1ref, D2ref, phases_test, p1_ref, p2_ref = test_data_ACU();

        n = U(8);
        dx = F(10.0e-3);
        aa = DistCtrl4DistMan.ActuatorArray(n, n, dx);

        # two agents at positions close to the center
        z0 = F(-50e-3);
        maxDist1 = 3*dx; #2
        maxDist2 = 3*dx; #3
        oa1_pos = ((n-1)*aa.dx/2-dx, (n-1)*aa.dx/2, z0);
        oa2_pos = ((n-1)*aa.dx/2+dx, (n-1)*aa.dx/2+1.6f0*dx, z0);
        Pdes = F(1000); # desired presure (same for both agents)

        # Generate list of used actuators for each object agent
        aL1, a_used1 = DistCtrl4DistMan.genActList(aa, oa1_pos, maxDist1, maxDist2);
        aL2, a_used2 = DistCtrl4DistMan.genActList(aa, oa2_pos, maxDist1, maxDist2);

        oa1 = DistCtrl4DistMan.ObjectAgent_ACU("Agent1", oa1_pos, Pdes, aa, aL1, a_used1);
        oa2 = DistCtrl4DistMan.ObjectAgent_ACU("Agent2", oa2_pos, Pdes, aa, aL2, a_used2);

        @testset "Test the function generating the C matrix" begin
            @test isapprox(oa1.zvec*oa1.zvec' + 1im*oa1.zvec*[0 1; -1 0]*oa1.zvec', C1ref, atol=1e-6)
            @test isapprox(oa2.zvec*oa2.zvec' + 1im*oa2.zvec*[0 1; -1 0]*oa2.zvec', C2ref, atol=1e-6)
        end
        @testset "Test the cost function" begin
            oa1.vk_r = cos.(th1);
            oa1.vk_i = sin.(th1);
            oa2.vk_r = cos.(th2);
            oa2.vk_i = sin.(th2);

            # Compute the values of the cost function (oa.fvk)
            DistCtrl4DistMan.costFun!(oa1)
            DistCtrl4DistMan.costFun!(oa2)

            @test isapprox(oa1.fvk[1],  objVal1_ref, atol=1e-5)
            @test isapprox(oa2.fvk[1],  objVal2_ref, atol=1e-5)
        end
        @testset "Test the cost function calculating the jacobian" begin
            # Compute the jacobians
            DistCtrl4DistMan.ObjectAgentModule.jac!(oa1);
            DistCtrl4DistMan.ObjectAgentModule.jac!(oa2);

            @test isapprox(oa1.J,  D1ref', atol=1e-6)
            @test isapprox(oa2.J,  D2ref', atol=1e-6)
        end
        @testset "Test the function calculating the pressure" begin
            @test isapprox(DistCtrl4DistMan.calcPressure(aa, oa1_pos, convert.(F, phases_test)), F(p1_ref), atol=1e-4)
            @test isapprox(DistCtrl4DistMan.calcPressure(aa, oa2_pos, convert.(F, phases_test)), F(p2_ref), atol=1e-4)
        end
        @testset "Test the function finding the shared actuators among agents (function from ActuatorArray module)" begin
            DistCtrl4DistMan.resolveNeighbrRelations!([oa1, oa2])
            @test oa1.actList[oa1.neighbors[1][2]] == oa2.actList[oa1.neighbors[1][3]]
            @test oa2.actList[oa2.neighbors[1][2]] == oa1.actList[oa2.neighbors[1][3]]
        end
    end
end

@testset "ActuatorArray tests" begin
    # Used types
    F = Float64;
    U = UInt64;

    # Load test data
    aL1_ref = Array{Tuple{UInt8,UInt8},1}([(1, 3), (1, 4), (1, 5), (1, 6), (2, 2), (2, 3), (2, 4), (2, 5), (2, 6), (2,7), (3, 2), (3, 3), (3, 4), (3, 5), (3, 6), (3, 7), (4, 2), (4, 3), (4, 4), (4, 5), (4, 6), (4, 7), (5, 2), (5, 3),(5, 4), (5, 5), (5, 6), (5, 7), (6, 3), (6, 4), (6, 5), (6, 6)]);
    aL2_ref = Array{Tuple{UInt8,UInt8},1}([(3, 5), (3, 6), (3, 7), (4, 4), (4, 5), (4, 6), (4, 7), (4, 8), (5, 4), (5, 5), (5, 6), (5, 7), (5, 8), (6, 4), (6, 5), (6, 6), (6, 7), (6, 8), (7, 4), (7, 5), (7, 6), (7, 7), (7, 8), (8, 5), (8, 6), (8, 7)]);

    n = U(8);
    dx = F(10.0e-3);
    aa = DistCtrl4DistMan.ActuatorArray(n, n, dx);

    # two agents at positions close to the center
    z0 = F(-50e-3);
    maxDist1 = 3*dx; #2
    maxDist2 = 3*dx; #3
    oa1_pos = ((n-1)*aa.dx/2-dx, (n-1)*aa.dx/2, z0);
    oa2_pos = ((n-1)*aa.dx/2+dx, (n-1)*aa.dx/2+1.6f0*dx, z0);

    # Generate list of used actuators for each object agent
    aL1, a_used1 = DistCtrl4DistMan.genActList(aa, oa1_pos, maxDist1, maxDist2);
    aL2, a_used2 = DistCtrl4DistMan.genActList(aa, oa2_pos, maxDist1, maxDist2);

    @testset "Test the function generating the actuator list" begin
        @test aL1 == aL1_ref
        @test aL2 == aL2_ref
    end
end

@testset "LinSolvers.jl package tests" begin
    n = 100;
    m = 5;

    @testset "Test backward and forward substitution for solution of R^T*Rx=b" begin
        for i=1:20
            R = Matrix(UpperTriangular(rand(n,n) + I));
            A = R'*R;
            x_test = rand(n);
            b_test = A*x_test;
            b = copy(b_test)

            DistCtrl4DistMan.LinSolvers.forward_substitution!(R, b)
            DistCtrl4DistMan.LinSolvers.backward_substitution!(R, b)
            x_sol = b;

            @test  A*x_sol ≈ b_test
        end
    end
    @testset "Test solving of (D^T*D + 1/λ*I)x=D^T*b by kernel_solve!()" begin
        x_sol = zeros(n);
        lsd = DistCtrl4DistMan.LinSystemData(n, m, x_sol);
        λ = 1.0e3;
        for i in 1:20
            Dt = rand(n,m);
            b  = rand(m);
            DistCtrl4DistMan.kernel_solve!(lsd, Dt, λ, b);

            @test  (Dt*Dt' + 1/λ*I)*x_sol ≈ Dt*b
        end
    end
    @testset "Test solving of (D^T*D + 1/λ*I)x=D^T*b by kernel_solve()" begin
        λ = 1.0e3;
        for i in 1:20
            Dt = rand(n,m);
            b  = rand(m);
            x_sol = DistCtrl4DistMan.kernel_solve(Dt, λ, b);

            @test  (Dt*Dt' + 1/λ*I)*x_sol ≈ Dt*b
        end
    end
    @testset "Test solving of (D^T*D + 1/λ*I)x=D^T*b by kernel_solve_prealloc!()" begin
        λ = 1.0e3;
        x_sol = zeros(n);
        A_prealloc = zeros(m+n,m);
        b_prealloc = zeros(m);
        R_prealloc = zeros(m,m);

        for i in 1:20
            Dt = rand(n,m);
            b  = rand(m);
            DistCtrl4DistMan.LinSolvers.kernel_solve_prealloc!(Dt, λ, b, x_sol, A_prealloc, b_prealloc, R_prealloc);

            @test  (Dt*Dt' + 1/λ*I)*x_sol ≈ Dt*b
        end
    end

    @testset "Test qr_r!() function" begin
        for i in 1:20
            A = rand(n,m);
            R  = zeros(m,m);

            # Compute the upper triangular R matrix from QR factorization by qr_r!()
            # Matrix A is modified in qr_r!(), that is why we use copy of it as an argument
            DistCtrl4DistMan.LinSolvers.qr_r!(copy(A), R);

            # Compute the cholesky factor of AᵀA which has to match to R
            C = cholesky(A'*A);

            @test  C.U ≈ R
        end
    end

end

@testset "Test - ADMM" begin
    F = Float64;
    U = UInt64;

    n = U(8);
    dx = F(25e-3);
    aa = DistCtrl4DistMan.ActuatorArray(n, n, dx, dx, :coil);

    # Set the maximum distance so that both agents use all the coils
    maxdist = n*dx;

    for i in 1:20
        # Generate random positions of the agents
        ball_pos1 = ( dx*(n-1)*rand(),  dx*(n-1)*rand());
        ball_pos2 = ( dx*(n-1)*rand(),  dx*(n-1)*rand());

        # Generate the desired forces so that one can actually generate them
        random_ctrls = rand(F, n, n);
        F_des1 = DistCtrl4DistMan.calcMAGForce(aa, ball_pos1, random_ctrls)
        F_des2 = DistCtrl4DistMan.calcMAGForce(aa, ball_pos2, random_ctrls)

        # Initialize the agents
        aL1, a_used1 = DistCtrl4DistMan.genActList(aa, (ball_pos1[1], ball_pos1[2], F(0)), maxdist, maxdist);
        aL2, a_used2 = DistCtrl4DistMan.genActList(aa, (ball_pos2[1], ball_pos2[2], F(0)), maxdist, maxdist);
        oa1 = DistCtrl4DistMan.ObjectAgent_MAG("Agent 1", ball_pos1, F_des1, aa, aL1, a_used1, F(1));
        oa2 = DistCtrl4DistMan.ObjectAgent_MAG("Agent 2", ball_pos2, F_des2, aa, aL2, a_used2, F(1));

        agents = [oa1, oa2];
        DistCtrl4DistMan.admm(agents,
            λ = 1.0, ρ = 1.0,
            log = false,
            maxiter = 1000,
            method = :freedir);

        F_dev1 = tuple((oa1.Gxy'*oa1.zk*oa1.Fdes_sc)...)
        F_dev2 = tuple((oa2.Gxy'*oa2.zk*oa2.Fdes_sc)...)

        @test all(isapprox.(F_dev1, F_des1, atol=1e-4))
        @test all(isapprox.(F_dev2, F_des2, atol=1e-4))
    end
end


#=
Sandia:
- Julia version: 1.6.2
- Author: samuel
- Date: 2021-09-10
=#

using SuiteSparseGraphBLAS
using SparseArrays
using LinearAlgebra

#A is the input matrix, Output: Triangle Count
function sandia(A)

    #extracting lower matrix from input matrix
    L = tril(A)

    C = mul(L, L, Semirings.PLUS_TIMES, mask=L)
    result = reduce(Monoids.PLUS_MONOID, C, dims=:)

    print("\n\n")
    print("RESULT ")
    print(result)
    return result

end


"""
$(TYPEDEF)

Houbolt's integration scheme with given step-size.

### Fields

- `Δt` -- step-size

### References

See [[HOU50]](@ref).
"""
struct Houbolt{N} <: AbstractSolver
    Δt::N
end

Houbolt(; Δt::N) where N = Houbolt(Δt)

function _init(alg::Houbolt, M, C, K)
    Δt = alg.Δt

    # compute integration constants
    a₀ = 2/Δt^2
    a₁ = 11/(6*Δt)
    a₂ = 5/Δt^2
    a₃ = 3/Δt
    a₄ = -2*a₀
    a₅ = -a₃/2
    a₆ = a₀/2
    a₇ = a₃/9

    K̂ = K + a₀ * M + a₁ * C

    return a₀, a₁, a₂, a₃, a₄, a₅, a₆, a₇, K̂
end

function _solve(alg::Houbolt{N},
                ivp::InitialValueProblem{ST, XT},
                NSTEPS::Int; kwargs...) where {N, VT, ST, XT<:Tuple{VT, VT}}

    sys = system(ivp)
    (U₀, U₀′) = initial_state(ivp)

    IMAX = NSTEPS + 1
    M, C, K, R = _unwrap(sys, IMAX)

    U₀′′ = M \ (R[1] - C * U₀′ - K * U₀)
    a₀, a₁, a₂, a₃, a₄, a₅, a₆, a₇, K̂ = _init(alg, M, C, K)
    K̂⁻¹ = factorize(K̂)

    U = Vector{VT}(undef, IMAX)
    U[1] = U₀

    # special starting procedure to calculate U(Δt) and U(2Δt)
    V = _solve(CentralDifference(Δt=alg.Δt), ivp, 2) |> displacements
    U[2] = V[2]
    U[3] = V[3]

    @inbounds for i in 3:NSTEPS
        mᵢ = M * (a₂ * U[i] + a₄ * U[i-1] + a₆ * U[i-2])
        cᵢ = C * (a₃ * U[i] + a₅ * U[i-1] + a₇ * U[i-2])
        R̂ᵢ₊₁ = R[i+1] + mᵢ + cᵢ
        U[i+1] = K̂⁻¹ \ R̂ᵢ₊₁
    end

    return _build_solution(alg, U, NSTEPS)
end

function _build_solution(alg::Houbolt{N}, U, NSTEPS) where {N}
    t = range(zero(N), step=alg.Δt, length=(NSTEPS+1))
    return Solution(alg, U, nothing, nothing, t)
end


using LinearAlgebra, SparseArrays, MAT
using Printf, Test
using BasicLU

include("utils.jl")

test_factorize("./data/", false)
test_factorize("./data/", true)
test_update("./data/")

m = 1000
this = BasicLU.initialize(m)
B = sprand(m, m, 5e-3) + I
err = BasicLU.factorize(this, B)
if err ≠ BasicLU.BASICLU_OK
    error("factorization failed or singular")
end
rhs = randn(m)
lhs = solve(this, rhs, 'N')
res = norm(B * lhs - rhs, Inf)
@test res ≤ √eps(cdbl)
col = sparsevec([1], [1.0], m)
lhs = solve4update(this, col, true)
(vmax,j) = findmax(abs.(lhs))
xtbl = lhs[j]
solve4update(this, j)
piverr = update(this, xtbl)
lhs = solve(this, rhs, 'N')
B[:,j] = col
res = norm(B * lhs - rhs, Inf)
@test res ≤ √eps(cdbl)


function  push!(Bop :: BFGSOperator{T, N, F1, F2, F3},
                sk:: Vector{T},
                yk:: Vector{T}) where {T, N, F1, F2, F3}

    scaling = Bop.data.scaling
    nsk2 =sk'*sk
    denom = yk'*sk
    B = Bop.data.M
    if (denom > 1.0e-20) # && (mod(iter,n)!=0)
        #self-scaled version aka Luenberger
        By = B*yk
        ytBy = (yk'*By)
        γ = 1.0
        if scaling
            γ = denom / (ytBy)
        end

        # math formula, involves matrix multiplications, thus O(n³) complexity
        M = I - sk*yk'/denom
        B1 = γ*M*B*M' + sk*sk'/denom
        B = 0.5*(B1+B1')  # make sure B is symmetric
    else
        @warn "No update, B=I, denom = ", denom
        B = I
    end

    Bop.data.M .= B

    return Bop
end


"""
The {2, 3, 5}-wheel is efficient because it compresses memory
perfectly (1 byte = 30 numbers) and removes 4/15th of all the
multiples already. We don't get that memory efficiency when
extending the wheel to {2, 3, 5, 7}, since we need to store
48 bits per 210 numbers, which is could be done with one 64-bit
integer per 210 numbers, which in fact compresses worse
(64 bits / 210 numbers) than the {2, 3, 5}-wheel
(8 bits / 30 numbers).

What we can do however, is compute the repeating pattern that
the first `n` primes create, and copy that pattern over. That
is, we look at a the numbers modulo p₁ * p₂ * ⋯ * pₙ.

For instance, when presieving all multiples of {2, 3, ..., 19}
we allocate a buffer for the range 1 : 2 * 3 * ... * 19 =
1:9_699_690. In a {2, 3, 5} wheel this means a buffer of
9_699_690 ÷ 30 = 323_323 bytes.
"""
function create_presieve_buffer()
    n_bytes = 7 * 11 * 13 * 17
    xs = fill(0xFF, n_bytes)

    @inbounds for p in (7, 11, 13, 17)
        p²        = p * p
        byte_idx  = p² ÷ 30 + 1
        wheel     = to_idx(p)
        wheel_idx = 8 * (wheel - 1) + wheel
        increment = 0
        @sieve_loop :unroll
    end

    @inbounds xs[1] = 0b11100001 # remove 7, 11, 13 and 17
    return xs
end

"""
When applying the presieve buffer, we have to compute the offset in
"""
function apply_presieve_buffer!(xs::Vector{UInt8}, buffer::Vector{UInt8}, byte_start, byte_stop)

    len = byte_stop - byte_start + 1

    # todo, clean this up a bit.
    from_idx = (byte_start - 1) % length(buffer) + 1
    to = min(len, length(buffer) - from_idx + 1)

    # First copy the remainder of buffer at the front
    copyto!(view(xs, Base.OneTo(to)), view(buffer, from_idx:from_idx + to - 1))
    from = to + 1

    # Then copy buffer multiple times
    while from + length(buffer) - 1 <= len
        copyto!(view(xs, from : from + length(buffer) - 1), buffer)
        from += length(buffer)
    end

    # And finally copy the remainder of buffer again
    copyto!(view(xs, from:len), view(buffer, Base.OneTo(length(from:len))))

    xs
end


function hull_reformulation{HR<:HRepresentation}(model::Model, regions::Vector{HR}, x::AbstractVector{Variable})
    n = length(x)
    R = length(regions)
    @boundscheck begin
        for region in regions
            # ensure that dimension of region matches number of variables
            @assert fulldim(region) == n

            # ensure that there are no rays, so that setting the bounds of the inequalities to zero will
            # result in only the origin being inside the polyhedron
            vr = vrep(polyhedron(region))
            @assert nrays(vr) == 0
        end
    end

    Hrs = @variable(model, [r = 1 : R], Bin) # indicator variables, 4.2c
    @constraint(model, sum(Hrs) == 1) # 4.2b

    xrs = [@variable(model, [k = 1 : n]) for r = 1 : R] # auxiliary variables
    @constraint(model, x .== sum(xrs)) # 4.4a

    for (region, Hr, xr) in zip(regions, Hrs, xrs)
        for ineq in ineqs(region)
            @constraint(model, ineq.a ⋅ xr <= Hr * ineq.β) # 4.4b
        end
        for eq in eqs(region)
            @constraint(model, eq.a ⋅ xr == Hr * eq.β) # 4.4b
        end
    end
    Hrs
end


"""
    surrogate_model(plan::AbstractArray{T,2}, samples::AbstractArray{T,2}; options=Options()) where T

Returns a surrogate model function based on an optimized Radial Basis Function
interpolant. Depending on the options, the kernel, kernel width and scaling of
input data is optimized.

...
# Arguments
- `plan::AbstractArray{T,2}`:
    sample locations where each column corresponds to the location of one point.
    `size(plan) = (num_dimensions,num_samples)`.
- `samples::AbstractArray{T,2}`:
    function value at each sample location. each column contains one value from
    the corresponding plan location. `size(samples) = (1,num_samples)`.
- `options=Options()`:
    all options available to customize the surrogate optimization.
...
"""
function surrogate_model(plan::AbstractArray{T,2}, samples::AbstractArray{T,2}; options::Options=Options()) where T

    @unpack num_interpolants, trace, parallel_surrogate = options

    trace && println("Creating optimized surrogate model ...")

    (length(samples) != size(plan,2)) && error("plan and samples do not have the correct sizes")

    #Optimize RBF hypers for the ensamble of interpolants
    if parallel_surrogate
        optres = pmap(  (x)->rbf_hypers_opt(samples, plan, options),
                        1:num_interpolants)
    else
        optres = map(  (x)->rbf_hypers_opt(samples, plan, options),
                        1:num_interpolants)
    end

    # Sample order for ScatteredInterpolation
    samples_vec = vec(samples)

    # Interpolation object based on the optimization results
    itp = Array{ScatteredInterpolation.RBFInterpolant,1}(undef,num_interpolants)
    for i = 1:num_interpolants
        # Preprocess the points based on the settings used.
        preprocessed_point = preprocess_point(plan,optres[i],base_scale=plan)

        itp[i] = interpolate(optres[i].kernelFunc, preprocessed_point,
                            samples_vec, smooth=optres[i].smooth)
    end

    # Pre allocate and calculate to reduce cost of evaluating surrogate model.
    preprocessed_est_point = Array{Float64,2}(undef,size(plan,1),1)
    old_min = minimum(plan,dims=2)
    old_max = maximum(plan,dims=2)

    sm_func = (estimation_point) ->  surrogate_evaluate.(Ref(preprocessed_est_point),
                                        Ref(estimation_point),Tuple(itp),Tuple(optres),
                                        Ref(old_min),Ref(old_max))

    return sm_func, optres
end


include("../src/Opt.jl")
include("../src/Model.jl")
h = Model.hill(k=0.5,n=3)

@show Opt.find_min(h,[0.,1.])

@show Opt.find_max(hill,[0.,30.])

@show Opt.find_MinMax(hill, [0.,30.])

using Optim

xx(w) = w

function t(x;functor)

return functor(x)

Optim.optimize(h, 0.1, 0.2)

##
using Optim
using Test

struct y
    p
end

function (s::y)(x)
    return x*s.p
end

model = y(1.0)
model2(x;p=1.0) = x*p

@show  model(2.0)
@show  model2(2.0)

@test typeof(y) <: Any # Because Optim.jl accepts :Any object

Optim.optimize(model, 0.1, 0.2);
Optim.optimize(model2, 0.1, 0.2);


"""
standard Armijo linesearch to determine step length
default parameters c = τ = 0.5 as in Armijo's paper (1966)
"""
function backtracking_linesearch(X, U, Z, dZ, dU, μ0 = 1, c = 0.5, τ = 0.5)
    μ = μ0
    val = norm(X - U*Z')^2
    m = -(norm(dU)^2 + norm(dZ)^2)
    A = X - U*Z'
    a = [dU*Z', U*dZ', dU*dZ']
    while norm(A + μ*(a[1] + a[2]) - μ^2*a[3])^2 > val + μ*c*m
        μ *= τ
    end
    return μ
end

"""
augments low rank representation with relevant directions from normal cone
"""
function augment_rank(N, U0, Z0, σ)
    V, Σ, P = svd(N)
    r = sum(Σ .>= σ)
    N_U = V[:, 1:r]
    N_Z = P[:, 1:r]*Diagonal(Σ[1:r])
    U = cat(U0, N_U, dims=2)
    Z = cat(Z0, N_Z, dims=2)
    project_to_low_rank!(U, Z, U0, Z0)
    return U, Z
end

"""
finds normal to low rank matrix manifold
"""
normal_direction(U, Z, L_U, L_Z) = (I - U*U')*L_U*(L_Z*(I - Z*pinv(Z'*Z)*Z'))
normal_direction(U, Z, L) = (I - U*U')*L*(I - Z*pinv(Z'*Z)*Z')

"""
finds orthonormal ̄U such that || (U-̄U)Z || is minimal (via gradient descent)
"""
function orthonormalize!(U, Z; ϵ = 1e-8, μ = 0.1, iter_max = 1000)
    iter = 0
    K = U'*U
    A = Matrix{Float64}(I, size(K))
    Ainv = A
    dA = A'*K*A - I
    while norm(dA) > ϵ && iter < iter_max
        A += μ*dA
        Ainv -= μ*Ainv*A*Ainv
        dA = - K * A * (A'*K*A - I)
        iter += 1
    end

    if iter == iter_max
        @warn "Maximum number of iterations exceeded in orthonormalization step"
    end

    U .= U*A
    Z .= Z*Ainv'
end

"""
finds orthonormal U and Z such that || UZ' - U0 Z0'|| is minimal (via gradient
descent)
"""
function project_to_low_rank!(U, Z, U0, Z0; μ = 0.1, ϵ = 1e-8, iter_max = 1000)
    iter = 0
    while iter < iter_max
        dZ = (Z - Z0*(U0'*U))
        dU = -(I - U*U')*U0*(Z0'*Z*pinv(Z'*Z)) # probably better to use a backsolve
        Z .-= μ*dZ
        U .-= μ*dU
        orthonormalize!(U, Z; μ = 0.1)
        if √(norm(dZ)^2 + norm(dU)^2) < ϵ
            break
        end
        iter += 1
    end
    if iter == iter_max
        @warn "Maximum number of iterations exceeded in projection step"
    end
end

"""
finds orthonormal U and Z such that || X - U Z'|| is minimal (via gradient
descent)
"""
function project_to_low_rank!(U, Z, X; μ = 0.1, ϵ = 1e-8, iter_max = 1000)
    iter = 0
    while iter < iter_max
        dZ = (Z - X*U)
        dU = -(I - U*U')*X*Z*pinv(Z'*Z) # probably better to use a backsolve
        Z .-= μ*dZ
        U .-= μ*dU
        orthonormalize!(U, Z; μ = 0.1)
        if √(norm(dZ)^2 + norm(dU)^2) < ϵ
            break
        end
        iter += 1
    end

    if iter == iter_max
        @warn "Maximum number of iterations exceeded in projection step"
    end
end


using QuantEcon
using PyPlot

include("jv.jl")

wp = JvWorker(grid_size=25)
v_init = wp.x_grid .* 0.5

f(x) = bellman_operator(wp, x)
V = compute_fixed_point(f, v_init, max_iter=40)

s_policy, phi_policy = bellman_operator(wp, V, return_policies=true)

# === plot policies === #
fig, ax = subplots()
ax[:set_xlim](0, maximum(wp.x_grid))
ax[:set_ylim](-0.1, 1.1)
ax[:plot](wp.x_grid, phi_policy, "b-", label="phi")
ax[:plot](wp.x_grid, s_policy, "g-", label="s")
ax[:set_xlabel]("x")
ax[:legend]()


"This script computes the Laguerre-Kravchuk functions, who are the Kravchuk functions on a angular momentum space."

using PixelsGT # Call the parent module

export  ϕd, ϕu, LKd, LKu, LK

"`Φd(q₁, q₂, n, m, j)` rotates π/2 the Kravchuk function modes at (n, |m| ≤ n) for n ∈ [0, 2j]. It needs to be evaluted at (q₁, q₂) ∈ [0, 2j + 1]."
function Φd(q₁, q₂, n, m, j)
    v₀ = zeros(Integer(2 * j + 1)) # (Sep 9, 2021) Added Integer() to meet the half-spin feature in j
    for n₁ in 0:n
        v₀[Integer(n₁ + 1)] = Φ2(n₁, n - n₁, q₁, q₂, j) * dd(n, j)[n₁ + 1, Integer((m + n)/2) + 1]
    end
    return sum(v₀)
end

"`Φu(q₁, q₂, n, m, j)` rotates π/2 the Kravchuk function modes at (n, |m| ≤ 4j - n) for n ∈ [4j, 2j]. It needs to be evaluted at (q₁, q₂) ∈ [0, 2j + 1]."
function Φu(q₁, q₂, n, m, j)
    v₁ = zeros(Integer(2 * j + 1))
    for n₁ in 0:Integer((4* j - n))
        v₁[n₁ + 1] = Φ2(2 * j - n₁, n + n₁ - 2 * j, q₁, q₂, j) * du(n, j)[n₁ + 1, Integer((m + 4 * j - n)/2) + 1]
    end
    return sum(v₁)
end

"Laguerre-Kravchuk modes functions are the core of the rotation of images, they encodes the original position information in a mode-angular momentum space, ready to apply an angle phase who rotates the image"

"`LKd(q₁, q₂, n, m, j)` gives Laguerre-Kravchuk modes at (n, |m| ≤ n) for n ∈ [0, 2j]. It needs to be evaluted at (q₁, q₂) ∈ [0, 2j + 1]."
function LKd(q₁, q₂, n, m, j)
    vlkd = zeros(ComplexF64, Integer(2 * j + 1))
    for n₁ in 0:n
        vlkd[n₁ + 1] = (-1.0im)^(n - n₁) * Φ2(n₁, n - n₁, q₁, q₂, j) * dd(n, j)[n₁ + 1, Integer((m + n)/2) + 1]
    end
    return (-1.0)^((0.5) * (abs(m) - m)) * sum(vlkd)
end

"`LKu(q₁, q₂, n, m, j)` gives the Laguerre-Kravchuk modes at (n, |m| ≤ 4j - n) for n ∈ [4j, 2j]. It needs to be evaluted at (q₁, q₂) ∈ [0, 2j + 1]."
function LKu(q₁, q₂, n, m, j)
    vlku = zeros(ComplexF64, Integer(2 * j + 1))
    for n₁ in 0:Integer(4* j - n)
        vlku[n₁ + 1] = (-1.0im)^(n - n₁) * Φ2(2 * j - n₁, n + n₁ - 2 * j, q₁, q₂, j) * du(n, j)[n₁ + 1, Integer((m + 4 * j - n)/2) + 1]
    end
    return (-1.0)^((0.5) * (abs(m) - m)) * sum(vlku)
end

"`LK(q₁, q₂, n, m, j)` gives the complete Laguerre-Kravchuk modes. Their index ranges are n ∈ [0, 4j] and m ∈ [-n,n] in steps of 2."
function LK(q₁, q₂, n, m, j)
    if 0 ≤ n ≤ 2 * j
        return LKd(q₁, q₂, n, m, j)
    elseif 2 * j < n ≤ 4 * j
        return LKu(q₁, q₂, n, m, j)
    end
end


module TestPaths
  using Test
  using EcologicalNetworks

  # Generate some data

  A = UnipartiteProbabilisticNetwork([0.0 0.1; 0.2 0.8])

  P = number_of_paths(A, n=2)

  @test P[1,1] ≈ 0.02
  @test P[1,2] ≈ 0.08
  @test P[2,2] ≈ 0.66

  m = rand([0,1], (5,5))
  U = UnipartiteQuantitativeNetwork(m)
  @test number_of_paths(U) == number_of_paths(UnipartiteNetwork(m.>0))
  @test shortest_path(U) == shortest_path(UnipartiteNetwork(m.>0))

  N = unipartitemotifs()[:S1]
  b = bellman_ford(N)
  @test length(b) == 3
  @test length(filter(x -> x.to == :c, b)) == 2

  @test sort(dijkstra(N, :a)) == sort(bellman_ford(N, :a))
  @test sort(dijkstra(N)) == sort(bellman_ford(N))
  d = dijkstra(N, :a)
  @test length(d) == 2
  @test length(filter(x -> x.from == :a, d)) == 2

end


#  Copyright 2016, Iain Dunning, Joey Huchette, Miles Lubin, and contributors
#  This Source Code Form is subject to the terms of the Mozilla Public
#  License, v. 2.0. If a copy of the MPL was not distributed with this
#  file, You can obtain one at http://mozilla.org/MPL/2.0/.
#############################################################################
# JuMP
# An algebraic modelling langauge for Julia
# See http://github.com/JuliaOpt/JuMP.jl
#############################################################################
# diet.jl
#
# Solve the classic "diet problem".
# Based on
#  http://www.gurobi.com/documentation/5.6/example-tour/diet_cpp_cpp
#############################################################################

using JuMP

function PrintSolution(status, foods, buy)
    println("RESULTS:")
    if status == :Optimal
        for i = 1:length(foods)
            println("  $(foods[i]) = $(getvalue(buy[i]))")
        end
    else
        println("  No solution")
    end
    println("")
end

function SolveDiet()

    # Nutrition guidelines
    numCategories = 4
    categories = ["calories", "protein", "fat", "sodium"]
    minNutrition = [1800, 91, 0, 0]
    maxNutrition = [2200, Inf, 65, 1779]

    # Foods
    numFoods = 9
    foods = ["hamburger", "chicken", "hot dog", "fries",
                     "macaroni", "pizza", "salad", "milk", "ice cream"]
    cost = [2.49, 2.89, 1.50, 1.89, 2.09, 1.99, 2.49, 0.89, 1.59]
    nutritionValues = [410 24 26 730;
                       420 32 10 1190;
                       560 20 32 1800;
                       380  4 19 270;
                       320 12 10 930;
                       320 15 12 820;
                       320 31 12 1230;
                       100  8 2.5 125;
                       330  8 10 180]

    # Build model
    m = Model()

    # Variables for nutrition info
    @variable(m, minNutrition[i] <= nutrition[i=1:numCategories] <= maxNutrition[i])
    # Variables for which foods to buy
    @variable(m, buy[i=1:numFoods] >= 0)

    # Objective - minimize cost
    @objective(m, Min, dot(cost, buy))

    # Nutrition constraints
    for j = 1:numCategories
        @constraint(m, sum{nutritionValues[i,j]*buy[i], i=1:numFoods} == nutrition[j])
    end

    # Solve
    println("Solving original problem...")
    status = solve(m)
    PrintSolution(status, foods, buy)

    # Limit dairy
    @constraint(m, buy[8] + buy[9] <= 6)
    println("Solving dairy-limited problem...")
    status = solve(m)
    PrintSolution(status, foods, buy)

end

SolveDiet()