using Base.Test
using QuantumOptics

shape1 = [5]
shape2 = [2, 3]

b1 = GenericBasis(shape1)
b2 = GenericBasis(shape2)

@test b1.shape == shape1
@test b2.shape == shape2
@test b1 != b2
@test b1 == b1

comp_b1 = tensor(b1, b2)
comp_b2 = tensor(b1, b1, b2)
@test comp_b1.shape == [prod(shape1), prod(shape2)]
@test comp_b2.shape == [prod(shape1), prod(shape1), prod(shape2)]

comp_b1_b2 = tensor(comp_b1, comp_b2)
@test comp_b1_b2.shape == [prod(shape1), prod(shape2), prod(shape1), prod(shape1), prod(shape2)]
@test comp_b1_b2 == CompositeBasis(b1, b2, b1, b1, b2)

@test comp_b2.shape == tensor(b1, comp_b1).shape
@test comp_b2 == tensor(b1, comp_b1)

@test ptrace(comp_b2, [1]) == ptrace(comp_b2, [2]) == comp_b1
@test ptrace(comp_b2, [1, 2]) == ptrace(comp_b1, [1])
@test ptrace(comp_b2, [2, 3]) == ptrace(comp_b1, [2])


using PQEcolaPoint
using DataFrames

# setup DTE -
T = 149.0
DTE = range(T, stop=0.0,step=-1.0) |> collect

# setup underlying range -
Sₒ = 69.41
IV = 0.4693
σ = Sₒ*(IV)*(√(T/365))
L = round(Sₒ - 1.96*σ)  # 95%
U = round(Sₒ + 1.96*σ)  # 95%
S = range(max.(L, 0.0),stop=U, step=0.01) |> collect

# setup IV array -
IV_array = range(0.1, stop = 1.2, step = 0.005) |> collect

# setup strike coordinates -
K = range(45.0, stop = 90.0, step = 5.0) |> collect

# build JIT lattice model -
JIT_lattice = build(CRRJITContractPremiumLatticeModel, PutContractModel, S, K, IV_array, DTE; number_of_levels = 80)

# origin version -
origin_point = encode(JIT_lattice; S = 70.12, DTE = 53.0, IV = 0.5129, K = 65.0)
start_point = encode(JIT_lattice; S = 68.41, DTE = 56.0, IV = 0.5957, K = 65.0)
end_point = encode(JIT_lattice; S = 68.41 , DTE = 84.0, IV = 0.5737, K = 55.0)
C = roll(JIT_lattice, origin_point, start_point=>end_point)

# build my points -
test = roll(JIT_lattice, start_point=>end_point)


using TailRiskScenGen: vec_isless, SurvivorApproximator, below_nonrisk_frontier, above_risk_frontier, decapitate!
using Distributions
using LinearAlgebra: I
using DataStructures

A = [[1.0, 2.0, 3.0] [2.0, 0.0, 6.0] [3.0, 8.0, 9.0]]
@test vec_isless([1.0, 2.0, 3.0], [2.0, 3.0, 4.0])
@test !vec_isless([1.0, 1.0], [1.0, 2.0])
@test !vec_isless(view(A,:,1), view(A,:,2))
@test vec_isless(view(A,:,1), view(A,:,3))
@test !vec_isless(view(A,:,3), [10.0, 9.0, 7.0])

d = 3
dist = MvNormal(Array{Float64}(I, d, d))
sample = rand(dist, 10000)
α = 0.99
surv = SurvivorApproximator(sample, α)
approx_prob, err = surv(zeros(d))
@test 0.5^d ≈ approx_prob atol=err

frontier=list([2.0, 0.0, 1.0], [1.0, 1.0, 1.0], [1.0, 0.0, 2.0])
@test above_risk_frontier([3.0, 3.0, 3.0], frontier)
@test above_risk_frontier([1.5,1.5,1.5], frontier)
@test !above_risk_frontier([3.0, 0.0, 3.0], frontier)
@test !above_risk_frontier([-2.0, -1.0, -3.0], frontier)
@test below_nonrisk_frontier([0.0, 0.0, 0.0], frontier)
@test below_nonrisk_frontier([0.5, -0.1, 1.5], frontier)
@test !below_nonrisk_frontier([1.0, 1.0, 1.0], frontier)
@test !below_nonrisk_frontier([5.0, 5.0, 1.0], frontier)

l = list(1,2,3,4)
l2 = tail(l)
decapitate!(l2)
@test l == list(1,3,4)
@test l2 == list(3,4)


abstract type Integrator{DT} end

struct ExplicitEuler{DT, ET <: Equation{DT}} <: Integrator{DT}
    equ::ET
    Δt::DT
    ẋ::Vector{DT}

    function ExplicitEuler(equ::ET, Δt::DT) where {DT, ET <: Equation{DT}}
        new{DT,ET}(equ, Δt, zeros(DT, ndims(equ)))
    end
end

function (int::ExplicitEuler)(x₀::AbstractVector, x₁::AbstractVector)
    int.equ.f(int.ẋ, x₀)
    x₁ .= x₀ .+ int.Δt .* int.ẋ
end


using PrimeFactorsKata
using Test

@testset "PrimeFactorsKata.jl" begin
    @test factorsOf(1) == []
    @test factorsOf(2) == [2]
    @test factorsOf(3) == [3]
    @test factorsOf(4) == [2, 2]
    @test factorsOf(5) == [5]
    @test factorsOf(6) == [2, 3]
    @test factorsOf(7) == [7]
    @test factorsOf(8) == [2, 2, 2]
    @test factorsOf(9) == [3, 3]
    @test factorsOf(2 * 2 * 3 * 5 * 7 * 7 * 11 * 13 * 13) == [2,2,3,5,7,7,11,13,13]
end


module Evaluation

using DataFrames
using DataStructures: OrderedDict
using MLDataPattern: nobs, getobs
using Statistics

using ..Transforms: ColTrans

export prob2class, probs2class, ids2classes, accuracy, cmatrix

# binary for single output
prob2class(data::AbstractMatrix{<:AbstractFloat}) = reshape(Int.(getobs(data)), :)

probs2class(data::AbstractMatrix{<:AbstractFloat}) = reshape(getindex.(argmax(data; dims=1), 1), :)

ids2classes(is::Symbol...) = begin pairs = split.(string.(is), "_"); Symbol(pairs[1][1]) => [string(v[2]) for v in pairs] end
ids2classes(i::ColTrans) = ids2classes(outids(i)...)

function accuracy(Y::AbstractVector{S}, T::AbstractVector{S}) where S
    size(Y) == size(T) || error("different dimensions size(Y)=$(size(Y)), size(T)=$(size(T))")
    mean(Y .== T)
end

accuracy(Y::AbstractMatrix{<:AbstractFloat}, T::AbstractMatrix{<:AbstractFloat}) = accuracy(probs2class(Y), probs2class(T))

function cmatrix(Y::AbstractVector{S}, T::AbstractVector{S}, classnames::Union{AbstractVector,Nothing} = nothing) where S
    size(Y) == size(T) || error("different dimensions size(Y)=$(size(Y)), size(T)=$(size(T))")
    if isnothing(classnames)
        classnames = unique(vcat(unique(Y), unique(T)))
        sort!(classnames)
    end
    cndict = OrderedDict((cn => i for (i, cn) in enumerate(classnames))...)

    cm = zeros(Int, length(classnames), length(classnames))
    for (gnd, pred) in zip(T, Y)
         cm[cndict[gnd], cndict[pred]] += 1
    end
    names = [Symbol("truth"), (Symbol("$n (pred)") for n in classnames)...]
    DataFrame([classnames, (col for col in eachcol(cm))...], names)
end

cmatrix(Y::AbstractMatrix{<:AbstractFloat}, T::AbstractMatrix{<:AbstractFloat}, classnames::Union{AbstractVector,Nothing} = nothing) =
    cmatrix(probs2class(Y), probs2class(T), classnames)

end


# Let's make the stroud stuff @generated.

function stroud_eval1d(deg, coeffs, result, qpts)
    nqp1d = size(qpts, 1)
    for i=1:nqp1d
        result[i] = 0.0
        xi = qpts[i, 1]
        s = 1-xi
        r = xi / s
        w = s^deg
        for alpha=0:deg
            result[i] += w * coeffs[alpha+1]
            w *= r * (deg-alpha) / (alpha+1.0)
        end
    end
    return
end

function stroud_eval2d(deg, coeffs, result, qpts)
    nqp1d = size(qpts, 1)
    result[:] = 0.0
    tmp = zeros(nqp1d, deg+1)
    for i2=1:nqp1d
        xi = qpts[i2, 2]
        s = 1 - xi
        r = xi / s
        for alpha1=0:deg
            w = s^(deg-alpha1)
            aind = polydim(2,deg) - polydim(2, deg-alpha1) + 1
            for alpha2=0:deg-alpha1
                tmp[i2, alpha1+1] += w * coeffs[aind+alpha2]
                w = w * r * (deg-alpha1-alpha2) / (1.0 + alpha2)
            end
        end
    end
    for i1=1:nqp1d
        xi = qpts[i1, 1]
        s = 1-xi
        r = xi / s
        w = s^deg
        for alpha1=0:deg
            for i2=1:nqp1d
                result[(i1-1)*nqp1d + i2] += w*tmp[i2, alpha1+1]
            end
            w = w * r * (deg-alpha1) / (1. + alpha1)
        end
    end

    return
end

val2val{d}(::Type{Val{d}}) = d

@generated function stroud_eval{sdim}(::Type{Val{sdim}}, deg, coeffs, result, qpts)
    return stroud_eval_impl(Val{sdim}, deg, coeffs, result, qpts)
end

function make_stroud_eval_constructor(sd)
    # return sdim = 1 code as a special case
    if sd == 1
        function makeit(deg, qpts)
            function doit(coeffs, result)
                nqp1d = size(qpts, 1)
                for i=1:nqp1d
                    result[i] = 0.0
                    xi = qpts[i, 1]
                    s = 1-xi
                    r = xi / s
                    w = s^deg
                    for alpha=0:deg
                        result[i] += w * coeffs[alpha+1]
                        w *= r * (deg-alpha) / (alpha+1.0)
                    end
                end
                return
            end
            return doit
        end
        return makeit
    else # sd >= 2
        # first we have to generate the function body.
        ai = [symbol("alpha",i) for i=1:sd]
        ii = [symbol("i",i) for i=1:sd]

        degminalphas = Dict()
        degminalphas[0] = :deg
        for i=1:sd
            degminalphas[i] = :($(degminalphas[i-1]) - $(ai[i]))
        end

        # work on the first stage, which is different since it
        # reads from the input

        tmp_inds = vcat([:($(ii[sd]))],
                        [:($(ai[i])+1) for i=(sd-1):-1:1],
                        [:1])
        tmp_ref = Expr(:ref,
                       :tmp)

        for ti in tmp_inds
            push!(tmp_ref.args, ti)
        end

        first_stage = Expr(:for,
                           :($(ai[sd])=0:$(degminalphas[sd-1])),
                           Expr(:block,
                                :($(tmp_ref) += w*coeffs[aind+$(ai[sd])]),
                                :(w = w * r * $(degminalphas[sd]) / (1.0 + $(ai[sd])))))

        aind = :(binomial($(sd)+deg, $(sd)) - binomial($(sd)+deg-alpha1, $(sd)))
        for i=(sd-1):-1:2
            aind = :($(aind) + binomial($(i)+ $(degminalphas[sd-i]),$(i))
                     - binomial($(i)+ $(degminalphas[sd-i+1]), $(i)))
        end
        aind = :($(aind)+1)
        first_stage = Expr(:block,
                           :(aind = $(aind)),
                           :(w=s^($(degminalphas[sd-1]))),
                           first_stage)

        for i=(sd-1):-1:1
            first_stage = Expr(:for,
                               :($(ai[i])=0:$(degminalphas[i-1])),
                               first_stage)
        end

        first_stage = Expr(:for,
                           :($(ii[sd])=1:nqp1d),
                           Expr(:block,
                                :(xi=qpts[$(ii[sd]), $(sd)]),
                                :(s=1-xi),
                                :(r=xi/s),
                                first_stage))

        # middle stages:
        # on stage k, we have k-1 innermost i loops over last i variables
        # tucked inside of (sd-k+1) alpha loops
        # tucked inside of a loop over ik

        middle_stages = []

        for k=2:(sd-1)
            # start wth write location for tmp
            iinds = [:($(ii[sd-i+1])) for i=1:k]
            ainds = [:($(ai[i])+1) for i=(sd-k):-1:1]
            tmp_inds = vcat(iinds, ainds, [:($k)])
            tmp_write_ref = Expr(:ref,
                                 :tmp)

            for ti in tmp_inds
                push!(tmp_write_ref.args, ti)
            end

            # read location for tmp
            iinds = [:($(ii[sd-i+1])) for i=1:(k-1)]
            ainds = [:($(ai[i])+1) for i=(sd-k+1):-1:1]
            tmp_inds = vcat(iinds, ainds, [:($(k-1))])
            tmp_read_ref = Expr(:ref,
                                :tmp)

            for ti in tmp_inds
                push!(tmp_read_ref.args, ti)
            end

            ms = :($(tmp_write_ref) += w * $(tmp_read_ref))

            # wrap ms into the i loops
            for i=1:(k-1)
                ms = Expr(:for,
                          :($(ii[sd-i+1])=1:nqp1d),
                          ms)
            end

            # the innermost alpha is sd - k + 1
            j = sd - k
            ms = Expr(:for,
                      :($(ai[j+1])=0:$(degminalphas[j])),
                      Expr(:block,
                           ms,
                           :(w=w*r*($(degminalphas[j+1])) / (1.0+$(ai[j+1])))))

            ms = Expr(:block,
                      :(w=s^($(degminalphas[j]))),
                      ms)

            for i=k-1:-1:1
                ms = Expr(:for,
                          :($(ai[i])=0: $(degminalphas[i-1])),
                          ms)
            end

            ms = Expr(:for,
                      :($(ii[sd-k+1])=1:nqp1d),
                      Expr(:block,
                           :(xi=qpts[$(ii[sd-k+1]), $(sd-k+1)]),
                           :(s=1-xi),
                           :(r=xi/s),
                           ms)
                      )

            push!(middle_stages, ms)

        end

        # last stage is different because we're writing into the result
        # we have an outer loop over i1
        write_idx = :($(ii[sd]))
        for i=(sd-1):-1:1
            write_idx = :(($(ii[i])-1)*nqp1d^($(sd-i)) + $(write_idx))
        end

        tmp_idx = Expr(:ref,
                       :tmp)
        for i=sd:(-1):2
            push!(tmp_idx.args, :($(ii[i])))
        end
        push!(tmp_idx.args, :(alpha1+1))
        push!(tmp_idx.args, :($(sd-1)))

        last_stage = Expr(:+=,
                          Expr(:ref,
                               :result,
                               write_idx),
                          :(w * $(tmp_idx)))

        for i=sd:-1:2
            last_stage = Expr(:for,
                              :($(ii[i])=1:nqp1d),
                              last_stage)
        end

        last_stage = Expr(:for,
                          :(i1=1:nqp1d),
                          Expr(:block,
                               :(xi=qpts[i1,1]),
                               :(s=1-xi),
                               :(r=xi/s),
                               :(w=s^deg),
                               Expr(:for,
                                    :(alpha1=0:deg),
                                    Expr(:block,
                                         last_stage,
                                         :(w=w*r*(deg-alpha1)/(1.+alpha1))))))

        make_tmp = Expr(:call, :zeros)
        push!(make_tmp.args, :nqp1d)
        for i=1:sd-1
           push!(make_tmp.args, :(deg+1))
        end
        push!(make_tmp.args, sd-1)

        doit_func = Expr(:function,
                         :(doit(coeffs, result)))

        doit_func_body = Expr(:block)

        push!(doit_func_body.args,
              :(for i=1:length(result) result[i] = 0.0 end))
        push!(doit_func_body.args,
              first_stage)
        for ms in middle_stages
            push!(doit_func_body.args, ms)
        end
        push!(doit_func_body.args, last_stage)

        push!(doit_func.args, doit_func_body)

        func_impl = quote
            function makeit(deg, qpts)
                nqp1d = size(qpts, 1)
                tmp = $(make_tmp)
                $(doit_func)
                return doit
            end
        end

        println(func_impl)

        return eval(func_impl)
