"""
Abstract types and methods for codes, error models and decoders.
"""
module Model

# imports
using ..Qecsim:QecsimError
using ..PauliTools:bsp
using LinearAlgebra:I

# exports
# - abstract model
export AbstractModel
export label
# - stabilizer code
export StabilizerCode
export stabilizers, logical_xs, logical_zs, logicals, nkd, validate
# - error model
export ErrorModel
export generate, probability_distribution
# - decoder
export Decoder
export DecodeResult
export decode

# AbstractModel

"""
Abstract supertype for models.
"""
abstract type AbstractModel end

"""
    label(model) -> String

Return a label suitable for use in plots and for grouping results.

!!! note "Abstract method"

    This method should be implemented for concrete subtypes or duck-typed implementations of
    [`AbstractModel`](@ref).
"""
function label end

# StabilizerCode

"""
    StabilizerCode <: AbstractModel

Abstract supertype for stabilizer codes.
"""
abstract type StabilizerCode <: AbstractModel end

"""
    stabilizers(code) -> BitMatrix

Return the stabilizers in binary symplectic form.

Each row is a stabilizer generator. An overcomplete set of generators can be included to
simplify decoding.

!!! note "Abstract method"

    This method should be implemented for concrete subtypes or duck-typed implementations of
    [`StabilizerCode`](@ref).
"""
function stabilizers end

"""
    logical_xs(code) -> BitMatrix

Return the logical X operators in binary symplectic form.

Each row is an operator. The order should match that of [`logical_zs`](@ref).

!!! note "Abstract method"

    This method should be implemented for concrete subtypes or duck-typed implementations of
    [`StabilizerCode`](@ref).
"""
function logical_xs end

"""
    logical_zs(code) -> BitMatrix

Return the logical Z operators in binary symplectic form.

Each row is an operator. The order should match that of [`logical_xs`](@ref).

!!! note "Abstract method"

    This method should be implemented for concrete subtypes or duck-typed implementations of
    [`StabilizerCode`](@ref).
"""
function logical_zs end

"""
    logicals(code) -> BitMatrix

Return the logical operators in binary symplectic form.

Each row is an operator. X operators are stacked above Z operators in the order given by
[`logical_xs`](@ref) and [`logical_zs`](@ref).
"""
function logicals(code)
    return vcat(logical_xs(code), logical_zs(code))
end

"""
    nkd(code) -> Tuple{Int,Int,Union{Int,Missing}}

Return a descriptor in the format `(n, k, d)`, where `n` is the number of physical qubits,
`k` is the number of logical qubits, and `d` is the distance of the code (or `missing` if
unknown).

!!! note "Abstract method"

    This method should be implemented for concrete subtypes or duck-typed implementations of
    [`StabilizerCode`](@ref).
"""
function nkd end

@doc raw"""
    validate(code)

Perform sanity checks.

If any of the following fail then a [`QecsimError`](@ref) is thrown:

* ``S ⊙ S^T = 0``
* ``S ⊙ L^T = 0``
* ``L ⊙ L^T = Λ``

where ``S`` and ``L`` are the code [`stabilizers`](@ref) and [`logicals`](@ref),
respectively, and ``⊙`` and ``Λ`` are defined in [`PauliTools.bsp`](@ref bsp).
"""
function validate(code)
    s, l = stabilizers(code), logicals(code)
    !any(bsp(s, transpose(s))) || throw(QecsimError("stabilizers do not mutually commute"))
    !any(bsp(s, transpose(l))) || throw(QecsimError(
        "stabilizers do not commute with logicals"))
    # twisted identity with same size as logicals
    nlogicals = size(l, 1)
    twistedI = circshift(Matrix(I, nlogicals, nlogicals), nlogicals / 2)
    (bsp(l, transpose(l)) == twistedI) || throw(QecsimError(
        "logicals do not mutually twist commute"))
    return nothing
end

# ErrorModel

"""
    ErrorModel <: AbstractModel

Abstract supertype for error models.
"""
abstract type ErrorModel <: AbstractModel end

"""
    generate(error_model, code, p::Real, [rng::AbstractRNG=GLOBAL_RNG]) -> BitVector

Generate a new error in binary symplectic form according to the `error_model` and `code`,
where `p` is typically the probability of an error on a single qubit.

!!! note "Abstract method"

    This method should be implemented for concrete subtypes or duck-typed implementations of
    [`ErrorModel`](@ref).
"""
function generate end

"""
    probability_distribution(error_model, p::Real) -> NTuple{4,Real}

Return the single-qubit probability distribution amongst Pauli I, X, Y and Z, where `p` is
the overall probability of an error on a single qubit.

!!! note "Abstract method [optional]"

    This method is **not** invoked by any core modules. Since it is often useful for
    decoders, it is provided as a template and concrete subtypes or duck-typed
    implementations of [`ErrorModel`](@ref) are encouraged to implement it when appropriate,
    particularly for IID error models.
"""
function probability_distribution end

# Decoder
"""
    Decoder <: AbstractModel

Abstract supertype for decoders.
"""
abstract type Decoder <: AbstractModel end

"""
    decode(decoder, code, syndrome::AbstractVector{Bool}; kwargs...) -> DecodeResult

Resolve a recovery operation for the given `code` and `syndrome`, or evaluate the success of
decoding, as encapsulated in the decode result.

The syndrome has length equal to the number of code stabilizers, and element values of 0
or 1 (false or true) indicate whether the corresponding stabilizer does or does not commute
with the error, respectively.

Keyword parameters `kwargs` may be provided by the client, e.g. [`App`](@ref App), with
context values such as `error_model`, `error_probability` and `error`. Most implementations
will ignore such parameters; however, if they are used, implementations should declare them
explicitly and treat them as optional.

See also [`DecodeResult`](@ref).

!!! note "Abstract method"

    This method should be implemented for concrete subtypes or duck-typed implementations of
    [`Decoder`](@ref).
"""
function decode end

"""
    DecodeResult(
        success::Union{Nothing,Bool},
        recovery::Union{Nothing,AbstractVector{Bool}},
        logical_commutations::Union{Nothing,AbstractVector{Bool}}
        custom_values::Union{Nothing,AbstractVector}
    )
    DecodeResult(;
        success::Union{Nothing,Bool}=nothing,
        recovery::Union{Nothing,AbstractVector{Bool}}=nothing,
        logical_commutations::Union{Nothing,AbstractVector{Bool}}=nothing
        custom_values::Union{Nothing,AbstractVector}=nothing
    )

Construct a decoding result as returned by [`decode`](@ref).

Typically decoders will provide a `recovery` operation and delegate the evaluation of
`success` and `logical_commutations` to the client, e.g. [`App`](@ref App). Optionally,
`success` and/or `logical_commutations` may be provided as overrides. At least one of
`recovery` or `success` must be specified to allow a success value to be resolved.
Additionally `custom_values` may be specified. If `logical_commutations` and/or
`custom_values` are provided then they should be of consistent type and size over
identically parameterized simulation runs. For example, [`App`](@ref App) will sum
`logical_commutations` across runs and, similarly, if `custom_values` are numbers they will
be summed across runs, and if they are arrays they will be concatenated using `vcat`.

See also [`decode`](@ref).

# Examples
```jldoctest
julia> DecodeResult(; recovery=BitVector([1, 0, 1, 0, 1, 1]))  # typical use-case
DecodeResult{Nothing}(nothing, Bool[1, 0, 1, 0, 1, 1], nothing, nothing)

julia> DecodeResult(; success=true)  # override success
DecodeResult{Nothing}(true, nothing, nothing, nothing)

julia> DecodeResult(; success=false, logical_commutations=BitVector([1, 0]))  # override all
DecodeResult{Nothing}(false, nothing, Bool[1, 0], nothing)

julia> DecodeResult(; success=true, custom_values=[2.3, 4.1])  # custom values (numbers)
DecodeResult{Vector{Float64}}(true, nothing, nothing, [2.3, 4.1])

julia> DecodeResult(; success=true, custom_values=[[2.3], [4.1]])  # custom values (arrays)
DecodeResult{Vector{Vector{Float64}}}(true, nothing, nothing, [[2.3], [4.1]])

julia> DecodeResult(true, nothing, nothing, [2.3, 4.1])  # positional parameters
DecodeResult{Vector{Float64}}(true, nothing, nothing, [2.3, 4.1])

julia> DecodeResult(; success=nothing, recovery=nothing)  # too few specified parameters
ERROR: QecsimError: at least one of 'success' or 'recovery' must be specified
Stacktrace:
[...]
```
"""
struct DecodeResult{T<:Union{Nothing,Vector}}
    # Parametric type to avoid abstract types in struct, see:
    # https://docs.julialang.org/en/v1/manual/performance-tips/#Type-declarations
    # https://discourse.julialang.org/t/union-of-nothing-and-parametric-types/11986/5
    success::Union{Nothing,Bool}
    recovery::Union{Nothing,BitVector}
    logical_commutations::Union{Nothing,BitVector}
    custom_values::T
    # lenient constructor that infers types
    function DecodeResult(success, recovery, logical_commutations,
                          custom_values::Union{Nothing,AbstractVector{S}}) where S
        if isnothing(success) && isnothing(recovery)
            throw(QecsimError("at least one of 'success' or 'recovery' must be specified"))
        end
        if isnothing(custom_values)
            new{Nothing}(success, recovery, logical_commutations, custom_values)
        else
            new{Vector{S}}(success, recovery, logical_commutations, collect(custom_values))
        end
    end
end
# keyword convenience constructor
function DecodeResult(; success=nothing, recovery=nothing, logical_commutations=nothing,
                      custom_values=nothing)
    DecodeResult(success, recovery, logical_commutations, custom_values)
end

end


# Portions translated from SLICOT-Reference distribution
# Copyright (c) 2002-2020 NICONET e.V.
function run_mb03qd(datfile, io=stdout)
    NIN = 5
    NOUT = 6
    NMAX = 10
    LDA = NMAX
    LDU = NMAX
    LDWORK = 3*NMAX
    A = Array{Float64,2}(undef, LDA,NMAX)
    U = Array{Float64,2}(undef, LDU,NMAX)
    DWORK = Array{Float64,1}(undef, LDWORK)
    WI = Array{Float64,1}(undef, NMAX)
    WR = Array{Float64,1}(undef, NMAX)
    BWORK = Array{Bool,1}(undef, NMAX)
    f = open(datfile,"r")
    readline(f)
    vs = split(readline(f))
    N = parse(BlasInt, vs[1])
    NLOW = parse(BlasInt, vs[2])
    NSUP = parse(BlasInt, vs[3])
    ALPHA = parse(Float64, replace(vs[4],'D'=>'E'))
    DICO =  vs[5][1]
    STDOM =  vs[6][1]
    JOBU =  vs[7][1]
    if ( N<0 || N>NMAX )
        @error "invalid N=$N"
    else

    vs = String[]
    _isz,_jsz = (N,N)
    while length(vs) < _isz*_jsz
        append!(vs, replace.(split(readline(f)),'D'=>'E'))
    end
    for i in 1:_isz
       _i0 = (i-1)*_jsz
       A[i,1:_jsz] .= parsex.(Float64, vs[_i0+1:_i0+_jsz])
    end
# interp call 1

    #FOREIGN.dgees!( 'Vectors', 'Not sorted', SELECT, N, A, LDA, NDIM, WR, WI, U, LDU, DWORK, LDWORK, BWORK, INFO )
    Aschur = schur(A[1:N,1:N])
    WR[1:N] .= real.(Aschur.values)
        WI[1:N] .= imag.(Aschur.values)
        U[1:N,1:N] .= Aschur.Z
        A[1:N,1:N] .= Aschur.T

    println(io, "orig T:")
    _nc = N
    _nr = N
    show(io, "text/plain", A[1:_nr,1:_nc])
    println(io)
# interp call 2

    NDIM, INFO = SLICOT.mb03qd!(DICO, STDOM, JOBU, N, NLOW, NSUP, ALPHA, A, U)
    println(io, "NDIM = $NDIM")
    @test INFO == 0
    INFO == 0 || return

if ( INFO!=0 )
else
# interp output 1
    println(io, "A:")
    _nc = N
    _nr = N
    show(io, "text/plain", A[1:_nr,1:_nc])
    println(io)

# interp output 2
    println(io, "U:")
    _nc = N
    _nr = N
    show(io, "text/plain", U[1:_nr,1:_nc])
    println(io)

end # if
end # if
    close(f)
end # run_mb03qd()


using Bridge
using StaticArrays
import Bridge: b, σ, B, β, a, constdiff
const ℝ = SVector{N,T} where {N,T}
import Base.display
"""
    LotkaVolterraDiffusion <: ContinuousTimeProcess{ℝ{2}}
Struct defining a stochastic Lotka Volterra model
"""
struct LotkaVolterraDiffusion{T} <: ContinuousTimeProcess{ℝ{2,T}}
    α::T
    β::T
    γ::T
    δ::T
    σ1::T
    σ2::T
    function LotkaVolterraDiffusion(α::T, β, γ, δ, σ1, σ2) where T
        new{T}(α, β, γ, δ, σ1, σ2)
    end
end
function b(t, x, P::LotkaVolterraDiffusion{T}) where T
    ℝ{2}(P.α*x[1] - P.β*x[1]*x[2], P.δ*x[1]*x[2] - P.γ*x[2])
end
function σ(t, x, P::LotkaVolterraDiffusion{T}) where T
    SDiagonal(P.σ1, P.σ2)
end
constdiff(::LotkaVolterraDiffusion) = true
clone(P::LotkaVolterraDiffusion, θ) = LotkaVolterraDiffusion(θ...)
params(P::LotkaVolterraDiffusion) = [P.α, P.β, P.γ, P.δ, P.σ1, P.σ2]
domain(P::LotkaVolterraDiffusion) = LowerBoundedDomain((0.0, 0.0), (1,2))

# <---------------------------------------------
# this is optional, needed for conjugate updates

nonhypo(P::LotkaVolterraDiffusion, x) = x
@inline hypo_a_inv(P::LotkaVolterraDiffusion, t, x) = SDiagonal(1.0/P.σ1^2, 1.0/P.σ2^2)
num_non_hypo(P::Type{<:LotkaVolterraDiffusion}) = 2

phi(::Val{0}, t, x, P::LotkaVolterraDiffusion) = (zero(x[1]), zero(x[2]))
phi(::Val{1}, t, x, P::LotkaVolterraDiffusion) = (x[1], zero(x[2]))
phi(::Val{2}, t, x, P::LotkaVolterraDiffusion) = (-x[1]*x[2], zero(x[2]))
phi(::Val{3}, t, x, P::LotkaVolterraDiffusion) = (zero(x[1]), x[1]*x[2])
phi(::Val{4}, t, x, P::LotkaVolterraDiffusion) = (zero(x[1]), -x[2])
phi(::Val{5}, t, x, P::LotkaVolterraDiffusion) = (zero(x[1]), zero(x[2]))
phi(::Val{6}, t, x, P::LotkaVolterraDiffusion) = (zero(x[1]), zero(x[2]))

#
# <---------------------------------------------

struct LotkaVolterraDiffusionAux{R,S1,S2} <: ContinuousTimeProcess{ℝ{2,R}}
    α::R
    β::R
    γ::R
    δ::R
    σ1::R
    σ2::R
    t::Float64
    u::S1
    T::Float64
    v::S2

    function LotkaVolterraDiffusionAux(α::R, β, γ, δ, σ1, σ2, t::Float64, u::S1,
                                       T::Float64, v::S2) where {R,S1,S2}
        new{R,S1,S2}(α, β, γ, δ, σ1, σ2, t, u, T, v)
    end
end

function B(t, P::LotkaVolterraDiffusionAux)
#    ℝ{2}(P.α*x[1] - P.β*x[1]*x[2], P.δ*x[1]*x[2] - P.γ*x[2])
    @SMatrix [-0.0 -P.β*P.γ/P.δ; P.α*P.δ/P.β -0.0]
end

# mean = ℝ{2}(P.γ/P.δ, P.α/P.β)
function β(t, P::LotkaVolterraDiffusionAux)
    ℝ{2}(P.γ/P.δ*P.α, -P.α/P.β*P.γ)
end

function σ(t, P::LotkaVolterraDiffusionAux)
    SDiagonal(P.σ1, P.σ2)
end

function σ(t, x, P::LotkaVolterraDiffusionAux)
    SDiagonal(P.σ1, P.σ2)
end
depends_on_params(::LotkaVolterraDiffusionAux) = (1, 2, 3, 4, 5, 6)

constdiff(::LotkaVolterraDiffusionAux) = true
b(t, x, P::LotkaVolterraDiffusionAux) = B(t,P) * x + β(t,P)
a(t, P::LotkaVolterraDiffusionAux) = σ(t,P) * σ(t, P)'

clone(P::LotkaVolterraDiffusionAux, θ) = LotkaVolterraDiffusionAux(θ..., P.t,
                                                                   P.u, P.T, P.v)

clone(P::LotkaVolterraDiffusionAux, θ, v) = LotkaVolterraDiffusionAux(θ..., P.t,
                                                                      v, P.T, v)
params(P::LotkaVolterraDiffusionAux) = [P.α, P.β, P.γ, P.δ, P.σ1, P.σ2]


using ARules

transactions = [["milk", "eggs", "bread"],
			    ["butter", "milk", "sugar", "flour", "eggs"],
			    ["bacon", "eggs", "milk", "beer"],
			    ["bread", "ham", "turkey"],
			    ["cheese", "ham", "bread", "ketchup"],
			    ["mustard", "hot dogs", "buns", "hamburger", "cheese", "beer"]]

# freq = frequent(transactions, 1, 8)
# display(freq)
#
# for x in freq[:itemset]
#     println(x)
# end

unq = unique_items(transactions)

@time node = ARules.frequent_item_tree(transactions, unq, 1, 8)

sup = gen_support_dict(node, 9)

for k in keys(sup)
    println(k)
end

rules = apriori(transactions, supp = 0.01, conf = 0.01, maxlen = 6)
showall(rules)


import Random
import Statistics: mean
import SelectivelyAdaptiveLasso as SAL
using Test

@testset "FeatureVector" begin for i in 1:100
    n = 100
    U = rand(n)
    X = SAL.FeatureVector(U)
    I = Random.shuffle(1:n)[1:10]

    @test all(U[I] == X.sorted[X.raw_to_sorted[I]])
    @test all(U[X.sorted_to_raw[I]] == X.sorted[I])
    @test X.sorted[sort(X.raw_to_sorted[I])] == U[I[SAL.sortperm_subset(X,I)]]
    @test issorted(X.sorted[X.raw_to_sorted[I[SAL.sortperm_subset(X,I)]]], rev=true)
end end

@testset "1-dim basis selection (right)" begin for i in 1:100
    n,c = 100, 0.4
    I = collect(1:2:n) # must be increasing to keep X sorted
    X = sort(rand(n), rev=true)
    Y = Vector{Float64}(X.≥c)

    Ŷ(X,Y) = X*mean(Y[X])
    sses_brute(X,Y) = [sum((Y.-Ŷ(X.≥xi,Y)).^2) for xi in X]
    sses_clever = SAL.sse_indicator_ols(Y)

    @test all(sses_brute(X,Y) .≈ sses_clever)
    @assert issorted(X[I], rev=true)

    Ĩ, Ỹ, _, x̃, l̃ = SAL.best_split_sorted(I,X[I],Y[I])
    @test Set(Ĩ) == Set((1:n)[X.≥c]) ∩ Set(I)
    @test all(Ỹ .== 1)
    @test x̃ == min(X[I][X[I].≥c]...)
    @test l̃ ≈ min(sses_brute(X[I], Y[I])...)
end end

@testset "changepoint inclusion" begin for i in 1:100
    X = [1.0, 0.9, 0.9, 0.5, 0.1]
    I = 1:length(X)
    Y = rand(length(X))
    Ĩ, Ỹ, _, x̃, l̃ = SAL.best_split_sorted(I,X[I],Y[I])
    @test ((2 ∉ Ĩ) & (3 ∉ Ĩ)) | ((2 ∈ Ĩ) & (3 ∈ Ĩ))
end end

@testset "multidim basis selection (right)" begin for i in 1:100
    c = 0.4
    n,p = 100,3
    j° = p
    U = rand(n,p)
    I = collect(1:2:n) # must be increasing to keep X sorted
    X = [SAL.FeatureVector(x) for x in eachcol(U)]
    Y = Vector{Float64}(U[:,j°] .≥ c)

    Ĩ, Ỹ, _, (j̃, x̃, d), l̃ = SAL.interact_basis(I, X, Y[I])
    @test Set(Ĩ) == Set((1:n)[U[:,j°].≥c]) ∩ Set(I)
    @test all(Ỹ .== 1)
    @test x̃ == min(U[I,j°][U[I,j°].≥c]...)
    @test d == :right
    @test j̃ == j°
end end

@testset "multidim basis selection (left)" begin for i in 1:100
    c = 0.4
    n,p = 100,3
    j° = p
    U = rand(n,p)
    I = collect(1:2:n) # must be increasing to keep X sorted
    X = [SAL.FeatureVector(x) for x in eachcol(U)]
    Y = Vector{Float64}(U[:,j°] .≤ c)

    Ĩ, Ỹ, _, (j̃, x̃, d), l̃ = SAL.interact_basis(I, X, Y[I])
    @test Set(Ĩ) == Set((1:n)[U[:,j°].≤c]) ∩ Set(I)
    @test all(Ỹ .== 1)
    @test x̃ == max(U[I,j°][U[I,j°].≤c]...)
    @test d == :left
    @test j̃ == j°
end end

@testset "build basis" begin for i in 1:100
    c = 0.4
    n,p = 100,3
    U = rand(n,p)
    X = [SAL.FeatureVector(x) for x in eachcol(U)]

    index, basis = SAL.basis_search_random(X)
    n_val = 20
    U_val = rand(n_val,p)
    X_val = [SAL.FeatureVector(x) for x in eachcol(U_val)]
    basis_val = SAL.build_basis(X_val, index)

    build_basis_brute(U, index) = reduce(.*, [d==:right ? U[:,j] .≥ c : U[:,j] .≤ c for (j,c,d) in index])
    bool_basis = build_basis_brute(U, index)
    bool_basis_val = build_basis_brute(U_val, index)
    @test Set(Vector{Int}(basis.nzind)) == Set((1:n)[bool_basis])
    @test Set(Vector{Int}(basis_val.nzind)) == Set((1:n_val)[bool_basis_val])
end end


# # Soil types

## load packages
using PlantHydraulics
using PlotPlants
FT = Float32;
#------------------------------------------------------------------------------

## prescribed soil types
_sh1 = BrooksCorey{FT}();
_sh2 = VanGenuchten{FT}();
_STS = ["Sand", "Loamy Sand", "Sandy Loam", "Loam", "Sandy Clay Loam",
        "Silt Loam", "Silt", "Clay Loam", "Silty Clay Loam", "Sandy Clay",
        "Silty Clay", "Clay"];

function compare_soil_vc(sn::Int)
    _bc = create_soil_VC(_sh1, _STS[sn]);
    _vg = create_soil_VC(_sh2, _STS[sn]);
    @show _bc.b, _bc.ϕs;

    _fig,_axes = create_canvas(_STS[sn]; ncol=2);
    _ax1,_ax2  = _axes;

    _Θs = range(_bc.Θr+FT(1e-3); stop=_bc.Θs, length=50);
    _P1 = -1 .* soil_p_25_swc.([_bc], _Θs);
    _P2 = -1 .* soil_p_25_swc.([_vg], _Θs);
    _K1 = soil_k_ratio_swc.([_bc], _Θs);
    _K2 = soil_k_ratio_swc.([_vg], _Θs);

    _ax1.plot(_Θs, _P1, "b-", label="Brooks Corey" );
    _ax1.plot(_Θs, _P2, "r-", label="van Genuchten");
    _ax2.plot(_Θs, _K1, "b-");
    _ax2.plot(_Θs, _K2, "r-");
    _ax1.set_yscale("log");
    _ax1.legend(loc="upper right");
    set_xlabels!(_axes, ["Soil Water Content (-)" for i in 1:2]);
    set_ylabels!(_axes, ["Metric potential (-MPa)", "Relative K"]);
    _fig.set_tight_layout(true);

    return _fig
end
#------------------------------------------------------------------------------

# ## Pre-defined soil types
# ### Sand
compare_soil_vc(1)
#------------------------------------------------------------------------------

# ### Loamy Sand
compare_soil_vc(2)
#------------------------------------------------------------------------------

# ### Sandy Loam
compare_soil_vc(3)
#------------------------------------------------------------------------------

# ### Loam
compare_soil_vc(4)
#------------------------------------------------------------------------------

# ### Sandy Clay Loam
compare_soil_vc(5)
#------------------------------------------------------------------------------

# ### Silt Loam
compare_soil_vc(6)
#------------------------------------------------------------------------------

# ### Silt
compare_soil_vc(7)
#------------------------------------------------------------------------------

# ### Clay Loam
compare_soil_vc(8)
#------------------------------------------------------------------------------

# ### Silty Clay Loam
compare_soil_vc(9)
#------------------------------------------------------------------------------

# ### Sandy Clay
compare_soil_vc(10)
#------------------------------------------------------------------------------

# ### Silty Clay
compare_soil_vc(11)
#------------------------------------------------------------------------------

# ### Clay
compare_soil_vc(12)
#------------------------------------------------------------------------------


include("utils.jl")

@testset "fixed Rank" begin
    M = FixedRankMatrices(3,2,2)
    Mc = FixedRankMatrices(3,2,2,ℂ)
    x = SVDMPoint([1.0 0.0; 0.0 1.0; 0.0 0.0])
    v = UMVTVector([0. 0.; 0. 0.; 1. 1.], [1. 0.; 0. 1.],zeros(2,2))
    @test repr(M) == "FixedRankMatrices(3, 2, 2, ℝ)"
    @test repr(Mc) == "FixedRankMatrices(3, 2, 2, ℂ)"
    @test sprint(show, "text/plain", x) == """
    SVDMPoint{Array{Float64,2},Array{Float64,1},Array{Float64,2}}
    U factor:
     3×2 Array{Float64,2}:
      1.0  0.0
      0.0  1.0
      0.0  0.0
    singular values:
     2-element Array{Float64,1}:
      1.0
      1.0
    Vt factor:
     2×2 Array{Float64,2}:
      1.0  0.0
      0.0  1.0"""
    @test sprint(show, "text/plain", v) == """
    UMVTVector{Array{Float64,2},Array{Float64,2},Array{Float64,2}}
    U factor:
     3×2 Array{Float64,2}:
      0.0  0.0
      0.0  0.0
      1.0  1.0
    M factor:
     2×2 Array{Float64,2}:
      1.0  0.0
      0.0  1.0
    Vt factor:
     2×2 Array{Float64,2}:
      0.0  0.0
      0.0  0.0"""
    @test inner(M,x,v,v) == norm(M,x,v)^2
    @test x == SVDMPoint(x.U,x.S,x.Vt)
    @test v == UMVTVector(v.U, v.M, v.Vt)
    @testset "Fixed Rank Matrices – Basics" begin
        @test representation_size(M) == (3,2)
        @test representation_size(Mc) == (3,2)
        @test manifold_dimension(M) == 6
        @test manifold_dimension(Mc) == 12
        @test !is_manifold_point(M, SVDMPoint([1. 0.; 0. 0.],2))
        @test_throws DomainError is_manifold_point(M, SVDMPoint([1. 0.; 0. 0.],2), true)

        @test !is_tangent_vector(M, SVDMPoint( [1. 0.; 0. 1.; 0. 0.] ), UMVTVector( zeros(2,1), zeros(1,2), zeros(2,2) )  )
        @test !is_tangent_vector(M, SVDMPoint([1. 0.; 0. 0.],2),v)
        @test_throws DomainError is_tangent_vector(M, SVDMPoint([1. 0.; 0. 0.],2), v, true)
        @test !is_tangent_vector(M, x, UMVTVector(x.U, v.M, x.Vt,2))
        @test_throws DomainError is_tangent_vector(M, x, UMVTVector(x.U, v.M, x.Vt,2), true)
        @test !is_tangent_vector(M, x, UMVTVector(v.U, v.M, x.Vt,2))
        @test_throws DomainError is_tangent_vector(M, x, UMVTVector(v.U, v.M, x.Vt,2), true)

        @test is_manifold_point(M,x)
        @test is_tangent_vector(M,x,v)
    end
    types = [
       [ Matrix{Float64}, Vector{Float64}, Matrix{Float64} ],
    ]
    TEST_FLOAT32 && push!(types, [ Matrix{Float32}, Vector{Float32}, Matrix{Float32} ] )

    for T in types
        @testset "Type $T" begin
            y = retract(M, x, v, PolarRetraction())
            z = SVDMPoint( [1/sqrt(2) 1/sqrt(2); 1/sqrt(2) -1/sqrt(2); 0. 0.] )
            pts = []
            for p in [x,y,z]
                push!(pts, SVDMPoint( convert.( T , [ p.U, p.S, p.Vt ] )... )   )
            end
            for p in pts
                @test is_manifold_point(M,p)
            end
            @testset "SVD MPoint Basics" begin
                s = svd(x.U*Diagonal(x.S)*x.Vt)
                x2 = SVDMPoint(s)
                x3 = SVDMPoint(s.U, s.S, s.Vt)
                @test SVDMPoint(x.U, x.S, x.Vt) == x
                @test x.S == x2.S
                @test x.U == x2.U
                @test x.Vt == x2.Vt
                @test x == x2
                @test x2.U == x3.U
                @test x2.S == x3.S
                @test x2.Vt == x3.Vt
                @test x2 == x3
                y = SVDMPoint([1. 0.; 0. 0.; 0. 0.], 1)
                s2 = svd([1. 0.; 0. 0.; 0. 0.])
                y2 = SVDMPoint(s2,1)
                y3 = SVDMPoint(s2.U, s2.S, s2.Vt,1)
                @test y.S == y2.S
                @test y.U == y2.U
                @test y.Vt == y2.Vt
                @test y == y2
                @test y2.U == y3.U
                @test y2.S == y3.S
                @test y2.Vt == y3.Vt
                @test y2==y3

                @test is_manifold_point(M,x)
                xM = x.U*Diagonal(x.S)*x.Vt
                @test is_manifold_point(M,xM)
                @test !is_manifold_point(M,xM[1:2,:])
                @test_throws DomainError is_manifold_point(M,xM[1:2,:], true)
                @test_throws DomainError is_manifold_point(FixedRankMatrices(3,2,1), x, true)
                @test_throws DomainError is_manifold_point(FixedRankMatrices(3,2,1), xM, true)
                xF1 = SVDMPoint(2*x.U, x.S, x.Vt)
                @test !is_manifold_point(M,xF1)
                @test_throws DomainError is_manifold_point(M,xF1,true)
                xF2 = SVDMPoint(x.U, x.S, 2*x.Vt)
                @test !is_manifold_point(M,xF2)
                @test_throws DomainError is_manifold_point(M,xF2,true)
            end
            @testset "UMV TVector Basics" begin
                w = UMVTVector(v.U, 2*v.M, v.Vt)
                @test v+w == UMVTVector(2*v.U, 3*v.M, 2*v.Vt)
                @test v-w == UMVTVector(0*v.U, -v.M, 0*v.Vt)
                @test 2*v == UMVTVector(2*v.U, 2*v.M, 2*v.Vt)
                @test v*2 == UMVTVector(v.U*2, v.M*2, v.Vt*2)
                @test 2\v == UMVTVector(2\v.U, 2\v.M, 2\v.Vt)
                @test v/2 == UMVTVector(v.U/2, v.M/2, v.Vt/2)
                @test +v == v
                @test -v == UMVTVector(-v.U, -v.M, -v.Vt)
                w = UMVTVector(v.U, v.M, v.Vt)
                @test v == w
                w = allocate(v, number_eltype(v))
                zero_tangent_vector!(M,w,x)
                oneP = SVDMPoint(one(zeros(3, 3)), ones(2), one(zeros(2, 2)), 2)
                @test oneP == one(x)
                oneV = UMVTVector(one(zeros(3,3)),one(zeros(2,2)),one(zeros(2,2)),2)
                @test oneV == one(v)
            end
            test_manifold(
                M,
                pts,
                test_exp_log = false,
                default_inverse_retraction_method = nothing,
                test_injectivity_radius = false,
                default_retraction_method = PolarRetraction(),
                test_is_tangent = false,
                test_project_tangent = true,
                test_vector_transport = false,
                test_forward_diff = false,
                test_reverse_diff = false,
                test_vector_spaces = false,
                test_vee_hat = false,
                test_tangent_vector_broadcasting = false, #broadcast not so easy for 3 matrix type
                projection_atol_multiplier = 15,
                retraction_methods = [PolarRetraction()]
            )
        end
    end