end

end


# using Plots; pyplot(); x = linspace(-2,2,20); y = linspace(-1,3,20); contourf(x, y, rosenbrock, zlim=(0,1000))
#
# julia> xy = vec([(xi,yi) for xi=x,yi=y]);
#
# julia> uv = map(xyi -> (g=rosenbrock_gradient(collect(xyi))/2000; (g[1],g[2])), xy);
#
# julia> quiver!(Plots.unzip(xy)..., quiver=uv, c=:cyan, alpha=0.4)
#
# julia> uv = map(xyi -> (g=rosenbrock_gradient(collect(xyi))/-2000; (g[1],g[2])), xy);

using Plots; begin
pyplot()
x = linspace(-2,2,20)
y = linspace(-1,3,20)
xy = vec([(xi,yi) for xi=x,yi=y])
uv = map(xyi -> (g=rosenbrock_gradient(collect(xyi))/-2000; (g[1],g[2])), xy)
contourf(x, y, rosenbrock, xlim=extrema(x), ylim=extrema(y))
quiver!(Plots.unzip(xy)..., quiver=uv, c=:cyan, alpha=0.4)
end


if get(ENV, "CI", nothing) !== nothing
    # copied and modified from https://github.com/tkf/IPython.jl/blob/master/test/install_dependencies.jl

    # Adding Pkg in test/REQUIRE would be an error in 0.6.  Using
    # Project.toml still has some gotchas.  So:
    Pkg = Base.require(Base.PkgId(Base.UUID(0x44cfe95a1eb252eab672e2afdf69b78f), "Pkg"))

    # Let PyCall.jl use Python interpreter from Conda.jl
    # See: https://github.com/JuliaPy/PyCall.jl
    ENV["PYTHON"] = ""
    Pkg.build("PyCall")
end

using Test
using SymPy
using Symbolics

@variables t p(t) x y(t)
expr = x * p + (x^2 - 1 + y) * (p + 2t)
sexpr = symbolics_to_sympy(expr)
sp = symbolics_to_sympy(p)

@test SymPy.simplify(symbolics_to_sympy(Symbolics.solve_for(expr, p))) == SymPy.solve(sexpr, sp)[1]


using Plots

f(x) = 100*x - 60
g(x) = 30 - 60*x
h(x) = 0
v(x) = max.(f(x),g(x),h(x))
x = 0:0.1:1

plot(x, f.(x),
    xaxis = ("prob. (mu)", [0.0, 0.5, 0.6, 1.0]),
    yaxis = ("expected value", (-10,40)),
    label = "outdoor")
plot!(x, g.(x),
    label = "indoor")
plot!(x, h.(x),
    label = "none")
plot!(x, v.(x),
    label = "max{o,i,n}")


# ---
# title: 1498. Number of Subsequences That Satisfy the Given Sum Condition
# id: problem1498
# author: Tian Jun
# date: 2020-10-31
# difficulty: Medium
# categories: Sort, Sliding Window
# link: <https://leetcode.com/problems/number-of-subsequences-that-satisfy-the-given-sum-condition/description/>
# hidden: true
# ---
#
# Given an array of integers `nums` and an integer `target`.
#
# Return the number of **non-empty** subsequences of `nums` such that the sum of
# the minimum and maximum element on it is less or equal to `target`. Since the
# answer may be too large, return it modulo `109 + 7`.
#
#
#
# **Example 1:**
#
#
#
#     Input: nums = [3,5,6,7], target = 9
#     Output: 4
#     Explanation: There are 4 subsequences that satisfy the condition.
#     [3] -> Min value + max value <= target (3 + 3 <= 9)
#     [3,5] -> (3 + 5 <= 9)
#     [3,5,6] -> (3 + 6 <= 9)
#     [3,6] -> (3 + 6 <= 9)
#
#
# **Example 2:**
#
#
#
#     Input: nums = [3,3,6,8], target = 10
#     Output: 6
#     Explanation: There are 6 subsequences that satisfy the condition. (nums can have repeated numbers).
#     [3] , [3] , [3,3], [3,6] , [3,6] , [3,3,6]
#
# **Example 3:**
#
#
#
#     Input: nums = [2,3,3,4,6,7], target = 12
#     Output: 61
#     Explanation: There are 63 non-empty subsequences, two of them don't satisfy the condition ([6,7], [7]).
#     Number of valid subsequences (63 - 2 = 61).
#
#
# **Example 4:**
#
#
#
#     Input: nums = [5,2,4,1,7,6,8], target = 16
#     Output: 127
#     Explanation: All non-empty subset satisfy the condition (2^7 - 1) = 127
#
#
#
# **Constraints:**
#
#   * `1 <= nums.length <= 105`
#   * `1 <= nums[i] <= 106`
#   * `1 <= target <= 106`
#
#
## @lc code=start
using LeetCode

## add your code here:
## @lc code=end


# ========== Single Step Stationary Wavelet Transform ==========
# ----- 1D SWT -----
"""
    sdwt_step(v, d, h, g)

Perform one level of the stationary discrete wavelet transform (SDWT) on the vector `v`,
which is the `d`-th level scaling coefficients (Note the 0th level scaling coefficients is
the raw signal). The vectors `h` and `g` are the detail and scaling filters.

# Arguments
- `v::AbstractArray{T} where T<:Number`: Array of coefficients from a node at level `d`.
- `d::Integer`: Depth level of `v`.
- `h::Vector{S} where S<:Number`: High pass filter.
- `g::Vector{S} where S<:Number`: Low pass filter.

# Returns
- `w₁::AbstractVector{T} where T<:Number` or `w₁::AbstractMatrix{T} where T<:Number`: Output
  from low pass filter (1D case); or output from low + low pass filter (2D case).
- `w₂::AbstractVector{T} where T<:Number` or `w₂::AbstractMatrix{T} where T<:Number`: Output
  from high pass filter (1D case); or output from low + high pass filter (2D case).
- `w₃::AbstractVector{T} where T<:Number` or `w₃::AbstractMatrix{T} where T<:Number`: Output
  from high + low pass filter (2D case).
- `w₄::AbstractVector{T} where T<:Number` or `w₄::AbstractMatrix{T} where T<:Number`: Output
  from high + high pass filter (2D case).

# Examples
```julia
using Wavelets, WaveletsExt

# Setup
v = randn(8)
wt = wavelet(WT.haar)
g, h = WT.makereverseqmfpair(wt, true)

# One step of SDWT
SWT.sdwt_step(v, 0, h, g)
```

**See also:** [`sdwt_step!`](@ref)
"""
function sdwt_step(v::AbstractVector{T}, d::Integer, h::Array{S,1}, g::Array{S,1}) where
                  {T<:Number, S<:Number}
    n = length(v)
    w₁ = zeros(T, n)
    w₂ = zeros(T, n)

    sdwt_step!(w₁, w₂, v, d, h, g)
    return w₁, w₂
end

"""
    sdwt_step!(w₁, w₂, v, d, h, g)

Same as `sdwt_step` but without array allocation.

# Arguments
- `w₁::AbstractVector{T} where T<:Number` or `w₁::AbstractMatrix{T} where T<:Number`: Vector
  allocation for output from low pass filter (1D case); or matrix allocation for output from
  low + low pass filter (2D case).
- `w₂::AbstractVector{T} where T<:Number` or `w₂::AbstractMatrix{T} where T<:Number`: Vector
  allocation for output from high pass filter (1D case); or matrix allocation for output
  from low + high pass filter (2D case).
- `w₃::AbstractVector{T} where T<:Number` or `w₃::AbstractMatrix{T} where T<:Number`: Matrix
  allocation for output from high + low pass filter (2D case).
- `w₄::AbstractVector{T} where T<:Number` or `w₄::AbstractMatrix{T} where T<:Number`: Matrix
  allocation for output from high + high pass filter (2D case).
- `v::AbstractArray{T} where T<:Number`: Array of coefficients from a node at level `d`.
- `d::Integer`: Depth level of `v`.
- `h::Vector{S} where S<:Number`: High pass filter.
- `g::Vector{S} where S<:Number`: Low pass filter.

# Returns
- `w₁::AbstractVector{T} where T<:Number` or `w₁::AbstractMatrix{T} where T<:Number`: Output
  from low pass filter (1D case); or output from low + low pass filter (2D case).
- `w₂::AbstractVector{T} where T<:Number` or `w₂::AbstractMatrix{T} where T<:Number`: Output
  from high pass filter (1D case); or output from low + high pass filter (2D case).
- `w₃::AbstractVector{T} where T<:Number` or `w₃::AbstractMatrix{T} where T<:Number`: Output
  from high + low pass filter (2D case).
- `w₄::AbstractVector{T} where T<:Number` or `w₄::AbstractMatrix{T} where T<:Number`: Output
  from high + high pass filter (2D case).

# Examples
```julia
using Wavelets, WaveletsExt

# Setup
v = randn(8)
wt = wavelet(WT.haar)
g, h = WT.makereverseqmfpair(wt, true)
w₁ = zeros(8)
w₂ = zeros(8)

# One step of SDWT
SWT.sdwt_step!(w₁, w₂, v, 0, h, g)
```

**See also:** [`sdwt_step`](@ref)
"""
function sdwt_step!(w₁::AbstractVector{T},
                    w₂::AbstractVector{T},
                    v::AbstractVector{T},
                    d::Integer,
                    h::Array{S,1},
                    g::Array{S,1}) where {T<:Number, S<:Number}
    # Sanity check
    @assert length(w₁) == length(w₂) == length(v)
    @assert length(h) == length(g)

    # Setup
    n = length(v)                   # Signal length
    filtlen = length(h)             # Filter length

    # One step of stationary transform
    for i in 1:n
        k₁ = mod1(i-(1<<d),n)       # Start index for low pass filtering
        k₂ = i                      # Start index for high pass filtering
        @inbounds w₁[i] = g[end] * v[k₁]
        @inbounds w₂[i] = h[1] * v[k₂]
        for j in 2:filtlen
            k₁ = k₁+(1<<d) |> k₁ -> k₁>n ? mod1(k₁,n) : k₁
            k₂ = k₂-(1<<d) |> k₂ -> k₂≤0 ? mod1(k₂,n) : k₂
            @inbounds w₁[i] += g[end-j+1] * v[k₁]
            @inbounds w₂[i] += h[j] * v[k₂]
        end
    end
    return w₁, w₂
end

"""
    isdwt_step(w₁, w₂, d, h, g)
    isdwt_step(w₁, w₂, d, sv, sw, h, g)

Perform one level of the inverse stationary discrete wavelet transform (ISDWT) on the
vectors `w₁` and `w₂`, which are the `d+1`-th level scaling coefficients (Note the 0th level
scaling coefficients is the raw signal). The vectors `h` and `g` are the detail and scaling
filters.

!!! note
    One can decide to choose the average-based inverse transform or the shift-based inverse
    transform. For shift based, one needs to specify the shifts of the parent and children
    nodes; for the average based, the output is the average of all possible shift-based
    inverse transform outputs.

# Arguments
- `w₁::AbstractVector{T} where T<:Number` or `w₁::AbstractMatrix{T} where T<:Number`:
  Coefficients of left child node (1D case); or coefficients from top left child node (2D
  case).
- `w₂::AbstractVector{T} where T<:Number` or `w₂::AbstractMatrix{T} where T<:Number`:
  Coefficients of right child node (1D case); or coefficients from top right child node (2D
  case).
- `w₃::AbstractVector{T} where T<:Number` or `w₃::AbstractMatrix{T} where T<:Number`:
  Coefficients from bottom left child node (2D case).
- `w₄::AbstractVector{T} where T<:Number` or `w₄::AbstractMatrix{T} where T<:Number`:
  Coefficients from bottom right child node (2D case).
- `d::Integer`: Depth level of parent node of `w₁` and `w₂`.
- `sv::Integer`: Shift of parent node `v`.
- `sw::Integer`: Shift of children nodes `w₁` and `w₂`.
- `h::Vector{S} where S<:Number`: High pass filter.
- `g::Vector{S} where S<:Number`: Low pass filter.

# Returns
- `v::Vector{T}`: Reconstructed coefficients.

# Examples
```julia
using Wavelets, WaveletsExt

# Setup
v = randn(8)
wt = wavelet(WT.haar)
g, h = WT.makereverseqmfpair(wt, true)

# One step of SDWT
w₁, w₂ = SWT.sdwt_step(v, 0, h, g)

# One step of ISDWT
v̂ = SWT.isdwt_step(w₁, w₂, 0, h, g)         # Average based
ṽ = SWT.isdwt_step(w₁, w₂, 0, 0, 1, h, g)   # Shift based
```

**See also:** [`isdwt_step!`](@ref)
"""
function isdwt_step(w₁::AbstractVector{T},
                    w₂::AbstractVector{T},
                    d::Integer,
                    h::Array{S,1},
                    g::Array{S,1}) where {T<:Number, S<:Number}
    v = similar(w₁)
    isdwt_step!(v, w₁, w₂, d, h, g)
    return v
end

function isdwt_step(w₁::AbstractVector{T},
                    w₂::AbstractVector{T},
                    d::Integer,
                    sv::Integer,
                    sw::Integer,
                    h::Array{S,1},
                    g::Array{S,1}) where {T<:Number, S<:Number}
    v = similar(w₁)
    isdwt_step!(v, w₁, w₂, d, sv, sw, h, g)
    return v
end

"""
    isdwt_step!(v, w₁, w₂, d, h, g)
    isdwt_step!(v, w₁, w₂, d, sv, sw, h, g[; add2out])

Same as `isdwt_step` but without array allocation.

# Arguments
- `v::AbstractArray{T} where T<:Number`: Array allocation for reconstructed coefficients.
- `w₁::AbstractVector{T} where T<:Number` or `w₁::AbstractMatrix{T} where T<:Number`:
  Coefficients of left child node (1D case); or coefficients from top left child node (2D
  case).
- `w₂::AbstractVector{T} where T<:Number` or `w₂::AbstractMatrix{T} where T<:Number`:
  Coefficients of right child node (1D case); or coefficients from top right child node (2D
  case).
- `w₃::AbstractVector{T} where T<:Number` or `w₃::AbstractMatrix{T} where T<:Number`:
  Coefficients from bottom left child node (2D case).
- `w₄::AbstractVector{T} where T<:Number` or `w₄::AbstractMatrix{T} where T<:Number`:
  Coefficients from bottom right child node (2D case).
- `d::Integer`: Depth level of parent node of `w₁` and `w₂`.
- `sv::Integer`: Shift of parent node `v`.
- `sw::Integer`: Shift of children nodes `w₁` and `w₂`.
- `h::Vector{S} where S<:Number`: High pass filter.
- `g::Vector{S} where S<:Number`: Low pass filter.

# Keyword Arguments
- `add2out::Bool`: (Default: `false`) Whether to add computed result directly to output `v`
  or rewrite computed result to `v`.

# Returns
- `v::Array{T}`: Reconstructed coefficients.

# Examples
```julia
using Wavelets, WaveletsExt

# Setup
v = randn(8)
v̂ = similar(v)
ṽ = similar(v)
wt = wavelet(WT.haar)
g, h = WT.makereverseqmfpair(wt, true)

# One step of SDWT
w₁, w₂ = SWT.sdwt_step(v, 0, h, g)

# One step of ISDWT
SWT.isdwt_step!(v̂, w₁, w₂, 0, h, g)          # Average based
SWT.isdwt_step!(ṽ, w₁, w₂, 0, 0, 1, h, g)    # Shift based
```

**See also:** [`isdwt_step`](@ref)
"""
function isdwt_step!(v::AbstractVector{T},
                     w₁::AbstractVector{T},
                     w₂::AbstractVector{T},
                     d::Integer,
                     h::Array{S,1},
                     g::Array{S,1}) where {T<:Number, S<:Number}
    nd = 1 << d             # Number of blocks for parent node

    # isdwt_step for each shift
    for sv in 0:(nd-1)
        sw₁ = sv            # Shift of w₁, w₂ with no additional shift
        sw₂ = sv + 1<<d     # Shift of w₁, w₂ with addtional shift
        isdwt_step!(v, w₁, w₂, d, sv, sw₁, h, g)                 # Without shift
        isdwt_step!(v, w₁, w₂, d, sv, sw₂, h, g, add2out=true)   # With shift
    end

    # Get average output of shift vs no shift
    for i in eachindex(v)
        @inbounds v[i] /= 2
    end