end


using RobustAdaptiveMetropolisSampler
using Distributions
using Statistics
using LinearAlgebra
using Random
using Test

@testset "RobustAdaptiveMetropolisSampler" begin

    Random.seed!(10)

    result = RAM_sample(p -> logpdf(Normal(3., 2.), p[1]), [0.], 0.5, 100_000, show_progress=false, output_log_probability_x=true)

    @test length(result.chain) == 100_000
    @test result.acceptance_rate ≈ 0.234 atol = 0.01
    @test mean(result.chain[:,1]) ≈ 3. atol = 0.1
    @test std(result.chain[:,1]) ≈ 2. atol = 0.01
    @test result.log_probabilities_x[1] == logpdf(Normal(3., 2.), result.chain[1])
    @test result.log_probabilities_x[length(result.chain)] == logpdf(Normal(3., 2.), result.chain[length(result.chain)])

    result = RAM_sample(p -> logpdf(Normal(3., 2.), p[1]), [0.], [0.5], 1_000, show_progress=false, output_log_probability_x=false)
    @test length(result.chain) == 1_000

    result = RAM_sample(p -> logpdf(Normal(3., 2.), p[1]), [0.], Diagonal([0.5]), 1_000, show_progress=false, output_log_probability_x=false)
    @test length(result.chain) == 1_000

    result = RAM_sample(p -> logpdf(Normal(3., 2.), p[1]), [0.], fill(0.5, 1, 1), 1_000, show_progress=false, output_log_probability_x=false)
    @test length(result.chain) == 1_000

end


"""
    test_constvel(; kwargs...)

Advects a gaussian concentration c0(x, y, t) with a constant velocity flow
u(x, y) = uvel and v(x, y) = vvel and compares the final state with
cfinal = c0(x-uvel*tfinal, y-vvel*tfinal)
"""
function test_constvel(stepper, dt, nsteps, dev::Device=CPU())

  nx, Lx = 128, 2π
  uvel, vvel = 0.2, 0.1
  u(x, y) = uvel
  v(x, y) = vvel

  prob = TracerAdvectionDiffusion.Problem(dev; nx=nx, Lx=Lx, κ=0.0, u=u, v=v, dt=dt, stepper=stepper, steadyflow=true)
  sol, cl, vs, pr, gr = prob.sol, prob.clock, prob.vars, prob.params, prob.grid

  x, y = gridpoints(gr)

  σ = 0.1
  c0ampl = 0.1
  c0func(x, y) = c0ampl * exp(-(x^2+y^2)/(2σ^2))

  c0 = c0func.(x, y)
  tfinal = nsteps*dt
  cfinal = @. c0func(x - uvel*tfinal, y - vvel*tfinal)

  TracerAdvectionDiffusion.set_c!(prob, c0)

  stepforward!(prob, nsteps)
  TracerAdvectionDiffusion.updatevars!(prob)

  isapprox(cfinal, vs.c, rtol=gr.nx*gr.ny*nsteps*1e-12)
end

"""
    test_timedependenttvel(; kwargs...)

Advects a gaussian concentration c0(x, y, t) with a time-varying velocity flow
u(x, y, t) = uvel and v(x, y, t) = vvel*sign(-t+tfinal/2) and compares the final
state with cfinal = c0(x-uvel*tfinal, y)
"""
function test_timedependentvel(stepper, dt, tfinal, dev::Device=CPU(); uvel=0.5, αv=0.5)

  nx, Lx = 128, 2π
  nsteps = round(Int, tfinal/dt)

  if !isapprox(tfinal, nsteps*dt, rtol=rtol_traceradvectiondiffusion)
    error("tfinal is not multiple of dt")
  end

  u(x, y, t) = uvel
  v(x, y, t) = αv*t + αv*dt/2

  prob = TracerAdvectionDiffusion.Problem(dev; nx=nx, Lx=Lx, κ=0.0, u=u, v=v, dt=dt, stepper=stepper)
  sol, cl, vs, pr, gr = prob.sol, prob.clock, prob.vars, prob.params, prob.grid
  x, y = gridpoints(gr)

  σ = 0.2
  c0func(x, y) = 0.1*exp(-(x^2+y^2)/(2σ^2))

  c0 = @. c0func(x, y)
  tfinal = nsteps*dt
  cfinal = @. c0func(x - uvel*tfinal, y - 0.5αv*tfinal^2)

  TracerAdvectionDiffusion.set_c!(prob, c0)

  stepforward!(prob, nsteps)
  TracerAdvectionDiffusion.updatevars!(prob)
  isapprox(cfinal, vs.c, rtol=gr.nx*gr.ny*nsteps*1e-12)
end

"""
    test_diffusion(; kwargs...)

Diffuses a gaussian concentration c0(x, y, t) and compares the final state with
the analytic solution of the heat equation, cfinal
"""
function test_diffusion(stepper, dt, tfinal, dev::Device=CPU(); steadyflow = true)

  nx = 128
  Lx = 2π
   κ = 0.01
  nsteps = round(Int, tfinal/dt)

  if !isapprox(tfinal, nsteps*dt, rtol=rtol_traceradvectiondiffusion)
    error("tfinal is not multiple of dt")
  end

  prob = TracerAdvectionDiffusion.Problem(dev; steadyflow=steadyflow, nx=nx,
    Lx=Lx, κ=κ, dt=dt, stepper=stepper)
  sol, cl, vs, pr, gr = prob.sol, prob.clock, prob.vars, prob.params, prob.grid
  x, y = gridpoints(gr)

  c0ampl, σ = 0.1, 0.1
  c0func(x, y) = c0ampl*exp(-(x^2+y^2)/(2σ^2))

  c0 = @. c0func(x, y)
  tfinal = nsteps*dt
  σt = sqrt(2*κ*tfinal + σ^2)
  cfinal = @. c0ampl*(σ^2/σt^2)*exp(-(x^2+y^2)/(2*σt^2))

  TracerAdvectionDiffusion.set_c!(prob, c0)

  stepforward!(prob, nsteps)
  TracerAdvectionDiffusion.updatevars!(prob)

  isapprox(cfinal, vs.c, rtol=gr.nx*gr.ny*nsteps*1e-12)
end

"""
    test_hyperdiffusion(; kwargs...)

Diffuses a gaussian concentration c0(x, y, t) using hyperdiffusivity and
compares the final state with the analytic solution of the heat equation, cfinal
"""
function test_hyperdiffusion(stepper, dt, tfinal, dev::Device=CPU(); steadyflow = true)

   nx = 128
   Lx = 2π
    κ = 0.0   # no diffusivity
    η = κ     # no diffusivity
   κh = 0.01  # hyperdiffusivity coeff
  nκh = 1     # nκh=1 converts hyperdiffusivity to plain diffusivity
              # so we can compare with the analytic solution of heat equation

  nsteps = round(Int, tfinal/dt)

  if !isapprox(tfinal, nsteps*dt, rtol=rtol_traceradvectiondiffusion)
    error("tfinal is not multiple of dt")
  end

  gr = TwoDGrid(dev, nx, Lx)
  x, y = gridpoints(gr)

  u, v = zero(x), zero(x) #0*x, 0*x

  vs = TracerAdvectionDiffusion.Vars(dev, gr)
  pr = TracerAdvectionDiffusion.ConstDiffSteadyFlowParams(η, κ, κh, nκh, u, v)
  eq = TracerAdvectionDiffusion.Equation(pr, gr)
  prob = FourierFlows.Problem(eq, stepper, dt, gr, vs, pr, dev)

  c0ampl, σ = 0.1, 0.1
  c0func(x, y) = c0ampl*exp(-(x^2+y^2)/(2σ^2))

  c0 = @. c0func(x, y)
  tfinal = nsteps*dt
  σt = sqrt(2*κh*tfinal + σ^2)
  cfinal = @. c0ampl*σ^2/σt^2 * exp(-(x^2+y^2)/(2*σt^2))

  TracerAdvectionDiffusion.set_c!(prob, c0)

  stepforward!(prob, nsteps)
  TracerAdvectionDiffusion.updatevars!(prob)

  isapprox(cfinal, vs.c, rtol=gr.nx*gr.ny*nsteps*1e-12)
end


@defcomp AbatementCosts begin
    region = Index()

   # From the AbatementCostParameters
    zc_zerocostemissions = Parameter(index=[time, region], unit= "%")
    q0_absolutecutbacksatnegativecost = Parameter(index=[time, region], unit= "Mtonne")
    blo = Parameter(index=[time, region], unit = "per Mtonne")
    alo = Parameter(index=[time, region], unit = "\$/tonne")
    bhi = Parameter(index=[time, region], unit = "per Mtonne")
    ahi = Parameter(index=[time, region], unit = "\$/tonne")

    # Driver of abatement costs
    e0_baselineemissions = Parameter(index=[region], unit= "Mtonne/year")
    er_emissionsgrowth = Parameter(index=[time, region], unit= "%")

    # Intermediate outputs
    cb_reductionsfromzerocostemissions = Variable(index=[time, region], unit= "%")
    cbe_absoluteemissionreductions = Variable(index=[time, region], unit="Mtonne") # Goes to AbatementCostParameters

    # Main costs results
    mc_marginalcost = Variable(index=[time, region], unit = "\$/tonne")
    tcq0 = Variable(index=[time, region], unit = "\$million")
    tc_totalcost = Variable(index=[time, region], unit= "\$million")

    function run_timestep(p, v, d, t)

        for r in d.region
            v.cb_reductionsfromzerocostemissions[t,r] = max(p.zc_zerocostemissions[t,r] - p.er_emissionsgrowth[t,r], 0)
            v.cbe_absoluteemissionreductions[t,r] = v.cb_reductionsfromzerocostemissions[t,r]* p.e0_baselineemissions[r]/100

            if v.cbe_absoluteemissionreductions[t,r]< p.q0_absolutecutbacksatnegativecost[t,r]
                v.mc_marginalcost[t,r] = p.alo[t,r]* (exp(p.blo[t,r]*(v.cbe_absoluteemissionreductions[t,r]- p.q0_absolutecutbacksatnegativecost[t,r]))-1)
            else
                v.mc_marginalcost[t,r] = p.ahi[t,r]*(exp(p.bhi[t,r]*(v.cbe_absoluteemissionreductions[t,r]- p.q0_absolutecutbacksatnegativecost[t,r]))-1)
            end

            if p.q0_absolutecutbacksatnegativecost[t,r] == 0.
                v.tcq0[t,r] = 0.
            else
                v.tcq0[t,r] = (p.alo[t,r]/p.blo[t,r])*(1-exp(-p.blo[t,r]* p.q0_absolutecutbacksatnegativecost[t,r]))- p.alo[t,r]*p.q0_absolutecutbacksatnegativecost[t,r]
            end

            if v.cbe_absoluteemissionreductions[t,r]<p.q0_absolutecutbacksatnegativecost[t,r]
                v.tc_totalcost[t,r] = (p.alo[t,r]/p.blo[t,r])*(exp(p.blo[t,r]*(v.cbe_absoluteemissionreductions[t,r]- p.q0_absolutecutbacksatnegativecost[t,r]))- exp(-p.blo[t,r]*p.q0_absolutecutbacksatnegativecost[t,r])) - p.alo[t,r]*v.cbe_absoluteemissionreductions[t,r]
            else
                v.tc_totalcost[t,r] = (p.ahi[t,r]/p.bhi[t,r])* (exp(p.bhi[t,r]*(v.cbe_absoluteemissionreductions[t,r]-p.q0_absolutecutbacksatnegativecost[t,r]))-1) - p.ahi[t,r]*(v.cbe_absoluteemissionreductions[t,r] - p.q0_absolutecutbacksatnegativecost[t,r]) + v.tcq0[t,r]
            end
        end
    end
end

function addabatementcosts(model::Model, class::Symbol)
    componentname = Symbol("AbatementCosts$class")
    abatementcostscomp = add_comp!(model, AbatementCosts, componentname)

    if class == :CO2
        setdistinctparameter(model, componentname, :e0_baselineemissions, readpagedata(model, "data/e0_baselineCO2emissions.csv"))
    elseif class == :CH4
        setdistinctparameter(model, componentname, :e0_baselineemissions, readpagedata(model, "data/e0_baselineCH4emissions.csv"))
    elseif class == :N2O
        setdistinctparameter(model, componentname, :e0_baselineemissions, readpagedata(model, "data/e0_baselineN2Oemissions.csv"))
    elseif class == :Lin
        setdistinctparameter(model, componentname, :e0_baselineemissions, readpagedata(model,"data/e0_baselineLGemissions.csv"))
    else
        error("Unknown class of abatement costs.")
    end

    return abatementcostscomp