end

function isdwt_step!(v::AbstractVector{T},
                     w₁::AbstractVector{T},
                     w₂::AbstractVector{T},
                     d::Integer,
                     sv::Integer,
                     sw::Integer,
                     h::Array{S,1},
                     g::Array{S,1};
                     add2out::Bool = false) where {T<:Number, S<:Number}
    # Sanity check
    @assert 0 ≤ sv < 1<<d
    @assert sv ≤ sw < 1<<(d+1)

    # Setup
    n = length(v)               # Signal length
    filtlen = length(h)         # filter length
    ip = sv + 1                 # Parent start index
    sp = 1 << d                 # Parent step size
    ic = sw + 1                 # Child start index
    sc = 1 << (d+1)             # Child step size

    # One step of inverse SDWT
    for (t, m) in enumerate(ip:sp:n)
        i₀ = mod1(t,2)                                  # Pivot point for filter
        i₁ = filtlen-i₀+1                               # Index for low pass filter g
        i₂ = mod1(t+1,2)                                # Index for high pass filter h
        j = sw==sv ? mod1(m-1<<d,n) : mod1(m+sp-1<<d,n) # Position of v, shift needed if sw≠sv
        k₁ = ((t-1)>>1) * sc + ic                       # Index for approx coefs w₁
        k₂ = ((t-1)>>1) * sc + ic                       # Index for detail coefs w₂
        @inbounds v[j] = add2out ? v[j]+g[i₁]*w₁[k₁]+h[i₂]*w₂[k₂] : g[i₁]*w₁[k₁]+h[i₂]*w₂[k₂]
        for i in (i₀+2):2:filtlen
            i₁ = filtlen-i+1
            i₂ = i + isodd(i) - iseven(i)
            k₁ = k₁-(1<<(d+1)) |> k₁ -> k₁≤0 ? mod1(k₁,n) : k₁
            k₂ = k₂+(1<<(d+1)) |> k₂ -> k₂>n ? mod1(k₂,n) : k₂
            @inbounds v[j] += g[i₁]*w₁[k₁]+h[i₂]*w₂[k₂]
        end
    end
    return v
end

# ----- 2D SWT -----
# Forward transform with allocation
function sdwt_step(v::AbstractMatrix{T}, d::Integer, h::Vector{S}, g::Vector{S}) where
                  {T<:Number, S<:Number}
    n, m = size(v)
    w₁ = Matrix{T}(undef, (n,m))
    w₂ = Matrix{T}(undef, (n,m))
    w₃ = Matrix{T}(undef, (n,m))
    w₄ = Matrix{T}(undef, (n,m))
    temp = Array{T,3}(undef, (n,m,2))
    sdwt_step!(w₁, w₂, w₃, w₄, v, d, h, g, temp)
    return w₁, w₂, w₃, w₄
end
# Forward transform without allocation
function sdwt_step!(w₁::AbstractMatrix{T}, w₂::AbstractMatrix{T},
                    w₃::AbstractMatrix{T}, w₄::AbstractMatrix{T},
                    v::AbstractMatrix{T},
                    d::Integer,
                    h::Vector{S}, g::Vector{S},
                    temp::AbstractArray{T,3}) where {T<:Number, S<:Number}
    # Sanity check
    @assert size(v) == size(w₁) == size(w₂) == size(w₃) == size(w₄)
    @assert ndims(temp) == 3
    @assert size(temp,1) == size(v,1)
    @assert size(temp,2) == size(v,2)

    # Setup
    n, m = size(w₁)

    # --- Transform ---
    # Compute sdwt for all columns
    for j in 1:m
        @inbounds temp₁ⱼ = @view temp[:,j,1]
        @inbounds temp₂ⱼ = @view temp[:,j,2]
        @inbounds vⱼ = @view v[:,j]
        @inbounds sdwt_step!(temp₁ⱼ, temp₂ⱼ, vⱼ, d, h, g)
    end
    # Compute sdwt for all rows
    for i in 1:n
        @inbounds temp₁ᵢ = @view temp[i,:,1]
        @inbounds w₁ᵢ = @view w₁[i,:]
        @inbounds w₂ᵢ = @view w₂[i,:]
        @inbounds sdwt_step!(w₁ᵢ, w₂ᵢ, temp₁ᵢ, d, h, g)

        @inbounds temp₂ᵢ = @view temp[i,:,2]
        @inbounds w₃ᵢ = @view w₃[i,:]
        @inbounds w₄ᵢ = @view w₄[i,:]
        @inbounds sdwt_step!(w₃ᵢ, w₄ᵢ, temp₂ᵢ, d, h, g)
    end
    return w₁, w₂, w₃, w₄
end

# Inverse transform (Average based) with allocation
function isdwt_step(w₁::AbstractMatrix{T}, w₂::AbstractMatrix{T},
                    w₃::AbstractMatrix{T}, w₄::AbstractMatrix{T},
                    d::Integer,
                    h::Array{S,1}, g::Array{S,1}) where {T<:Number, S<:Number}
    n, m = size(w₁)
    v = Matrix{T}(undef, (n,m))
    temp = Array{T,3}(undef, (n,m,2))
    isdwt_step!(v, w₁, w₂, w₃, w₄, d, h, g, temp)
    return v
end
# Inverse transform (Shift based) with allocation
function isdwt_step(w₁::AbstractMatrix{T}, w₂::AbstractMatrix{T},
                    w₃::AbstractMatrix{T}, w₄::AbstractMatrix{T},
                    d::Integer, sv::Integer, sw::Integer,
                    h::Vector{S}, g::Vector{S}) where {T<:Number, S<:Number}
    n, m = size(w₁)
    v = Matrix{T}(undef, (n,m))
    temp = Array{T,3}(undef, (n,m,2))
    isdwt_step!(v, w₁, w₂, w₃, w₄, d, sv, sw, h, g, temp)
    return v
end
# Inverse transform (Average based) without allocation
function isdwt_step!(v::AbstractMatrix{T},
                     w₁::AbstractMatrix{T}, w₂::AbstractMatrix{T},
                     w₃::AbstractMatrix{T}, w₄::AbstractMatrix{T},
                     d::Integer,
                     h::Vector{S}, g::Vector{S},
                     temp::AbstractArray{T,3}) where {T<:Number, S<:Number}
    # Sanity check
    @assert size(v) == size(w₁) == size(w₂) == size(w₃) == size(w₄)
    @assert ndims(temp) == 3
    @assert size(temp,1) == size(v,1)
    @assert size(temp,2) == size(v,2)

    # Setup
    n, m = size(w₁)

    # --- Transform ---
    # Compute isdwt for all rows
    for i in 1:n
        @inbounds temp₁ᵢ = @view temp[i,:,1]
        @inbounds w₁ᵢ = @view w₁[i,:]
        @inbounds w₂ᵢ = @view w₂[i,:]
        @inbounds isdwt_step!(temp₁ᵢ, w₁ᵢ, w₂ᵢ, d, h, g)

        @inbounds temp₂ᵢ = @view temp[i,:,2]
        @inbounds w₃ᵢ = @view w₃[i,:]
        @inbounds w₄ᵢ = @view w₄[i,:]
        @inbounds isdwt_step!(temp₂ᵢ, w₃ᵢ, w₄ᵢ, d, h, g)
    end
    # Compute isdwt for all columns
    for j in 1:m
        @inbounds temp₁ⱼ = @view temp[:,j,1]
        @inbounds temp₂ⱼ = @view temp[:,j,2]
        @inbounds vⱼ = @view v[:,j]
        @inbounds isdwt_step!(vⱼ, temp₁ⱼ, temp₂ⱼ, d, h, g)
    end
    return v
end
# Inverse transform (Shift based) without array allocation
function isdwt_step!(v::AbstractMatrix{T},
                     w₁::AbstractMatrix{T}, w₂::AbstractMatrix{T},
                     w₃::AbstractMatrix{T}, w₄::AbstractMatrix{T},
                     d::Integer, sv::Integer, sw::Integer,
                     h::Vector{S}, g::Vector{S},
                     temp::AbstractArray{T,3}) where {T<:Number, S<:Number}
    # Sanity check
    @assert size(v) == size(w₁) == size(w₂) == size(w₃) == size(w₄)
    @assert ndims(temp) == 3
    @assert size(temp,1) == size(v,1)
    @assert size(temp,2) == size(v,2)

    # Setup
    n, m = size(w₁)

    # --- Transform ---
    # Compute isdwt for all rows
    for i in 1:n
        @inbounds temp₁ᵢ = @view temp[i,:,1]
        @inbounds w₁ᵢ = @view w₁[i,:]
        @inbounds w₂ᵢ = @view w₂[i,:]
        @inbounds isdwt_step!(temp₁ᵢ, w₁ᵢ, w₂ᵢ, d, sv, sw, h, g)

        @inbounds temp₂ᵢ = @view temp[i,:,2]
        @inbounds w₃ᵢ = @view w₃[i,:]
        @inbounds w₄ᵢ = @view w₄[i,:]
        @inbounds isdwt_step!(temp₂ᵢ, w₃ᵢ, w₄ᵢ, d, sv, sw, h, g)
    end
    # Compute isdwt for all columns
    for j in 1:m
        @inbounds temp₁ⱼ = @view temp[:,j,1]
        @inbounds temp₂ⱼ = @view temp[:,j,2]
        @inbounds vⱼ = @view v[:,j]
        @inbounds isdwt_step!(vⱼ, temp₁ⱼ, temp₂ⱼ, d, sv, sw, h, g)
    end
    return v
end


## Test nonlinear function
function nonlinear_function(y, u, v)
    return ((0.8 - 0.5*exp(-y[1]^2))*y[1] - (0.3 + 0.9*exp(-y[1]^2))*y[2] + u[1]
            + 0.2*u[2] + 0.1*u[1]*u[2] + 0.1*v[1] + 0.3*v[2])
end

@testset "Test simulation for arbitrary nonlinear function" begin
    u = NarmaxLasso.generate_random_input(10, 5)
    y = NarmaxLasso.simulate(u, nonlinear_function, [1, 2], [1, 2], [1, 2], σv=0.1)
    @test y ≈ [0.0, 0.0, 0.318563, 0.480975, 0.134545, -0.139843, 0.458256,
                   0.603922, 0.422271, -0.0208464] rtol=1e-2
    @test length(y) == length(u)
end

## Test one step ahead
m1 = NarmaxLasso.Monomial([1, 0, 0])
m2 = NarmaxLasso.Monomial([0, 0, 1])
m3 = NarmaxLasso.Monomial([0, 2, 0])

@testset "Test one-step-ahead" begin
    @test NarmaxLasso.one_step_ahead([1.0], [2.0], [3.0],
                                     [m1, m2, m3],
                                     [1.0, 2.0, 3.0]) ≈ 19.0
    @test NarmaxLasso.one_step_ahead([1.0], [2.0, 3.0],
                                     [m1, m2, m3],
                                     [1.0, 2.0, 3.0]) ≈ 19.0
end

# simulate
@testset "Test simulate" begin
    mdl1 = NarmaxLasso.NarmaxRegressors([m1, m2, m3], [1], [1], [1])
    mdl2 = NarmaxLasso.NarmaxRegressors([m1, m2, m3], [1], [1, 2])
    u = NarmaxLasso.generate_random_input(20, 20)
    y = NarmaxLasso.simulate(u, mdl1, [0.5, 1, 0])
    @test (1-0.5)*y[end]/(u[end]^2) ≈ 1.0 rtol=1e-2

    y = NarmaxLasso.simulate(u, mdl2, [0.5, 1, 0.2])
    @test (1-0.5)*y[end]/(0.2*u[end]^2+u[end]) ≈ 1.0 rtol=1e-2
end


export rotate, rotate!

"""
    rotate(arr, θ, rotation_plane=(1,2))

Rotate an `arr` in the plane `rotation_plane` with an angle `θ` in degree
around the center pixel.

`rotate!` is also available.
"""
function rotate(arr, θ, rotation_plane=(1, 2))
    return rotate!(copy(arr), θ, rotation_plane)
end

"""
    rotate!(arr, θ, rotation_plane=(1,2))

Rotate an `arr` in the plane `rotation_plane` with an angle `θ` in degree
around the center pixel.
"""
function rotate!(arr, θ, rotation_plane=(1, 2))
    # convert to radiants
    θ = deg2rad(θ)
    a,b = rotation_plane

    # parameters for shearing
    α = -tan(θ/2) * size(arr, a)
	β = sin(θ) * size(arr, b)

    # do the three step shearing
    shear!(arr, α, a, b)
    shear!(arr, β, b, a)
    shear!(arr, α, a, b)

    return arr
end


struct Verlet{TGV, TF<:AbstractFloat, TM} <: SecondOrderNonMetropolisSampler
    ∇V!::TGV
    M::TM
    Δt::TF
end

"""
    Veret(∇V!, M, Δt)

Set up the Verlet integrator.

### Fields

* ∇V!   - In place gradient of the potential
* M     - Mass (either scalar or vector)
* Δt    - Time step
"""
function Veret(∇V!::TGV, M::TM, Δt::TF) where {TGV, TF<:AbstractFloat,TM}
    return Verlet(∇V!,M, Δt);
end

mutable struct VerletState{Tq, Tx<:AbstractVector{Tq}, TF<:AbstractFloat} <:SecondOrderNonMetropolisSamplerState
    x::Tx
    ∇V::Tq
    p_mid::Tq
end

function InitState!(x₀, sampler::Verlet)
    ∇V = similar(x₀[1])
    sampler.∇V!(∇V , x₀);
    return VerletState(x₀, copy(∇V) , similar(x₀[1]));
end

function InitState(x₀, sampler::Verlet)
    ∇V = similar(x₀[1])
    sampler.∇V!(∇V , x₀);
    return VerletState(copy(x₀), copy(∇V) , similar(x₀[1]));
end

function UpdateState!(state::VerletState, sampler::Verlet)
    @. state.p_mid = state.p - 0.5 * sampler.Δt * state.∇V;
    @. state.x = state.x + sampler.Δt * state.p_mid/sampler.M;
    sampler.∇V!(state.∇V, state.x);
    @. state.p = state.p_mid - 0.5 * sampler.Δt * state.∇V;
    state
end


struct DyadicRational <: AbstractFloat
    numerator::Int64
    R::Int64

    function DyadicRational(numerator, R)
        if R < 0
            throw(DomainError(R, "Scale must be non-negative"))
        end

        largest_power_of_two_dividing_numerator = trailing_zeros(numerator)
        common_power_of_two = min(R, largest_power_of_two_dividing_numerator)

        new(numerator >> common_power_of_two, R - common_power_of_two)
    end
end

numerator(x::DyadicRational) = x.numerator
resolution(x::DyadicRational) = x.R

function Base.show(io::IO, x::DyadicRational)
    if resolution(x) == 0
        print(io, numerator(x))
    elseif resolution(x) == 1
        print(io, numerator(x), "/2")
    else
        print(io, numerator(x), "/2^", resolution(x))
    end
end

Base.convert(::Type{DyadicRational}, x::Integer) = DyadicRational(x, 0)
Base.promote_rule(::Type{DyadicRational}, ::Type{T}) where {T <: Integer} = DyadicRational

Base.:*(x::DyadicRational, y::DyadicRational) = DyadicRational(numerator(x) * numerator(y), resolution(x) + resolution(y))
Base.:*(x::DyadicRational, k::Integer) = DyadicRational(k * numerator(x), resolution(x))
Base.:*(k::Integer, x::DyadicRational) = x * k

for op in [:+, :-]
    @eval begin
        Base.$op(k::Integer, x::DyadicRational) = $op(x, k)

        function Base.$op(x::DyadicRational, k::Integer)
            new_numerator = $op(numerator(x), k << resolution(x))
            DyadicRational(new_numerator, resolution(x))
        end

        function Base.$op(a::DyadicRational, b::DyadicRational)
            common_numerator = $op(numerator(a) << resolution(b), numerator(b) << resolution(a))
            common_resolution = resolution(a) + resolution(b)

            DyadicRational(common_numerator, common_resolution)
        end
    end
end

for op in [:<, :<=]
    @eval begin
        Base.$op(x::DyadicRational, y::DyadicRational) = $op(numerator(x) << resolution(y), numerator(y) << resolution(x))
    end
end

function Base.Integer(x::DyadicRational)
    if resolution(x) != 0
        throw(InexactError(:Integer, Int64, x))
    end

    numerator(x)
end

function all_dyadic_rationals(left, right, R)
    all_dyadic_rationals(convert(DyadicRational, left), convert(DyadicRational, right), R)
end

function all_dyadic_rationals(left::DyadicRational, right::DyadicRational, R::Integer)
    if R < resolution(left) || R < resolution(right)
        throw(DomainError(R, "Resolution must be at least that of the endpoints"))
    end

    common_resolution = max(resolution(left), resolution(right), R)
    left_numerator = numerator(left) << (common_resolution - resolution(left))
    right_numerator = numerator(right) << (common_resolution - resolution(right))

    DyadicRational.(left_numerator:right_numerator, R)
end

function AbstractFloat(x::DyadicRational)
    numerator(x) / 2.0^resolution(x)
end


using Distributions
using Test

@testset "Dirac tests" begin
    for val in (3, 3.0, -3.5)
        d = Dirac(val)

        @test minimum(d) == val
        @test maximum(d) == val
        @test !insupport(d, prevfloat(float(val)))
        @test insupport(d, val)
        @test !insupport(d, nextfloat(float(val)))
        @test support(d) == (val,)

        @test iszero(pdf(d, prevfloat(float(val))))
        @test isone(pdf(d, val))
        @test iszero(pdf(d, nextfloat(float(val))))

        @test logpdf(d, prevfloat(float(val))) == -Inf
        @test iszero(logpdf(d, val))
        @test logpdf(d, nextfloat(float(val))) == -Inf

        @test iszero(cdf(d, prevfloat(float(val))))
        @test isone(cdf(d, val))
        @test isone(cdf(d, nextfloat(float(val))))

        @test quantile(d, 0) == val
        @test quantile(d, 0.5) == val
        @test quantile(d, 1) == val

        @test rand(d) == val

        @test mean(d) == val
        @test iszero(var(d))

        @test mode(d) == val

        @test iszero(entropy(d))

        t = rand()
        @test mgf(d, t) == exp(t * val)
        @test cf(d, t) == cis(t * val)
    end
end


using ControlSystemIdentification, ControlSystems, Statistics, Random
Random.seed!(0)
wtest = exp10.(LinRange(-3, log10(pi), 30))
freqresptest(G, model) = quantile(
    vec(log10.(abs2.(freqresp(model, wtest))) - log10.(abs2.(freqresp(G, wtest)))),
    0.9,
)

freqresptest(G, model, tol) = freqresptest(G, model) < tol

N = 3000
h = 0.05
t = range(0, length = N, step = h)

ω1, ω2 = 10, 20
f1, f2 = ω1 / 2π, ω2 / 2π
ωvec = exp10.(LinRange(-0, log10(pi / h), 300))

G1 = c2d(tf(ω1, [1, 2 * 0.02 * ω1, ω1^2]), h)
G2 = c2d(tf(ω2, [1, 2 * 0.02 * ω2, ω2^2]), h)
# bodeplot([G1, G2], ωvec)

e = randn(N)
y1, y2 = lsim(G1, e, t)[1], lsim(G2, e, t)[1]
d = iddata(y1 + y2, e, h)

# Gh = arx(d,2,2)
# bodeplot([G1, G2, Gh], ωvec)

##

# P = mt_pgram(vec(d.y), fs=1/h)

H1 = Bandstop(0.9 * f1, 1.1 * f1, fs = d.fs)
H2 = Bandstop(0.9 * f2, 1.1 * f2, fs = d.fs)

Gf1 = arx(d, 2, 2, estimator = weighted_estimator(H1))
Gf2 = arx(d, 2, 2, estimator = weighted_estimator(H2))

@test norm(pole(G1) - pole(Gf1)) / norm(pole(G1)) < 0.05
@test norm(pole(G2) - pole(Gf2)) / norm(pole(G1)) < 0.05

@test freqresptest(G1, Gf1) < 0.22
@test freqresptest(G2, Gf2) < 0.22

isinteractive() && bodeplot(
    [G1, G2, Gf1, Gf2],
    ωvec,
    plotphase = false,
    lab = ["G1" "G2" "Focus f1" "Focus f2"],
)
# plot!(2pi .* P.freq, sqrt.(P.power), l=(:black, 2, :dash), xlims=extrema(ωvec))

Random.seed!(0)
N = 2000
e = randn(N)
t = range(0, length = N, step = h)
y1, y2 = lsim(G1, e, t)[1], lsim(G2, e, t)[1]
d = iddata(y1 + y2, e, h)

# d.y .+= 0.1 .* randn.()

H1 = Bandstop(0.5 * f1, 1.2 * f1, fs = 1 / h)
H2 = Bandstop(0.9 * f2, 1.5 * f2, fs = 1 / h)

Gf1 = n4sid(d, 2, Wf = H1)
Gf2 = n4sid(d, 2, Wf = H2)
Gf1 = Gf1.sys
Gf2 = Gf2.sys

@test norm(pole(G1) - pole(Gf1)) / norm(pole(G1)) < 0.05
@test norm(pole(G2) - pole(Gf2)) / norm(pole(G1)) < 0.05

@test freqresptest(G1, Gf1) < 1
@test freqresptest(G2, Gf2) < 1

# isinteractive() && bodeplot([G1, G2, Gf1, Gf2], ωvec, plotphase=false, lab=["G1" "G2" "Focus f1" "Focus f2"])


using Revise, Test, ForwardDiff, Parameters, Setfield, Plots, LinearAlgebra
using BifurcationKit, Test
const BK = BifurcationKit

norminf(x) = norm(x, Inf)
####################################################################################################
function COm(u, p)
	@unpack q1,q2,q3,q4,q5,q6,k = p
	x, y, s = u
	z = 1-x-y-s
	[
		2q1 * z^2 - 2q5 * x^2 - q3 * x * y,
		q2 * z - q6 * y - q3 * x * y,
		q4 * (z - k * s)
	]
end
jet = BK.getJet(COm; matrixfree=false)

par_com = (q1 = 2.5, q2 = 2.0, q3 = 10., q4 = 0.0675, q5 = 1., q6 = 0.1, k = 0.4)

z0 = [0.001137, 0.891483, 0.062345]

opts_br = ContinuationPar(pMin = 0.5, pMax = 2.0, ds = 0.002, dsmax = 0.01, nInversion = 6, detectBifurcation = 3, maxBisectionSteps = 25, nev = 3, maxSteps = 20000)
	@set! opts_br.newtonOptions.verbose = true
	br, = @time continuation(jet[1], jet[2], z0, par_com, (@lens _.q2), opts_br;
		recordFromSolution = (x, p) -> (x = x[1], y = x[2], s = x[3]),
		plot = false, verbosity = 3, normC = norminf,
	bothside = true)
	show(br)

plot(br, plotfold=false, markersize=4, legend=:topright, ylims=(0,0.16))
####################################################################################################
@set! opts_br.newtonOptions.verbose = true
@set! opts_br.newtonOptions.maxIter = 10
opts_br = @set opts_br.newtonOptions.tol = 1e-12

sn, = newton(jet[1:2]..., br, 2; options = opts_br.newtonOptions, bdlinsolver = MatrixBLS())

hp, = newton(jet[1:2]..., br, 1; options = NewtonPar( opts_br.newtonOptions; maxIter = 10),startWithEigen=true, d2F = jet[3])

BK.hopfNormalForm(jet..., hp, opts_br.newtonOptions.linsolver)
hpnf = computeNormalForm(jet..., br, 1)

sn_codim2, = continuationFold(jet[1:2]..., br, 2, (@lens _.k), ContinuationPar(opts_br, pMax = 3.2, pMin = 0., detectBifurcation = 0, dsmin=1e-5, ds = -0.001, dsmax = 0.05, nInversion = 6, detectEvent = 2, detectFold = false) ; plot = true,
	verbosity = 3,
	normC = norminf,
	updateMinAugEveryStep = 1,
	startWithEigen = true,
	recordFromSolution = (u,p; kw...) -> (x = u.u[1] ),
	bothside=true,
	bdlinsolver = MatrixBLS()
	)

using Test
@test sn_codim2.specialpoint[1].printsol.k 	≈ 0.971397 rtol = 1e-4
@test sn_codim2.specialpoint[1].printsol.q2 ≈ 1.417628 rtol = 1e-4
@test sn_codim2.specialpoint[3].printsol.k 	≈ 0.722339 rtol = 1e-4
@test sn_codim2.specialpoint[3].printsol.q2 ≈ 1.161199 rtol = 1e-4

plot(sn_codim2)#, real.(sn_codim2.BT), ylims = (-1,1), xlims=(0,2))

plot(sn_codim2, vars=(:q2, :x), branchlabel = "Fold", plotstability = false);plot!(br,xlims=(0.8,1.8))

hp_codim2, = continuation(jet[1:2]..., br, 1, (@lens _.k), ContinuationPar(opts_br, pMin = 0., pMax = 2.8, detectBifurcation = 0, ds = -0.0001, dsmax = 0.08, dsmin = 1e-4, nInversion = 6, detectEvent = 2, detectLoop = true, maxSteps = 50, detectFold=false) ; plot = true,
	verbosity = 3,
	normC = norminf,
	tangentAlgo = BorderedPred(),
	detectCodim2Bifurcation = 2,
	updateMinAugEveryStep = 1,
	startWithEigen = true,
	recordFromSolution = (u,p; kw...) -> (x = u.u[1] ),
	d2F = jet[3], d3F = jet[4],
	bothside = true,
	bdlinsolver = MatrixBLS())

@test hp_codim2.branch[5].l1 |> real 		≈ 33.15920 rtol = 1e-1
@test hp_codim2.specialpoint[1].printsol.k 	≈ 0.305879 rtol = 1e-3
@test hp_codim2.specialpoint[1].printsol.q2 ≈ 0.924255 rtol = 1e-3
@test hp_codim2.specialpoint[2].printsol.k 	≈ 0.235550 rtol = 1e-4
@test hp_codim2.specialpoint[2].printsol.q2 ≈ 0.896099 rtol = 1e-4

plot(sn_codim2, vars=(:q2, :x), branchlabel = "Fold", plotcirclesbif = true)
	plot!(hp_codim2, vars=(:q2, :x), branchlabel = "Hopf",plotcirclesbif = true)
	plot!(br,xlims=(0.6,1.5))

plot(sn_codim2, vars=(:k, :q2), branchlabel = "Fold")
	plot!(hp_codim2, vars=(:k, :q2), branchlabel = "Hopf",)

plot(hp_codim2, vars=(:q2, :x), branchlabel = "Hopf")
####################################################################################################