end


function makecinfo(ecut,Ω,basis="Hermite")
cinfo = Cinfo()
if basis=="Hermite"
  dim=length(Ω)
  if dim==1
    ωx = Ω[1]
    e0 = 0.5*ωx
    ecut < e0 && error("ecut must exceed the zero point energy.")
    Mx,nx,en = nenergy(ecut,e0,ωx,basis)
    Mmax = Mx
    P  = en .< ecut; Mult = sum(P)
    en = P.*en
    M  = [Mx]
  elseif dim==2
    ωx,ωy = Ω
    e0 = 0.5(ωx+ωy)
    ecut < e0 && error("ecut must exceed the zero point energy.")
    Mx,nx,ex = nenergy(ecut,e0,ωx,basis)
    My,ny,ey = nenergy(ecut,e0,ωy,basis)
    M  = [Mx,My]; Mmax = maximum(M)
    en = [ex[i+1]+ey[j+1] for i=nx, j=ny]
    P  = en .< ecut; Mult = sum(P)
    en = P.*en
    M  = [Mx,My]
  elseif dim==3
    ωx,ωy,ωz = Ω
    e0 = 0.5(ωx+ωy+ωz)
    ecut < e0 && error("ecut must exceed the zero point energy.")
    Mx,nx,ex = nenergy(ecut,e0,ωx,basis)
    My,ny,ey = nenergy(ecut,e0,ωy,basis)
    Mz,nz,ez = nenergy(ecut,e0,ωz,basis)
    M  = [Mx,My,Mz]; Mmax = maximum(M)
    en = [ex[i+1]+ey[j+1]+ez[k+1] for i=nx, j=ny, k=nz]
    P  = en .< ecut; Mult = sum(P)
    en = P.*en
else error("Spatial dimension must be 1, 2, or 3.")
  end
end
  @pack cinfo = basis, Ω, ecut, e0, Mmax, M, en, P, Mult
  return cinfo
end


#based off https://github.com/Roboneet/CartPole.jl/blob/master/src/Envs/classic_control/CartPole.jl
@with_kw struct CartPoleMDP <: MDP{Array{Float32}, Array{Float32}}
    gravity::Float32 = 98f-1
    masscart::Float32 = 1f0
    masspole::Float32 = 1f-1
    total_mass::Float32 = masspole + masscart
    length::Float32 = 5f-1 # actually half the pole's length
    polemass_length::Float32 = (masspole * length)
    force_mag::Float32 = 1f1
    τ::Float32 = 2f-2 # seconds between state updates
    maxT::Float32 = τ*100
    kinematics_integrator::AbstractString = "euler"

    # Angle at which to fail the episode
    θ_threshold_radians::Float32 = Float32(12 * 2 * π / 360)
    x_threshold::Float32 = 24f-1
    γ::Float32 = 0.99f0
end

function POMDPs.gen(env::CartPoleMDP, s, a, rng::AbstractRNG = Random.GLOBAL_RNG)
    x, ẋ, θ, θ̇ , t = s

    failed =  !(all(vcat(-env.x_threshold .≤ x .≤ env.x_threshold,
            -env.θ_threshold_radians .≤ θ .≤ env.θ_threshold_radians)))
    timeout = t > env.maxT || t ≈ env.maxT

    if failed
        return (sp=fill(-100, size(s)), r=0)
    elseif timeout
        return (sp=fill(100, size(s)), r=0)
    end

     force = a .* env.force_mag  # action is +1 or -1
     cosθ = cos.(θ)
     sinθ = sin.(θ)
     temp = (force .+ env.polemass_length * θ̇  .^ 2 .* sinθ) / env.total_mass
     θacc = (env.gravity*sinθ .- cosθ.*temp) ./
            (env.length * (4f0/3 .- env.masspole * cosθ .^ 2 / env.total_mass))
     xacc  = temp .- env.polemass_length * θacc .* cosθ / env.total_mass
     if env.kinematics_integrator == "euler"
         x_ = x  .+ env.τ * ẋ
         ẋ_ = ẋ  .+ env.τ * xacc
         θ_ = θ  .+ env.τ * θ̇
         θ̇_ = θ̇  .+ env.τ * θacc
     else # semi-implicit euler
         ẋ_ = ẋ  .+ env.τ * xacc
         x_ = x  .+ env.τ * ẋ_
         θ̇_ = θ̇  .+ env.τ * θacc
         θ_ = θ  .+ env.τ * θ̇_
     end

     sp = vcat(x_, ẋ_, θ_, θ̇_, t+env.τ)
     (sp=sp, r=1f0)
end

function POMDPs.initialstate(mdp::CartPoleMDP)
    ImplicitDistribution((rng) -> Float32[(rand(Float32, 4) * 1f-1 .- 5f-2)..., 0f0])
end

isfailure(mdp::CartPoleMDP, s) = s[1] == -100
# Terminate if there is a failure or if we exceed the maximum time (in the case that we are tracking such things)
POMDPs.isterminal(mdp::CartPoleMDP, s) = isfailure(mdp, s) || s[1] == 100
POMDPs.discount(mdp::CartPoleMDP) = mdp.γ


@testset "intrinsiccoregion" begin
    rng = MersenneTwister(123)

    dims = (in=3, out=2, obs=3)
    rank = 1

    A = randn(dims.out, rank)
    B = A * transpose(A) + Diagonal(rand(dims.out))

    X = [(rand(dims.in), rand(1:(dims.out))) for i in 1:(dims.obs)]

    kernel = SqExponentialKernel()
    icoregionkernel = IntrinsicCoregionMOKernel(; kernel=kernel, B=B)

    @test icoregionkernel.B == B
    @test icoregionkernel.kernel == kernel
    @test icoregionkernel(X[1], X[1]) ≈ B[X[1][2], X[1][2]] * kernel(X[1][1], X[1][1])
    @test icoregionkernel(X[1], X[end]) ≈ B[X[1][2], X[end][2]] * kernel(X[1][1], X[end][1])

    KernelFunctions.TestUtils.test_interface(
        icoregionkernel, Vector{Tuple{Float64,Int}}; dim_out=dims.out
    )

    test_ADs(icoregionkernel; dims=dims)

    @test string(icoregionkernel) ==
          string("Intrinsic Coregion Kernel: ", kernel, " with ", dims.out, " outputs")
end


abstract type FunctionalResponse end

struct HollingTypeI{AT} <: FunctionalResponse
    A::AT
    HollingTypeI(a::A) where {A} = new{A}(a)
end
(fr::HollingTypeI)(c::T, r::T) where {T<:Biomass} = @fastmath fr.A * r

# aliases
LinearFunctionalResponse(a) = HollingTypeI(a)
LotkaVolterra(a) = HollingTypeI(a)

struct HollingTypeII{AT,HT} <: FunctionalResponse
    a::AT  ## attack rate
    h::HT ## handling time
    HollingTypeII(a::A, h::H) where {A,H} = new{A,H}(a, h)
end
(afr::HollingTypeII)(c, r) = @fastmath (afr.a * r) / (afr.a * afr.h * r + 1)
# aliases
MichaelisMenten(a, h) = HollingTypeII(a, h)

struct HollingTypeIII{AT,HT,BT} <: FunctionalResponse
    a::AT  ## attack rate
    h::HT ## handling time
    β::BT ## scaling
    HollingTypeIII(a::A, h::H, β::B) where {A,H,B} = new{A,H,B}(a, h, β)
end
(afr::HollingTypeIII)(c, r) = @fastmath (afr.a * r) / ((afr.a * afr.h * r)^β + 1)

struct CrowleyMartin{WT,UT,VT} <: FunctionalResponse
    w::WT
    u::UT
    v::VT
    CrowleyMartin(w::W, u::U, v::V) where {W,U,V} = new{W,U,V}(w, u, v)
end
(cmfr::CrowleyMartin)(c, r) =
    @fastmath (cmfr.w * c * r) / (1 + c * cmfr.u + r * cmfr.v + c * r * cmfr.u * cmfr.v)

struct Eating{R,W,FR,D} <: CellRule{R,W}
    functionalresponse::FR
    dt::D
end

Eating{R,W}(; functionalresponse::F = LotkaVolterra(1.5), dt::D = 0.1) where {R,W,F,D} =
    Eating{R,W}(functionalresponse, dt)

function DynamicGrids.applyrule(data, rule::Eating, (C, R), index)
    C > zero(C) || return zero(C), R
    R > zero(R) || return C, zero(R)

    fr = rule.functionalresponse(C, R)
    C += C * fr * rule.dt
    R -= C * fr * rule.dt

    # TODO
    # in the case this causes R to go extinct,
    # this should only grow C as much as the remaining value of R

    C > zero(C) || return zero(C), R
    R > zero(R) || return C, zero(R)
    return C, R
end

"""
    Shortcut to make many C->R layers. Instead of having all of the eating between pairs
    of species in a single rule, builds a ruleset where each rule is a single consumer-resource interaction.
"""
function FoodWebEating(
    consumernames::Vector{Symbol},
    resourcenames::Vector{Symbol},
    fr::FunctionalResponse,
    metaweb::Matrix;
    dt = 0.1,
)
    nspecies = size(metaweb)[1]

    allrules = Ruleset()
    for r = 1:length(resourcenames), c = 1:length(consumernames)
        if metaweb[r, c] == 1
            Csym = consumernames[c]
            Rsym = resourcenames[r]
            thisrule = Eating{Tuple{Csym,Rsym}}(functionalresponse = fr, dt = dt)
            allrules += thisrule
        end
    end

    return allrules
end


using JLD
using Plots
using Images
using LsqFit

plotlyjs()