"""
$(TYPEDEF)

Bathe's integration scheme with sub-steps of equal size.

### Fields

- `Δt` -- step-size
- `α`  -- parameter α of the method
- `δ`  -- parameter δ of the method

### References

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

Bathe(; Δt::N) where {N} = Bathe(Δt)

step_size(alg::Bathe) = alg.Δt

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

    # compute integration constants
    ξ = 1/Δt
    ξ² = ξ*ξ
    a₀ = 16*ξ²
    a₁ = 4*ξ
    a₂ = 9*ξ²
    a₃ = 3*ξ
    a₄ = 2*a₁
    a₅ = 12*ξ²
    a₆ = -3*ξ²
    a₇ = -ξ

    K̂₁ = K + a₀ * M + a₁ * C
    K̂₂ = K + a₂ * M + a₃ * C
    return a₀, a₁, a₂, a₃, a₄, a₅, a₆, a₇, K̂₁, K̂₂
end

# a comment about notation: variables with superscript `+` correspond to
# evaluation at intermediate times, i.e. t + Δt/2
function _solve(alg::Bathe{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̂₁, K̂₂ = _init(alg, M, C, K)
    K̂₁⁻¹ = factorize(K̂₁)
    K̂₂⁻¹ = factorize(K̂₂)

    # initialize displacements, velocities and accelerations
    U = Vector{VT}(undef, IMAX)
    U′ = Vector{VT}(undef, IMAX)
    U′′ = Vector{VT}(undef, IMAX)
    U[1] = U₀
    U′[1] = U₀′
    U′′[1] = U₀′′

    @inbounds for i in 1:NSTEPS
        # ---------------------------------------------
        # First sub-step (Newmark trapezoidal rule)
        # ---------------------------------------------

        # calculate effective loads at time t + Δt/2
        mᵢ = M * (a₀ * U[i] + a₄ * U′[i] + U′′[i])
        cᵢ = C * (a₁ * U[i] + U′[i])
        Rᵢ⁺ = (R[i] + R[i+1])/2
        R̂ᵢ⁺ = Rᵢ⁺ + mᵢ + cᵢ

        # solve for displacements at time t + Δt/2
        Uᵢ⁺ = K̂₁⁻¹ \ R̂ᵢ⁺

        # calculate velocities at time t + Δt/2
        U′ᵢ⁺ = a₁ * (Uᵢ⁺ - U[i]) - U′[i]

        # ---------------------------------------------
        # Second sub-step (3-point Euler backward method)
        # ---------------------------------------------

        # calculate effective loads at time t + Δt
        mᵢ = M * (a₅ * Uᵢ⁺ + a₆ * U[i] + a₁ * U′ᵢ⁺ + a₇ * U′[i])
        cᵢ = C * (a₁ * Uᵢ⁺ + a₇ * U[i])
        R̂ᵢ₊₁ = R[i+1] + mᵢ + cᵢ

        # solve for displacements at time t + Δt
        U[i+1] = K̂₂⁻¹ \ R̂ᵢ₊₁

        # calculate velocities and accelerations at time t + Δt
        U′[i+1] = -a₇ * U[i] - a₁ * Uᵢ⁺ + a₃ * U[i+1]
        U′′[i+1] = -a₇ * U′[i] - a₁ * U′ᵢ⁺ + a₃ * U′[i+1]
    end

    return _build_solution(alg, U, U′, U′′, NSTEPS)
end

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


using Distributions

m = constructdice()
run(m)

# Uncertain parameters defined in "Projections and Uncertainties about Climate
# Change in an Era of Minimal Climate Policies" (Nordhaus, 2018) EViews file
# provided by the AEJ : https://www.aeaweb.org/articles?id=10.1257/pol.20170046

mcs = @defsim begin
    t2xco2 = LogNormal(1.106, 0.2646)
	ga0 = Normal(0.076, 0.056)
	gsigma1 = Normal(-0.0152, 0.0032)
	mueq = LogNormal(5.851, 0.2649)
	a2 = Normal(0.00227, 0.001135)

    save(damages.DAMAGES)
	save(grosseconomy.YGROSS)
	save(emissions.E)
	save(co2cycle.MU)
	save(damages.DAMFRAC)
end

generate_trials!(mcs, 10000, filename="/tmp/dice-2016/trialdata.csv")

# Run trials 1:4, and save results to the indicated directory
set_models!(mcs, m)
run_sim(mcs, output_dir="/tmp/dice-2016")


using OrdinaryDiffEq, BenchmarkTools, StaticArrays
abstract type ContinuousDynamicalSystem end

function lorenz1(u0=[0.0, 10.0, 0.0]; σ = 10.0, ρ = 28.0, β = 8//3)
    @inline @inbounds function eom_lorenz!(du, u)
        du[1] = σ*(u[2]-u[1])
        du[2] = u[1]*(ρ-u[3]) - u[2]
        du[3] = u[1]*u[2] - β*u[3]
    end
    @inline @inbounds function jacob_lorenz(u)
        i = one(eltype(u))
        o = zero(eltype(u))
        @SMatrix [-σ*i           σ*i    zero(i);
                  (ρ*i - u[3])   (-i)   (-u[1]);
                  u[2]           u[1]   (-β*i) ]
    end# should give exponents [0.9056, 0, -14.5723]
    return ContinuousDS1(u0, eom_lorenz!, jacob_lorenz)
end
mutable struct ContinuousDS1{T<:AbstractVector, F, J} <: ContinuousDynamicalSystem
    state::T
    eom!::F
    jacob::J
end

function lorenz2(u0=[0.0, 10.0, 0.0]; σ = 10.0, ρ = 28.0, β = 8//3)
    @inline @inbounds function eom_lorenz!(du, u)
        du[1] = σ*(u[2]-u[1])
        du[2] = u[1]*(ρ-u[3]) - u[2]
        du[3] = u[1]*u[2] - β*u[3]
    end
    i = one(eltype(u0))
    o = zero(eltype(u0))
    J = zeros(eltype(u0), 3, 3)
    J[1,:] .= [-σ*i,        σ*i,    o]
    J[2,:] .= [ρ*i - u0[3],   -i,   -u0[1]]
    J[3,:] .= [u0[2],        u0[1],   (-β*i)]

    @inline @inbounds function jacob_lorenz!(J, u)
        J[2,1] = ρ - u[3]
        J[2,3] = -u[1]
        J[3,1] = u[2]; J[3,2] = u[1]
    end
    return ContinuousDS2(u0, eom_lorenz!, jacob_lorenz!, J)
end
mutable struct ContinuousDS2{T<:Real, F, J} <: ContinuousDynamicalSystem
    state::Vector{T}
    eom!::F
    jacob::J
    J::Matrix{T}
end

### Test jacobian function speed
function bjac()
    ds1 = lorenz1()
    b1 = @benchmark $ds1.jacob($ds1.state)
    println("Jacobian of 1: ")
    Juno.render(ds1.jacob(ds1.state))
    Juno.render(b1)

    ds2 = lorenz2()
    b2 = @benchmark $ds2.jacob($ds2.J, $ds2.state)
    println("Jacobian of 2: ")
    Juno.render(ds2.J)
    Juno.render(b2)
end
bjac()

### TANGENT BUNDLE!!!
function tangentbundle_setup_integrator(ds::ContinuousDS1, t_final;
  diff_eq_kwargs=Dict())

    D = dimension(ds)
    f! = ds.eom!
    jac = ds.jacob

    tbeom! = (t, u, du) -> begin
        f!(view(du, :, 1), u)
        A_mul_B!(
            view(du, :, 2:D+1),
            jac(view(u, :, 1)),
            view(u, :, 2:D+1)
        )
    end

    S = [ds.state eye(eltype(ds.state), D)]
    tbprob = ODEProblem(tbeom!, S, (zero(t_final), t_final))
    tb_integ = init(tbprob, Tsit5(); diff_eq_kwargs..., save_everystep=false)
    return tb_integ
end

function tangentbundle_setup_integrator(ds::ContinuousDS2, t_final;
  diff_eq_kwargs=Dict())

    D = dimension(ds)
    f! = ds.eom!

    tbeom! = (t, u, du) -> begin
        us = view(u, :, 1)
        f!(view(du, :, 1), us)
        ds.jacob(ds.J, us)
        A_mul_B!(
            view(du, :, 2:D+1),
            ds.J,
            view(u, :, 2:D+1)
        )
    end

    S = [ds.state eye(eltype(ds.state), D)]
    tbprob = ODEProblem(tbeom!, S, (zero(t_final), t_final))
    tb_integ = init(tbprob, Tsit5(); diff_eq_kwargs..., save_everystep=false)
    return tb_integ
end

# Benchmark it

function lyapunovspectrum1(ds::ContinuousDynamicalSystem, N::Real=1000;
    Ttr::Real = 0.0, diff_eq_kwargs::Dict = Dict(), dt::Real = 0.1)
    # Initialize
    tstops = dt:dt:N*dt
    D = dimension(ds)
    λ = zeros(eltype(ds.state), D)
    Q = eye(eltype(ds.state), D)
    # Transient evolution:
    Ttr != 0 && evolve!(ds, Ttr; diff_eq_kwargs = diff_eq_kwargs)
    # Create integrator for dynamics and tangent space:
    integ = tangentbundle_setup_integrator(
    ds, tstops[end]; diff_eq_kwargs = diff_eq_kwargs)

    # Main algorithm
    for τ in tstops
        integ.u[:, 2:end] .= Q # update tangent dynamics state (super important!)
        u_modified!(integ, true)
        # Integrate
        while integ.t < τ
            step!(integ)
        end
        # Perform QR (on the tangent flow):
        Q, R = qr(view(integ.u, :, 2:D+1))
        # Add correct (positive) numbers to Lyapunov spectrum
        for j in 1:D
            λ[j] += log(abs(R[j,j]))
        end
    end
    λ./(integ.t) #return spectrum
end

function lyapunovspectrum1(ds::ContinuousDynamicalSystem, N::Real=1000;
    Ttr::Real = 0.0, diff_eq_kwargs::Dict = Dict(), dt::Real = 0.1)
    # Initialize
    tstops = dt:dt:N*dt
    D = dimension(ds)
    λ = zeros(eltype(ds.state), D)
    Q = eye(eltype(ds.state), D)
    # Transient evolution:
    Ttr != 0 && evolve!(ds, Ttr; diff_eq_kwargs = diff_eq_kwargs)
    # Create integrator for dynamics and tangent space:
    integ = tangentbundle_setup_integrator(
    ds, tstops[end]; diff_eq_kwargs = diff_eq_kwargs)

    # Main algorithm
    for τ in tstops
        integ.u[:, 2:end] .= Q # update tangent dynamics state (super important!)
        u_modified!(integ, true)
        # Integrate
        while integ.t < τ
            step!(integ)
        end
        # Perform QR (on the tangent flow):
        Q, R = qr(view(integ.u, :, 2:D+1))
        # Add correct (positive) numbers to Lyapunov spectrum
        for j in 1:D
            λ[j] += log(abs(R[j,j]))
        end
    end
    λ./(integ.t) #return spectrum
end
dimension(ds) = 3

function binteg()
    A = Dict(:abstol=> 1e-6, :reltol=>1e-6)
    ds1 = lorenz1()
    ls1 = lyapunovspectrum1(ds1, 5000.0, diff_eq_kwargs=A)
    b1 = @benchmark lyapunovspectrum1($ds1, 500.0, diff_eq_kwargs=$A)
    println("Lyapunovs 1: $(ls1)")
    Juno.render(b1)

    ds2 = lorenz2()
    ls2 = lyapunovspectrum1(ds2, 5000.0, diff_eq_kwargs=A)
    b2 = @benchmark lyapunovspectrum1($ds2, 500.0, diff_eq_kwargs=$A)
    println("Lyapunovs 2: $(ls2)")

    Juno.render(b2)

end
binteg()

#
# b1 = @benchmark ds1.jacob(ds1.state)
# println("Jacobian of 1: ")
# Juno.render(b1)
#
#
# ds2 = lorenz2()
# b2 = @benchmark ds2.jacob(ds2.J, ds2.state)
# println("Jacobian of 2: ")
# Juno.render(b2)


function distance(x,y)
    return sqrt((x[1] - y[1])^2  + (x[2] - y[2])^2)
end

function read_data()
    try
        count = 0
        g = Dict()
        for line in eachline(STDIN)
            cur = split(strip(line))
            if !in(cur[1], Set(["NAME", "COMMENT", "TYPE", "DIMENSION:",
                                "EDGE_WEIGHT_TYPE", "NODE_COORD_SECTION", "EOF"]))
                g[cur[1]] = [[parse(Int,cur[2]), parse(Int,cur[3])], Any[]]
                count += 1
            end
        end
        for i in keys(g)
            for j in keys(g)
                if i != j
                    dist = distance(g[i][1], g[j][1])
                    push!(g[i][2], (dist, j))
                end
            end
        end
        for i in keys(g)
            sort!(g[i][2], lt=(x,y)->isless(x[1],y[1]))
        end
        return (count, g)
    catch err
        showerror(STDOUT, err, backtrace())
        return (0, Dict())
    end
end

function tsp(g, n)
    cur = "1"
    seen = Set()
    path = []
    count = 0
    dist = 0
    show = false
    for i in 1:n-1
        push!(seen, cur)
        push!(path, cur)
        count += 1
        d, cur = get_nearest_unseen(g[cur][2], seen, count, n)
        dist += d
        show = false
    end
    println(dist)
    println(path)
end

function get_nearest_unseen(ls, seen, num_seen, n)
    for entry in ls
        if  num_seen == n - 1 && entry[2] == "1"
            return entry
        elseif !in(entry[2], seen)
            return  entry
        end
    end
    println("EXITED")
    exit(1)
end

function main()
    n, graph = read_data()
    tsp(graph, n)
end

main()


export diffusionReg, wdiffusionReg, wTVReg,anisoTVReg, wdiffusionRegNodal,wTVRegNodal,computeRegularizer,smallnessReg,logBarrier,logBarrierSquared

function computeRegularizer(regFun::Function,mc::Vector,mref::Vector,MInv::AbstractMesh,alpha)
	R,dR,d2R = regFun(mc,mref,MInv)
	return alpha*R, alpha*dR, alpha*d2R
end

function computeRegularizer(regFun::Array{Function},mc::Vector,mref::Array,MInv::AbstractMesh,alpha::Array{Float64})
	numReg = length(regFun)
	if size(mref,2)!=numReg;
		error("computeRegularizer: number of regularizer (=$numReg) does not match number of reference models (=$(size(mref,2))).")
	end
	if length(alpha)!=numReg;
		error("computeRegularizer: number of regularizer (=$numReg) does not match number of alphas (=$(length(alpha))).")
	end

	R = 0.0; dR = zeros(length(mc)); d2R = spzeros(length(mc),length(mc))
	for k=1:numReg
		Rt,dRt,d2Rt = regFun[k](mc,mref[:,k],MInv)
		R   += alpha[k]*Rt
		dR  += alpha[k]*dRt
		d2R += alpha[k]*d2Rt
	end
	return R, dR, d2R
end

"""
	Rc,dR,d2R = diffusionReg(m,mref,M,Iact=1.0)

	Compute diffusion regularizer
		0.5*||GRAD*(m-mref)||_V^2

	Input:
		m     - model
		mref  - reference model
		M     - Mesh
		Iact  - projector on active cells

	Output
		Rc    - value of regularizer
		dR    - gradient w.r.t. m
		d2R   - Hessian
"""
function diffusionReg(m::Vector,mref,M::AbstractMesh;Iact=1.0)
	# Rc = .5* || Grad*m ||^2
	dm   = m .- mref
	Div  = getDivergenceMatrix(M)
	Div  = Iact'*Div   # project to the active cells
	V    = getVolume(M)
	Af   = getFaceAverageMatrix(M)
	d2R  = Div * sdiag(Af'*vec(Vector(diag(V)))) * Div'
	dR   = d2R*dm
	Rc   = 0.5*dot(dm,dR)
	return Rc,dR,d2R
end

"""
	Rc,dR,d2R = smallnessReg(m,mref,M,Iact=1.0)

	Compute smallness regularizer (L2 difference to reference model)

		R(m) = 0.5*||m-mref||_V^2

	Input:
		m     - model
		mref  - reference model
		M     - Mesh

	Output
		Rc    - value of regularizer
		dR    - gradient w.r.t. m
		d2R   - Hessian
"""
function smallnessReg(m::Vector,mref,M::AbstractMesh)
	dm   = m .- mref
	dR   = dm
	Rc   = 0.5*dot(dm,dm)
	return Rc,dR,sparse(1.0I,length(m),length(m))
end

function wdiffusionReg(m::Vector, mref::Vector, M::AbstractMesh; Iact=1.0, C=[])
   # Rc = a1*\\Dx(m-mref)||^2 + .. + a3*\\Dz(m-mref)||^2 + a4*|| m -mref ||^2

   if (M.dim==3) && (length(C) == 0)
      alpha1 = 1; alpha2 = 1; alpha3 = 1; alpha4 = 1e-4;
      Wt   = sdiag([alpha1*ones(M.nf[1]);alpha2*ones(M.nf[2]);alpha3*ones(M.nf[3])])
   elseif (M.dim==2) && (length(C) == 0)
	  alpha1 = 1; alpha2 = 1;alpha4 = 1e-4;
	  Wt   = sdiag([alpha1*ones(M.nf[1]);alpha2*ones(M.nf[2])])
   elseif (M.dim==3) && (length(C) == 4)
      # C = alphax, alphay, alphaz, alphas
      alpha1 = C[1]; alpha2 = C[2]; alpha3 = C[3]; alpha4 = C[4]
      Wt   = sdiag([alpha1*ones(M.nf[1]);alpha2*ones(M.nf[2]);alpha3*ones(M.nf[3])])
   elseif length(C) == 5
      #Af  = getFaceAverageMatrix(M)
      #wt  = Af'*C[5]
      alpha1 = C[1]; alpha2 = C[2]; alpha3 = C[3]; alpha4 = C[4]
      wt   = C[5]
      Wt   = sdiag([alpha1*ones(M.nf[1]);alpha2*ones(M.nf[2]);alpha3*ones(M.nf[3])]) *
             sdiag(wt)
      #Wt = sdiag(C[5])
   elseif length(C) == sum(M.nf) + 1
      # C contains alphas followed by interface weights.
      alpha4 = C[1]
      Wt = sdiag( C[2:end] )
      #println("inteface weights used.")
   else
      error("bad regparams in wdiffusionReg")
   end

   dm   = m .- mref
   Div  = getDivergenceMatrix(M)

   V    = getVolume(M)
   #Div  = Iact'*(Div*Wt)   # project to the active cells
   Div  = Iact'*(V*Div*Wt)   # project to the active cells

   Af   = getFaceAverageMatrix(M)

   #d2R  = Div * sdiag(Af'*Vector(diag(V))) * Div' + alpha4*Iact'*V*Iact
   d2R  = Div * sdiag(1 ./(Af'*Vector(diag(V)))) * Div' + alpha4*Iact'*V*Iact

   dR   = d2R*dm
   Rc   = 0.5*dot(dm,dR)

   return Rc,dR,d2R
end # function wdiffusionReg

"""
	Rc,dR,d2R = wTVReg(m,mref,M,Iact,C=[])

	Compute weighted total variation regularizer

	Input:
		m     - model
		mref  - reference model
		M     - Mesh
		Iact  - projector on active cells
		C     - anisotropy parameters (default: [1 1 1])
		eps   - conditioning parameter for TV norm (default: 1e-3)

	Output
		Rc    - value of regularizer
		dR    - gradient w.r.t. m
		d2R   - Hessian
"""
function wTVReg(m::Vector,mref,M::AbstractMesh; Iact=1.0, C=[],eps=1e-3)
	# Rc = sqrt(a1*\\Dx(m-mref)||^2 + .. + a3*\\Dz(m-mref)||^2) + a4*|| m -mref ||^2
	if length(C) == 0
		alpha1 = 1; alpha2 = 1; alpha3 = 1;
	elseif length(C) == 1
		alpha1 = 1; alpha2 = 1; alpha3 = 1;
	else
		alpha1 = C[1]; alpha2 = C[2]; alpha3 = C[3];
	end
	dm   = m .- mref
	Div  = getDivergenceMatrix(M)
	if M.dim==3
		Wt   = sdiag([alpha1*ones(M.nf[1]);alpha2*ones(M.nf[2]);alpha3*ones(M.nf[3])])
	elseif M.dim==2
		Wt   = sdiag([alpha1*ones(M.nf[1]);alpha2*ones(M.nf[2])])
	end
	Div  = Iact'*(Div*Wt)   # project to the active cells
	V    = getVolume(M); v = Vector(diag(V))
	Af   = getFaceAverageMatrix(M)

	wTV  = sqrt.(Af*(Div'*dm).^2 .+eps);

	Rc   = dot(v,wTV);
	d2R  = Div*(sdiag(Af'*(v./wTV))*Div')
	dR   = d2R*dm;
	return Rc,dR,d2R
end

"""
	Rc,dR,d2R = anisoTVReg(m,mref,M,Iact,eps)

	Compute anisotropic total variation regularizer

	Input:
		m     - model
		mref  - reference model
		M     - Mesh
		Iact  - projector on active cells
		eps   - conditioning parameter for TV norm (default: 1e-3)

	Output
		Rc    - value of regularizer
		dR    - gradient w.r.t. m
		d2R   - Hessian
"""
function anisoTVReg(m::Vector,mref,M::AbstractMesh; Iact=1.0, eps=1e-3)
	# Rc = sqrt(a1*\\Dx(m-mref)||^2 + .. + a3*\\Dz(m-mref)||^2) + a4*|| m -mref ||^2
	dm   = m .- mref
	Div  = getVolume(M)*getDivergenceMatrix(M)
	Div  = Iact'*(Div)   # project to the active cells
	v    = ones(size(Div,2))

	wTV  = sqrt.((Div'*dm).^2 .+eps);
	Rc   = dot(v,wTV);
	d2R  = Div*(sdiag(v./wTV)*Div')
	dR   = d2R*dm;
	return Rc,dR,d2R
end

"""
	Rc,dR,d2R = wdiffusionRegNodal(m::Vector, mref::Vector, M::AbstractMesh; Iact=1.0, C=[])

	Computes weighted diffusion regularizer for nodal model

	Input:
		m     - model
		mref  - reference model
		M     - Mesh
		Iact  - projector on active cells
		C     - optional parameters

	Output
		Rc    - value of regularizer
		dR    - gradient w.r.t. m
		d2R   - Hessian
"""
function wdiffusionRegNodal(m::Vector, mref::Vector, M::AbstractMesh; Iact=1.0, C=[])
	dm = m.-mref;
	if isempty(C)
		C = [1,1,1,1,1e-5];
	end
	if M.dim==3
		Wt   = [C[1]*ones(M.ne[1]);C[2]*ones(M.ne[2]);C[3]*ones(M.ne[3])];
	else
		Wt   = [C[1]*ones(M.ne[1]);C[3]*ones(M.ne[2])];
	end

	V    = getVolume(M);
	Af   = getEdgeAverageMatrix(M)
	v    = (Af'*Vector(diag(V)))
	Wt   = sdiag(Wt.*(v.*Wt));

	Av = getNodalAverageMatrix(M);

	Av = Av*Iact;
	mass = Av'*V*Av;

	Grad   = getNodalGradientMatrix(M)*Iact;

	d2R = Grad'*Wt*Grad;
	d2R += C[4]*mass;
	dR  = d2R*dm;
	Rc  = 0.5*dot(dm,dR);

	if isnan(Rc)
		dump(m)
	end
	clear!(M);
   return Rc,dR,d2R
end

function wTVRegNodal(m::Vector, mref::Vector, M::AbstractMesh; Iact=1.0, C=[])
	dm = m.-mref;
	if isempty(C)
		C = [1,1,1,1,1e-5];
	end
	eps = 1e-3;
	if M.dim==3
		Wt   = sdiag([C[1]*ones(M.ne[1]);C[2]*ones(M.ne[2]);C[3]*ones(M.ne[3])]);
	else
		Wt   = sdiag([C[1]*ones(M.ne[1]);C[3]*ones(M.ne[2])]);
	end

	V    = getVolume(M);
	Av = getNodalAverageMatrix(M);
	mass = Iact'*Av'*V*Av*Iact;
	d2R  = C[4]*mass;
	Rc   = 0.5*dot(dm,d2R*dm);

	v      = Vector(diag(V))
	Af     = getEdgeAverageMatrix(M)
	Grad   = Wt*getNodalGradientMatrix(M)*Iact;
	wTV    = sqrt.(Af*(((Grad*dm).^2).+eps));
	Rc     += dot(v,wTV);
	d2R    += Grad'*sdiag(Af'*(v./wTV))*Grad;
	dR     = d2R*dm
	if isnan(Rc)
		dump(m)
	end
	clear!(M);
   return Rc,dR,d2R
end

"""
	Rc,dR,d2R = logBarrier(m::Vector, z::Vector, M::AbstractMesh,low::Vector,high::Vector, epsilon)

	Computes logBarrier regularizer

	R = -log(1 - ((m-high)/epsilon).^2) if high-epsilon <  m < high
					0					if low+epsilon  <= m <= high-epsilon
		-log(1 - ((m-low)/epsilon).^2)	if low          <  m < low+epsilon

	Input:
		m     	- model
		z     	- not being used. Here for compatibility.
		M     	- Mesh. not being used. Here for compatibility.
		low   	- low bound for each coordinate.
		high  	- high bound for each coordinate.
		epsilon - layer width of the barier.

	Output
		g    - value of regularizer
		dg    - gradient w.r.t. m
		d2g   - Hessian (diagonal matrix). Second derivative is not continous.
"""
function logBarrier(m::Vector,z::Vector,M::AbstractMesh, low::Vector, high::Vector ,epsilon = min(0.1*abs(low),0.1*abs(high)))

	if length(findall(m.>=high)) + length(findall(m.<=low))>0
		return Inf, zeros(length(m)), spzeros(length(m),length(m));
	end

	z = copy(m);
	indProj = zeros(length(m));

	low = low + epsilon;
	high = high - epsilon;
	t = z.<low;
	indProj[t]  .= 1.0;
	z[t]  = low[t];
	t = z.>high;
	indProj[t] .= 1.0;
	z[t] = high[t];

	e 	= (epsilon.+1e-10);
	dm  = (m-z)./e;
	f   = 1.0 .- dm.^2; # > 0
	if minimum(f) < 0.0
		error("negative f");
	end
	df  = (-2.0).*indProj.*(dm./e);
	d2f = (-2.0).*(indProj./(e.^2));

	g   = -sum(log.(f));
	dg  = -df./f;
	d2g = dg.^2 - d2f./f  ;
	return g,dg,sdiag(d2g);
end

"""
	Rc,dR,d2R = logBarrierSquared(m::Vector, z::Vector, M::AbstractMesh,low::Vector,high::Vector, epsilon)

	Computes logBarrier regularizer

	R = (log(1 - ((m-high)/epsilon).^2))^2 if high-epsilon <  m < high
					0					   if low+epsilon  <= m <= high-epsilon
		(log(1 - ((m-low)/epsilon).^2))^2  if low          <  m < low+epsilon

	Input:
		m     	- model
		z     	- not being used. Here for compatibility.
		M     	- Mesh. not being used. Here for compatibility.
		low   	- low bound for each coordinate.
		high  	- high bound for each coordinate.
		epsilon - layer width of the barier.

	Output
		g    - value of regularizer
		dg    - gradient w.r.t. m
		d2g   - Gauss Newton Hessian approximation (diagonal matrix). Second derivative approx is continous.