d = load("save.jld")
zs, κs, σs = d["zs"], d["ks"], d["ss"]
nzs = imfilter(zs, Kernel.gaussian(0.7))
plot(κs, σs, nzs', st=:surface, zlims=(0, 1), xlabel="κ" , ylabel="σ", zlabel="r_inf", title="R∞(σ, κ)")
display(current())

k = nzs[2:end, :] - nzs[1:end-1, :]

plot(κs, σs, k', zlims=(0,0.5), st=:surface, xlabel="κ" , ylabel="σ", zlabel="r_inf", title="R∞(σ, κ)")
display(current())

j = findall(k .> 0.09)

ys = map(j) do i; κs[i[1]]; end
xs = map(j) do i; σs[i[2]]; end

@. model(x, m) = m*x
fit = curve_fit(model, xs, ys, [0.])

plot(0:0.5:2.5, model(0:0.5:2.5, fit.param), yaxis="κ", xaxis="σ", xlims=(0,2.5), ylims=(0, 2.5))
plot!(xs, ys, st=:scatter, legend=:bottomright)

display(current())


@testset "solve_adbc.jl" begin
    a = 0.0
    b = 1.0
    domain = ClosedIntervalDomain(a, b)
    g_a(t) = 0.0
    g_b(t) = 0.0
    f(ξ, t) = (1.0 - ξ)ξ
    u₀(ξ) = 0.0
    T = 1.0
    N = 100
    Δx = (b - a) / N
    h(x) = 3.1Δx
    σ = Δx / 2
    filter = GaussianFilter(h, σ)
    equation = DiffusionEquation(domain, filter, f, g_a, g_b)
    ū_adbc = solve_adbc(equation, u₀, (0.0, T), N, T / 1000)
end


using LinearAlgebra
using Distributions

"""
    multinomial_hyper(α::AbstractArray{Float32,1})
[Dirichlet Distribution](https://en.wikipedia.org/wiki/Dirichlet_distribution)
"""
struct multinomial_dist <: distibution_sample
    α::AbstractArray{Float32,1}
end

function log_likelihood!(r::AbstractArray,x::AbstractArray, distibution_sample::multinomial_dist , group::Int64 = -1)
    r .= (distibution_sample.α' * x)[1,:]
end


using Test

using EquationsOfState.FiniteStrains
using EquationsOfState.LinearFitting

@testset "Test strain-volume-derivative" begin
    volumes = [
        25.987454833,
        26.9045702104,
        27.8430241908,
        28.8029649591,
        29.7848370694,
        30.7887887064,
        31.814968055,
        32.8638196693,
        33.9353435494,
        35.0299842495,
        36.1477417695,
        37.2892088485,
        38.4543854865,
        39.6437162376,
        40.857201102,
        42.095136449,
        43.3579668329,
        44.6456922537,
        45.9587572656,
        47.2973100535,
        48.6614988019,
        50.0517680652,
        51.4682660281,
        52.9112890601,
        54.3808371612,
        55.8775030703,
        57.4014349722,
        58.9526328669,
    ]
    v0 = 40.98926572870838
    @test [strain_volume_derivative(EulerianStrain(), v0, v, 3) for v in volumes] ≈ [
        -0.00025163164656036914,
        -0.00022158317701314234,
        -0.0001954064296869249,
        -0.00017256893337511754,
        -0.00015261005948744114,
        -0.00013514276136071226,
        -0.00011983495896811362,
        -0.00010639785139417004,
        -9.458902516425933e-5,
        -8.419491823560175e-5,
        -7.503571748024374e-5,
        -6.695156194310254e-5,
        -5.980866702052554e-5,
        -5.348840953851233e-5,
        -4.789028506457302e-5,
        -4.2925492613058147e-5,
        -3.85164573279167e-5,
        -3.459717265567486e-5,
        -3.110870879523787e-5,
        -2.8000473828556818e-5,
        -2.5228159545528874e-5,
        -2.275242969867225e-5,
        -2.0539381634855355e-5,
        -1.855902166896191e-5,
        -1.6785363112484185e-5,
        -1.5194859884358581e-5,
        -1.3767304534537318e-5,
        -1.2484962604556655e-5,
    ]
    @test [strain_volume_derivative(LagrangianStrain(), v0, v, 3) for v in volumes] ≈ [
        -4.360792585446927e-5,
        -4.021795998862133e-5,
        -3.7125806529670504e-5,
        -3.430262724447532e-5,
        -3.172185978064063e-5,
        -2.9360589247189284e-5,
        -2.719823329453535e-5,
        -2.5215761270023045e-5,
        -2.3396932620829496e-5,
        -2.172639361674028e-5,
        -2.0191013020817355e-5,
        -1.8778177794469247e-5,
        -1.7477278166117428e-5,
        -1.6278235883744374e-5,
        -1.5172392867363944e-5,
        -1.4151626497343523e-5,
        -1.3208496564809248e-5,
        -1.2336591903572137e-5,
        -1.1529793498552908e-5,
        -1.0782740265865238e-5,
        -1.009054720444801e-5,
        -9.448630234274912e-6,
        -8.852961760583228e-6,
        -8.299807942932397e-6,
        -7.78586914649821e-6,
        -7.307932527332951e-6,
        -6.863217457745042e-6,
        -6.449212613211634e-6,
    ]
    @test [strain_volume_derivative(NaturalStrain(), v0, v, 3) for v in volumes] ≈ [
        0.00011395640842831919,
        0.00010269559592471744,
        9.265754335710983e-5,
        8.369864119123017e-5,
        7.569101157907677e-5,
        6.852549573283301e-5,
        6.210628060176279e-5,
        5.634765917870933e-5,
        5.117681395143497e-5,
        4.6527548395465794e-5,
        4.234347362622052e-5,
        3.857273983853629e-5,
        3.517161972472347e-5,
        3.2100139446683775e-5,
        2.9324064018820016e-5,
        2.681230993011352e-5,
        2.4537102053587615e-5,
        2.2474565886168472e-5,
        2.0602748703699993e-5,
        1.890256196526239e-5,
        1.7356955567635e-5,
        1.5950405389513548e-5,
        1.466936735556791e-5,
        1.3501591655278468e-5,
        1.2436332606671682e-5,
        1.1463547171440715e-5,
        1.0574546405738849e-5,
        9.761585855581402e-6,
    ]
    @test [strain_volume_derivative(InfinitesimalStrain(), v0, v, 3) for v in volumes] ≈ [
        0.00017195382081596378,
        0.00015318072693119346,
        0.0001366374157858279,
        0.00012203947625819057,
        0.00010913738901171863,
        9.77197419230021e-5,
        8.76030749059566e-5,
        7.862563981068296e-5,
        7.065075730232455e-5,
        6.355618910218268e-5,
        5.723834723338108e-5,
        5.160365044520063e-5,
        4.657341353283599e-5,
        4.207684089417924e-5,
        3.80535882418961e-5,
        3.444963171076881e-5,
        3.1217248819476426e-5,
        2.8315601024208774e-5,
        2.5707710435800747e-5,
        2.3361614304347095e-5,
        2.1249043437552616e-5,
        1.9344592130116945e-5,
        1.7626220489559587e-5,
        1.6074217776270604e-5,
        1.4671393271712412e-5,
        1.3401940487277108e-5,
        1.2252230465391096e-5,
        1.121020710562245e-5,
    ]
end


@doc raw"""
    energy_balance(psi_container::PSIContainer,
                        initial_conditions::Vector{InitialCondition},
                        efficiency_data::Tuple{Vector{String}, Vector{InOut}},
                        cons_name::Symbol,
                        var_names::Tuple{Symbol, Symbol, Symbol})

Constructs multi-timestep constraint from initial condition, efficiency data, and variable tuple

# Constraints

If t = 1:

``` varenergy[name, 1] == initial_conditions[ix].value + varin[name, 1]*eff_in*fraction_of_hour - varout[name, 1]*fraction_of_hour/eff_out ```

If t > 1:

``` varenergy[name, t] == varenergy[name, t-1] + varin[name, t]*eff_in*fraction_of_hour - varout[name, t]*fraction_of_hour/eff_out ```

# LaTeX

`` x^{energy}_1 == x^{energy}_{init} + frhr \eta^{in} x^{in}_1 - \frac{frhr}{\eta^{out}} x^{out}_1, \text{ for } t = 1 ``

`` x^{energy}_t == x^{energy}_{t-1} + frhr \eta^{in} x^{in}_t - \frac{frhr}{\eta^{out}} x^{out}_t, \forall t \geq 2 ``

# Arguments
* psi_container::PSIContainer : the psi_container model built in PowerSimulations
* initial_conditions::Vector{InitialCondition} : for time zero 'varenergy'
* efficiency_data::Tuple{Vector{String}, Vector{InOut}} :: charging/discharging efficiencies
* cons_name::Symbol : name of the constraint
* var_names::Tuple{Symbol, Symbol, Symbol} : the names of the variables
- : var_names[1] : varin
- : var_names[2] : varout
- : var_names[3] : varenergy

"""
function energy_balance(
    psi_container::PSIContainer,
    initial_conditions::Vector{InitialCondition},
    efficiency_data::Tuple{Vector{String}, Vector{InOut}},
    cons_name::Symbol,
    var_names::Tuple{Symbol, Symbol, Symbol},
)
    time_steps = model_time_steps(psi_container)
    resolution = model_resolution(psi_container)
    fraction_of_hour = Dates.value(Dates.Minute(resolution)) / 60
    name_index = efficiency_data[1]

    varin = get_variable(psi_container, var_names[1])
    varout = get_variable(psi_container, var_names[2])
    varenergy = get_variable(psi_container, var_names[3])

    constraint = add_cons_container!(psi_container, cons_name, name_index, time_steps)

    for (ix, name) in enumerate(name_index)
        eff_in = efficiency_data[2][ix].in
        eff_out = efficiency_data[2][ix].out

        constraint[name, 1] = JuMP.@constraint(
            psi_container.JuMPmodel,
            varenergy[name, 1] ==
            initial_conditions[ix].value + varin[name, 1] * eff_in * fraction_of_hour -
            (varout[name, 1]) * fraction_of_hour / eff_out
        )

    end

    for t in time_steps[2:end], (ix, name) in enumerate(name_index)
        eff_in = efficiency_data[2][ix].in
        eff_out = efficiency_data[2][ix].out

        constraint[name, t] = JuMP.@constraint(
            psi_container.JuMPmodel,
            varenergy[name, t] ==
            varenergy[name, t - 1] + varin[name, t] * eff_in * fraction_of_hour -
            (varout[name, t]) * fraction_of_hour / eff_out
        )
    end

    return
end

@doc raw"""
    energy_balance_external_input(psi_container::PSIContainer,
                        initial_conditions::Vector{InitialCondition},
                        efficiency_data::Tuple{Vector{String}, Vector{InOut}},
                        cons_name::Symbol,
                        var_names::Tuple{Symbol, Symbol, Symbol})
Constructs multi-timestep constraint from initial condition, efficiency data, and variable tuple
# Constraints
If t = 1:
``` varenergy[name, 1] == initial_conditions[ix].value + (paraminflow[name, t] - varspill[name, 1] - varout[name, 1])*fraction_of_hour ```
If t > 1:
``` varenergy[name, t] == varenergy[name, t-1] + (paraminflow[name, t] - varspill[name, t] - varout[name, t])*fraction_of_hour ```
# LaTeX
`` x^{energy}_1 == x^{energy}_{init} + frhr  (x^{in}_1 - x^{spillage}_1 -  x^{out}_1), \text{ for } t = 1 ``
`` x^{energy}_t == x^{energy}_{t-1} + frhr (x^{in}_t - x^{spillage}_t - x^{out}_t), \forall t \geq 2 ``
# Arguments
* psi_container::PSIContainer : the psi_container model built in PowerSimulations
* initial_conditions::Vector{InitialCondition} : for time zero 'varenergy'
* inflow_data::Vector{DeviceTimeSeries} :: Inflow energy forecast information
* cons_name::Symbol : name of the constraint
* var_names::Tuple{Symbol, Symbol, Symbol} : the names of the variables
- : var_names[1] : varspill
- : var_names[2] : varout
- : var_names[3] : varenergy
* param_reference::UpdateRef : UpdateRef to access the inflow parameter
"""
function energy_balance_external_input_param(
    psi_container::PSIContainer,
    initial_conditions::Vector{InitialCondition},
    inflow_data::Vector{DeviceTimeSeries},
    cons_name::Symbol,
    var_names::Tuple{Symbol, Symbol, Symbol},
    param_reference::UpdateRef,
)
    time_steps = model_time_steps(psi_container)
    resolution = model_resolution(psi_container)
    fraction_of_hour = Dates.value(Dates.Minute(resolution)) / 60
    name_index = (d.name for d in inflow_data)

    varspill = get_variable(psi_container, var_names[1])
    varout = get_variable(psi_container, var_names[2])
    varenergy = get_variable(psi_container, var_names[3])
    container = add_param_container!(psi_container, param_reference, name_index, time_steps)
    paraminflow = get_parameter_array(container)
    multiplier = get_multiplier_array(container)
    constraint = add_cons_container!(psi_container, cons_name, name_index, time_steps)

    for (ix, d) in enumerate(inflow_data)
        multiplier[d.name, 1] = d.multiplier
        paraminflow[d.name, 1] = PJ.add_parameter(psi_container.JuMPmodel, d.timeseries[1])
        exp =
            initial_conditions[ix].value +
            (
                multiplier[d.name, 1] * paraminflow[d.name, 1] - varspill[d.name, 1] -
                varout[d.name, 1]
            ) * fraction_of_hour
        constraint[d.name, 1] =
            JuMP.@constraint(psi_container.JuMPmodel, varenergy[d.name, 1] == exp)

        for t in time_steps[2:end]
            paraminflow[d.name, t] =
                PJ.add_parameter(psi_container.JuMPmodel, d.timeseries[t])
            exp =
                varenergy[d.name, t - 1] +
                (
                    d.multiplier * paraminflow[d.name, 1] - varspill[d.name, t] -
                    varout[d.name, t]
                ) * fraction_of_hour
            constraint[d.name, t] =
                JuMP.@constraint(psi_container.JuMPmodel, varenergy[d.name, t] == exp)
        end
    end
    return
end

@doc raw"""
    energy_balance_external_input(psi_container::PSIContainer,
                        initial_conditions::Vector{InitialCondition},
                        efficiency_data::Tuple{Vector{String}, Vector{InOut}},
                        cons_name::Symbol,
                        var_names::Tuple{Symbol, Symbol, Symbol})
Constructs multi-timestep constraint from initial condition, efficiency data, and variable tuple
# Constraints
If t = 1:
``` varenergy[name, 1] == initial_conditions[ix].value + (paraminflow[name, t] - varspill[name, 1] - varout[name, 1])*fraction_of_hour ```
If t > 1:
``` varenergy[name, t] == varenergy[name, t-1] + (paraminflow[name, t] - varspill[name, t] - varout[name, t])*fraction_of_hour ```
# LaTeX
`` x^{energy}_1 == x^{energy}_{init} + frhr  (x^{in}_1 - x^{spillage}_1 -  x^{out}_1), \text{ for } t = 1 ``
`` x^{energy}_t == x^{energy}_{t-1} + frhr (x^{in}_t - x^{spillage}_t - x^{out}_t), \forall t \geq 2 ``
# Arguments
* psi_container::PSIContainer : the psi_container model built in PowerSimulations
* initial_conditions::Vector{InitialCondition} : for time zero 'varenergy'
* inflow_data::TVector{DeviceTimeSeries} :: Inflow energy forecast information
* cons_name::Symbol : name of the constraint
* var_names::Tuple{Symbol, Symbol, Symbol} : the names of the variables
- : var_names[1] : varspill
- : var_names[2] : varout
- : var_names[3] : varenergy
"""
function energy_balance_external_input(
    psi_container::PSIContainer,
    initial_conditions::Vector{InitialCondition},
    inflow_data::Vector{DeviceTimeSeries},
    cons_name::Symbol,
    var_names::Tuple{Symbol, Symbol, Symbol},
)
    time_steps = model_time_steps(psi_container)
    resolution = model_resolution(psi_container)
    fraction_of_hour = Dates.value(Dates.Minute(resolution)) / 60
    name_index = (d.name for d in inflow_data)

    varspill = get_variable(psi_container, var_names[1])
    varout = get_variable(psi_container, var_names[2])
    varenergy = get_variable(psi_container, var_names[3])

    constraint = add_cons_container!(psi_container, cons_name, name_index, time_steps)

    for (ix, d) in enumerate(inflow_data)
        constraint[d.name, 1] = JuMP.@constraint(
            psi_container.JuMPmodel,
            varenergy[d.name, 1] ==
            initial_conditions[ix].value +
            (d.multiplier * d.timeseries[1] - varspill[d.name, 1] - varout[d.name, 1]) *
            fraction_of_hour
        )

        for t in time_steps[2:end]
            constraint[d.name, t] = JuMP.@constraint(
                psi_container.JuMPmodel,
                varenergy[d.name, t] ==
                varenergy[d.name, t - 1] +
                (d.multiplier * d.timeseries[t] - varspill[d.name, t] - varout[d.name, t]) * fraction_of_hour
            )
        end
    end
    return
end


@testset "interpolate.interp1d" begin
    x = collect(0:10)
    y = exp.(-x./3.0)
    f = interpolate.interp1d(x, y)
    @test f([0.5, 0.6]) == [0.8582656552868946, 0.8299187863442736]
    @test f(0.5)[1] == 0.8582656552868946
end


fcalls = 0
function f(dx,x)
  global fcalls += 1
  for i in 2:length(x)-1
    dx[i] = x[i-1] - 2x[i] + x[i+1]
  end
  dx[1] = -2x[1] + x[2]
  dx[end] = x[end-1] - 2x[end]
  nothing
end

using SparsityDetection, SparseArrays
input = rand(10)
output = similar(input)
sparsity_pattern = sparsity!(f,output,input)

@test nnz(sparse(sparsity_pattern)) == 28


# --------------------------------------------------------------------------
# ACE.jl and SHIPs.jl: Julia implementation of the Atomic Cluster Expansion
# Copyright (c) 2019 Christoph Ortner <christophortner0@gmail.com>
# All rights reserved.
# --------------------------------------------------------------------------

using LinearAlgebra: cross, norm, dot
using StaticArrays

struct CylindricalCoordinateSystem{T}
   ez::SVector{3, T}
   ex::SVector{3, T}
   ey::SVector{3, T}
   o::SVector{3, T}
end

struct CylindricalCoordinates{T}
   cosθ::T
   sinθ::T
   r::T
   z::T
end

function CylindricalCoordinateSystem(R::AbstractVector{T},
                                     o=zero(SVector{3,T})) where {T}
   @assert length(R) == 3
   # first coordinate
   ez = R / norm(R)
   # second coordinate
   if abs(ez[1] - 1) > 0.33
      ex = @SVector T[1, 0, 0]
   else
      ex = @SVector T[0, 1, 0]
   end
   ex = ex - dot(ez, ex) * ez
   ex /= norm(ex)
   # third coordinate
   ey = cross(ex, ez)
   return CylindricalCoordinateSystem(ez, ex, ey, o)
end

(C::CylindricalCoordinateSystem)(R) = cylindrical(C, R)

function cylindrical(C::CylindricalCoordinateSystem{T},
                     R::AbstractVector{T}) where {T}
   @assert length(R) == 3
   Rc = R - C.o
   x = dot(Rc, C.ex)
   y = dot(Rc, C.ey)
   z = dot(Rc, C.ez)
   r = sqrt(x^2 + y^2)
   cosθ = x / r
   sinθ = y / r
   return CylindricalCoordinates(cosθ, sinθ, r, z)
end

cartesian(C::CylindricalCoordinateSystem{T},
          c::CylindricalCoordinates{T}) where {T} =
      C.o + (c.r * c.cosθ * C.ex + c.r * c.sinθ * C.ey + c.z * C.ez)


using Test
using StaticArrays
using IntervalArithmetic
# using Revise
using HDGElasticity

function allapprox(v1,v2)
    return all(v1 .≈ v2)
end

@test HDGElasticity.linear_map_slope(1.,2.,[1.,2.],[3.,4.]) ≈ [2.,2.]

map = HDGElasticity.LineMap(0.,1.,[0.,0.],[1.,1.])
@test map.xiL ≈ 0.0
@test map.xiR ≈ 1.0
@test allapprox(map.xL,[0.,0.])
@test allapprox(map.xR,[1.,1.])
@test allapprox(map.slope,[1.,1.])
@test HDGElasticity.dimension(map) == 2
@test allapprox(map(0.),[0.,0.])
@test allapprox(map(0.5),[0.5,0.5])

map = HDGElasticity.LineMap([0.5,0.5],[1.,0.5])
@test map.xiL ≈ -1.0
@test map.xiR ≈ 1.0
@test allapprox(map.xL,[0.5,0.5])
@test allapprox(map.xR,[1.0,0.5])
@test allapprox(map.slope,[0.5,0.0]/2)

map = HDGElasticity.LineMap([0.,0.],[1.,1.]*sqrt(2))
@test allapprox(map(-1),[0.,0.])
@test allapprox(map(1),[1,1]*sqrt(2))
@test allapprox(HDGElasticity.jacobian(map),[1,1]/sqrt(2))
@test allapprox(HDGElasticity.determinant_jacobian(map),1.)
@test allapprox(HDGElasticity.inverse_jacobian(map),[1,1]*sqrt(2))

xiL = [0.,0.]
xiR = [1.,1.]
xL = [0.,0.]
xR = [2.,1.]
map = HDGElasticity.CellMap(xiL,xiR,xL,xR)
@test allapprox(map.xiL,xiL)
@test allapprox(map.xiR,xiR)
@test allapprox(map.xL,xL)
@test allapprox(map.xR,xR)
@test HDGElasticity.dimension(map) == 2
@test allapprox(map([0.5,0.5]),[1.,0.5])
@test allapprox(map([0.25,0.75]),[0.5,0.75])
@test allapprox(HDGElasticity.jacobian(map),[2.,1.])
@test allapprox(HDGElasticity.inverse_jacobian(map),[0.5,1.])
@test HDGElasticity.determinant_jacobian(map) ≈ 2.0
@test allapprox(HDGElasticity.face_determinant_jacobian(map),[2.,1.,2.,1.])

HDGElasticity.update_range!(map,[1.,1.],[2.,2.])
@test allapprox(map([0.5,0.5]),[1.5,1.5])
@test allapprox(map([0.75,0.25]),[1.75,1.25])
@test allapprox(HDGElasticity.jacobian(map),[1.,1.])
@test HDGElasticity.determinant_jacobian(map) ≈ 1.0

map = HDGElasticity.CellMap(2)
@test allapprox(map.xiL,[-1.,-1.])
@test allapprox(map.xiR,[+1.,+1.])
@test allapprox(map.xL,[-1.,-1.])
@test allapprox(map.xR,[+1.,+1.])
@test allapprox(map.slope,[1.,1.])
HDGElasticity.update_range!(map,[0.,0.],[1.,1.])
@test allapprox(map.xL,[0.,0.])
@test allapprox(map.xR,[1.,1.])
@test allapprox(map.slope,[0.5,0.5])

box = IntervalBox(1..2,3..5)
HDGElasticity.update_range!(map,box)
@test allapprox(map.xL,[1.,3.])
@test allapprox(map.xR,[2.,5.])
@test allapprox(map.slope,[0.5,1.])

map = HDGElasticity.CellMap([1.,2.],[2.,4.])
@test allapprox(map.xiL,[-1.,-1.])
@test allapprox(map.xiR,[1.,1.])
@test allapprox(map.xL,[1.,2.])
@test allapprox(map.xR,[2.,4.])
@test allapprox(HDGElasticity.jacobian(map),[0.5,1.])
@test HDGElasticity.determinant_jacobian(map) ≈ 0.5

xi = IntervalBox(-1..1,2)
x = IntervalBox(0..1,2)
map = HDGElasticity.CellMap(xi,x)
@test allapprox(map.xiL,[-1.,-1.])
@test allapprox(map.xiR,[1.,1.])
@test allapprox(map.xL,[0.,0.])
@test allapprox(map.xR,[1.,1.])

map = HDGElasticity.CellMap(IntervalBox(0..2,1..2))
@test allapprox(map.xiL,[-1.,-1.])
@test allapprox(map.xiR,[1.,1.])
@test allapprox(map.xL,[0.,1.])
@test allapprox(map.xR,[2.,2.])

fmaps = HDGElasticity.reference_cell_facemaps(2)
@test allapprox(fmaps[1](-1.0),[-1.,-1.])
@test allapprox(fmaps[1](+1.0),[+1.,-1.])

@test allapprox(fmaps[2](-1.0),[+1.,-1.])
@test allapprox(fmaps[2](+1.0),[+1.,+1.])

@test allapprox(fmaps[3](-1.0),[-1.,+1.])
@test allapprox(fmaps[3](+1.0),[+1.,+1.])

@test allapprox(fmaps[4](-1.0),[-1.,-1.])
@test allapprox(fmaps[4](+1.0),[-1.,+1.])


# Unit Tests
using Test
using MLKernels

using MLKernels:
    BaseFunction,
        InnerProduct,
            ScalarProduct,
        PreMetric,
            Metric,
                SquaredEuclidean

using SpecialFunctions:
    besselk, gamma

import LinearAlgebra
import Statistics

const FloatingPointTypes = (Float32, Float64)
const IntegerTypes = (Int32, Int64)
const MLK = MLKernels

include("reference.jl")

@testset "Testing utility functions" begin
    include("utils.jl")
end

@testset "Testing BaseFunction types" begin
    include("basefunctions.jl")
end

@testset "Testing base function evaluation" begin
    include("basematrix.jl")
end

@testset "Testing Kernel types" begin
    include("kernelfunctions.jl")
end

@testset "Testing kernel function evaluation" begin
    include("kernelmatrix.jl")
end

@testset "Testing nystrom" begin
    include("nystrom.jl")
end


using MAT
using Dates
using GLPK
using JuMP
using Plots
using Plots.PlotMeasures
include("../utils.jl")

function get_wave_t_in_range(param, min, max)
    indices_to_keep = findall(>=(min), param)
    param = param[indices_to_keep]

    indices_to_keep = findall(<(max), param)
    param = param[indices_to_keep]
    param, indices_to_keep
end

function fit_poly_inf_norm(x,y,n)
	m = Model(GLPK.Optimizer)

    @variable(m, alpha[1:n+1])
	@variable(m, t >= 0.)
	@constraint(m, [i=1:length(y)], sum(alpha[j]*x[i]^(j-1) for j in 1:n+1) - y[i] <= t);
	@constraint(m, [i=1:length(y)], -t <= sum(alpha[j]*x[i]^(j-1) for j in 1:n+1) - y[i]);
	@objective(m, Min, t)

	optimize!(m) # added new 2020-04-08

    return JuMP.value.(alpha)
end

f2(x, args) = (args[2] * x) .+ args[1]

const MATLAB_EPOCH = Dates.DateTime(-1,12,31)

date2num(d::Dates.DateTime) = Dates.value(d-MATLAB_EPOCH)/(1000*60*60*24)
num2date(n::Number) =  MATLAB_EPOCH + Dates.Millisecond(round(Int64, n*1000*60*60*24))

function plot_bathy(xs, profile; show_points=false, x_last_positive=NaN, x_first_negative=NaN, value_last_positive=NaN, value_first_negative=NaN, x_at_zero=NaN, title=NaN)
    p1 = plot(xs, profile, label="profile_t0")
    if show_points
        scatter!([x_last_positive, x_first_negative], [value_last_positive, value_first_negative], label="A, B")
        scatter!([x_at_zero], [0], label="Shoreline")
    end
    display(plot(p1, title=title))
end

function find_x_at_zero(profile, xs, i)
    t0_index_last_positive = findlast(profile.>0)
    if t0_index_last_positive == nothing
        println(i, ": t0_index_last_positive")
        t0_index_last_positive = 1
    end
    # if t0_index_last_positive == nothing t0_index_last_positive = 1 end

    t0_index_first_negative = findfirst(profile.<0)
    if t0_index_first_negative == nothing
        println(i, ": t0_index_first_negative")
        t0_index_first_negative = 1
    end

    t0_x_last_positive = xs[t0_index_last_positive]
    t0_x_first_negative = xs[t0_index_first_negative]
    t0_value_last_positive = profile[t0_index_last_positive]
    t0_value_first_negative = profile[t0_index_first_negative]

    t0_coeffs = fit_poly_inf_norm([t0_value_last_positive, t0_value_first_negative], [t0_x_last_positive, t0_x_first_negative], 1)
    t0_x_at_zero = f2([0], t0_coeffs)[1]
    t0_x_at_zero
end

dir_path = get_dataset_path(;dataset_name="grand_popo")
kalman_bathy_vars = matread(string(dir_path, "Kalman_bathy.mat"))
bathy_xb = kalman_bathy_vars["xb"][1, :]# 1×213 Matrix{Float64}
# bathy_tn = kalman_bathy_vars["tn"]      # 1×1214 Matrix{Float64}
bathy_hbK = kalman_bathy_vars["hbK"]    # 213×1214 Matrix{Float64}

min_target_matdate = minimum(filter(!isnan, kalman_bathy_vars["tn"]))
max_target_matdate = maximum(filter(!isnan, kalman_bathy_vars["tn"]))

wave_vars = matread(string(dir_path, "Wave_GPP_ERA.mat"))
wave_t = wave_vars["t"]
wave_Tm = wave_vars["Tm"]
wave_hs = wave_vars["hs"]
wave_dir = wave_vars["dir"]

wave_t, indices = get_wave_t_in_range(wave_t, min_target_matdate, max_target_matdate)
wave_Tm = wave_Tm[indices]
wave_hs = wave_hs[indices]
wave_dir = wave_dir[indices]

bathy_dates = map(num2date, collect(min_target_matdate:max_target_matdate))  # collect() instead of using bathy_tn because bathy_tn includes NaN's, even though bathy_hbK has no NaN
wave_dates = map(num2date, wave_t)

changes = []

for i in 2:length(bathy_dates)
    _d0 = bathy_dates[i-1]
    _d1 = bathy_dates[i]

    _t0 = date2num(_d0)
    _t1 = date2num(_d1)

    _bathy_t0 = bathy_hbK[:, i-1]
    _bathy_t1 = bathy_hbK[:, i]

    _wave_t = deepcopy(wave_t)
    _wave_Tm = deepcopy(wave_Tm)
    _wave_dir = deepcopy(wave_dir)
    _wave_hs = deepcopy(wave_hs)

    _wave_t, _indices = get_wave_t_in_range(_wave_t, _t0, _t1)
    _wave_Tm = _wave_Tm[_indices]
    _wave_hs = _wave_hs[_indices]
    _wave_dir = _wave_dir[_indices]

    t0_x_at_zero = find_x_at_zero(_bathy_t0, bathy_xb, i)
    t1_x_at_zero = find_x_at_zero(_bathy_t1, bathy_xb, i)
    push!(changes, t1_x_at_zero - t0_x_at_zero)

end

p1 = plot(bathy_dates[2:end], changes, ylabel="Shoreline change", legend=false)
p2 = plot(wave_dates, wave_Tm, ylabel="Tm", legend=false)
p3 = plot(wave_dates, wave_hs, ylabel="Hs", legend=false)
p4 = plot(wave_dates, wave_dir, ylabel="Dir", legend=false)

full_plot = plot(p1, p2, p3, p4, layout=(4, 1), legend=false)#, title="Grand Popo shoreline change")
plot!(size=(1920, 1080), left_margin=10mm) # 1920, 1080  # 1280, 720
savefig("/Users/mn/JuliaProjects/cgp-morphcast/output/grand_popo/wavesandshorelinechange.pdf")


"""
    STL{T <: Real}
"""
struct STL{T <: Real}
    trend::Vector
    seasonal::Vector
    remainder::Vector
end

function nextodd(x::T)::Integer where T <: Real
    cx = ceil(Integer, x)
    return isodd(cx) ? cx+1 : cx
end

function sma(x, n::Integer; center::Bool = true)
    N = length(x)
    @assert 1 <= n & n <= N
    if n == 1
        return x
    end
    res = Vector{Any}(missing,N)

    # initial and final value positions
    ivp = findfirst(!ismissing, x)
    fvp = findlast(!ismissing, x)

    # using missing values to center ma
    a = center ? div(n,2) : 0

    # initial moving average value
    ma = sum(x[ivp:n+ivp-1])/n
    res[a+ivp] = ma
    for i in 1:(N-n-ivp-(N-fvp)+1)
        resai = ma + (x[n+ivp+i-1] - x[ivp+i-1])/n
        # missing values are imputed
        res[a+ivp+i] = ismissing(resai) ? ma : resai
        ma = res[a+ivp+i]
    end
    return res
end

"""
    stl(Yv, np; robust=false,
                nl=nextodd(np),
                ns=10*length(Yv)+1,
                nt=nextodd(1.5*np/(1-1.5/ns)),
                ni=robust ? 1 : 2,
                no=0,
                spm=false,
                qsmp=max(div(np,7),2))

Decompose a time series into trend, seasonal, and remainder components.
"STL has a simple design that consists of a sequence of applications of the loess smoother; the simplicity
allows analysis of the properties of the procedure and allows fast computation, even for very long time series
and large amounts of trend and seasonal smoothing. Other features of STL  are specification of amounts of seasonal
and trend smoothing that range, in a nearly continous way, from very small amount of smoothing to a very large
amount; robust estimates of the trend and seasonal components that are not distorted by aberrant behavior in
the data; specification of the period of the seasonal component to any intenger multiple of the time sampling
interval greater than one; and the ability to decompose time series with missing values."*

All default values are chosen following the recommendations of the original paper when those were recommended.
`ns` is recommended to be chosen of the basis of knowledge of the time series and on the basis of diagnostic
methods; it must nonethelessbe always odd and at least 7. A default value is not advised on the original
paper, instead the same default value used in the stl implementation in R in usere here.
for `no` the authors advise 5 ("safe value") or 10 ("near certainty of convergence") cycles or a
convergence criterion when robustness is required, in this case when `robust` is true computations
stop when convergence is achieved in trend and seasonality.
for `qsmp` the authors do not adivise a default but they use a value close to div(`np`,7).

# Arguments
- `np`: Seasonality.
- `robust`: Robust stl.
- `nl`: Smoothing parameter of the low-pass filter.
- `ns`: Smoothing parameter for the seasonal component.
- `nt`: Smoothing parameter for the trend decomposition.
- `ni`: Number of inner loop cycles.
- `no`: Number of outer loop cycles.
- `spm`: Seasonal post-smoothing.
- `qsmp`: Loess q window for Seasonal post-smoothing.
- `verbose`: If true shows updates for the Seasonal and Trend convergence.
- `cth`: Corvengence threshold for Seasonal and Trend.

# Returns
An `STL` object with the seasonal, trend and remainder components.
* STL: A Seasonal, Trend Decomposition Procedure Based on Loess" Robert B. Cleveland, William S. Cleveland, Jean E. McRae, and Irma Terpenning. Journal of Official Statistics Vol. 6. No. 1, 1990, pp. 3-73 (c) Statistics Sweden.

# Examples
```julia-repl
julia> stl_co2 = stl(co2(),365; robust=true, spm=true)
[ Info: Dataset used in Cleveland et al. paper
[ Info: Corvengence achieved (< 0.01); Stopping computation...
STL{TimeSeries.TimeArray{Union{Missing, Float64},2,
Dates.Date,Array{Union{Missing, Float64},2}}}
(4609×3 TimeSeries.TimeArray{Union{Missing, Float64},2,
Dates.Date,Array{Union{Missing, Float64},2}} 1974-05-17 to 1986-12-31,
"stl(Yn, np=365; nl=365, ns=46091, nt=549, ni=1, no=0, spm=true, qsmp=52)")
julia> plot(stl_co2)
```
"""
function stl(Yv::Vector{T},
             np::Integer;
             robust::Bool=false,
             nl::Integer=nextodd(np),
             ns::Integer=10*length(Yv)+1,
             nt::Integer=nextodd(1.5*np/(1-1.5/ns)),
             ni::Integer=robust ? 1 : 2,
             no::Integer=robust ? 15 : 0,
             spm::Bool=false,
             qsmp::Integer=max(div(np,7),2),
             verbose::Bool=false,
             cth::Float64 = 0.01) where T

    @assert isodd(ns) & (ns>=7) "`ns` is chosen of the basis of knowledge of the time series and on the basis of diagnostic methods; must always be odd and at least 7"

    function B(u)
        (u < 1) & !ismissing(u) ? (1.0-u^2)^2 : 0.0
    end

    N = length(Yv)
    # initial robustness weights
    rhov = ones(N)
    # intial trend
    Tv = Tv0 = zeros(N)
    Sv = Sv0 = zeros(N)
    Rv = Array{Float64,1}(undef,N)
    Cv = Array{Float64,1}(undef,N+2*np)
    scnv = false # seasonal convergence flag
    tcnv = false # trend convergence flag
    o = 0
    while robust | (o <= no)
        for k in 1:ni
            # Updating sesonal and trend components
            ## 1. Detrending (Yv = Tv + Sv)
            Sv = Yv - Tv

            ### Seasonal convergence criterion
            Md = maximum(abs.(skipmissing(Sv-Sv0)))
            M0 = maximum(skipmissing(Sv0))
            m0 = minimum(skipmissing(Sv0))
            scnv = (Md/(M0-m0) < cth)
            if verbose
                println("Outer loop: " * string(o) * " - " * "Inner loop: " * string(k))
                println("Seasonal Convergence: " * string(Md/(M0-m0)))
            end
            Sv0 = Sv

            # Seasonal Smoothing 2,3,4
            ## 2. Cycle-subseries Smoothing
            for csi in 1:np
                Cv[csi:np:N+2*np] = loess(csi:np:N,
                                          Sv[csi:np:N];
                                          q=ns,d=1,rho=rhov[csi:np:N],
                                          predict=(1.0*csi-np):np:(N+np))
            end
            ## 3. Low-Pass Filtering of Smoothed Cycle-Subseries
            ### centered support instead 1:N to balance out machine error
            Lv = loess(1-ceil(N/2):N-ceil(N/2),
                       collect(skipmissing(sma(sma(sma(Cv,np),np),3))),
                       d=1,q=nl,rho=rhov)
            ## 4. Detreending of Smoothed Cycle-Subseries
            ### Lv is substracted to prevent low-frenquency power
            ### from entering the seasonal component.
            Sv = Cv[np+1:end-np] - Lv

            ## 5. Deseasonalizing
            Dv = Yv - Sv

            # Trend Smoothing
            ## 6. Trend Smoothing
            ### centered support instead 1:N to balance out machine error
            ### (floor isntead ceil like in Lv to balance out even lengths)
            Tv = loess(1-floor(N/2):N-floor(N/2),Dv,q=nt,d=1,rho=rhov)

            ### Trend convergence criterion
            Md = maximum(abs.(skipmissing(Tv-Tv0)))
            M0 = maximum(skipmissing(Tv0))
            m0 = minimum(skipmissing(Tv0))
            tcnv = (Md/(M0-m0) < cth)
            if verbose
                println("Trend    Convergence: " * string(Md/(M0-m0)) * "\n")
            end
            Tv0 = Tv

            scnv & tcnv ? break : nothing
        end
        # Computation of robustness weights
        ## These new weights will reduce the influence of transient behaviour
        ## on the trend and seasonal components.

        # Tv and Sv defined everywhere but Rv is not defined
        # where Yv has missing values

        Rv = Yv - Tv - Sv

        if scnv & tcnv
            verbose &&  println("Corvengence achieved (< " * string(cth) * "); Stopping computation...")
            break
        end

        if 0 < o <= no
            smRv = skipmissing(Rv)
            h = 6*median(abs.(smRv))
            rhov = B.(abs.(Rv)/h)
        end
        o += 1
    end

    if spm
        # Using div(np,7) as default to approximate
        # the q=51 for a np=365 chosen in the original paper
        Sv = loess(1-ceil(N/2):N-ceil(N/2),Sv,q=qsmp)
        Rv = Yv - Tv - Sv
    end

    if !scnv
        @warn "Seasonal convergence not achieved (>= " * string(cth) * ")
         Consider a robust estimation"
    end

    if !tcnv
        @warn "Trend convergence not achieved (>= " * string(cth) * ")
         Consider a robust estimation"
    end

    call = "stl(Yn, np="*string(np)*
        "; nl="*string(nl)*
        ", ns="*string(ns)*
        ", nt="*string(nt)*
        ", ni="*string(ni)*
        ", no="*string(no)*
        ", spm="*string(spm)*
        ", qsmp="*string(qsmp)*")"

    return STL{T}(Tv, Sv, Rv)
end


using Gadfly

set_default_plot_size(20cm, 15cm)

#     A spiking neuron model based on:
#     Izhikevich, E. M.
#     Simple Model of Spiking Neurons
#     IEEE TRANSACTIONS ON NEURAL NETWORKS, VOL. 14, NO. 6, NOVEMBER 2003

#         a, b, c, d are the parameters of this model.
#         a: the time scale of the recovery variable.
#         b: the sensitivity of the recovery variable.
#         c: the after-spike reset value of the membrane potential.
#         d: after-spike reset of the recovery variable.

type IZNeuron
    a::Float64
    b::Float64
    c::Float64
    d::Float64
    v::Float64 # membrane potential
    u::Float64 # membrane recovery variable
    has_fired::Bool
    bias::Float64
    current::Float64
    function IZNeuron(bias::Float64=0; a::Float64=0.02,b::Float64=0.2, c::Float64=-65.0, d::Float64=8.0)
       new(a, b, c, d, c, b*c, false, bias, bias)
    end
end

function update!(n::IZNeuron)
#         Advances time in 1 ms.
    n.v += 0.5 * (0.04 * n.v ^ 2. + 5. * n.v + 140. - n.u + n.current)
    n.v += 0.5 * (0.04 * n.v ^ 2. + 5. * n.v + 140. - n.u + n.current)
    n.u += n.a * (n.b * n.v - n.u)
    if n.v > 30
        n.has_fired = true
        n.v = n.c
        n.u += n.d
    else
        n.has_fired = false
        n.current = n.bias
    end
    return n.v
end

function reset!(n::IZNeuron)
#         'Resets all state variables.'
    n.v = n.c
    n.u = n.b * n.v
    n.has_fired = false
    n.current = n.bias
end

function getResponse(V::Vector{Float64}, s::IZNeuron)
    n = length(V)
    data = zeros(n)
    tm = zeros(Int64,n)
    for i = 1:n
        s.current = V[i]
        v = update!(s)
        data[i] = v
        tm[i] = i == 1? 1 : tm[i] = tm[i-1] + 1
    end
    return data,tm
end

# The following parameters produce some known spiking behaviors:
# Regular spiking:
s = IZNeuron(0.,a = 0.02, b = 0.2, c = -65., d = 8.)
# Intrinsically bursting:
s = IZNeuron(0.,a = 0.02, b = 0.2, c = -55., d = 4.)
# Chattering:
s = IZNeuron(0.,a = 0.02, b = 0.2, c = -50., d = 2.)
# Fast spiking:
s = IZNeuron(0.,a = 0.1, b = 0.2, c = -65., d = 2.)
# Thalamo-/cortical:
s = IZNeuron(0.,a = 0.02, b = 0.25, c = -65., d = 0.05)
# Resonator:
s = IZNeuron(0.,a = 0.1, b = 0.25, c = -65., d = 2.)
# Low-threshold spiking:
s = IZNeuron(0.,a = 0.02, b = 0.25, c = -65., d = 2.)

n = 2000
V = ones(n) * 5
response, tm = getResponse(V, s)
vstack(plot(x = tm, y=V, Geom.line), plot(x=tm, y=response, Geom.line))

response[1:10]
tm[1:10]

Ne = 800;                     Ni = 200;
re = rand(Ne);                ri = rand(Ni);
a = [0.02 * ones(Ne),         0.02 + 0.08 * ri]
b = [0.2  * ones(Ne),         0.25 - 0.05 * ri]
c = [-65. + 15.*re.^2,        -65. * ones(Ni)]
d = [  8. - 6. *re.^2,        2. * ones(Ni)]
S = hcat(0.5*rand(Ne+Ni,Ne),  -rand(Ne+Ni,Ni))

v =  -65*ones(Ne+Ni)
u =  b .* v

firings = zeros(Ne+Ni,1000)

for t =  1:1000
    I = [5.*randn(Ne), 2.*randn(Ni)]
    fired = find(x -> x >= 30,v)
    firings[fired,t] = v[fired]
    v[fired] = c[fired]
    u[fired] = u[fired] + d[fired]
    I = I + sum(S[:,fired],2)
    v = v .+ 0.5 * (0.04 * v .^ 2. + 5. .* v + 140. .- u .+ I)
    v = v .+ 0.5 * (0.04 * v .^ 2. + 5. .* v + 140. .- u .+ I)
    u = u .+ a .* (b .* v .- u)
end
firings
# plot(x=[1:1000], y=firings[3,:], Geom.line)
spy(firings)

# class Synapse:
#     """ A synapse indicates the connection strength between two neurons (or itself) """
#     def __init__(self, source, dest, weight):
#         self.__weight = weight
#         self.__source = source
#         self.__dest = dest

#     def advance(self):
#         'Advances time in 1 ms.'
#         if self.__source.has_fired:
#             self.__dest.current += self.__weight # dest.current or dest.__v ?


#-----------------------------------------------------------------------# Counter
"""
    Counter(T=Number)

Count the number of items in a data stream with elements of type `T`.

# Example

    fit!(Counter(Int), 1:100)
"""
mutable struct Counter{T} <: OnlineStat{T}
    n::Int
    Counter{T}() where {T} = new{T}(0)
end
Counter(T = Number) = Counter{T}()
_fit!(o::Counter{T}, y) where {T} = (o.n += 1)
_merge!(a::Counter, b::Counter) = (a.n += b.n)

#-----------------------------------------------------------------------# CountMap
"""
    CountMap(T::Type)
    CountMap(dict::AbstractDict{T, Int})

Track a dictionary that maps unique values to its number of occurrences.  Similar to
`StatsBase.countmap`.

# Example

    o = fit!(CountMap(Int), rand(1:10, 1000))
    value(o)
    probs(o)
    OnlineStats.pdf(o, 1)
    collect(keys(o))
"""
mutable struct CountMap{T, A <: AbstractDict{T, Int}} <: OnlineStat{T}
    value::A  # OrderedDict by default
    n::Int
end
CountMap{T}() where {T} = CountMap{T, OrderedDict{T,Int}}(OrderedDict{T,Int}(), 0)
CountMap(T::Type = Any) = CountMap{T, OrderedDict{T,Int}}(OrderedDict{T, Int}(), 0)
CountMap(d::D) where {T,D<:AbstractDict{T, Int}} = CountMap{T, D}(d, 0)
function _fit!(o::CountMap, x)
    o.n += 1
    o.value[x] = get!(o.value, x, 0) + 1
end
_merge!(o::CountMap, o2::CountMap) = (merge!(+, o.value, o2.value); o.n += o2.n)
function probs(o::CountMap, kys = keys(o.value))
    out = zeros(Int, length(kys))
    valkeys = keys(o.value)
    for (i, k) in enumerate(kys)
        out[i] = k in valkeys ? o.value[k] : 0
    end
    sum(out) == 0 ? Float64.(out) : out ./ sum(out)
end
pdf(o::CountMap, y) = y in keys(o.value) ? o.value[y] / nobs(o) : 0.0
Base.keys(o::CountMap) = keys(o.value)
nkeys(o::CountMap) = length(o.value)
Base.values(o::CountMap) = values(o.value)
Base.getindex(o::CountMap, i) = o.value[i]

#-----------------------------------------------------------------------------# CountMissing
"""
    CountMissing(stat)

Calculate a `stat` along with the count of `missing` values.

# Example

    o = CountMissing(Mean())
    fit!(o, [1, missing, 3])
"""
mutable struct CountMissing{T, O<:OnlineStat{T}} <: OnlineStat{Union{Missing,T}}
    stat::O
    nmissing::Int
end
CountMissing(stat::OnlineStat) = CountMissing(stat, 0)
value(o::CountMissing) = (nmissing=o.nmissing, stat=o.stat)
nobs(o::CountMissing) = nobs(o.stat) + o.nmissing

_fit!(o::CountMissing, x) = _fit!(o.stat, x)
_fit!(o::CountMissing, ::Missing) = (o.nmissing += 1)

_merge!(a::CountMissing, b::CountMissing) = (merge!(a.stat, b.stat); a.nmissing += b.nmissing)

#-----------------------------------------------------------------------# CovMatrix
"""
    CovMatrix(p=0; weight=EqualWeight())
    CovMatrix(::Type{T}, p=0; weight=EqualWeight())

Calculate a covariance/correlation matrix of `p` variables.  If the number of variables is
unknown, leave the default `p=0`.

# Example

    o = fit!(CovMatrix(), randn(100, 4) |> eachrow)
    cor(o)
    cov(o)
    mean(o)
    var(o)
"""
mutable struct CovMatrix{T,W} <: OnlineStat{Union{Tuple, NamedTuple, AbstractVector}} where T<:Number
    value::Matrix{T}
    A::Matrix{T}  # x'x/n
    b::Vector{T}  # 1'x/n
    weight::W
    n::Int
end
function CovMatrix(::Type{T}, p::Int=0; weight = EqualWeight()) where T<:Number
    CovMatrix(zeros(T,p,p), zeros(T,p,p), zeros(T,p), weight, 0)
end
CovMatrix(p::Int=0; weight = EqualWeight()) = CovMatrix(zeros(p,p), zeros(p,p), zeros(p), weight, 0)
function _fit!(o::CovMatrix{T}, x) where {T}
    γ = o.weight(o.n += 1)
    if isempty(o.A)
        p = length(x)
        o.b = zeros(T, p)
        o.A = zeros(T, p, p)
        o.value = zeros(T, p, p)
    end
    smooth!(o.b, x, γ)
    smooth_syr!(o.A, x, γ)
end
nvars(o::CovMatrix) = size(o.A, 1)
function value(o::CovMatrix; corrected::Bool = true)
    o.value[:] = Matrix(Hermitian((o.A - o.b * o.b')))
    corrected && rmul!(o.value, bessel(o))
    o.value
end
function _merge!(o::CovMatrix, o2::CovMatrix)
    γ = o2.n / (o.n += o2.n)
    smooth!(o.A, o2.A, γ)
    smooth!(o.b, o2.b, γ)
end
Statistics.cov(o::CovMatrix; corrected::Bool = true) = value(o; corrected=corrected)
Statistics.mean(o::CovMatrix) = o.b
Statistics.var(o::CovMatrix; kw...) = diag(value(o; kw...))
function Statistics.cor(o::CovMatrix; kw...)
    value(o; kw...)
    v = 1.0 ./ sqrt.(diag(o.value))
    rmul!(o.value, Diagonal(v))
    lmul!(Diagonal(v), o.value)
    o.value
end

#-----------------------------------------------------------------------# Extrema
"""
    Extrema(T::Type = Float64)

Maximum and minimum (and number of occurrences for each) for a data stream of type `T`.

# Example

    o = fit!(Extrema(), rand(10^5))
    extrema(o)
    maximum(o)
    minimum(o)
"""
# T is type to store data, S is type of single observation.
# E.g. you may want to accept any Number even if you are storing values as Float64
mutable struct Extrema{T,S} <: OnlineStat{S}
    min::T
    max::T
    nmin::Int
    nmax::Int
    n::Int
end
function Extrema(T::Type = Float64)
    a, b, S = extrema_init(T)
    Extrema{T,S}(a, b, 0, 0, 0)
end
extrema_init(T::Type{<:Number}) = typemax(T), typemin(T), Number
extrema_init(T::Type{String}) = "", "", String
extrema_init(T::Type{Date}) = typemax(Date), typemin(Date), Union{Date, Dates.AbstractDateTime}
extrema_init(T::Type) = rand(T), rand(T), T
function _fit!(o::Extrema, y)
    (o.n += 1) == 1 && (o.min = o.max = y)
    if y < o.min
        o.min = y
        o.nmin = 0
    elseif y > o.max
        o.max = y
        o.nmax = 0
    end
    y == o.min && (o.nmin += 1)
    y == o.max && (o.nmax += 1)
end
function _merge!(a::Extrema, b::Extrema)
    if a.min == b.min
        a.nmin += b.nmin
    elseif b.min < a.min
        a.min = b.min
        a.nmin = b.nmin
    end
    if a.max == b.max
        a.nmax += b.nmax
    elseif b.max > a.max
        a.max = b.max
        a.nmax = b.nmax
    end
    a.n += b.n
end
value(o::Extrema) = (min=o.min, max=o.max, nmin=o.nmin, nmax=o.nmax)
Base.extrema(o::Extrema) = (o.min, o.max)
Base.maximum(o::Extrema) = o.max
Base.minimum(o::Extrema) = o.min

#-----------------------------------------------------------------------# Group
"""
    Group(stats::OnlineStat...)
    Group(; stats...)
    Group(collection)

Create a vector-input stat from several scalar-input stats.  For a new
observation `y`, `y[i]` is sent to `stats[i]`.

# Examples

    x = randn(100, 2)

    fit!(Group(Mean(), Mean()), eachrow(x))
    fit!(Group(Mean(), Variance()), eachrow(x))

    o = fit!(Group(m1 = Mean(), m2 = Mean()), eachrow(x))
    o.stats.m1
    o.stats.m2
"""
struct Group{T, S} <: StatCollection{S}
    stats::T
    function Group(stats::T) where {T}
        inputs = map(input, stats)
        tup = Tuple{inputs...}
        S = Union{tup, NamedTuple{names, R} where R<:tup, AbstractVector{<: promote_type(inputs...)}} where names
        new{T,S}(stats)
    end
end
Group(o::OnlineStat...) = Group(o)
Group(;o...) = Group(o.data)
nobs(o::Group) = nobs(first(o.stats))
Base.:(==)(a::Group, b::Group) = all(a.stats .== b.stats)

Base.getindex(o::Group, i) = o.stats[i]
Base.first(o::Group) = first(o.stats)
Base.last(o::Group) = last(o.stats)
Base.lastindex(o::Group) = length(o)
Base.length(o::Group) = length(o.stats)
Base.values(o::Group) = map(value, o.stats)

Base.iterate(o::Group) = (o.stats[1], 2)
Base.iterate(o::Group, i) = i > length(o) ? nothing : (o.stats[i], i + 1)

@generated function _fit!(o::Group{T}, y) where {T}
    N = fieldcount(T)
    :(Base.Cartesian.@nexprs $N i -> @inbounds(_fit!(o.stats[i], y[i])))
end
function _fit!(o::Group{T}, y) where {T<:AbstractVector}
    for (i,yi) in enumerate(y)
        _fit!(o.stats[i], yi)
    end
end

_merge!(o::Group, o2::Group) = map(merge!, o.stats, o2.stats)

Base.:*(n::Integer, o::OnlineStat) = Group([copy(o) for i in 1:n]...)

#-----------------------------------------------------------------------# GroupBy
"""
    GroupBy(T, stat)

Update `stat` for each group (of type `T`).  A single observation is either a (named)
tuple with two elements or a Pair.

# Example

    x = rand(Bool, 10^5)
    y = x .+ randn(10^5)
    fit!(GroupBy(Bool, Series(Mean(), Extrema())), zip(x,y))
"""
mutable struct GroupBy{T, S, O <: OnlineStat{S}} <: StatCollection{TwoThings{T,S}}
    value::OrderedDict{T, O}
    init::O
    n::Int
    function GroupBy(value::OrderedDict{T,O}, init::O, n::Int) where {T,S,O<:OnlineStat{S}}
        new{T,S,O}(value, init, n)
    end
end
GroupBy(T::Type, stat::O) where {O<:OnlineStat} = GroupBy(OrderedDict{T, O}(), stat, 0)
function _fit!(o::GroupBy, xy)
    o.n += 1
    x, y = xy
    x in keys(o.value) ? fit!(o.value[x], y) : (o.value[x] = fit!(copy(o.init), y))
end
Base.getindex(o::GroupBy{T}, i::T) where {T} = o.value[i]
AbstractTrees.children(o::GroupBy) = collect(o.value)
AbstractTrees.printnode(io::IO, o::GroupBy{T,S,O}) where {T,S,O} = print(io, "GroupBy: $T => $(name(O,false,false))")
Base.sort!(o::GroupBy) = (sort!(o.value); o)
function _merge!(a::GroupBy{T,O}, b::GroupBy{T,O}) where {T,O}
    a.init == b.init || error("Cannot merge GroupBy objects with different inits")
    a.n += b.n
    merge!((o1, o2) -> merge!(o1, o2), a.value, b.value)
end

#-----------------------------------------------------------------------# Mean
"""
    Mean(T = Float64; weight=EqualWeight())

Track a univariate mean, stored as type `T`.

# Example

    @time fit!(Mean(), randn(10^6))
"""
mutable struct Mean{T,W} <: OnlineStat{Number}
    μ::T
    weight::W
    n::Int
end
Mean(T::Type{<:Number} = Float64; weight = EqualWeight()) = Mean(zero(T), weight, 0)
_fit!(o::Mean{T}, x) where {T} = (o.μ = smooth(o.μ, x, T(o.weight(o.n += 1))))
function _merge!(o::Mean, o2::Mean)
    o.n += o2.n
    o.μ = smooth(o.μ, o2.μ, o2.n / o.n)
end
Statistics.mean(o::Mean) = o.μ
Base.copy(o::Mean) = Mean(o.μ, o.weight, o.n)

#-----------------------------------------------------------------------# Moments
"""
    Moments(; weight=EqualWeight())

First four non-central moments.

# Example

    o = fit!(Moments(), randn(1000))
    mean(o)
    var(o)
    std(o)
    skewness(o)
    kurtosis(o)
"""
mutable struct Moments{W} <: OnlineStat{Number}
    m::Vector{Float64}
    weight::W
    n::Int
end
Moments(;weight = EqualWeight()) = Moments(zeros(4), weight, 0)
function _fit!(o::Moments, y::Real)
    γ = o.weight(o.n += 1)
    y2 = y * y
    @inbounds o.m[1] = smooth(o.m[1], y, γ)
    @inbounds o.m[2] = smooth(o.m[2], y2, γ)
    @inbounds o.m[3] = smooth(o.m[3], y * y2, γ)
    @inbounds o.m[4] = smooth(o.m[4], y2 * y2, γ)
end
Statistics.mean(o::Moments) = o.m[1]
function Statistics.var(o::Moments; corrected=true)
    out = (o.m[2] - o.m[1] ^ 2)
    corrected ? bessel(o) * out : out
end
function StatsBase.skewness(o::Moments)
    v = value(o)
    vr = o.m[2] - o.m[1]^2
    (v[3] - 3.0 * v[1] * vr - v[1] ^ 3) / vr ^ 1.5
end
function StatsBase.kurtosis(o::Moments)
    # v = value(o)
    # (v[4] - 4.0 * v[1] * v[3] + 6.0 * v[1] ^ 2 * v[2] - 3.0 * v[1] ^ 4) / var(o) ^ 2 - 3.0
    m1, m2, m3, m4 = value(o)
    (m4 - 4.0 * m1 * m3 + 6.0 * m1^2 * m2 - 3.0 * m1 ^ 4) / var(o; corrected=false) ^ 2 - 3.0
end
function _merge!(o::Moments, o2::Moments)
    γ = o2.n / (o.n += o2.n)
    smooth!(o.m, o2.m, γ)
end

#-----------------------------------------------------------------------# Sum
"""
    Sum(T::Type = Float64)

Track the overall sum.

# Example

    fit!(Sum(Int), fill(1, 100))
"""
mutable struct Sum{T} <: OnlineStat{Number}
    sum::T
    n::Int
end
Sum(T::Type = Float64) = Sum(T(0), 0)
Base.sum(o::Sum) = o.sum
_fit!(o::Sum{T}, x::Real) where {T<:AbstractFloat} = (o.sum += convert(T, x); o.n += 1)
_fit!(o::Sum{T}, x::Real) where {T<:Integer} =       (o.sum += round(T, x); o.n += 1)
_merge!(o::T, o2::T) where {T <: Sum} = (o.sum += o2.sum; o.n += o2.n; o)

#-----------------------------------------------------------------------# Variance
"""
    Variance(T = Float64; weight=EqualWeight())

Univariate variance, tracked as type `T`.

# Example

    o = fit!(Variance(), randn(10^6))
    mean(o)
    var(o)
    std(o)
"""
mutable struct Variance{T, W} <: OnlineStat{Number}
    σ2::T
    μ::T
    weight::W
    n::Int
end
function Variance(T::Type{<:Number} = Float64; weight = EqualWeight())
    Variance(zero(T), zero(T), weight, 0)
end
Base.copy(o::Variance) = Variance(o.σ2, o.μ, o.weight, o.n)
function _fit!(o::Variance{T}, x) where {T}
    μ = o.μ
    γ = T(o.weight(o.n += 1))
    o.μ = smooth(o.μ, T(x), γ)
    o.σ2 = smooth(o.σ2, (T(x) - o.μ) * (T(x) - μ), γ)
end
function _merge!(o::Variance, o2::Variance)
    γ = o2.n / (o.n += o2.n)
    δ = o2.μ - o.μ
    o.σ2 = smooth(o.σ2, o2.σ2, γ) + δ ^ 2 * γ * (1.0 - γ)
    o.μ = smooth(o.μ, o2.μ, γ)
end
value(o::Variance) = o.n > 1 ? o.σ2 * bessel(o) : 1.0
Statistics.var(o::Variance) = value(o)
Statistics.mean(o::Variance) = o.μ

#-----------------------------------------------------------------------# Series
"""
    Series(stats)
    Series(stats...)
    Series(; stats...)

Track a collection stats for one data stream.

# Example

    s = Series(Mean(), Variance())
    fit!(s, randn(1000))
"""
struct Series{IN, T} <: StatCollection{IN}
    stats::T
    Series(stats::T) where {T} = new{Union{map(input, stats)...}, T}(stats)
end
Series(t::OnlineStat...) = Series(t)
Series(; t...) = Series(t.data)

value(o::Series) = map(value, o.stats)
nobs(o::Series) = nobs(o.stats[1])
@generated function _fit!(o::Series{IN, T}, y) where {IN, T}
    n = length(fieldnames(T))
    :(Base.Cartesian.@nexprs $n i -> _fit!(o.stats[i], y))
end
_merge!(o::Series, o2::Series) = map(_merge!, o.stats, o2.stats)

#-----------------------------------------------------------------------# FTSeries
"""
    FTSeries(stats...; filter=x->true, transform=identity)

Track multiple stats for one data stream that is filtered and transformed before being
fitted.

    FTSeries(T, stats...; filter, transform)

Create an FTSeries and specify the type `T` of the pre-transformed values.

# Example

    o = FTSeries(Mean(), Variance(); transform=abs)
    fit!(o, -rand(1000))

    # Remove missing values represented as DataValues
    using DataValues
    y = DataValueArray(randn(100), rand(Bool, 100))
    o = FTSeries(DataValue, Mean(); transform=get, filter=!isna)
    fit!(o, y)

    # Remove missing values represented as Missing
    y = [rand(Bool) ? rand() : missing for i in 1:100]
    o = FTSeries(Mean(); filter=!ismissing)
    fit!(o, y)

    # Alternatively for Missing:
    fit!(Mean(), skipmissing(y))
"""
mutable struct FTSeries{IN, OS, F, T} <: StatCollection{Union{IN,Missing}}
    stats::OS
    filter::F
    transform::T
    nfiltered::Int
end
function FTSeries(stats::OnlineStat...; filter=x->true, transform=identity)
    IN, OS = Union{map(input, stats)...}, typeof(stats)
    FTSeries{IN, OS, typeof(filter), typeof(transform)}(stats, filter, transform, 0)
end
function FTSeries(T::Type, stats::OnlineStat...; filter=x->true, transform=identity)
    FTSeries{T, typeof(stats), typeof(filter), typeof(transform)}(stats, filter, transform, 0)
end
value(o::FTSeries) = value.(o.stats)
nobs(o::FTSeries) = nobs(o.stats[1])
@generated function _fit!(o::FTSeries{N, OS}, y) where {N, OS}
    n = length(fieldnames(OS))
    quote
        if o.filter(y)
            yt = o.transform(y)
            Base.Cartesian.@nexprs $n i -> @inbounds begin
                _fit!(o.stats[i], yt)
            end
        else
            o.nfiltered += 1
        end
    end
end
function _merge!(o::FTSeries, o2::FTSeries)
    o.nfiltered += o2.nfiltered
    _merge!.(o.stats, o2.stats)
end