"""
function logBarrierSquared(m::Vector,z::Vector,M::AbstractMesh, low::Vector, high::Vector ,epsilon = min(0.1*abs(low),0.1*abs(high)))
	if length(findall(m.>=high)) + length(findall(m.<=low))>0
		return Inf, zeros(length(m)), spzeros(length(m),length(m));
	end
	z = copy(m);
	indProj = zeros(length(m));
	low = low + epsilon;
	high = high - epsilon;
	t = z.<low;
	indProj[t]  .= 1.0;
	z[t]  = low[t];
	t = z.>high;
	indProj[t] .= 1.0;
	z[t] = high[t];

	e 	= (epsilon.+1e-10);
	dm  = (m-z)./e;
	f   = 1.0 .- dm.^2; # > 0
	if minimum(f) < 0.0
		error("negative f");
	end
	df  = (-2.0).*indProj.*(dm./e);

	r   = log.(f);
	dr  = df./f;
	g   = 0.5*dot(r,r);
	dg  = dr.*r;
	d2g = dr.^2;
	return g,dg,sdiag(d2g);
end

export TikhonovReg
function TikhonovReg(m::Vector,mref::Vector,M::AbstractMesh,sigma::SparseMatrixCSC)
d = (m-mref);
df = sigma*d;
f = 0.5*dot(d,df);
d2f = sigma;
return f,df,d2f;
end


""" Find the Best Matching Unit (BMU) given an array dist of distances."""
function bestMatchingUnit(dist::Array{Float64,1})
    return argmin(abs.(dist))
end

function neighbourhood(bmu::Int,s::Int,smax::Int,nodes::Array{Float64,1},w::Float64,t::Float64,r::Float64)

    decay = αexp(s,t)
    shrink = αexp(s,r)

    # Gaussian neighbourhood function
    d = decay*exp.(-(((nodes .- nodes[bmu])/(w*shrink)).^2)./2)
    return d
end


import LinearAlgebra: I, mul!

struct MESH
    Δx::Float64
    Δy::Float64
    x::Array{Float64,1}
    y::Array{Float64,1}
    X::Array{Float64,2}
    Y::Array{Float64,2}
    xsize::Int64
    ysize::Int64
end

#Eulerian Mesh
function EulerianMesh(xsize=Int(3e5), ysize=Int(2e5), xnum=61, ynum=41)
    Δx = xsize/(xnum-1)
    Δy = ysize/(ynum-1)
    x = 0:Δx:xsize
    y = 0:Δy:ysize
    X = x'.*ones(length(y))
    Y = y.*ones(1,length(x))
    return MESH(Δx,Δy,x,y,X,Y,Int(xsize),Int(ysize))
end

mutable struct Particles
    x::Array{Float64,1}
    y::Array{Float64,1}
    n::Int64 # number of particles
end

#Lagrangian Particles
function LagrangianParticles(Δx,Δy,mxcell=2, mycell=2, xnum=61, ynum=41)
    mΔx = Δx/mxcell
    mΔy = Δy/mycell
    mxnum = (xnum-1)*mxcell
    mynum = (ynum-1)*mycell
    marknum = mxnum*mynum
    mx = Vector{Float64}(undef,marknum)
    my = similar(mx)
    @inbounds for i=1:mxnum, j=1:mynum
        m = i + (j-1)*mxnum
        mx[m] = mΔx*(i-0.5)
        my[m] = mΔy*(j-0.5)
    end
    return Particles(mx,my,marknum)
end

struct Velocity
    x::Array{Float64,2}
    y::Array{Float64,2}
end

# Eulerian velocity field
function EulerianVelocity(M)
    Vscale = 1
    xnum, ynum = size(M.X,2), size(M.X,1)
    Vx = Array{Float64,2}(undef,ynum,xnum)
    Vy = similar(Vx)
    kx, ky = 1, 1 # wavenumber
    vx0 = 1e-9/Vscale
    vy0 = vx0*kx/ky*M.ysize/M.xsize # Incompressibility
    cte₁ = kx*π/M.xsize
    cte₂ = ky*π/M.ysize
    @inbounds for i=1:xnum, j=1:ynum
        Vx[j,i]=-vx0*sin(cte₁*M.x[i])*cos(cte₂*M.y[j])
        Vy[j,i]= vy0*cos(cte₁*M.x[i])*sin(cte₂*M.y[j])
    end
    return Velocity(Vx, Vy)
end

struct Gradient
    ∂x∂x::Array{Float64,2}
    ∂x∂y::Array{Float64,2}
    ∂y∂y::Array{Float64,2}
    ∂y∂x::Array{Float64,2}
end

# Velocity gradient
function VelocityGradient(M)
    xnum, ynum = size(M.X,2), size(M.X,1)
    dVxdx = Array{Float64,2}(undef,ynum,xnum)
    dVxdy, dVydx, dVydy = similar(dVxdx), similar(dVxdx), similar(dVxdx)
    Vscale = 1
    vx0 = 1e-9/Vscale
    kx, ky = 1, 1 # wavenumber
    vy0 = vx0*kx/ky*M.ysize/M.xsize # Incompressibility
    cte₁ = kx*π/M.xsize
    cte₂ = ky*π/M.ysize
    @inbounds for i=1:xnum, j=1:ynum
        dVxdx[j,i] =-vx0*cos(cte₁*M.x[i])*cte₁*cos(cte₂*M.y[j]);
        dVxdy[j,i] = vx0*sin(cte₁*M.x[i])*sin(cte₂*M.y[j])*cte₂;
        dVydx[j,i] =-vy0*sin(cte₁*M.x[i])*cte₁*sin(cte₂*M.y[j]);
        dVydy[j,i] = vy0*cos(cte₁*M.x[i])*cos(cte₂*M.y[j])*cte₂;
    end
    return Gradient(dVxdx,dVxdy,dVydy,dVydx)
end

function interp∇U!(l,m,∇U,P,M)
    xnum, ynum = size(M.X,2), size(M.X,1)
    i = floor(Int,P.x[m]/M.Δx) + 1
    j = floor(Int,P.y[m]/M.Δy) + 1

    if i<1
        i = 1
        Δxx = 0.0
    elseif i > xnum-1
        i = xnum-1
        Δxx = 1.0
    else
        Δxx = (P.x[m]-M.x[i])/M.Δx
    end
    if j<1
        j = 1
        Δyy = 0.0
    elseif j > ynum-1
        j = ynum-1
        Δyy = 1.0
    else
        Δyy = (P.y[m]-M.y[j])/M.Δy
    end

    w1 = (1.0-Δxx)*(1.0-Δyy)
    w2 = (1.0-Δxx)*Δyy
    w3 = Δxx*(1.0-Δyy)
    w4 = Δxx*Δyy

    @inbounds l[1,1] = ∇U.∂x∂x[j,i]*w1 + ∇U.∂x∂x[j+1,i]*w2 + ∇U.∂x∂x[j,i+1]*w3 + ∇U.∂x∂x[j+1,i+1]*w4
    @inbounds l[1,2] = ∇U.∂x∂y[j,i]*w1 + ∇U.∂x∂y[j+1,i]*w2 + ∇U.∂x∂y[j,i+1]*w3 + ∇U.∂x∂y[j+1,i+1]*w4
    @inbounds l[2,1] = ∇U.∂y∂y[j,i]*w1 + ∇U.∂y∂y[j+1,i]*w2 + ∇U.∂y∂y[j,i+1]*w3 + ∇U.∂y∂y[j+1,i+1]*w4
    @inbounds l[2,2] = ∇U.∂y∂x[j,i]*w1 + ∇U.∂y∂x[j+1,i]*w2 + ∇U.∂y∂x[j,i+1]*w3 + ∇U.∂y∂x[j+1,i+1]*w4

end

function interpVel(mxA,myA,M,U)
    xnum, ynum = size(M.X,2), size(M.X,1)

    i = floor(Int, mxA/M.Δx)+1
    j = floor(Int, myA/M.Δy)+1

    if i<1
        i = 1
        Δxx = 0.0
    elseif i > xnum-1
        i = xnum-1
        Δxx = 1.0
    else
        Δxx = (mxA-M.x[i])/M.Δx
    end
    if j<1
        j = 1
        Δyy = 0.0
    elseif j > ynum-1
        j = ynum-1
        Δyy = 1.0
    else
        Δyy = (myA-M.y[j])/M.Δy
    end

    w1 = (1.0-Δxx)*(1.0-Δyy)
    w2 = (1.0-Δxx)*Δyy
    w3 = Δxx*(1.0-Δyy)
    w4 = Δxx*Δyy

    @inbounds Ux = U.x[j,i]*w1 + U.x[j+1,i]*w2 + U.x[j,i+1]*w3 + U.x[j+1,i+1]*w4
    @inbounds Uz = U.y[j,i]*w1 + U.y[j+1,i]*w2 + U.y[j,i+1]*w3 + U.y[j+1,i+1]*w4
    return Ux, Uz
end

function advectparticle!(P,M,U,Δt,m)
    mxA = P.x[m]
    myA = P.y[m]
    # 4th order RK
    mVxA, mVyA = interpVel(mxA,myA,M,U)
    mxB = mxA + 0.5*mVxA*Δt; myB = myA + 0.5*mVyA*Δt
    mVxB, mVyB = interpVel(mxB,myB,M,U)
    mxC = mxA + 0.5*mVxB*Δt; myC = myA + 0.5*mVyB*Δt
    mVxC, mVyC = interpVel(mxC,myC,M,U)
    mxD = mxA + mVxC*Δt; myD = myA + mVyC*Δt
    mVxD, mVyD = interpVel(mxD,myD,M,U)
    P.x[m] += (mVxA + 2*(mVxB + mVxC) + mVxD)/6*Δt
    P.y[m] += (mVyA + 2*(mVyB + mVyC) + mVyD)/6*Δt
end

mutable struct RK
    Δ1::Array{Float64,2}
    Δ2::Array{Float64,2}
    Δ3::Array{Float64,2}
    Δ4::Array{Float64,2}
end

function updateFSE(F0,l,Δt,RK4)
    Fᵢ = deepcopy(F0)
    # # 4th-order Runge-Kutta
    # RK4.Δ1 .= l*Fᵢ*Δt
    # Fᵢ = F0 + 0.5*RK4.Δ1
    # RK4.Δ2 .= l*Fᵢ*Δt
    # Fᵢ = F0 + 0.5*RK4.Δ2
    # RK4.Δ3 .= l*Fᵢ*Δt
    # Fᵢ = F0 + RK4.Δ3
    # RK4.Δ4 .= l*Fᵢ*Δt

    # 4th-order Runge-Kutta
    mul!(RK4.Δ1, l,Fᵢ*Δt)
    Fᵢ = F0 + 0.5*RK4.Δ1
    mul!(RK4.Δ2, l,Fᵢ*Δt)
    Fᵢ = F0 + 0.5*RK4.Δ2
    mul!(RK4.Δ3, l,Fᵢ*Δt)
    Fᵢ = F0 + RK4.Δ3
    mul!(RK4.Δ4, l,Fᵢ*Δt)

    return F0 += (RK4.Δ1 + 2*(RK4.Δ2 + RK4.Δ3) + RK4.Δ4)/6
end


function calc_grad( LF, Potentials, Focc, v::Array{Float64,2} )
  #
  Npoints = size(v)[1]
  Ncol    = size(v)[2]
  ΔV = LF.LFx.h * LF.LFy.h * LF.LFz.h
  grad = zeros( Float64,Npoints,Ncol )
  #
  for ic = 1:Ncol
    Hv = op_H( LF, Potentials, v[:,ic] )
    grad[:,ic] = Hv # copy?
    for icc = 1:Ncol
      grad[:,ic] = grad[:,ic] - dot( v[:,icc], Hv ) * v[:,icc] * ΔV
    end
    grad[:,ic] = Focc[ic] * grad[:,ic]
  end
  return grad
end

function calc_grad( LF, ∇2::SparseMatrixCSC{Float64,Int64},
                Potentials, Focc, v::Array{Float64,2} )
  #
  Npoints = size(v)[1]
  Ncol    = size(v)[2]
  ΔV = LF.LFx.h * LF.LFy.h * LF.LFz.h
  grad = zeros(Float64,Npoints,Ncol)
  for ic = 1:Ncol
    Hv = op_H( LF, ∇2, Potentials, v[:,ic] )
    grad[:,ic] = Hv # copy?
    for icc = 1:Ncol
      grad[:,ic] = grad[:,ic] - dot( v[:,icc], Hv ) * v[:,icc] * ΔV
    end
    grad[:,ic] = Focc[ic] * grad[:,ic]
  end
  return grad
end


using Random
using LinearAlgebra
using Statistics

mutable struct CMAES
    N::Int
    μ::Int
    λ::Int
    τ::Float64
    τ_c::Float64
    τ_σ::Float64
    population::Array{Array{Float64}}
    offspring::Array{Array{Float64}}
    F_μ::Array{Float64}
    F_λ::Array{Float64}
    C::Array{Float64}
    s::Array{Float64}
    s_σ::Array{Float64}
    σ::Float64
    E::Array{Float64}
    W::Array{Float64}
    x::Array{Float64}
end

function CMAES(;N=2, μ=1, λ=10, τ=sqrt(N), τ_c=N^2, τ_σ=sqrt(N))
    x = randn(N)
    population = fill(x, µ)
    offspring = Array{Array{Float64}}(undef, λ)
    F_µ = Inf .* ones(µ)
    F_λ = Inf .* ones(λ)
    C = Array(Diagonal{Float64}(I, N))
    s = zeros(N)
    s_σ = zeros(N)
    σ = 1.0
    E = zeros(N, λ)
    W = zeros(N, λ);
    CMAES(N, μ, λ, τ, τ_c, τ_σ, population, offspring, F_µ, F_λ, C, s, s_σ, σ, E, W, x)
end

function step!(c::CMAES, obj::Function)
    # L1
    sqrt_c = cholesky((c.C + c.C') / 2.0).U
    for i in 1:c.λ
        c.E[:,i] = randn(c.N)
        c.W[:,i] = c.σ * (sqrt_c * c.E[:,i])
        c.offspring[i] = c.x + c.W[:,i]
        c.F_λ[i] = obj(c.offspring[i])
    end
    # Select new parent population
    idx = sortperm(c.F_λ)[1:c.μ]
    for i in 1:c.μ
        c.population[i] = c.offspring[idx[i]]
        c.F_μ[i] = c.F_λ[idx[i]]
    end
    # L2
    w = vec(mean(c.W[:,idx], dims=2))
    c.x += w
    # L3
    c.s = (1.0 - 1.0/c.τ)*c.s + (sqrt(c.μ/c.τ * (2.0 - 1.0/c.τ))/c.σ)*w
    # L4
    c.C = (1.0 - 1.0/c.τ_c).*c.C + (c.s./c.τ_c)*c.s'
    # L5
    ɛ = vec(mean(c.E[:,idx], dims=2))
    c.s_σ = (1.0 - 1.0/c.τ_σ)*c.s_σ + sqrt(c.μ/c.τ_σ*(2.0 - 1.0/c.τ_σ))*ɛ
    # L6
    c.σ = c.σ*exp(((c.s_σ'*c.s_σ)[1] - c.N)/(2*c.N*sqrt(c.N)))
    c
end


# Alex Hoffman
#cd("/Users/hoffm/.julia/dev")
using astrofunk
using Test

#variables for testing
Earth_a, Mars_a, Sun_mu = 149600000, 227920000, 1.3271*10^11
pos = [-0.270; -0.420; 0] #nondimensional
vel = [0.300; -1.00; 0] #nondimensional
μ₃ = 0.012150664267189

#test loop
@testset "astrofunk.jl" begin
    # Write your tests here.
    @test DCM(pi,'x',false) ≈ [1 0 0; 0 -1 0; 0 0 -1]
    @test circHohmann(Earth_a, Mars_a, Sun_mu) ≈ [2.944004579970198, 2.648345397475325, 44729392.19214477]
    @test JC([pos pos], [vel vel], μ₃) ≈ permutedims([3.186469735193683; 3.186469735193683])
    @test CR3BP_EOM([pos;vel],μ₃,'x') ≈ -0.135541250504792
    @test U_star(pos,μ₃,"yy") ≈ 10.724399149271521
end


# "Dirichlet distribution"
# abstract type Dirichlet <: Dist end

# function dirichlet(ω::Ω, α)
#   gammas = [gammarv(ω[@uid][i], αi, 1.0) for (i, αi) in enumerate(α)]
#   Σ = sum(gammas)
#   [gamma/Σ for gamma in gammas]
# end
# # FIXME: Type
# dirichlet(α::MaybeRV{T}) where T = RandVar{T, Dirichlet}(dirichlet, (α,))

# _rand!(rng::AbstractRNG, d::Djl.MvNormal, x::VecOrMat) = add!(unwhiten!(d.Σ, randn!(rng, x)), d.μ)

# mvchol(x, μ, Σ::PDMat) = Distributions.unwhiten(Σ, x) .+ μ

# "Multivariate Normal Distribution with mean vector `μ` and covariance `Σ`"
# abstract type MvNormal <: Dist end

# mvnormal(ω::Ω, μ::Vector, Σ::PDMat) = mvchol(normal(ω, 0.0, 1.0, size(μ)), μ, Σ)
# # mvnormal(ω::Ω, μ, Σ) = rand(ω, MvNormal(μ, Σ))
# mvnormal(μ::MaybeRV{T1}, Σ::MaybeRV{T2}) where {T1, T2} =
#   RandVar{T1, MvNormal}(mvnormal, (μ, PDMat(Σ)))

# lift(:PDMat, 1)


# this module is currently only used for setup and testing covariance properties
# the spherically covariant ACE basis. There are few/no optimisations for
# performance.
module Wigner

using StaticArrays
using ACE: rand_rot, rand_refl

__L2syms = [:s, :p, :d, :f, :g, :h, :i, :k]
__syms2L = Dict( [sym => L-1 for (L, sym) in enumerate(__L2syms)]... )
get_orbsym(L::Integer)  = __L2syms[L+1]

"""
Index of entries in D matrix (sign included)
"""
struct D_Index
	sign::Int64
	μ::Int64
	m::Int64
end

"""
auxiliary matrix - indices for D matrix
"""
wigner_D_indices(L::Integer) = (   @assert L >= 0;
		[ D_Index(1, i - 1 - L, j - 1 - L) for j = 1:2*L+1, i = 1:2*L+1] )

Base.adjoint(idx::D_Index) = D_Index( (-1)^(idx.μ+idx.m), - idx.μ, - idx.m)

"""
One entry of the Wigner-big-D matrix, `[D^l]_{mu, m}`
"""
wigner_D(μ,m,l,α,β,γ) = (exp(-im*α*m) * wigner_d(m,μ,l,β)  * exp(-im*γ*μ))'

"""
One entry of the Wigner-small-d matrix,

Wigner small d, modified from
```
https://github.com/cortner/SlaterKoster.jl/blob/
8dceecb073709e6448a7a219ed9d3a010fa06724/src/code_generation.jl#L73
```
"""
function wigner_d(μ, m, l, β)
    fc1 = factorial(l+m)
    fc2 = factorial(l-m)
    fc3 = factorial(l+μ)
    fc4 = factorial(l-μ)
    fcm1 = sqrt(fc1 * fc2 * fc3 * fc4)

    cosb = cos(β / 2.0)
    sinb = sin(β / 2.0)

    p = m - μ
    low  = max(0,p)
    high = min(l+m,l-μ)

    temp = 0.0
    for s = low:high
       fc5 = factorial(s)
       fc6 = factorial(l+m-s)
       fc7 = factorial(l-μ-s)
       fc8 = factorial(s-p)
       fcm2 = fc5 * fc6 * fc7 * fc8
       pow1 = 2 * l - 2 * s + p
       pow2 = 2 * s - p
       temp += (-1)^(s+p) * cosb^pow1 * sinb^pow2 / fcm2
    end
    temp *= fcm1

    return temp
end

mat2ang(Q) = mod(atan(Q[2,3],Q[1,3]),2pi), acos(Q[3,3]), mod(atan(Q[3,2],-Q[3,1]),2pi);

function wigner_D(L::Integer, Q::AbstractMatrix)
	D = wigner_D_indices(L);
	α, β, γ = mat2ang(Q);
	Mat_D = [ wigner_D(D[i,j].μ, D[i,j].m, L, α, β, γ)
			    for i = 1:2*L+1, j = 1:2*L+1 ]
   # NB: type instability here, but performance is not important.
	return SMatrix{2L+1, 2L+1, ComplexF64}(Mat_D)
end

function rand_QD(L::Integer)
	# random rotation matrix, reflection
	Q, σ = rand_rot(), rand_refl()
	# return resulting D matrix
	return σ * Q, σ^L * wigner_D(L, Q)
end

function rand_QD(L1::Integer, L2::Integer)
	# random rotation matrix, reflection
	Q, σ = rand_rot(), rand_refl()
	# return Q, D_L1, D_L2
	return σ * Q, σ^(L1) * wigner_D(L1, Q), σ^(L2) * wigner_D(L2, Q)
end

end


kernel = "../src/pca/logreg.jl"

function logreg_loss(v, X, b, λ)
    rv = 0.0
    for i in 1:length(b)
        rv += log(1 + exp(-b[i]*(v[1]+dot(X[:, i], view(v, 2:length(v))))))
    end
    rv / length(b) + λ/2 * norm(v)^2
end

# input data set with known optimal solution
X = [1.444786643000158 0.49236792885913283 -0.53258473265429 0.05476455630673194 -1.3473893605265843; 0.48932299731783646 2.0708445447107926 1.2414596020757043 0.9131934117095984 -0.15692043560721075; 0.7774625331093794 0.7234405608945721 -0.037446104354257874 -1.1104987697394342 1.354975413199728]
b = [-1, 1, -1, -1, 1]
v_opt = [-0.3423591553493419, -0.41317049965033387, 0.007294166575956451, 0.6763846515861628]
m, n = size(X)
λ = 1 / n
opt = logreg_loss(v_opt, X, b, λ)

# write test problem to disk
inputfile = tempname()
inputdataset = "X"
outputdataset = "V"
labeldataset = "b"
h5open(inputfile, "w") do file
    file[inputdataset] = X
    file[labeldataset] = b
end

# GD
nworkers = 2
niterations = 200
stepsize = 0.1
outputfile = tempname()
mpiexec(cmd -> run(```
    $cmd -n $(nworkers+1) julia --project $kernel $inputfile $outputfile
    --inputdataset $inputdataset
    --outputdataset $outputdataset
    --niterations $niterations
    --saveiterates
    --lambda $λ
    ```))
vs = load_logreg_iterates(outputfile, outputdataset)
v = vs[end]
f = logreg_loss(v, X, b, λ)
@test f < opt * (1+1e-2)

# same as previous, but with nslow > 0
nworkers = 2
nslow = 1
niterations = 100
stepsize = 0.1
outputfile = tempname()
mpiexec(cmd -> run(```
    $cmd -n $(nworkers+1) julia --project $kernel $inputfile $outputfile
    --inputdataset $inputdataset
    --outputdataset $outputdataset
    --niterations $niterations
    --saveiterates
    --lambda $λ
    --nslow $nslow
    ```))
vs = load_logreg_iterates(outputfile, outputdataset)
v = vs[end]
f = logreg_loss(v, X, b, λ)
@test f < opt * (1+1e-2)

# same as previous, but with slowprob. > 0
nworkers = 2
slowprob = 0.5
niterations = 100
stepsize = 0.1
outputfile = tempname()
mpiexec(cmd -> run(```
    $cmd -n $(nworkers+1) julia --project $kernel $inputfile $outputfile
    --inputdataset $inputdataset
    --outputdataset $outputdataset
    --niterations $niterations
    --saveiterates
    --lambda $λ
    --slowprob $slowprob
    ```))
vs = load_logreg_iterates(outputfile, outputdataset)
v = vs[end]
f = logreg_loss(v, X, b, λ)
@test f < opt * (1+1e-2)

# DSAG
vralgo = "tree"
nworkers = 2
nwait = 1
niterations = 100
stepsize = 0.1
nsubpartitions = 2
outputfile = tempname()
mpiexec(cmd -> run(```
    $cmd -n $(nworkers+1) julia --project $kernel $inputfile $outputfile
    --inputdataset $inputdataset
    --nwait $nwait
    --variancereduced
    --vralgo $vralgo
    --nsubpartitions $nsubpartitions
    --outputdataset $outputdataset
    --niterations $niterations
    --saveiterates
    --lambda $λ
    ```))
vs = load_logreg_iterates(outputfile, outputdataset)
v = vs[end]
f = logreg_loss(v, X, b, λ)
@test f < opt * (1+1e-2)

# DSAG w. nwaitschedule < 1.0
nworkers = 2
nwait = 2
niterations = 100
stepsize = 0.1
nsubpartitions = 2
nwaitschedule = 0.9
outputfile = tempname()
mpiexec(cmd -> run(```
    $cmd -n $(nworkers+1) julia --project $kernel $inputfile $outputfile
    --inputdataset $inputdataset
    --nwait $nwait
    --variancereduced
    --vralgo $vralgo
    --nsubpartitions $nsubpartitions
    --outputdataset $outputdataset
    --niterations $niterations
    --saveiterates
    --lambda $λ
    --nwaitschedule $nwaitschedule
    ```))
vs = load_logreg_iterates(outputfile, outputdataset)
v = vs[end]
f = logreg_loss(v, X, b, λ)
@test f < opt * (1+1e-2)

# DSAG w. sparse input data
inputfile = tempname()
h5open(inputfile, "cw") do file
    H5SparseMatrixCSC(file, inputdataset, sparse(X))
    file[labeldataset] = b
    flush(file)
end

nworkers = 2
nwait = 1
niterations = 100
stepsize = 0.1
nsubpartitions = 2
outputfile = tempname()
mpiexec(cmd -> run(```
    $cmd -n $(nworkers+1) julia --project $kernel $inputfile $outputfile
    --inputdataset $inputdataset
    --nwait $nwait
    --variancereduced
    --vralgo $vralgo
    --nsubpartitions $nsubpartitions
    --outputdataset $outputdataset
    --niterations $niterations
    --lambda $λ
    ```))
vs = load_logreg_iterates(outputfile, outputdataset)
v = vs[end]
f = logreg_loss(v, X, b, λ)
@test f < opt * (1+1e-2)


using LinearAlgebra
using SparseArrays

function expmv(t::StepRangeLen, A, b::AbstractVecOrMat;
                M = nothing, precision = "double", shift = false)

    t0 = Float64(t.ref)
    tmax = Float64(t.ref + (t.len - 1.) * t.step)
    q = t.len - 1
    n = size(A, 1)

    if shift == true && !hasmethod(tr, typeof(A))
        shift = false
        @warn "Shift set to false as $(typeof(A)) doesn't support tr"
    end

    force_estm = !hasmethod(opnorm, Tuple{typeof(A), typeof(1)})
    temp = (tmax - t0) * normAm(A, 1);

    if (precision == "single" || precision == "half" && temp > 85.496) || ( precision == "double" && temp > 63.152)
       force_estm = true;
    end

    if M == nothing
        (M, alpha, unA) = select_taylor_degree(A, b, precision=precision, shift=shift, force_estm=force_estm)
    end

    tol =
      if precision == "double"
          2.0^(-53)
      elseif precision == "single"
          2.0^(-24)
      elseif precision == "half"
          2.0^(-10)
      end

    X = zeros(eltype(A), n, q+1);
    (m_max, p) = size(M);
    U = diagm(0 => 1:m_max);

    temp, s = degree_selector(tmax - t0, M, U, p)
    h = (tmax - t0)/q;

    X[:,1] = expmv(t0,A,b,M=M,precision=precision,shift=shift);

    mu = 0.
    if shift
        mu = tr(A)/n
        A = A - mu * I
    end

    d = max(1, Integer(floor(q/s)))
    j = Integer(floor(q/d))
    r = q - d * j
    z = X[:,1]
    m_opt, = degree_selector(d, M, U, p)
    dr = d
    K = zeros(eltype(A), n, m_opt+1)

    temp = similar(b)
    f = similar(b)

    for i = 1:j+1
        if i > j
            dr = r
        end
        K[:,1] = z
        m = 0
        for k = 1:dr
            f = copy(z)
            c1 = norm(z, Inf)

            p = 1
            for outer p = 1:m_opt
                if p > m
                    @views K[:,p+1] .= (h/p).*(A*K[:,p])
                end

                @views temp .= (Float64(k)^p).*K[:,p+1]
                f .+= temp
                c2 = norm(temp, Inf)
                if c1 + c2 <= tol * norm(f, Inf)
                    break
                end
                c1 = c2
            end

            m = max(m,p)
            @views X[:, k + (i - 1) * d + 1] .= exp(k * h * mu) * f
        end

        if i <= j
            z = X[:,i*d+1]
        end
    end
    return X
end


# m6_11s.jl

using Pkg, DrWatson

using MonteCarloMeasurements
using StanSample
using StatisticalRethinking

N = 200
b_GP = 1                               # Direct effect of G on P
b_GC = 0                               # Direct effect of G on C
b_PC = 1                               # Direct effect of P on C
b_U = 2                                # Direct effect of U on P and C

df = DataFrame(
  :u => 2 * rand(Bernoulli(0.5), N) .- 1,
  :g => rand(Normal(), N)
)

df[!, :p] = [rand(Normal(b_GP * df[i, :g] + b_U * df[i, :u]), 1)[1] for i in 1:N]
df[!, :c] = [rand(Normal(b_PC * df[i, :p] + b_GC * df[i, :g] + b_U * df[i, :u]), 1)[1] for i in 1:N]

precis(df)

stan6_11 = "
data {
    int <lower=0> N;
    vector[N] C;
    vector[N] P;
    vector[N] G;
}
parameters {
    real <lower=0> sigma;
    real a;
    real b_PC;
    real b_GC;
}
model {
    vector[N] mu;
    sigma ~ exponential(1);
    a ~ normal(0, 1);
    b_PC ~ normal(0, 1);
    b_GC ~ normal(0, 1);
    mu = a + b_PC * P + b_GC * G;
    C ~ normal(mu, sigma);
}
";

m6_11s = SampleModel("m6.11s", stan6_11)
m6_11_data = Dict(:N => nrow(df), :C => df.c, :P => df.p, :G => df.g)
rc6_11s = stan_sample(m6_11s, data=m6_11_data)

if success(rc6_11s)
    part6_11s = read_samples(m6_11s, :particles)
    part6_11s |> display
end

# End of m6.11s.jl


const rₑ  = 6371220.0 # Radius of the Earth (meters)
const Ωₑ  = 7.292e-5  # Rotational frequency of the Earth (1/second)
const g   = 9.80616   # Gravitational constant (meters/second^2)
function f(x)         # Coriolis term
  2.0*Ωₑ*x[3]/rₑ
end


function wave_forward_operater(f, x, t, R, N_ϕ)
    N_b = size(f)[1]
    x_end = x[end]
    x2 = 2x_end
    N_t = length(t)        # Number of elements of vector t
    N_x = length(x)        # Number of elements of vector x
    N_x2 = 2N_x
    half_N_x = Int64(N_x/2)
    dt = t[2]-t[1]

    #outer larger grid for solution of wave equation
    g_x = range(-x2, stop=x2 - (2x2)/N_x2, length=N_x2)

    #frequency vectors
    k = range(-N_x, stop=N_x-1, length=N_x2) .* ( pi / x2)
    K = k.^2 .+ (k.^2)'
    KM2 = ifftshift(K)
    KM = sqrt.(KM2)

    #initial pressure field
    #P = zeros(N_b, N_x2, N_x2)
    #P[:, half_N_x:3half_N_x-1, half_N_x:3half_N_x-1] = f
    P = zeros(N_x2, N_x2)
    P[half_N_x:3half_N_x-1, half_N_x:3half_N_x-1] = f

    # k-space
    #W = zeros(N_b, N_t+1, N_x2, N_x2)
    #W[:, 2, :, :] = P
    #W[:, 1, :, :] = P   #t=-1일 때의 W
    W = zeros(N_x2, N_x2, N_t+1)
    W[:, :, 2] = P
    W[:, :, 1] = P   #t=-1일 때의 W

    detector_angle = 0: 2pi/N_ϕ : 2pi-2pi/N_ϕ
    detector_x = R * cos.(detector_angle)
    detector_y = R * sin.(detector_angle)

    #Wf = zeros(N_b, N_t, N_ϕ)
    Wf = zeros(N_t, N_ϕ)
    sin_term = @. -4(sin(dt*KM/2))^2
    #sin_term = reshape(sin_term, (1,size(sin_term)[1], size(sin_term)[2]))

    #for i ∈ 2:N_t
    #    LW = ifft(sin_term.*fft(W[:, i, :, :]))
    #    real_LW  = real(LW)
    #    W[:, i+1, :, :] = 2W[:, i, :, :] - W[:, i-1,: ,:] + real_LW
    #    #W[:, i+1, :, :] = 2*W[:, i, :, :] - W[:, i-1,: ,:] + cx*LW1
    #
    #    Wf[:, i, : ] = interp(g_x, W[:, i+1, :, :], detector_x, detector_y)
    #end
    for i ∈ 2:N_t
        LW = ifft(sin_term.*fft(W[:, :, i]))
        real_LW  = real(LW)
        W[:, :, i+1] = 2W[:, :, i] - W[:, :, i-1] + real_LW
        #W[:, i+1, :, :] = 2*W[:, i, :, :] - W[:, i-1,: ,:] + cx*LW1

        Wf[i, :] = interp(g_x, W[:, :, i+1], detector_x, detector_y)
    end
    return Wf, W
end


using ModelingToolkit

@parameters t x
domains = [t ∈ IntervalDomain(0.0,1.0),
           x ∈ IntervalDomain(0.0,1.0)]

@parameters z
z ∈ (IntervalDomain(0.0,1.0) ⊗ IntervalDomain(0.0,1.0))

@parameters y
(x,y) ∈ CircleDomain()
@parameters r θ
(r,θ) ∈ CircleDomain(true)


using MannKendall
using Test
# NoTrendData = ones(360)*rand(Int, 1)
TrendData = 1:100;
@testset "MannKendall.jl" begin
    # Write your tests here.
    TrendRes = mk_original_test(TrendData);

    @test TrendRes.trend == "increasing";
    @test TrendRes.h == true;
    @test TrendRes.p == 0.0;
    @test TrendRes.z == 14.73869998208331;
    @test TrendRes.τ == 1.0;
    @test TrendRes.s == 4950;
    @test TrendRes.var_s ==112750.0;

end


function adaLASSO{F<:AbstractFloat}(Y::Array{F})

    ## Initialize Parameters
    n      = length(Y)
    BICold = 10000000
    β      = Float64[]

    ## LASSO
    tic()
    B0TM, ~, bestB0TM, X, array_σ = CoordinateDescentTM(Y, n)
    TimeLASSO = toq();

    # println("LASSO took ", TimeLASSO, "s")

    β_LASSO = B0TM[:, bestB0TM]

    ## Making real copies of data, so we can change and it'll be okay!
    brTM = copy(B0TM[:, bestB0TM])
    br   = copy(brTM)

    temp = 1:size(X,2)
    idx = abs(br ./ array_σ) .<= 1e-3
    temp = temp[!idx]

    if size(temp,1) >= 1
        if temp[1] != 1
            temp = [1; temp]
        end
    else
        temp = []
    end

    adaX = X[:, temp]
    adaSigma = array_σ[temp]

    ## OLS normalization
    adabr = br[temp] ./ adaSigma
    adabr = \(adaX'*adaX,adaX'*Y) ./ adaSigma

    ## Path
    λ_max, λ_max_var = findmax(adaSigma .* abs(adaX[:,1:end]' * Y) )
    λ_max = λ_max / adaSigma[λ_max_var] / n

    MaxIterλ = 100
    passo = Float64(1 / MaxIterλ)
    grid_linear = Array{Float64}(collect(1+0.001:passo*10:10))
    LASSOpath = Array{Float64}(1 - log10(grid_linear))

    println("Initializing adaLASSO path.")

    ## Gamma Grid
    BetaGamma = []
    for gamma = 1*sort(1-log10(1:1/4:10))+0.001
        wg = 1 ./ ( abs(adabr).^gamma )
        Blasso, bestB = CoordinateDescent3(adaX, Y, LASSOpath, λ_max, wg, adaSigma)

        Best = Blasso[:,bestB]
        BetaGamma = [BetaGamma; Best]
        BICnew = n * log(var( Y - adaX*Best )) + sum( abs(Best./adaSigma) .> 1e-3) * log(n)
        if BICnew < BICold
            BICold = BICnew
            β = Best
            lastgamma = gamma
        end
    end

    println("adaLASSO path finished.")

    ## TODO

    ## Output
    temp2 = 1:size(adaSigma,1)
    idx = abs(β ./ adaSigma) .> 1e-3
    temp2 = temp2[idx]

    if size(temp2,1) >= 1
        if temp2[1] != 1
            temp2 = [1; temp2]
        end
    else
        temp2 = []
    end

    selected = temp[temp2]
    β_adaLASSO = zeros(size(X,2),1)
    β_adaLASSO[selected] = β[temp2]

    return β_adaLASSO, β_LASSO, X, array_σ, TimeLASSO
end


"""
$(TYPEDEF)

A simple feedback linearizable unicycle model.

# Fields

$(TYPEDFIELDS)
"""
struct Unicycle4D{ΔT} <: ControlSystem{ΔT, 4, 2} end

function dx(cs::Unicycle4D, x::SVector{4}, u::SVector{2}, t::AbstractFloat)
    # position: x
    dx1 = x[4] * cos(x[3])
    # position: y
    dx2 = x[4] * sin(x[3])
    # orientation
    dx3 = u[1]
    # velocity
    dx4 = u[2]

    return @SVector[dx1, dx2, dx3, dx4]
end

function linearize_discrete(cs::Unicycle4D, x::SVector{4}, u::SVector{2},
                            t::AbstractFloat)
    ΔT = samplingtime(cs)
    cθ = cos(x[3]) * ΔT
    sθ = sin(x[3]) * ΔT

    A = @SMatrix [1 0 -x[4]*sθ cθ;
                  0 1 x[4]*cθ  sθ;
                  0 0 1        0;
                  0 0 0        1]

    B = @SMatrix [0  0;
                  0  0;
                  ΔT 0;
                  0  ΔT];
    return LinearSystem{ΔT}(A, B)
end

xyindex(cs::Unicycle4D) = SVector(1, 2)

"------------------- Implement Feedback Linearization Interface -------------------"

n_linstates(cs::Unicycle4D) = 4
ξxyindex(cs::Unicycle4D) = SVector(1, 3)

# TODO: Maybe use StaticArrays.FieldVector here to have a type safe desctiction
# between ξ and x coordinates.

# implement holy trait
LinearizationStyle(::Unicycle4D) = FeedbackLinearization()

@inline function feedbacklin(cs::Unicycle4D)
    ΔT = samplingtime(cs)
    # ξ = (px, pẋ, py, pẏ)
    A = @SMatrix [1. ΔT 0. 0.;
                  0. 1. 0. 0.;
                  0. 0. 1. ΔT;
                  0. 0. 0. 1.;]
    B = @SMatrix [0. 0.;
                  ΔT 0.;
                  0. 0.;
                  0. ΔT];

    dyn = LinearSystem{ΔT}(A, B)
    return LTISystem(dyn, ξxyindex(cs))
end

function x_from(cs::Unicycle4D, ξ::SVector{4})
    # px, py, θ, v
    return @SVector [ξ[1], ξ[3], atan(ξ[4], ξ[2]), sqrt(ξ[2]^2 + ξ[4]^2)]
end

function ξ_from(cs::Unicycle4D, x::SVector{4})
    # px, pẋ, py, pẏ
    return @SVector [x[1], cos(x[3])*x[4], x[2], sin(x[3])*x[4]]
end

function λ_issingular(cs::Unicycle4D, ξ::SVector{4})
    vx = ξ[2]; vy = ξ[4]
    ϵ = 0.01
    # cant invert for velocities near 0
    return isnan(vx) || isnan(vy) || (abs(vx) < ϵ && abs(vy) < ϵ)
end

function inverse_decoupling_matrix(cs::Unicycle4D, x::SVector{4})
    cθ = cos(x[3]); sθ = sin(x[3]); v = x[4]
    return @SMatrix[cθ    sθ;
                   -sθ/v -cθ/v]
end

@inline decoupling_drift_term(cs::Unicycle4D, x::SVector{4}) = @SVector(zeros(n_controls(cs)))